allow secrets in the connector config

PabloPardoGarcia · PabloPardoGarcia · commit b0c90449ec2d · 2025-02-20T10:50:26.000+01:00
diff --git a/src/glassflow/client.py b/src/glassflow/client.py
@@ -2,7 +2,7 @@
 
 from .api_client import APIClient
 from .models import errors, responses
-from .pipeline import Pipeline
+from .pipeline import ConnectorConfiguration, Pipeline
 from .secret import Secret
 from .space import Space
 
@@ -65,9 +65,9 @@ def create_pipeline(
         transformation_file: str = None,
         requirements: str = None,
         source_kind: str = None,
-        source_config: dict = None,
+        source_config: ConnectorConfiguration = None,
         sink_kind: str = None,
-        sink_config: dict = None,
+        sink_config: ConnectorConfiguration = None,
         env_vars: list[dict[str, str]] = None,
         state: str = "running",
         metadata: dict = None,
diff --git a/src/glassflow/models/api/__init__.py b/src/glassflow/models/api/__init__.py
@@ -1,4 +1,6 @@
 from .api import (
+    ConnectorValueSecretRef,
+    ConnectorValueValue,
     ConsumeOutputEvent,
     CreatePipeline,
     CreateSecret,
@@ -32,4 +34,6 @@
     "ListSpaceScopes",
     "ListPipelines",
     "FunctionEnvironments",
+    "ConnectorValueSecretRef",
+    "ConnectorValueValue",
 ]
diff --git a/src/glassflow/pipeline.py b/src/glassflow/pipeline.py
@@ -1,9 +1,16 @@
 from __future__ import annotations
 
+from typing import Union
+
+from typing_extensions import TypeAlias
+
 from .client import APIClient
 from .models import api, errors, operations, responses
 from .models.responses.pipeline import AccessToken
 from .pipeline_data import PipelineDataSink, PipelineDataSource
+from .secret import Secret
+
+ConnectorConfiguration: TypeAlias = dict[str, Union[str, list, Secret]]
 
 
 class Pipeline(APIClient):
@@ -14,9 +21,9 @@ def __init__(
         space_id: str | None = None,
         id: str | None = None,
         source_kind: str | None = None,
-        source_config: dict | None = None,
+        source_config: ConnectorConfiguration | None = None,
         sink_kind: str | None = None,
-        sink_config: dict | None = None,
+        sink_config: ConnectorConfiguration | None = None,
         requirements: str | None = None,
         transformation_file: str | None = None,
         env_vars: list[dict[str, str]] | None = None,
@@ -182,9 +189,9 @@ def update(
         requirements: str | None = None,
         metadata: dict | None = None,
         source_kind: str | None = None,
-        source_config: dict | None = None,
+        source_config: ConnectorConfiguration | None = None,
         sink_kind: str | None = None,
-        sink_config: dict | None = None,
+        sink_config: ConnectorConfiguration | None = None,
         env_vars: list[dict[str, str]] | None = None,
     ) -> Pipeline:
         """
@@ -405,15 +412,19 @@ def _request(
                 ) from e
             raise e
 
-    @staticmethod
     def _fill_connector(
-        connector_type: str, kind: str, config: dict
+        self, connector_type: str, kind: str, config: dict
     ) -> api.SourceConnector | api.SinkConnector:
         """Format connector input"""
         if not kind and not config:
             connector = None
         elif kind and config:
-            connector = dict(kind=kind, config=config)
+            connector = dict(
+                kind=kind,
+                configuration={
+                    k: self._format_connector_config_value(v) for k, v in config.items()
+                },
+            )
         else:
             raise errors.ConnectorConfigValueError(connector_type)
 
@@ -424,6 +435,21 @@ def _fill_connector(
         else:
             raise ValueError("connector_type must be 'source' or 'sink'")
 
+    @staticmethod
+    def _format_connector_config_value(value: str | list | Secret):
+        """Formats configuration values to match API expectations"""
+        if isinstance(value, Secret):
+            config_value = api.ConnectorValueSecretRef(
+                **{"secret_ref": {"type": "organization", "key": value.key}}
+            )
+        elif isinstance(value, list):
+            config_value = value
+        else:
+            config_value = api.ConnectorValueValue(
+                value=value,
+            )
+        return config_value
+
     def _list_access_tokens(self) -> Pipeline:
         endpoint = f"/pipelines/{self.id}/access_tokens"
         http_res = self._request(method="GET", endpoint=endpoint)
diff --git a/tests/glassflow/integration_tests/conftest.py b/tests/glassflow/integration_tests/conftest.py
@@ -49,7 +49,38 @@ def space_with_random_id_and_invalid_token():
 
 
 @pytest.fixture
-def pipeline(client, creating_space):
+def secret(client):
+    return Secret(
+        key="SecretKey",
+        value="SecretValue",
+        personal_access_token=client.personal_access_token,
+    )
+
+
+@pytest.fixture
+def creating_secret(secret):
+    secret.create()
+    yield secret
+    secret.delete()
+
+
+@pytest.fixture
+def secret_with_invalid_key_and_token():
+    return Secret(
+        key="InvalidSecretKey",
+        personal_access_token="invalid-token",
+    )
+
+
+@pytest.fixture
+def secret_with_invalid_key(client):
+    return Secret(
+        key="InvalidSecretKey", personal_access_token=client.personal_access_token
+    )
+
+
+@pytest.fixture
+def pipeline(client, creating_space, creating_secret):
     return Pipeline(
         name="test_pipeline",
         space_id=creating_space.id,
@@ -60,7 +91,7 @@ def pipeline(client, creating_space):
         source_config={
             "project_id": "my-project-id",
             "subscription_id": "my-subscription-id",
-            "credentials_json": "my-credentials.json",
+            "credentials_json": creating_secret,
         },
     )
 
@@ -123,34 +154,3 @@ def sink(source_with_published_events):
         pipeline_id=source_with_published_events.pipeline_id,
         pipeline_access_token=source_with_published_events.pipeline_access_token,
     )
-
-
-@pytest.fixture
-def secret(client):
-    return Secret(
-        key="SecretKey",
-        value="SecretValue",
-        personal_access_token=client.personal_access_token,
-    )
-
-
-@pytest.fixture
-def creating_secret(secret):
-    secret.create()
-    yield secret
-    secret.delete()
-
-
-@pytest.fixture
-def secret_with_invalid_key_and_token():
-    return Secret(
-        key="InvalidSecretKey",
-        personal_access_token="invalid-token",
-    )
-
-
-@pytest.fixture
-def secret_with_invalid_key(client):
-    return Secret(
-        key="InvalidSecretKey", personal_access_token=client.personal_access_token
-    )
diff --git a/tests/glassflow/unit_tests/conftest.py b/tests/glassflow/unit_tests/conftest.py
@@ -70,10 +70,12 @@ def fetch_pipeline_response():
             "space_name": "test-space-name",
             "source_connector": {
                 "kind": "google_pubsub",
-                "config": {
-                    "project_id": "test-project",
-                    "subscription_id": "test-subscription",
-                    "credentials_json": "credentials.json",
+                "configuration": {
+                    "project_id": {"value": "test-project"},
+                    "subscription_id": {"value": "test-subscription"},
+                    "credentials_json": {
+                        "secret_ref": {"type": "organization", "key": "credentialsJson"}
+                    },
                 },
             },
             "sink_connector": {
diff --git a/tests/glassflow/unit_tests/pipeline_test.py b/tests/glassflow/unit_tests/pipeline_test.py
@@ -1,7 +1,7 @@
 import pytest
 
+from glassflow import Pipeline, Secret
 from glassflow.models import errors
-from glassflow.pipeline import Pipeline
 
 
 def test_pipeline_with_transformation_file():
@@ -47,6 +47,54 @@ def test_pipeline_fail_with_connection_config_value_error():
         )
 
 
+def test_pipeline_with_connector_secret_ok():
+    p = Pipeline(
+        transformation_file="tests/data/transformation.py",
+        personal_access_token="test-token",
+        sink_kind="webhook",
+        sink_config={
+            "url": Secret(key="testKey", personal_access_token="test-token"),
+            "method": "POST",
+            "headers": [{"name": "Content-Type", "value": "application/json"}],
+        },
+        source_kind="google_pubsub",
+        source_config={
+            "project_id": "test-project-id",
+            "subscription_id": "test-subscription-id",
+            "credentials_json": Secret(
+                key="testCredentials", personal_access_token="test-token"
+            ),
+        },
+    )
+
+    expected_source_connector = {
+        "kind": "google_pubsub",
+        "configuration": {
+            "project_id": {"value": "test-project-id"},
+            "subscription_id": {"value": "test-subscription-id"},
+            "credentials_json": {
+                "secret_ref": {"type": "organization", "key": "testCredentials"}
+            },
+        },
+    }
+    expected_sink_connector = {
+        "kind": "webhook",
+        "configuration": {
+            "url": {"secret_ref": {"type": "organization", "key": "testKey"}},
+            "method": {"value": "POST"},
+            "headers": [{"name": "Content-Type", "value": "application/json"}],
+        },
+    }
+    assert (
+        p.source_connector.model_dump(mode="json", exclude_none=True)
+        == expected_source_connector
+    )
+    assert (
+        p.sink_connector.model_dump(mode="json", exclude_none=True)
+        == expected_sink_connector
+    )
+
+
 def test_fetch_pipeline_ok(
     get_pipeline_request_mock,
     get_access_token_request_mock,
