fix(Parameters): add support for custom connection (#116)

openhexa/sdk/__init__.py

@@ -2,7 +2,14 @@
 
 from .pipelines import current_run, parameter, pipeline
 from .workspaces import workspace
-from .workspaces.connection import DHIS2Connection, GCSConnection, IASOConnection, PostgreSQLConnection, S3Connection
+from .workspaces.connection import (
+    CustomConnection,
+    DHIS2Connection,
+    GCSConnection,
+    IASOConnection,
+    PostgreSQLConnection,
+    S3Connection,
+)
 
 __all__ = [
     "workspace",
@@ -14,4 +21,5 @@
     "PostgreSQLConnection",
     "GCSConnection",
     "S3Connection",
+    "CustomConnection",
 ]

openhexa/sdk/pipelines/parameter.py

@@ -300,7 +300,7 @@ def spec_type(self) -> str:
     @property
     def expected_type(self) -> type:
         """Returns the python type expected for values."""
-        return str
+        return CustomConnectionType
 
     def to_connection(self, value: str) -> CustomConnection:
         """Build a custom connection instance from the provided value (which should be a connection identifier)."""
@@ -317,6 +317,7 @@ def to_connection(self, value: str) -> CustomConnection:
     IASOConnection: IASOConnectionType,
     S3Connection: S3ConnectionType,
     GCSConnection: GCSConnectionType,
+    CustomConnection: CustomConnectionType,
 }
 
 
@@ -468,6 +469,7 @@ def parameter(
         type[PostgreSQLConnection],
         type[GCSConnection],
         type[S3Connection],
+        type[CustomConnection],
     ],
     name: typing.Optional[str] = None,
     choices: typing.Optional[typing.Sequence] = None,

tests/test_parameter.py

@@ -12,9 +12,11 @@
     IASOConnection,
     PostgreSQLConnection,
     S3Connection,
+    workspace,
 )
 from openhexa.sdk.pipelines.parameter import (
     Boolean,
+    CustomConnectionType,
     DHIS2ConnectionType,
     Float,
     FunctionWithParameter,
@@ -198,6 +200,30 @@ def test_validate_s3_connection():
             s3_parameter_type.validate(86)
 
 
+def test_validate_custom_connection():
+    """Check Custom connection validation."""
+    identifier = "custom-connection-id"
+    env_variable_prefix = stringcase.constcase(identifier)
+    field_1 = "field_1"
+    field_2 = "field_2"
+
+    with mock.patch.dict(
+        os.environ,
+        {
+            f"{env_variable_prefix}_FIELD_1": field_1,
+            f"{env_variable_prefix}_FIELD_2": field_2,
+        },
+    ):
+        custom_co_type = CustomConnectionType()
+
+        custom_co = custom_co_type.validate(identifier)
+        _custom_co = workspace.custom_connection(identifier)
+
+        assert str(custom_co) == str(_custom_co)
+        with pytest.raises(ParameterValueError):
+            custom_co_type.validate(86)
+
+
 def test_parameter_init():
     """Sanity checks for parameter initialization."""
     # Wrong type