Repository: TobikoData/sqlmesh
Pull request #2596 (2024-05-11 13:42:14)

#2335 Add support for session properties in model defaults
This should allow you to set session_properties on the project configuration which will then be propagated on the models. If custom properties are also set on the specific model, those will not be overwritten by the default ones. It might need some more work when adding the default props to handle quoted/unquoted identifiers better; right now, it sets them by default as quoted=False. 

CLAassistant: [![CLA assistant check](https://cla-assistant.io/pull/badge/not_signed)](https://cla-assistant.io/TobikoData/sqlmesh?pullRequest=2596) <br/>Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our [Contributor License Agreement](https://cla-assistant.io/TobikoData/sqlmesh?pullRequest=2596) before we can accept your contribution.<br/><sub>You have signed the CLA already but the status is still pending? Let us [recheck](https://cla-assistant.io/check/TobikoData/sqlmesh?pullRequest=2596) it.</sub>
izeigerman: Make sure to run `make style` to run linters and code formatter

Patch:
diff --git a/sqlmesh/core/config/model.py b/sqlmesh/core/config/model.py
index 2fd586b8..ef527e5e 100644
--- a/sqlmesh/core/config/model.py
+++ b/sqlmesh/core/config/model.py
@@ -29,5 +29,6 @@ class ModelDefaultsConfig(BaseConfig):
     owner: t.Optional[str] = None
     start: t.Optional[TimeLike] = None
     storage_format: t.Optional[str] = None
+    session_properties: t.Optional[t.Dict[str, t.Any]] = None
 
     _model_kind_validator = model_kind_validator
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index a86d13bd..fea6d749 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -1813,6 +1813,10 @@ def _create_model(
 ) -> Model:
     _validate_model_fields(klass, {"name", *kwargs} - {"grain", "table_properties"}, path)
 
+    kwargs["session_properties"] = _resolve_custom_session_properties(
+        defaults, kwargs.get("session_properties")
+    )
+
     dialect = dialect or ""
     physical_schema_override = physical_schema_override or {}
 
@@ -1876,6 +1880,28 @@ def _split_sql_model_statements(
     return query, expressions[:pos], expressions[pos + 1 :]
 
 
+def _resolve_custom_session_properties(
+    defaults: t.Optional[t.Dict[str, t.Any]],
+    provided: t.Optional[exp.Expression] | t.Optional[t.Dict[str, t.Any]],
+) -> t.Optional[exp.Expression]:
+    if isinstance(provided, dict):
+        session_properties = {k: exp.Literal.string(k).eq(v) for k, v in provided.items()}
+    else:
+        session_properties = (
+            {expr.this.name: expr for expr in provided} if provided is not None else {}
+        )
+
+    if defaults and defaults.get("session_properties"):
+        for k, v in defaults["session_properties"].items():
+            if k not in session_properties:
+                session_properties[k] = exp.Literal.string(k).eq(v)
+
+    if session_properties:
+        return exp.Tuple(expressions=list(session_properties.values()))
+
+    return None
+
+
 def _validate_model_fields(klass: t.Type[_Model], provided_fields: t.Set[str], path: Path) -> None:
     missing_required_fields = klass.missing_required_fields(provided_fields)
     if missing_required_fields:
diff --git a/sqlmesh/dbt/adapter.py b/sqlmesh/dbt/adapter.py
index d955317c..b4ef2d60 100644
--- a/sqlmesh/dbt/adapter.py
+++ b/sqlmesh/dbt/adapter.py
@@ -278,9 +278,8 @@ class RuntimeAdapter(BaseAdapter):
         self, sql: str, auto_begin: bool = False, fetch: bool = False
     ) -> t.Tuple[AdapterResponse, agate.Table]:
         from dbt.adapters.base.impl import AdapterResponse
-        from dbt.clients.agate_helper import empty_table
 
-        from sqlmesh.dbt.util import pandas_to_agate
+        from sqlmesh.dbt.util import pandas_to_agate, empty_table
 
         # mypy bug: https://github.com/python/mypy/issues/10740
         exec_func: t.Callable[..., None | pd.DataFrame] = (
diff --git a/sqlmesh/dbt/builtin.py b/sqlmesh/dbt/builtin.py
index d95bd09d..27e8df5d 100644
--- a/sqlmesh/dbt/builtin.py
+++ b/sqlmesh/dbt/builtin.py
@@ -92,15 +92,15 @@ class SQLExecution:
         self._results: t.Dict[str, AttributeDict] = {}
 
     def store_result(self, name: str, response: t.Any, agate_table: t.Optional[agate.Table]) -> str:
-        from dbt.clients import agate_helper
+        from sqlmesh.dbt.util import empty_table, as_matrix
 
         if agate_table is None:
-            agate_table = agate_helper.empty_table()
+            agate_table = empty_table()
 
         self._results[name] = AttributeDict(
             {
                 "response": response,
-                "data": agate_helper.as_matrix(agate_table),
+                "data": as_matrix(agate_table),
                 "table": agate_table,
             }
         )
diff --git a/sqlmesh/dbt/manifest.py b/sqlmesh/dbt/manifest.py
index 75690103..83662b07 100644
--- a/sqlmesh/dbt/manifest.py
+++ b/sqlmesh/dbt/manifest.py
@@ -256,6 +256,7 @@ class ManifestHelper:
             profiles_dir=str(self.profiles_path),
             target=self.target.name,
             macro_debugging=False,
+            REQUIRE_RESOURCE_NAMES_WITHOUT_SPACES=True,
         )
         flags.set_from_args(args, None)
 
@@ -271,7 +272,15 @@ class ManifestHelper:
 
         self._project_name = project.project_name
 
-        register_adapter(runtime_config)
+        if DBT_VERSION >= (1, 8):
+            from dbt.mp_context import get_mp_context  # type: ignore
+            from dbt_common.context import set_invocation_context  # type: ignore
+
+            register_adapter(runtime_config, get_mp_context())  # type: ignore
+            set_invocation_context({})
+        else:
+            register_adapter(runtime_config)  # type: ignore
+
         manifest = ManifestLoader.get_full_manifest(runtime_config)
         reset_adapters()
         return manifest
diff --git a/sqlmesh/dbt/relation.py b/sqlmesh/dbt/relation.py
index 11708bc4..cc28e075 100644
--- a/sqlmesh/dbt/relation.py
+++ b/sqlmesh/dbt/relation.py
@@ -2,6 +2,6 @@ from sqlmesh.dbt.util import DBT_VERSION
 
 
 if DBT_VERSION < (1, 8):
-    from dbt.contracts.relation import *  # noqa: F403
+    from dbt.contracts.relation import *  # type: ignore  # noqa: F403
 else:
     from dbt.adapters.contracts.relation import *  # noqa: F403
diff --git a/sqlmesh/dbt/util.py b/sqlmesh/dbt/util.py
index 9edca3fc..a06f098d 100644
--- a/sqlmesh/dbt/util.py
+++ b/sqlmesh/dbt/util.py
@@ -5,18 +5,22 @@ import pandas as pd
 from dbt.version import get_installed_version
 
 
-def pandas_to_agate(df: pd.DataFrame) -> agate.Table:
-    """
-    Converts a Pandas DataFrame to an Agate Table
-    """
-    from dbt.clients.agate_helper import table_from_data
-
-    return table_from_data(df.to_dict(orient="records"), df.columns.tolist())
-
-
 def _get_dbt_version() -> t.Tuple[int, int]:
     dbt_version = get_installed_version()
     return (int(dbt_version.major or "0"), int(dbt_version.minor or "0"))
 
 
 DBT_VERSION = _get_dbt_version()
+
+if DBT_VERSION < (1, 8):
+    from dbt.clients.agate_helper import table_from_data, empty_table, as_matrix  # type: ignore  # noqa: F401
+else:
+    from dbt_common.clients.agate_helper import table_from_data, empty_table, as_matrix  # type: ignore  # noqa: F401
+
+
+def pandas_to_agate(df: pd.DataFrame) -> agate.Table:
+    """
+    Converts a Pandas DataFrame to an Agate Table
+    """
+
+    return table_from_data(df.to_dict(orient="records"), df.columns.tolist())


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 4bdd72b6..aa912f48 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -1554,6 +1554,28 @@ def test_python_model_depends_on() -> None:
     assert m.depends_on == {'"foo"."bar"'}
 
 
+def test_python_model_with_session_props():
+    @model(
+        name="python_model_prop",
+        kind="full",
+        columns={"some_col": "int"},
+        session_properties={"some_string": "string_prop", "some_bool": True, "some_float": 1.0},
+    )
+    def python_model_prop(context, **kwargs):
+        context.table("foo")
+
+    m = model.get_registry()["python_model_prop"].model(
+        module_path=Path("."),
+        path=Path("."),
+        dialect="duckdb",
+    )
+    assert m.session_properties == {
+        "some_string": "string_prop",
+        "some_bool": True,
+        "some_float": 1.0,
+    }
+
+
 def test_python_models_returning_sql(assert_exp_eq) -> None:
     config = Config(model_defaults=ModelDefaultsConfig(dialect="snowflake"))
     context = Context(config=config)
@@ -2698,6 +2720,76 @@ def test_model_table_properties_conflicts() -> None:
         sql_model.physical_properties
 
 
+def test_session_properties_on_model_and_project(sushi_context):
+    model_defaults = ModelDefaultsConfig(
+        session_properties={
+            "some_bool": False,
+            "quoted_identifier": "value_you_wont_see",
+            "project_level_property": "project_property",
+        }
+    )
+
+    model = load_sql_based_model(
+        d.parse(
+            """
+        MODEL (
+            name test_schema.test_model,
+            session_properties (
+                'spark.executor.cores' = 2,
+                'spark.executor.memory' = '1G',
+                some_bool = True,
+                some_float = 0.1,
+                quoted_identifier = "quoted identifier",
+                unquoted_identifier = unquoted_identifier,
+            )
+        );
+        SELECT a FROM tbl;
+        """,
+            default_dialect="snowflake",
+        ),
+        defaults=model_defaults.dict(),
+    )
+
+    assert model.session_properties == {
+        "spark.executor.cores": 2,
+        "spark.executor.memory": "1G",
+        "some_bool": True,
+        "some_float": 0.1,
+        "quoted_identifier": exp.column("quoted identifier", quoted=True),
+        "unquoted_identifier": exp.column("unquoted_identifier", quoted=False),
+        "project_level_property": "project_property",
+    }
+
+
+def test_project_level_session_properties(sushi_context):
+    model_defaults = ModelDefaultsConfig(
+        session_properties={
+            "some_bool": False,
+            "some_float": 0.1,
+            "project_level_property": "project_property",
+        }
+    )
+
+    model = load_sql_based_model(
+        d.parse(
+            """
+        MODEL (
+            name test_schema.test_model,
+        );
+        SELECT a FROM tbl;
+        """,
+            default_dialect="snowflake",
+        ),
+        defaults=model_defaults.dict(),
+    )
+
+    assert model.session_properties == {
+        "some_bool": False,
+        "some_float": 0.1,
+        "project_level_property": "project_property",
+    }
+
+
 def test_model_session_properties(sushi_context):
     assert sushi_context.models['"memory"."sushi"."items"'].session_properties == {
         "string_prop": "some_value",