Repository: TobikoData/sqlmesh
Pull request #1923 (2024-01-03 22:31:19)

fix: set airflow default catalog


Patch:
diff --git a/sqlmesh/engines/spark/app.py b/sqlmesh/engines/spark/app.py
index 6d03770c..a8709361 100644
--- a/sqlmesh/engines/spark/app.py
+++ b/sqlmesh/engines/spark/app.py
@@ -30,7 +30,11 @@ def get_or_create_spark_session(dialect: str) -> SparkSession:
 
 
 def main(
-    dialect: str, command_type: commands.CommandType, ddl_concurrent_tasks: int, payload_path: str
+    dialect: str,
+    default_catalog: str,
+    command_type: commands.CommandType,
+    ddl_concurrent_tasks: int,
+    payload_path: str,
 ) -> None:
     if dialect not in ("databricks", "spark"):
         raise NotSupportedError(
@@ -50,6 +54,7 @@ def main(
         create_engine_adapter(
             lambda: spark_session_db.connection(spark),
             dialect,
+            default_catalog=default_catalog,
             multithreaded=ddl_concurrent_tasks > 1,
             execute_log_level=logging.INFO,
         ),
@@ -79,6 +84,10 @@ if __name__ == "__main__":
         "--dialect",
         help="The dialect to use when creating the engine adapter.",
     )
+    parser.add_argument(
+        "--default_catalog",
+        help="The default catalog to use when creating the engine adapter.",
+    )
     parser.add_argument(
         "--command_type",
         type=commands.CommandType,
@@ -96,4 +105,10 @@ if __name__ == "__main__":
         help="Path to the payload object. Can be a local or remote path.",
     )
     args = parser.parse_args()
-    main(args.dialect, args.command_type, args.ddl_concurrent_tasks, args.payload_path)
+    main(
+        args.dialect,
+        args.default_catalog,
+        args.command_type,
+        args.ddl_concurrent_tasks,
+        args.payload_path,
+    )
diff --git a/sqlmesh/migrations/v0036_delete_plan_dags_bug_fix.py b/sqlmesh/migrations/v0036_delete_plan_dags_bug_fix.py
new file mode 100644
index 00000000..7f9f49d6
--- /dev/null
+++ b/sqlmesh/migrations/v0036_delete_plan_dags_bug_fix.py
@@ -0,0 +1,14 @@
+"""Add missing delete from migration #34."""
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    plan_dags_table = "_plan_dags"
+    if state_sync.schema:
+        plan_dags_table = f"{schema}.{plan_dags_table}"
+
+    # At the time of migration plan_dags table is only needed for in-flight DAGs and therefore we can safely
+    # just delete it instead of migrating it
+    # If reusing this code verify that this is still the case
+    engine_adapter.delete_from(plan_dags_table, "TRUE")
diff --git a/sqlmesh/schedulers/airflow/operators/databricks.py b/sqlmesh/schedulers/airflow/operators/databricks.py
index 4495dbe1..c023ad81 100644
--- a/sqlmesh/schedulers/airflow/operators/databricks.py
+++ b/sqlmesh/schedulers/airflow/operators/databricks.py
@@ -101,6 +101,7 @@ class SQLMeshDatabricksSubmitOperator(DatabricksSubmitRunOperator):
             )
         task_arguments = {
             "dialect": "databricks",
+            "default_catalog": self._target.default_catalog,
             "command_type": self._target.command_type.value if self._target.command_type else None,
             "ddl_concurrent_tasks": self._target.ddl_concurrent_tasks,
             "payload_path": remote_payload_path,
diff --git a/sqlmesh/schedulers/airflow/operators/spark_submit.py b/sqlmesh/schedulers/airflow/operators/spark_submit.py
index f2cbf913..7a0664fd 100644
--- a/sqlmesh/schedulers/airflow/operators/spark_submit.py
+++ b/sqlmesh/schedulers/airflow/operators/spark_submit.py
@@ -136,6 +136,7 @@ class SQLMeshSparkSubmitOperator(BaseOperator):
     ) -> SparkSubmitHook:
         application_args = {
             "dialect": "spark",
+            "default_catalog": self._target.default_catalog,
             "command_type": command_type.value if command_type else None,
             "ddl_concurrent_tasks": ddl_concurrent_tasks,
             "payload_path": command_payload_file_path.split("/")[-1]
diff --git a/sqlmesh/schedulers/airflow/operators/targets.py b/sqlmesh/schedulers/airflow/operators/targets.py
index bfeb2acb..2f25c743 100644
--- a/sqlmesh/schedulers/airflow/operators/targets.py
+++ b/sqlmesh/schedulers/airflow/operators/targets.py
@@ -3,6 +3,7 @@ import logging
 import typing as t
 
 from airflow.exceptions import AirflowSkipException
+from airflow.models import Variable
 from airflow.utils.context import Context
 from airflow.utils.session import provide_session
 from sqlalchemy.orm import Session
@@ -29,6 +30,10 @@ class BaseTarget(abc.ABC, t.Generic[CP]):
     command_handler: t.Callable[[SnapshotEvaluator, CP], None]
     ddl_concurrent_tasks: int
 
+    @property
+    def default_catalog(self) -> str:
+        return Variable.get(common.DEFAULT_CATALOG_VARIABLE_NAME)
+
     def serialized_command_payload(self, context: Context) -> str:
         """Returns the serialized command payload for the Spark application.
 
@@ -62,6 +67,7 @@ class BaseTarget(abc.ABC, t.Generic[CP]):
                 dialect,
                 multithreaded=self.ddl_concurrent_tasks > 1,
                 execute_log_level=logging.INFO,
+                default_catalog=self.default_catalog,
                 **kwargs,
             ),
             ddl_concurrent_tasks=self.ddl_concurrent_tasks,


Test patch:
diff --git a/tests/schedulers/airflow/operators/test_targets.py b/tests/schedulers/airflow/operators/test_targets.py
index 7c650e92..6aa37274 100644
--- a/tests/schedulers/airflow/operators/test_targets.py
+++ b/tests/schedulers/airflow/operators/test_targets.py
@@ -48,6 +48,10 @@ def test_evaluation_target_execute(mocker: MockerFixture, make_snapshot: t.Calla
         "sqlmesh.core.state_sync.engine_adapter.EngineAdapterStateSync.add_interval"
     )
 
+    variable_get_mock = mocker.patch("sqlmesh.schedulers.airflow.operators.targets.Variable.get")
+
+    variable_get_mock.return_value = "default_catalog"
+
     snapshot = make_snapshot(model)
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
     parent_snapshots = {snapshot.name: snapshot}
@@ -83,6 +87,10 @@ def test_evaluation_target_execute_seed_model(mocker: MockerFixture, make_snapsh
     dag_run_mock.data_interval_end = interval_ds
     dag_run_mock.logical_date = logical_ds
 
+    variable_get_mock = mocker.patch("sqlmesh.schedulers.airflow.operators.targets.Variable.get")
+
+    variable_get_mock.return_value = "default_catalog"
+
     context = Context(dag_run=dag_run_mock)  # type: ignore
 
     snapshot = make_snapshot(
@@ -150,6 +158,10 @@ def test_cleanup_target_execute(mocker: MockerFixture, make_snapshot: t.Callable
     task_instance_mock = mocker.Mock()
     task_instance_mock.xcom_pull.return_value = command.json()
 
+    variable_get_mock = mocker.patch("sqlmesh.schedulers.airflow.operators.targets.Variable.get")
+
+    variable_get_mock.return_value = "default_catalog"
+
     context = Context(ti=task_instance_mock)  # type: ignore
 
     evaluator_cleanup_mock = mocker.patch(