fix: set airflow default catalog (#1923)

Author: eakmanrq

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 @@ def 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 @@ def 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,
+    )

sqlmesh/schedulers/airflow/operators/databricks.py

@@ -101,6 +101,7 @@ def execute(self, context: Context) -> None:
             )
         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,

sqlmesh/schedulers/airflow/operators/spark_submit.py

@@ -136,6 +136,7 @@ def _get_hook(
     ) -> 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]

sqlmesh/schedulers/airflow/operators/targets.py

@@ -3,6 +3,7 @@
 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 @@ def execute(
                 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,

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(