From 3e879b648132f0f09e7400697d291a683df093a3 Mon Sep 17 00:00:00 2001 From: Shaan Satsangi Date: Thu, 14 May 2026 01:08:35 +0530 Subject: [PATCH 1/6] feat(providers/databricks): Inject Airflow context as query tags in SQL operators This PR injects the Airflow dag_id, task_id, and run_id into the session_configuration parameter of DatabricksSqlOperator and DatabricksCopyIntoOperator. This enhances observability on the Databricks side. The tags are safely escaped and user-defined query tags are preserved. Closes #66839 --- .../databricks/operators/databricks_sql.py | 49 +++++++++++++++++++ .../operators/test_databricks_copy.py | 31 ++++++++++++ .../operators/test_databricks_sql.py | 29 +++++++++++ 3 files changed, 109 insertions(+) diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py index b50c434d04cb9..03904b16bd571 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py @@ -46,6 +46,23 @@ _DISALLOWED_SQL_TOKENS = (";", "--", "/*", "*/") +def _format_query_tags(context: Context) -> str: + """Format Airflow context metadata into databricks-sql-connector query tags.""" + + def escape(val: str) -> str: + return str(val).replace("\\", "\\\\").replace(",", "\\,").replace(":", "\\:") + + tags = [] + if "dag" in context and getattr(context["dag"], "dag_id", None): + tags.append(f"airflow_dag_id:{escape(context['dag'].dag_id)}") + if "task" in context and getattr(context["task"], "task_id", None): + tags.append(f"airflow_task_id:{escape(context['task'].task_id)}") + if "run_id" in context and context["run_id"]: + tags.append(f"airflow_run_id:{escape(context['run_id'])}") + + return ",".join(tags) + + class DatabricksSqlOperator(SQLExecuteQueryOperator): """ Executes SQL code in a Databricks SQL endpoint or a Databricks cluster. @@ -153,6 +170,24 @@ def _hook(self) -> DatabricksSqlHook: def get_db_hook(self) -> DatabricksSqlHook: return self._hook + def execute(self, context: Context) -> Any: + hook = self.get_db_hook() + query_tags = _format_query_tags(context) + if query_tags: + if hook.session_config is None: + conn_extra = hook.databricks_conn.extra_dejson + hook.session_config = conn_extra.get("session_configuration", {}) + + if isinstance(hook.session_config, dict): + hook.session_config = hook.session_config.copy() + existing_tags = hook.session_config.get("query_tags", "") + if existing_tags: + hook.session_config["query_tags"] = f"{existing_tags},{query_tags}" + else: + hook.session_config["query_tags"] = query_tags + + return super().execute(context) + def _should_run_output_processing(self) -> bool: return self.do_xcom_push or bool(self._output_path) @@ -518,6 +553,20 @@ def execute(self, context: Context) -> Any: self._sql = self._create_sql_query() self.log.info("Executing: %s", self._sql) hook = self._get_hook() + + query_tags = _format_query_tags(context) + if query_tags: + if hook.session_config is None: + conn_extra = hook.databricks_conn.extra_dejson + hook.session_config = conn_extra.get("session_configuration", {}) + if isinstance(hook.session_config, dict): + hook.session_config = hook.session_config.copy() + existing_tags = hook.session_config.get("query_tags", "") + if existing_tags: + hook.session_config["query_tags"] = f"{existing_tags},{query_tags}" + else: + hook.session_config["query_tags"] = query_tags + hook.run(self._sql) def on_kill(self) -> None: diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py b/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py index f00653f8b22c8..d070b5055a0f9 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py @@ -522,3 +522,34 @@ def test_get_openlineage_facets(): "externalQuery": ExternalQueryRunFacet(externalQueryId="query_id", source="scheme://host") } assert result.job_facets == {"sql": SQLJobFacet(query=op._sql)} + + +def test_query_tags_injection(): + """Test that Airflow context is correctly injected into query_tags in session_configuration.""" + from unittest.mock import MagicMock + + with mock.patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as db_mock_class: + op = DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + ) + db_mock = db_mock_class.return_value + db_mock.session_config = {"query_tags": "user_tag:value"} + + class MockConn: + extra_dejson = {} + db_mock.databricks_conn = MockConn() + + context = { + "dag": MagicMock(dag_id="test_dag"), + "task": MagicMock(task_id="test_task"), + "run_id": "test_run_123", + } + + op.execute(context) + + expected_tags = "user_tag:value,airflow_dag_id:test_dag,airflow_task_id:test_task,airflow_run_id:test_run_123" + assert db_mock.session_config["query_tags"] == expected_tags + diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py b/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py index e216c56bea2b4..bb43dd86812a8 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py @@ -453,3 +453,32 @@ def test_parse_gcs_path(): bucket, object_name = op._parse_gcs_path("gs://my-bucket/path/to/file.parquet") assert bucket == "my-bucket" assert object_name == "path/to/file.parquet" + + +def test_query_tags_injection(): + """Test that Airflow context is correctly injected into query_tags in session_configuration.""" + from unittest.mock import MagicMock + + with patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as db_mock_class: + op = DatabricksSqlOperator( + task_id=TASK_ID, + sql="SELECT 1", + session_configuration={"query_tags": "user_tag:value"}, + ) + db_mock = db_mock_class.return_value + db_mock.session_config = {"query_tags": "user_tag:value"} + + class MockConn: + extra_dejson = {} + db_mock.databricks_conn = MockConn() + + context = { + "dag": MagicMock(dag_id="test_dag"), + "task": MagicMock(task_id="test_task"), + "run_id": "test_run_123", + } + + op.execute(context) + + expected_tags = "user_tag:value,airflow_dag_id:test_dag,airflow_task_id:test_task,airflow_run_id:test_run_123" + assert db_mock.session_config["query_tags"] == expected_tags From f38f4bf7ce6ec3b0b82dd8e04ec8186832f66e3b Mon Sep 17 00:00:00 2001 From: Shaan Satsangi <146729792+Shaan-alpha@users.noreply.github.com> Date: Thu, 14 May 2026 01:36:27 +0530 Subject: [PATCH 2/6] Refactor Databricks query tag injection --- .../databricks/operators/databricks_sql.py | 68 +++++++++---------- 1 file changed, 32 insertions(+), 36 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py index 03904b16bd571..65810480d1c92 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py @@ -46,23 +46,46 @@ _DISALLOWED_SQL_TOKENS = (";", "--", "/*", "*/") -def _format_query_tags(context: Context) -> str: - """Format Airflow context metadata into databricks-sql-connector query tags.""" +def _escape_query_tag_value(value: str) -> str: + """Escape Databricks query-tag separators in a tag value.""" + return str(value).replace("\\", "\\\\").replace(",", "\\,").replace(":", "\\:") - def escape(val: str) -> str: - return str(val).replace("\\", "\\\\").replace(",", "\\,").replace(":", "\\:") +def _format_query_tags(context: Context) -> str: + """Format Airflow context metadata into databricks-sql-connector query tags.""" tags = [] if "dag" in context and getattr(context["dag"], "dag_id", None): - tags.append(f"airflow_dag_id:{escape(context['dag'].dag_id)}") + tags.append(f"airflow_dag_id:{_escape_query_tag_value(context['dag'].dag_id)}") if "task" in context and getattr(context["task"], "task_id", None): - tags.append(f"airflow_task_id:{escape(context['task'].task_id)}") + tags.append(f"airflow_task_id:{_escape_query_tag_value(context['task'].task_id)}") if "run_id" in context and context["run_id"]: - tags.append(f"airflow_run_id:{escape(context['run_id'])}") + tags.append(f"airflow_run_id:{_escape_query_tag_value(context['run_id'])}") return ",".join(tags) +def _merge_query_tags(session_config: dict[str, Any], query_tags: str) -> dict[str, Any]: + """Return a copied session config with Airflow query tags appended.""" + updated_config = session_config.copy() + existing_tags = updated_config.get("query_tags", "") + updated_config["query_tags"] = f"{existing_tags},{query_tags}" if existing_tags else query_tags + return updated_config + + +def _inject_query_tags(hook: DatabricksSqlHook, context: Context) -> None: + """Inject Airflow context metadata into Databricks query tags.""" + query_tags = _format_query_tags(context) + if not query_tags: + return + + if hook.session_config is None: + conn_extra = hook.databricks_conn.extra_dejson + hook.session_config = conn_extra.get("session_configuration", {}) + + if isinstance(hook.session_config, dict): + hook.session_config = _merge_query_tags(hook.session_config, query_tags) + + class DatabricksSqlOperator(SQLExecuteQueryOperator): """ Executes SQL code in a Databricks SQL endpoint or a Databricks cluster. @@ -171,21 +194,7 @@ def get_db_hook(self) -> DatabricksSqlHook: return self._hook def execute(self, context: Context) -> Any: - hook = self.get_db_hook() - query_tags = _format_query_tags(context) - if query_tags: - if hook.session_config is None: - conn_extra = hook.databricks_conn.extra_dejson - hook.session_config = conn_extra.get("session_configuration", {}) - - if isinstance(hook.session_config, dict): - hook.session_config = hook.session_config.copy() - existing_tags = hook.session_config.get("query_tags", "") - if existing_tags: - hook.session_config["query_tags"] = f"{existing_tags},{query_tags}" - else: - hook.session_config["query_tags"] = query_tags - + _inject_query_tags(self.get_db_hook(), context) return super().execute(context) def _should_run_output_processing(self) -> bool: @@ -553,20 +562,7 @@ def execute(self, context: Context) -> Any: self._sql = self._create_sql_query() self.log.info("Executing: %s", self._sql) hook = self._get_hook() - - query_tags = _format_query_tags(context) - if query_tags: - if hook.session_config is None: - conn_extra = hook.databricks_conn.extra_dejson - hook.session_config = conn_extra.get("session_configuration", {}) - if isinstance(hook.session_config, dict): - hook.session_config = hook.session_config.copy() - existing_tags = hook.session_config.get("query_tags", "") - if existing_tags: - hook.session_config["query_tags"] = f"{existing_tags},{query_tags}" - else: - hook.session_config["query_tags"] = query_tags - + _inject_query_tags(hook, context) hook.run(self._sql) def on_kill(self) -> None: From e1c8b2d4b64c72f7f1256b6c70ec25897e398917 Mon Sep 17 00:00:00 2001 From: Shaan Satsangi <146729792+Shaan-alpha@users.noreply.github.com> Date: Thu, 14 May 2026 07:21:24 +0530 Subject: [PATCH 3/6] Refactor Databricks query tag helper utilities Refactor query tag formatting and escaping into mapping-driven helper utilities for maintainability and extensibility. --- .../databricks/operators/databricks_sql.py | 36 ++++++++++++++----- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py index 65810480d1c92..29a85ee91b831 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py @@ -45,21 +45,39 @@ _IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") _DISALLOWED_SQL_TOKENS = (";", "--", "/*", "*/") +_QUERY_TAG_FIELDS = { + "airflow_dag_id": ("dag", "dag_id"), + "airflow_task_id": ("task", "task_id"), + "airflow_run_id": ("run_id", None), +} + +_QUERY_TAG_ESCAPE_SEQUENCES = { + "\\": "\\\\", + ",": "\\,", + ":": "\\:", +} + def _escape_query_tag_value(value: str) -> str: - """Escape Databricks query-tag separators in a tag value.""" - return str(value).replace("\\", "\\\\").replace(",", "\\,").replace(":", "\\:") + escaped = str(value) + + for char, replacement in _QUERY_TAG_ESCAPE_SEQUENCES.items(): + escaped = escaped.replace(char, replacement) + + return escaped def _format_query_tags(context: Context) -> str: - """Format Airflow context metadata into databricks-sql-connector query tags.""" tags = [] - if "dag" in context and getattr(context["dag"], "dag_id", None): - tags.append(f"airflow_dag_id:{_escape_query_tag_value(context['dag'].dag_id)}") - if "task" in context and getattr(context["task"], "task_id", None): - tags.append(f"airflow_task_id:{_escape_query_tag_value(context['task'].task_id)}") - if "run_id" in context and context["run_id"]: - tags.append(f"airflow_run_id:{_escape_query_tag_value(context['run_id'])}") + + for tag_name, (context_key, attr) in _QUERY_TAG_FIELDS.items(): + value = context.get(context_key) + + if attr: + value = getattr(value, attr, None) + + if value: + tags.append(f"{tag_name}:{_escape_query_tag_value(value)}") return ",".join(tags) From f1178c5cf9228e452f5ed112cccb67050bef1c46 Mon Sep 17 00:00:00 2001 From: Shaan Satsangi <146729792+Shaan-alpha@users.noreply.github.com> Date: Thu, 14 May 2026 08:04:58 +0530 Subject: [PATCH 4/6] Add inject_query_tags opt-out and expand query-tag test coverage Address review feedback on #66886: - Add `inject_query_tags: bool = True` parameter to DatabricksSqlOperator and DatabricksCopyIntoOperator, allowing users to opt out of the automatic session_configuration mutation while preserving the default observability benefit. Documented in both operator docstrings. - Move query-tag tests above the OpenLineage block (test_databricks_copy) and after test_exec_write_gcs_parquet_output (test_databricks_sql) so they sit closer to core execution logic. - Expand coverage in both test files: empty/partial context, empty existing query_tags, end-to-end exercise of escape sequences for commas/colons/backslashes, preservation of unrelated session_configuration keys, fallback to connection extras when session_config is None, and verification of the opt-out path. --- .../databricks/operators/databricks_sql.py | 20 +- .../operators/test_databricks_copy.py | 186 +++++++++++++++--- .../operators/test_databricks_sql.py | 148 +++++++++++--- 3 files changed, 292 insertions(+), 62 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py index 29a85ee91b831..bb577d7b7b3a4 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py @@ -141,6 +141,11 @@ class DatabricksSqlOperator(SQLExecuteQueryOperator): :param gcs_impersonation_chain: Optional service account to impersonate using short-term credentials for GCS upload, or chained list of accounts required to get the access_token of the last account in the list, which will be impersonated in the request. (templated) + :param inject_query_tags: If ``True`` (default), Airflow context metadata + (``airflow_dag_id``, ``airflow_task_id``, ``airflow_run_id``) is injected into the + Databricks session ``query_tags`` at execution time, preserving any user-defined + ``query_tags`` already present in ``session_configuration``. Set to ``False`` to + retain full control over ``session_configuration`` and skip the automatic injection. """ template_fields: Sequence[str] = tuple( @@ -175,6 +180,7 @@ def __init__( client_parameters: dict[str, Any] | None = None, gcp_conn_id: str = "google_cloud_default", gcs_impersonation_chain: str | Sequence[str] | None = None, + inject_query_tags: bool = True, **kwargs, ) -> None: super().__init__(conn_id=databricks_conn_id, **kwargs) @@ -192,6 +198,7 @@ def __init__( self.schema = schema self._gcp_conn_id = gcp_conn_id self._gcs_impersonation_chain = gcs_impersonation_chain + self.inject_query_tags = inject_query_tags @cached_property def _hook(self) -> DatabricksSqlHook: @@ -212,7 +219,8 @@ def get_db_hook(self) -> DatabricksSqlHook: return self._hook def execute(self, context: Context) -> Any: - _inject_query_tags(self.get_db_hook(), context) + if self.inject_query_tags: + _inject_query_tags(self.get_db_hook(), context) return super().execute(context) def _should_run_output_processing(self) -> bool: @@ -410,6 +418,11 @@ class DatabricksCopyIntoOperator(BaseOperator): :param validate: optional configuration for schema & data validation. ``True`` forces validation of all rows, integer number - validate only N first rows :param copy_options: optional dictionary of copy options. Right now only ``force`` option is supported. + :param inject_query_tags: If ``True`` (default), Airflow context metadata + (``airflow_dag_id``, ``airflow_task_id``, ``airflow_run_id``) is injected into the + Databricks session ``query_tags`` at execution time, preserving any user-defined + ``query_tags`` already present in ``session_configuration``. Set to ``False`` to + retain full control over ``session_configuration`` and skip the automatic injection. """ template_fields: Sequence[str] = ( @@ -443,6 +456,7 @@ def __init__( force_copy: bool | None = None, copy_options: dict[str, str] | None = None, validate: bool | int | None = None, + inject_query_tags: bool = True, **kwargs, ) -> None: """Create a new ``DatabricksSqlOperator``.""" @@ -477,6 +491,7 @@ def __init__( self._client_parameters = client_parameters or {} if force_copy is not None: self._copy_options["force"] = "true" if force_copy else "false" + self.inject_query_tags = inject_query_tags self._sql: str | None = None def _get_hook(self) -> DatabricksSqlHook: @@ -580,7 +595,8 @@ def execute(self, context: Context) -> Any: self._sql = self._create_sql_query() self.log.info("Executing: %s", self._sql) hook = self._get_hook() - _inject_query_tags(hook, context) + if self.inject_query_tags: + _inject_query_tags(hook, context) hook.run(self._sql) def on_kill(self) -> None: diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py b/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py index d070b5055a0f9..1055bd5da2ed4 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py @@ -367,6 +367,161 @@ def test_hook_is_cached(): assert hook is hook2 +def _make_context(*, dag_id=None, task_id=None, run_id=None): + context: dict = {} + if dag_id is not None: + context["dag"] = mock.MagicMock(dag_id=dag_id) + if task_id is not None: + context["task"] = mock.MagicMock(task_id=task_id) + if run_id is not None: + context["run_id"] = run_id + return context + + +def _run_with_mocked_hook(op, context, initial_session_config, conn_extra=None): + """Execute the operator with a mocked hook and return the resulting session_config.""" + with mock.patch( + "airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook" + ) as db_mock_class: + db_mock = db_mock_class.return_value + db_mock.session_config = initial_session_config + db_mock.databricks_conn = mock.MagicMock(extra_dejson=conn_extra or {}) + op.execute(context) + return db_mock.session_config + + +def test_query_tags_injection_appends_to_existing_tags(): + op = DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + ) + context = _make_context(dag_id="test_dag", task_id="test_task", run_id="test_run_123") + + result = _run_with_mocked_hook(op, context, {"query_tags": "user_tag:value"}) + + assert result["query_tags"] == ( + "user_tag:value,airflow_dag_id:test_dag," + "airflow_task_id:test_task,airflow_run_id:test_run_123" + ) + + +def test_query_tags_injection_with_no_existing_tags(): + op = DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + ) + context = _make_context(dag_id="d", task_id="t", run_id="r") + + result = _run_with_mocked_hook(op, context, {}) + + assert result["query_tags"] == "airflow_dag_id:d,airflow_task_id:t,airflow_run_id:r" + + +def test_query_tags_injection_with_partial_context(): + op = DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + ) + context = _make_context(task_id="only_task") + + result = _run_with_mocked_hook(op, context, {}) + + assert result["query_tags"] == "airflow_task_id:only_task" + + +def test_query_tags_injection_with_empty_context(): + op = DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + ) + + result = _run_with_mocked_hook(op, {}, {"unrelated": "keep"}) + + assert result == {"unrelated": "keep"} + + +def test_query_tags_injection_escapes_special_chars(): + op = DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + ) + context = _make_context( + dag_id="dag,with,commas", + task_id="task:with:colons", + run_id="run\\with\\backslashes", + ) + + result = _run_with_mocked_hook(op, context, {}) + + assert result["query_tags"] == ( + "airflow_dag_id:dag\\,with\\,commas," + "airflow_task_id:task\\:with\\:colons," + "airflow_run_id:run\\\\with\\\\backslashes" + ) + + +def test_query_tags_injection_preserves_unrelated_session_config(): + op = DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + ) + context = _make_context(dag_id="d", task_id="t", run_id="r") + initial = {"spark.sql.shuffle.partitions": "200", "query_tags": "x:y"} + + result = _run_with_mocked_hook(op, context, initial) + + assert result["spark.sql.shuffle.partitions"] == "200" + assert result["query_tags"] == "x:y,airflow_dag_id:d,airflow_task_id:t,airflow_run_id:r" + + +def test_query_tags_injection_falls_back_to_conn_extra_when_session_config_none(): + op = DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + ) + context = _make_context(dag_id="d", task_id="t", run_id="r") + + result = _run_with_mocked_hook( + op, + context, + initial_session_config=None, + conn_extra={"session_configuration": {"query_tags": "conn_tag:1"}}, + ) + + assert result["query_tags"] == ( + "conn_tag:1,airflow_dag_id:d,airflow_task_id:t,airflow_run_id:r" + ) + + +def test_query_tags_injection_disabled(): + op = DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + inject_query_tags=False, + ) + context = _make_context(dag_id="d", task_id="t", run_id="r") + + result = _run_with_mocked_hook(op, context, {"query_tags": "user_tag:value"}) + + assert result == {"query_tags": "user_tag:value"} + + @pytest.mark.parametrize( ("file_location", "expected_namespace", "expected_name"), ( @@ -522,34 +677,3 @@ def test_get_openlineage_facets(): "externalQuery": ExternalQueryRunFacet(externalQueryId="query_id", source="scheme://host") } assert result.job_facets == {"sql": SQLJobFacet(query=op._sql)} - - -def test_query_tags_injection(): - """Test that Airflow context is correctly injected into query_tags in session_configuration.""" - from unittest.mock import MagicMock - - with mock.patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as db_mock_class: - op = DatabricksCopyIntoOperator( - task_id=TASK_ID, - file_location=COPY_FILE_LOCATION, - file_format="JSON", - table_name="test", - ) - db_mock = db_mock_class.return_value - db_mock.session_config = {"query_tags": "user_tag:value"} - - class MockConn: - extra_dejson = {} - db_mock.databricks_conn = MockConn() - - context = { - "dag": MagicMock(dag_id="test_dag"), - "task": MagicMock(task_id="test_task"), - "run_id": "test_run_123", - } - - op.execute(context) - - expected_tags = "user_tag:value,airflow_dag_id:test_dag,airflow_task_id:test_task,airflow_run_id:test_run_123" - assert db_mock.session_config["query_tags"] == expected_tags - diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py b/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py index bb43dd86812a8..721cbfbc4a24a 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py @@ -432,6 +432,125 @@ def test_exec_write_gcs_parquet_output(tmp_path): assert call_kwargs["object_name"] == "data/results.parquet" +def _make_context(*, dag_id=None, task_id=None, run_id=None): + from unittest.mock import MagicMock + + context: dict = {} + if dag_id is not None: + context["dag"] = MagicMock(dag_id=dag_id) + if task_id is not None: + context["task"] = MagicMock(task_id=task_id) + if run_id is not None: + context["run_id"] = run_id + return context + + +def _run_with_mocked_hook(op, context, initial_session_config, conn_extra=None): + """Execute the operator with a mocked hook and return the resulting session_config.""" + from unittest.mock import MagicMock + + op.do_xcom_push = False + with patch( + "airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook" + ) as db_mock_class: + db_mock = db_mock_class.return_value + db_mock.session_config = initial_session_config + db_mock.databricks_conn = MagicMock(extra_dejson=conn_extra or {}) + op.execute(context) + return db_mock.session_config + + +def test_query_tags_injection_appends_to_existing_tags(): + op = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1") + context = _make_context(dag_id="test_dag", task_id="test_task", run_id="test_run_123") + + result = _run_with_mocked_hook(op, context, {"query_tags": "user_tag:value"}) + + assert result["query_tags"] == ( + "user_tag:value,airflow_dag_id:test_dag," + "airflow_task_id:test_task,airflow_run_id:test_run_123" + ) + + +def test_query_tags_injection_with_no_existing_tags(): + op = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1") + context = _make_context(dag_id="d", task_id="t", run_id="r") + + result = _run_with_mocked_hook(op, context, {}) + + assert result["query_tags"] == "airflow_dag_id:d,airflow_task_id:t,airflow_run_id:r" + + +def test_query_tags_injection_with_partial_context(): + op = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1") + context = _make_context(task_id="only_task") + + result = _run_with_mocked_hook(op, context, {}) + + assert result["query_tags"] == "airflow_task_id:only_task" + + +def test_query_tags_injection_with_empty_context(): + op = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1") + + result = _run_with_mocked_hook(op, {}, {"unrelated": "keep"}) + + assert result == {"unrelated": "keep"} + + +def test_query_tags_injection_escapes_special_chars(): + op = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1") + context = _make_context( + dag_id="dag,with,commas", + task_id="task:with:colons", + run_id="run\\with\\backslashes", + ) + + result = _run_with_mocked_hook(op, context, {}) + + assert result["query_tags"] == ( + "airflow_dag_id:dag\\,with\\,commas," + "airflow_task_id:task\\:with\\:colons," + "airflow_run_id:run\\\\with\\\\backslashes" + ) + + +def test_query_tags_injection_preserves_unrelated_session_config(): + op = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1") + context = _make_context(dag_id="d", task_id="t", run_id="r") + initial = {"spark.sql.shuffle.partitions": "200", "query_tags": "x:y"} + + result = _run_with_mocked_hook(op, context, initial) + + assert result["spark.sql.shuffle.partitions"] == "200" + assert result["query_tags"] == "x:y,airflow_dag_id:d,airflow_task_id:t,airflow_run_id:r" + + +def test_query_tags_injection_falls_back_to_conn_extra_when_session_config_none(): + op = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1") + context = _make_context(dag_id="d", task_id="t", run_id="r") + + result = _run_with_mocked_hook( + op, + context, + initial_session_config=None, + conn_extra={"session_configuration": {"query_tags": "conn_tag:1"}}, + ) + + assert result["query_tags"] == ( + "conn_tag:1,airflow_dag_id:d,airflow_task_id:t,airflow_run_id:r" + ) + + +def test_query_tags_injection_disabled(): + op = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1", inject_query_tags=False) + context = _make_context(dag_id="d", task_id="t", run_id="r") + + result = _run_with_mocked_hook(op, context, {"query_tags": "user_tag:value"}) + + assert result == {"query_tags": "user_tag:value"} + + def test_is_gcs_output(): """Test _is_gcs_output property.""" op_gcs = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1", output_path="gs://bucket/path") @@ -453,32 +572,3 @@ def test_parse_gcs_path(): bucket, object_name = op._parse_gcs_path("gs://my-bucket/path/to/file.parquet") assert bucket == "my-bucket" assert object_name == "path/to/file.parquet" - - -def test_query_tags_injection(): - """Test that Airflow context is correctly injected into query_tags in session_configuration.""" - from unittest.mock import MagicMock - - with patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as db_mock_class: - op = DatabricksSqlOperator( - task_id=TASK_ID, - sql="SELECT 1", - session_configuration={"query_tags": "user_tag:value"}, - ) - db_mock = db_mock_class.return_value - db_mock.session_config = {"query_tags": "user_tag:value"} - - class MockConn: - extra_dejson = {} - db_mock.databricks_conn = MockConn() - - context = { - "dag": MagicMock(dag_id="test_dag"), - "task": MagicMock(task_id="test_task"), - "run_id": "test_run_123", - } - - op.execute(context) - - expected_tags = "user_tag:value,airflow_dag_id:test_dag,airflow_task_id:test_task,airflow_run_id:test_run_123" - assert db_mock.session_config["query_tags"] == expected_tags From c81da97b3bf4f584de83f96ba5146522d27c54b4 Mon Sep 17 00:00:00 2001 From: Shaan Satsangi Date: Mon, 18 May 2026 03:40:10 +0530 Subject: [PATCH 5/6] docs(providers/databricks): add inject_query_tags usage examples (#66886) --- providers/databricks/docs/operators/sql.rst | 10 ++++++++++ .../system/databricks/example_databricks_sql.py | 13 ++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/providers/databricks/docs/operators/sql.rst b/providers/databricks/docs/operators/sql.rst index 448fce7bbf8f0..06b395ac0e796 100644 --- a/providers/databricks/docs/operators/sql.rst +++ b/providers/databricks/docs/operators/sql.rst @@ -64,6 +64,16 @@ An example usage of the DatabricksSqlOperator to select data from a table and st :start-after: [START howto_operator_databricks_sql_select_file] :end-before: [END howto_operator_databricks_sql_select_file] +Executing with automatic query tags +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +An example usage of the DatabricksSqlOperator with automatic Airflow context query tags injection is as follows: + +.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_sql.py + :language: python + :start-after: [START howto_operator_databricks_sql_query_tags] + :end-before: [END howto_operator_databricks_sql_query_tags] + Executing multiple statements ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/providers/databricks/tests/system/databricks/example_databricks_sql.py b/providers/databricks/tests/system/databricks/example_databricks_sql.py index 2782525a55ccc..170c8f25de9f1 100644 --- a/providers/databricks/tests/system/databricks/example_databricks_sql.py +++ b/providers/databricks/tests/system/databricks/example_databricks_sql.py @@ -86,6 +86,17 @@ ) # [END howto_operator_databricks_sql_select_file] + # [START howto_operator_databricks_sql_query_tags] + # Example of using the Databricks SQL Operator with automatic Airflow query tags injection. + select_with_query_tags = DatabricksSqlOperator( + databricks_conn_id=connection_id, + sql_endpoint_name=sql_endpoint_name, + task_id="select_with_query_tags", + sql="select * from default.my_airflow_table", + inject_query_tags=True, + ) + # [END howto_operator_databricks_sql_query_tags] + # [START howto_operator_databricks_sql_multiple_file] # Example of using the Databricks SQL Operator to select data. # SQL statements should be in the file with name test.sql @@ -111,7 +122,7 @@ ) # [END howto_operator_databricks_copy_into] - (create >> create_file >> import_csv >> select >> select_into_file) + (create >> create_file >> import_csv >> select >> select_into_file >> select_with_query_tags) from tests_common.test_utils.watcher import watcher From 7441d5a11943a91a6e7b3bd95767bb85345a96b4 Mon Sep 17 00:00:00 2001 From: Shaan Satsangi Date: Tue, 19 May 2026 13:48:00 +0530 Subject: [PATCH 6/6] fix(providers/databricks): guard query-tag injection against None context Existing operator tests call op.execute(None), and the new _inject_query_tags helper unconditionally called context.get(...), raising AttributeError: 'NoneType' object has no attribute 'get' across ~30 tests in the compat, low-dep, and non-DB provider jobs. - Early-return in _inject_query_tags when context is falsy. - Broaden _escape_query_tag_value parameter to object (already calls str(value) internally), resolving the MyPy [arg-type] error at databricks_sql.py:80. - Apply ruff-format to the new test files for ruff==0.15.12. Co-Authored-By: Claude Opus 4.7 --- .../providers/databricks/operators/databricks_sql.py | 6 ++++-- .../unit/databricks/operators/test_databricks_copy.py | 7 ++----- .../unit/databricks/operators/test_databricks_sql.py | 11 +++-------- 3 files changed, 9 insertions(+), 15 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py index bb577d7b7b3a4..e972545cd03ae 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py @@ -58,7 +58,7 @@ } -def _escape_query_tag_value(value: str) -> str: +def _escape_query_tag_value(value: object) -> str: escaped = str(value) for char, replacement in _QUERY_TAG_ESCAPE_SEQUENCES.items(): @@ -90,8 +90,10 @@ def _merge_query_tags(session_config: dict[str, Any], query_tags: str) -> dict[s return updated_config -def _inject_query_tags(hook: DatabricksSqlHook, context: Context) -> None: +def _inject_query_tags(hook: DatabricksSqlHook, context: Context | None) -> None: """Inject Airflow context metadata into Databricks query tags.""" + if not context: + return query_tags = _format_query_tags(context) if not query_tags: return diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py b/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py index 1055bd5da2ed4..e1845cd357238 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks_copy.py @@ -402,8 +402,7 @@ def test_query_tags_injection_appends_to_existing_tags(): result = _run_with_mocked_hook(op, context, {"query_tags": "user_tag:value"}) assert result["query_tags"] == ( - "user_tag:value,airflow_dag_id:test_dag," - "airflow_task_id:test_task,airflow_run_id:test_run_123" + "user_tag:value,airflow_dag_id:test_dag,airflow_task_id:test_task,airflow_run_id:test_run_123" ) @@ -502,9 +501,7 @@ def test_query_tags_injection_falls_back_to_conn_extra_when_session_config_none( conn_extra={"session_configuration": {"query_tags": "conn_tag:1"}}, ) - assert result["query_tags"] == ( - "conn_tag:1,airflow_dag_id:d,airflow_task_id:t,airflow_run_id:r" - ) + assert result["query_tags"] == ("conn_tag:1,airflow_dag_id:d,airflow_task_id:t,airflow_run_id:r") def test_query_tags_injection_disabled(): diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py b/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py index 721cbfbc4a24a..06a7882c70874 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py @@ -450,9 +450,7 @@ def _run_with_mocked_hook(op, context, initial_session_config, conn_extra=None): from unittest.mock import MagicMock op.do_xcom_push = False - with patch( - "airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook" - ) as db_mock_class: + with patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as db_mock_class: db_mock = db_mock_class.return_value db_mock.session_config = initial_session_config db_mock.databricks_conn = MagicMock(extra_dejson=conn_extra or {}) @@ -467,8 +465,7 @@ def test_query_tags_injection_appends_to_existing_tags(): result = _run_with_mocked_hook(op, context, {"query_tags": "user_tag:value"}) assert result["query_tags"] == ( - "user_tag:value,airflow_dag_id:test_dag," - "airflow_task_id:test_task,airflow_run_id:test_run_123" + "user_tag:value,airflow_dag_id:test_dag,airflow_task_id:test_task,airflow_run_id:test_run_123" ) @@ -537,9 +534,7 @@ def test_query_tags_injection_falls_back_to_conn_extra_when_session_config_none( conn_extra={"session_configuration": {"query_tags": "conn_tag:1"}}, ) - assert result["query_tags"] == ( - "conn_tag:1,airflow_dag_id:d,airflow_task_id:t,airflow_run_id:r" - ) + assert result["query_tags"] == ("conn_tag:1,airflow_dag_id:d,airflow_task_id:t,airflow_run_id:r") def test_query_tags_injection_disabled():