Update CLI connection loading logic

GitOrigin-RevId: 77091e056e8658071423ce2699b5fdeb15e2f91f

src/gretel_client/cli/connections.py

@@ -1,11 +1,12 @@
+import json
+
+from json import JSONDecodeError
+from pathlib import Path
+
 import click
+import yaml
 
-from gretel_client.cli.common import (
-    parse_file,
-    pass_session,
-    project_option,
-    SessionContext,
-)
+from gretel_client.cli.common import pass_session, SessionContext
 from gretel_client.config import get_session_config
 from gretel_client.rest_v1.api.connections_api import ConnectionsApi
 from gretel_client.rest_v1.models import CreateConnectionRequest
@@ -23,18 +24,34 @@ def get_connections_api() -> ConnectionsApi:
     return get_session_config().get_v1_api(ConnectionsApi)
 
 
+def _read_connection_file(file: str) -> dict:
+    fp = Path(file).resolve()
+    try:
+        with open(fp) as fd:
+            return json.load(fd)
+    except JSONDecodeError:
+        with open(fp) as fd:
+            return yaml.safe_load(fd)
+
+
 @connections.command(help="Create a new connection.")
 @click.option(
     "--from-file",
     metavar="PATH",
     help="Path to the file containing Gretel connection.",
     required=True,
 )
-@project_option
+@click.option(
+    "--project",
+    metavar="PROJECT-ID",
+    help="Specify the project to create the connection in.",
+    required=False,
+)
 @pass_session
 def create(sc: SessionContext, from_file: str, project: str):
     connection_api = get_connections_api()
-    conn = parse_file(from_file)
+
+    conn = _read_connection_file(from_file)
 
     # we try and configure a project in the following order
     #
@@ -53,7 +70,7 @@ def create(sc: SessionContext, from_file: str, project: str):
     if project_id is None:
         project_id = sc.project.project_guid
 
-    conn["project_id"] = sc.project.project_guid
+    conn["project_id"] = project_id
 
     connection = connection_api.create_connection(CreateConnectionRequest(**conn))
 
@@ -71,7 +88,7 @@ def create(sc: SessionContext, from_file: str, project: str):
 @pass_session
 def update(sc: SessionContext, id: str, from_file: str):
     connection_api = get_connections_api()
-    conn = parse_file(from_file)
+    conn = _read_connection_file(from_file)
     connection = connection_api.update_connection(connection_id=id, connection=conn)
 
     sc.log.info("Updated connection:")

tests/gretel_client/fixtures/connections/test_connection.yaml

@@ -0,0 +1,5 @@
+type: test
+name: integration_test
+credentials:
+    name: test
+    secret: secret

tests/gretel_client/integration/test_cli_connections.py

@@ -2,9 +2,12 @@
 
 from typing import Callable
 
+import pytest
+
 from click.testing import CliRunner
 
 from gretel_client.cli.cli import cli
+from gretel_client.cli.connections import _read_connection_file
 from gretel_client.projects.projects import Project
 
 
@@ -85,3 +88,14 @@ def test_connection_crud_from_cli(get_fixture: Callable, project: Project):
     )
     assert "Deleted connection " + connection_result["id"] in cmd.output
     assert cmd.exit_code == 0
+
+
+def test_read_yaml(get_fixture: Callable):
+    json_conn = _read_connection_file(get_fixture("connections/test_connection.json"))
+    yaml_conn = _read_connection_file(get_fixture("connections/test_connection.yaml"))
+    assert json_conn == yaml_conn
+
+
+def test_file_not_found():
+    with pytest.raises(OSError):
+        _read_connection_file("bad/path")