Feat: Allow supplying complex connection config fields as JSON strings (#4519)

Author: erindru

sqlmesh/core/config/connection.py

@@ -13,6 +13,7 @@
 
 import pydantic
 from pydantic import Field
+from pydantic_core import from_json
 from packaging import version
 from sqlglot import exp
 from sqlglot.helper import subclasses
@@ -33,6 +34,7 @@
     field_validator,
     model_validator,
     validation_error_message,
+    get_concrete_types_from_typehint,
 )
 from sqlmesh.utils.aws import validate_s3_uri
 
@@ -177,6 +179,42 @@ def get_catalog(self) -> t.Optional[str]:
             return self.db
         return None
 
+    @model_validator(mode="before")
+    @classmethod
+    def _expand_json_strings_to_concrete_types(cls, data: t.Any) -> t.Any:
+        """
+        There are situations where a connection config class has a field that is some kind of complex type
+        (eg a list of strings or a dict) but the value is being supplied from a source such as an environment variable
+
+        When this happens, the value is supplied as a string rather than a Python object. We need some way
+        of turning this string into the corresponding Python list or dict.
+
+        Rather than doing this piecemeal on every config subclass, this provides a generic implementatation
+        to identify fields that may be be supplied as JSON strings and handle them transparently
+        """
+        if data and isinstance(data, dict):
+            for maybe_json_field_name in cls._get_list_and_dict_field_names():
+                if (value := data.get(maybe_json_field_name)) and isinstance(value, str):
+                    # crude JSON check as we dont want to try and parse every string we get
+                    value = value.strip()
+                    if value.startswith("{") or value.startswith("["):
+                        data[maybe_json_field_name] = from_json(value)
+
+        return data
+
+    @classmethod
+    def _get_list_and_dict_field_names(cls) -> t.Set[str]:
+        field_names = set()
+        for name, field in cls.model_fields.items():
+            if field.annotation:
+                field_types = get_concrete_types_from_typehint(field.annotation)
+
+                # check if the field type is something that could concievably be supplied as a json string
+                if any(ft is t for t in (list, tuple, set, dict) for ft in field_types):
+                    field_names.add(name)
+
+        return field_names
+
 
 class DuckDBAttachOptions(BaseConfig):
     type: str

sqlmesh/utils/pydantic.py

@@ -306,6 +306,24 @@ def cron_validator(v: t.Any) -> str:
     return v
 
 
+def get_concrete_types_from_typehint(typehint: type[t.Any]) -> set[type[t.Any]]:
+    concrete_types = set()
+    unpacked = t.get_origin(typehint)
+    if unpacked is None:
+        if type(typehint) == type(type):
+            return {typehint}
+    elif unpacked is t.Union:
+        for item in t.get_args(typehint):
+            if str(item).startswith("typing."):
+                concrete_types |= get_concrete_types_from_typehint(item)
+            else:
+                concrete_types.add(item)
+    else:
+        concrete_types.add(unpacked)
+
+    return concrete_types
+
+
 if t.TYPE_CHECKING:
     SQLGlotListOfStrings = t.List[str]
     SQLGlotString = str

tests/core/test_config.py

@@ -985,3 +985,37 @@ def test_config_subclassing() -> None:
     class ConfigSubclass(Config): ...
 
     ConfigSubclass()
+
+
+def test_config_complex_types_supplied_as_json_strings_from_env(tmp_path: Path) -> None:
+    config_path = tmp_path / "config.yaml"
+    config_path.write_text("""
+    gateways:
+      bigquery:
+        connection:
+          type: bigquery
+          project: unit-test
+
+    default_gateway: bigquery
+
+    model_defaults:
+      dialect: bigquery
+""")
+    with mock.patch.dict(
+        os.environ,
+        {
+            "SQLMESH__GATEWAYS__BIGQUERY__CONNECTION__SCOPES": '     ["a","b","c"]',  # note: leading whitespace is deliberate
+            "SQLMESH__GATEWAYS__BIGQUERY__CONNECTION__KEYFILE_JSON": '{ "foo": "bar" }',
+        },
+    ):
+        config = load_config_from_paths(
+            Config,
+            project_paths=[config_path],
+        )
+
+        conn = config.gateways["bigquery"].connection
+        assert isinstance(conn, BigQueryConnectionConfig)
+
+        assert conn.project == "unit-test"
+        assert conn.scopes == ("a", "b", "c")
+        assert conn.keyfile_json == {"foo": "bar"}

tests/core/test_connection_config.py

@@ -617,6 +617,27 @@ def test_duckdb_attach_options():
     assert options.to_sql(alias="db") == "ATTACH IF NOT EXISTS 'test.db' AS db"
 
 
+def test_duckdb_config_json_strings(make_config):
+    config = make_config(
+        type="duckdb",
+        extensions='["foo","bar"]',
+        catalogs="""{
+            "test1": "test1.duckdb",
+            "test2": {
+                "type": "duckdb",
+                "path": "test2.duckdb"
+            }
+        }""",
+    )
+    assert isinstance(config, DuckDBConnectionConfig)
+
+    assert config.extensions == ["foo", "bar"]
+
+    assert config.get_catalog() == "test1"
+    assert config.catalogs.get("test1") == "test1.duckdb"
+    assert config.catalogs.get("test2").path == "test2.duckdb"
+
+
 def test_motherduck_attach_catalog(make_config):
     config = make_config(
         type="motherduck",
@@ -779,6 +800,21 @@ def test_bigquery(make_config):
         make_config(type="bigquery", quota_project="quota_project", check_import=False)
 
 
+def test_bigquery_config_json_string(make_config):
+    config = make_config(
+        type="bigquery",
+        project="project",
+        # these can be present as strings if they came from env vars
+        scopes='["a","b","c"]',
+        keyfile_json='{"foo":"bar"}',
+    )
+
+    assert isinstance(config, BigQueryConnectionConfig)
+
+    assert config.scopes == ("a", "b", "c")
+    assert config.keyfile_json == {"foo": "bar"}
+
+
 def test_postgres(make_config):
     config = make_config(
         type="postgres",

tests/utils/test_pydantic.py

@@ -1,7 +1,9 @@
+import typing as t
+import pytest
 from functools import cached_property
 
 from sqlmesh.utils.date import TimeLike, to_date, to_datetime
-from sqlmesh.utils.pydantic import PydanticModel
+from sqlmesh.utils.pydantic import PydanticModel, get_concrete_types_from_typehint
 
 
 def test_datetime_date_serialization() -> None:
@@ -62,3 +64,26 @@ class TestModel(PydanticModel):
         name: str
 
     assert TestModel(name="foo").dict(by_alias=True)
+
+
+@pytest.mark.parametrize(
+    "input,output",
+    [
+        (t.Dict[str, t.Any], {dict}),
+        (dict, {dict}),
+        (t.List[str], {list}),
+        (list, {list}),
+        (t.Tuple[str, ...], {tuple}),
+        (tuple, {tuple}),
+        (t.Set[str], {set}),
+        (set, {set}),
+        (t.Optional[t.Dict[str, t.Any]], {dict, type(None)}),
+        (t.Optional[t.List[str]], {list, type(None)}),
+        (
+            t.Union[str, t.List[str], t.Dict[str, t.Any], t.Optional[t.Set[str]]],
+            {str, list, dict, set, type(None)},
+        ),
+    ],
+)
+def test_get_concrete_types_from_typehint(input: t.Any, output: set[type]) -> None:
+    assert get_concrete_types_from_typehint(input) == output