support Databricks query tags in session properties

Author: WhammyLeaf

sqlmesh/core/engine_adapter/databricks.py

@@ -30,6 +30,25 @@
 logger = logging.getLogger(__name__)
 
 
+def _query_tags(query_tags: t.Any) -> t.Optional[t.Dict[str, t.Optional[str]]]:
+    if not query_tags:
+        return None
+
+    if not isinstance(query_tags, dict):
+        raise SQLMeshError("Invalid value for `session_properties.query_tags`. Must be a dict.")
+
+    if not all(
+        isinstance(key, str) and (value is None or isinstance(value, str))
+        for key, value in query_tags.items()
+    ):
+        raise SQLMeshError(
+            "Invalid value for `session_properties.query_tags`. Keys must be strings "
+            "and values must be strings or None."
+        )
+
+    return query_tags
+
+
 class DatabricksEngineAdapter(SparkEngineAdapter, GrantsFromInfoSchemaMixin):
     DIALECT = "databricks"
     INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.REPLACE_WHERE
@@ -98,6 +117,13 @@ def _use_spark_session(self) -> bool:
     def is_spark_session_connection(self) -> bool:
         return isinstance(self.connection, SparkSessionConnection)
 
+    @property
+    def _is_databricks_sql_connector_connection(self) -> bool:
+        return (
+            not self.is_spark_session_connection
+            and not self._connection_pool.get_attribute("use_spark_engine_adapter")
+        )
+
     def _set_spark_engine_adapter_if_needed(self) -> None:
         self._spark_engine_adapter = None
 
@@ -181,10 +207,25 @@ def _begin_session(self, properties: SessionProperties) -> t.Any:
         """Begin a new session."""
         # Align the different possible connectors to a single catalog
         self.set_current_catalog(self.default_catalog)  # type: ignore
+        self._connection_pool.set_attribute(
+            "query_tags", _query_tags(properties.get("query_tags"))
+        )
 
     def _end_session(self) -> None:
+        self._connection_pool.set_attribute("query_tags", None)
         self._connection_pool.set_attribute("use_spark_engine_adapter", False)
 
+    def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any) -> None:
+        query_tags = self._connection_pool.get_attribute("query_tags")
+        if (
+            query_tags
+            and "query_tags" not in kwargs
+            and self._is_databricks_sql_connector_connection
+        ):
+            kwargs["query_tags"] = query_tags
+
+        return super()._execute(sql, track_rows_processed, **kwargs)
+
     def _df_to_source_queries(
         self,
         df: DF,

tests/core/engine_adapter/test_databricks.py

@@ -10,6 +10,7 @@
 from sqlmesh.core.engine_adapter import DatabricksEngineAdapter
 from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
 from sqlmesh.core.node import IntervalUnit
+from sqlmesh.utils.errors import SQLMeshError
 from tests.core.engine_adapter import to_sql_calls
 
 pytestmark = [pytest.mark.databricks, pytest.mark.engine]
@@ -117,6 +118,92 @@ def test_set_current_catalog(mocker: MockFixture, make_mocked_engine_adapter: t.
     assert to_sql_calls(adapter) == ["USE CATALOG `test_catalog2`"]
 
 
+def test_session_query_tags(mocker: MockFixture, make_mocked_engine_adapter: t.Callable):
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.set_current_catalog"
+    )
+    adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog")
+
+    with adapter.session({"query_tags": {"team": "data-eng", "app": "sqlmesh"}}):
+        adapter.execute("SELECT 1")
+
+    adapter.cursor.execute.assert_called_with(
+        "SELECT 1", query_tags={"team": "data-eng", "app": "sqlmesh"}
+    )
+
+    adapter.execute("SELECT 2")
+
+    adapter.cursor.execute.assert_called_with("SELECT 2")
+
+
+def test_session_query_tags_do_not_override_explicit_query_tags(
+    mocker: MockFixture, make_mocked_engine_adapter: t.Callable
+):
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.set_current_catalog"
+    )
+    adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog")
+
+    with adapter.session({"query_tags": {"team": "data-eng"}}):
+        adapter.execute("SELECT 1", query_tags={"team": "analytics"})
+
+    adapter.cursor.execute.assert_called_with("SELECT 1", query_tags={"team": "analytics"})
+
+
+def test_session_query_tags_not_applied_to_spark_session_connection(
+    mocker: MockFixture, make_mocked_engine_adapter: t.Callable
+):
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.set_current_catalog"
+    )
+    adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog")
+    mocker.patch.object(
+        DatabricksEngineAdapter,
+        "is_spark_session_connection",
+        new_callable=mocker.PropertyMock,
+        return_value=True,
+    )
+
+    with adapter.session({"query_tags": {"team": "data-eng"}}):
+        adapter.execute("SELECT 1")
+
+    adapter.cursor.execute.assert_called_with("SELECT 1")
+
+
+def test_session_query_tags_not_applied_to_spark_engine_adapter(
+    mocker: MockFixture, make_mocked_engine_adapter: t.Callable
+):
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.set_current_catalog"
+    )
+    adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog")
+    spark_cursor = mocker.Mock()
+    adapter._spark_engine_adapter = mocker.Mock(cursor=spark_cursor)
+    adapter._connection_pool.set_attribute("use_spark_engine_adapter", True)
+
+    with adapter.session({"query_tags": {"team": "data-eng"}}):
+        adapter._connection_pool.set_attribute("use_spark_engine_adapter", True)
+        adapter.execute("SELECT 1")
+
+    spark_cursor.execute.assert_called_with("SELECT 1")
+
+
+@pytest.mark.parametrize(
+    "query_tags",
+    [
+        "team:data-eng",
+        {"team": 1},
+        {1: "data-eng"},
+    ],
+)
+def test_session_query_tags_invalid(query_tags, make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog")
+
+    with pytest.raises(SQLMeshError, match="session_properties.query_tags"):
+        with adapter.session({"query_tags": query_tags}):
+            pass
+
+
 def test_get_current_catalog(mocker: MockFixture, make_mocked_engine_adapter: t.Callable):
     mocker.patch(
         "sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.set_current_catalog"