Feat: Support custom session properties in the Spark Submit Airflow operator (#1482)

Author: izeigerman

examples/sushi/models/items.py

@@ -69,6 +69,12 @@
         "float_prop": 1.0,
         "bool_prop": True,
     },
+    session_properties={
+        "string_prop": "some_value",
+        "int_prop": 1,
+        "float_prop": 1.0,
+        "bool_prop": True,
+    },
 )
 def execute(
     context: ExecutionContext,

sqlmesh/core/model/common.py

@@ -111,6 +111,7 @@ def parse_properties(cls: t.Type, v: t.Any, values: t.Dict[str, t.Any]) -> t.Opt
 
 properties_validator = field_validator(
     "table_properties_",
+    "session_properties_",
     mode="before",
     check_fields=False,
 )(parse_properties)

sqlmesh/core/model/definition.py

@@ -744,6 +744,7 @@ def metadata_hash(self, audits: t.Dict[str, ModelAudit]) -> str:
             str(self.disable_restatement),
             self.project,
             str(self.allow_partials),
+            self.session_properties_.sql() if self.session_properties_ else None,
         ]
 
         for audit_name, audit_args in sorted(self.audits, key=lambda a: a[0]):
@@ -1926,4 +1927,5 @@ def _refs_to_sql(values: t.Any) -> exp.Expression:
     "references": _refs_to_sql,
     "hash_raw_query": exp.convert,
     "table_properties_": lambda value: value,
+    "session_properties_": lambda value: value,
 }

sqlmesh/core/model/meta.py

@@ -51,13 +51,14 @@ class ModelMeta(_Node, extra="allow"):
     hash_raw_query: bool = False
     physical_schema_override: t.Optional[str] = None
     table_properties_: t.Optional[exp.Tuple] = Field(default=None, alias="table_properties")
+    session_properties_: t.Optional[exp.Tuple] = Field(default=None, alias="session_properties")
     allow_partials: bool = False
 
     _table_properties: t.Dict[str, exp.Expression] = {}
 
     _bool_validator = bool_validator
     _model_kind_validator = model_kind_validator
-    _table_properties_validator = properties_validator
+    _properties_validator = properties_validator
 
     @field_validator("audits", mode="before")
     @classmethod
@@ -332,6 +333,27 @@ def table_properties(self) -> t.Dict[str, exp.Expression]:
                 self._table_properties[expression.this.name] = expression.expression
         return self._table_properties
 
+    @property
+    def session_properties(self) -> t.Dict[str, t.Union[exp.Expression | str | int | float | bool]]:
+        """A dictionary of session properties."""
+        if not self.session_properties_:
+            return {}
+
+        def _interpret_expr(
+            e: exp.Expression,
+        ) -> t.Union[exp.Expression | str | int | float | bool]:
+            if e.is_int:
+                return int(e.this)
+            if e.is_number:
+                return float(e.this)
+            if isinstance(e, (exp.Literal, exp.Boolean)):
+                return e.this
+            return e
+
+        return {
+            e.this.name: _interpret_expr(e.expression) for e in self.session_properties_.expressions
+        }
+
     @property
     def all_references(self) -> t.List[Reference]:
         """All references including grains."""

sqlmesh/schedulers/airflow/operators/spark_submit.py

@@ -8,7 +8,10 @@
 
 import sqlmesh
 from sqlmesh.engines import commands
-from sqlmesh.schedulers.airflow.operators.targets import BaseTarget
+from sqlmesh.schedulers.airflow.operators.targets import (
+    BaseTarget,
+    SnapshotEvaluationTarget,
+)
 
 
 class SQLMeshSparkSubmitOperator(BaseOperator):
@@ -54,7 +57,7 @@ def __init__(
         super().__init__(**kwargs)
         self._target = target
         self._application_name = application_name
-        self._spark_conf = spark_conf
+        self._spark_conf = spark_conf or {}
         self._total_executor_cores = total_executor_cores
         self._executor_cores = executor_cores
         self._executor_memory = executor_memory
@@ -77,24 +80,59 @@ def execute(self, context: Context) -> None:
                 payload_fd.write(command_payload)
 
             if self._hook is None:
+                if (
+                    isinstance(self._target, SnapshotEvaluationTarget)
+                    and self._target.snapshot.is_model
+                ):
+                    session_properties = self._target.snapshot.model.session_properties
+                    executor_cores: t.Optional[int] = session_properties.pop(  # type: ignore
+                        "spark.executor.cores", self._executor_cores
+                    )
+                    executor_memory: t.Optional[str] = session_properties.pop(  # type: ignore
+                        "spark.executor.memory", self._executor_memory
+                    )
+                    driver_memory: t.Optional[str] = session_properties.pop(  # type: ignore
+                        "spark.driver.memory", self._driver_memory
+                    )
+                    num_executors: t.Optional[int] = session_properties.pop(  # type: ignore
+                        "spark.executor.instances", self._num_executors
+                    )
+                    spark_conf: t.Dict[str, t.Any] = {**self._spark_conf, **session_properties}
+                else:
+                    executor_cores = self._executor_cores
+                    executor_memory = self._executor_memory
+                    driver_memory = self._driver_memory
+                    num_executors = self._num_executors
+                    spark_conf = self._spark_conf
+
                 self._hook = self._get_hook(
                     self._target.command_type,
                     payload_file_path,
                     self._target.ddl_concurrent_tasks,
+                    spark_conf,
+                    executor_cores,
+                    executor_memory,
+                    driver_memory,
+                    num_executors,
                 )
             self._hook.submit(self._application)
         self._target.post_hook(context)
 
     def on_kill(self) -> None:
         if self._hook is None:
-            self._hook = self._get_hook(None, None, None)
+            self._hook = self._get_hook(None, None, None, None, None, None, None, None)
         self._hook.on_kill()
 
     def _get_hook(
         self,
         command_type: t.Optional[commands.CommandType],
         command_payload_file_path: t.Optional[str],
         ddl_concurrent_tasks: t.Optional[int],
+        spark_conf: t.Optional[t.Dict[str, t.Any]],
+        executor_cores: t.Optional[int],
+        executor_memory: t.Optional[str],
+        driver_memory: t.Optional[str],
+        num_executors: t.Optional[int],
     ) -> SparkSubmitHook:
         application_args = {
             "dialect": "spark",
@@ -105,17 +143,17 @@ def _get_hook(
             else None,
         }
         return SparkSubmitHook(
-            conf=self._spark_conf,
+            conf=spark_conf,
             conn_id=self._connection_id,
             total_executor_cores=self._total_executor_cores,
-            executor_cores=self._executor_cores,
-            executor_memory=self._executor_memory,
-            driver_memory=self._driver_memory,
+            executor_cores=executor_cores,
+            executor_memory=executor_memory,
+            driver_memory=driver_memory,
             keytab=self._keytab,
             principal=self._principal,
             proxy_user=self._proxy_user,
             name=self._application_name,
-            num_executors=self._num_executors,
+            num_executors=num_executors,
             application_args=[f"--{k}={v}" for k, v in application_args.items() if v is not None],
             files=command_payload_file_path,
         )

tests/core/test_model.py

@@ -1997,6 +1997,38 @@ def test_model_table_properties(sushi_context):
         )
 
 
+def test_model_session_properties(sushi_context):
+    assert sushi_context.models["sushi.items"].session_properties == {
+        "string_prop": "some_value",
+        "int_prop": 1,
+        "float_prop": 1.0,
+        "bool_prop": True,
+    }
+    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,
+            )
+        );
+        SELECT a FROM tbl;
+        """
+        )
+    )
+
+    assert model.session_properties == {
+        "spark.executor.cores": 2,
+        "spark.executor.memory": "1G",
+        "some_bool": True,
+        "some_float": 0.1,
+    }
+
+
 def test_model_jinja_macro_rendering():
     expressions = d.parse(
         """

tests/core/test_snapshot.py

@@ -437,7 +437,7 @@ def test_fingerprint(model: Model, parent_model: Model):
 
     original_fingerprint = SnapshotFingerprint(
         data_hash="3811098861",
-        metadata_hash="541992912",
+        metadata_hash="3858405978",
     )
 
     assert fingerprint == original_fingerprint
@@ -484,7 +484,7 @@ def test_fingerprint_seed_model():
 
     expected_fingerprint = SnapshotFingerprint(
         data_hash="3270932819",
-        metadata_hash="2823924537",
+        metadata_hash="1017437962",
     )
 
     model = load_sql_based_model(expressions, path=Path("./examples/sushi/models/test_model.sql"))
@@ -524,7 +524,7 @@ def test_fingerprint_jinja_macros(model: Model):
     )
     original_fingerprint = SnapshotFingerprint(
         data_hash="2864998504",
-        metadata_hash="541992912",
+        metadata_hash="3858405978",
     )
 
     fingerprint = fingerprint_from_node(model, nodes={})