diff --git a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py index 021142395b299..ffc089607c33a 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import logging import threading from collections import namedtuple from collections.abc import Callable, Iterable, Mapping, Sequence @@ -51,6 +52,8 @@ T = TypeVar("T") +log = logging.getLogger(__name__) + def create_timeout_thread( cur, execution_timeout: timedelta | None @@ -71,6 +74,35 @@ def _cancel(): return timer, timeout_event +def _format_query_tag_value(value: str) -> str: + """ + Escape special characters and truncate a single query tag value. + + Databricks ``QUERY_TAGS`` uses ``key:value`` pairs delimited by commas, so + backslash, comma and colon inside *values* must be escaped. Values are also + capped at 128 characters before escaping to keep the overall tag string + within reasonable bounds. + """ + raw = str(value) + if len(raw) > 128: + log.warning( + "Query tag value truncated to 128 characters (original length %d): %r", len(raw), raw[:128] + ) + value = raw[:128] + return value.replace("\\", "\\\\").replace(",", "\\,").replace(":", "\\:") + + +def _format_query_tags(tags: dict[str, str | None]) -> str: + """ + Serialize a query-tags dict to the ``key:value,key:value`` string expected by ``QUERY_TAGS``. + + Entries whose value is ``None`` are omitted. + """ + return ",".join( + f"{key}:{_format_query_tag_value(value)}" for key, value in tags.items() if value is not None + ) + + class DatabricksSqlHook(BaseDatabricksHook, DbApiHook): """ Hook to interact with Databricks SQL. @@ -88,6 +120,10 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook): on every request :param catalog: An optional initial catalog to use. Requires DBR version 9.0+ :param schema: An optional initial schema to use. Requires DBR version 9.0+ + :param query_tags: An optional dict of query tags to attach to every SQL statement executed by + this hook. Tags are injected via the ``QUERY_TAGS`` Databricks session parameter so they + appear in ``system.query.history``. Any existing ``QUERY_TAGS`` already present in + *session_configuration* are preserved and the new tags are appended. :param kwargs: Additional parameters internal to Databricks SQL Connector parameters """ @@ -104,6 +140,7 @@ def __init__( catalog: str | None = None, schema: str | None = None, caller: str = "DatabricksSqlHook", + query_tags: dict[str, str | None] | None = None, **kwargs, ) -> None: super().__init__(databricks_conn_id, caller=caller) @@ -118,6 +155,7 @@ def __init__( self.schema = schema self.additional_params = kwargs self.query_ids: list[str] = [] + self.query_tags = query_tags def _get_extra_config(self) -> dict[str, Any | None]: extra_params = copy(self.databricks_conn.extra_dejson) @@ -169,20 +207,32 @@ def get_conn(self) -> AirflowConnection: if not self.session_config: self.session_config = self.databricks_conn.extra_dejson.get("session_configuration") + # session_configuration (including QUERY_TAGS) is applied only when opening a new + # connection; changing query_tags after the first get_conn() call has no effect. if not self._sql_conn or prev_token != new_token: if self._sql_conn: # close already existing connection self._sql_conn.close() + session_config: dict[str, str] = dict(self.session_config) if self.session_config else {} + if self.query_tags: + tags_str = _format_query_tags(self.query_tags) + existing = session_config.get("QUERY_TAGS", "") + session_config["QUERY_TAGS"] = f"{existing},{tags_str}" if existing else tags_str + + connect_kwargs = { + "schema": self.schema, + "catalog": self.catalog, + "session_configuration": session_config or None, + "http_headers": self.http_headers, + "_user_agent_entry": self.user_agent_value, + **self._get_extra_config(), + **self.additional_params, + } + self._sql_conn = sql.connect( self.host, self._http_path, self._token, - schema=self.schema, - catalog=self.catalog, - session_configuration=self.session_config, - http_headers=self.http_headers, - _user_agent_entry=self.user_agent_value, - **self._get_extra_config(), - **self.additional_params, + **connect_kwargs, ) if self._sql_conn is None: 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..f72514f2488f9 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,24 @@ _DISALLOWED_SQL_TOKENS = (";", "--", "/*", "*/") +def _get_airflow_query_tags(context: Context) -> dict[str, str | None]: + """Return Airflow context metadata as a query-tags dict.""" + task_instance = context.get("ti") + if task_instance is None: + return {} + + def _as_str(value: Any) -> str | None: + return None if value is None else str(value) + + return { + "airflow_dag_id": _as_str(task_instance.dag_id), + "airflow_task_id": _as_str(task_instance.task_id), + "airflow_run_id": _as_str(task_instance.run_id), + "airflow_try_number": _as_str(task_instance.try_number), + "airflow_map_index": _as_str(task_instance.map_index), + } + + class DatabricksSqlOperator(SQLExecuteQueryOperator): """ Executes SQL code in a Databricks SQL endpoint or a Databricks cluster. @@ -68,6 +86,8 @@ class DatabricksSqlOperator(SQLExecuteQueryOperator): :param session_configuration: An optional dictionary of Spark session parameters. Defaults to None. If not specified, it could be specified in the Databricks connection's extra parameters. :param client_parameters: Additional parameters internal to Databricks SQL Connector parameters + :param query_tags: Optional dictionary of query tags to attach to Databricks SQL queries. + :param include_airflow_query_tags: If True, add Airflow DAG/task/run metadata as query tags. :param http_headers: An optional list of (k, v) pairs that will be set as HTTP headers on every request. (templated) :param catalog: An optional initial catalog to use. Requires DBR version 9.0+ (templated) @@ -93,6 +113,7 @@ class DatabricksSqlOperator(SQLExecuteQueryOperator): "http_headers", "databricks_conn_id", "_gcs_impersonation_chain", + "query_tags", } | set(SQLExecuteQueryOperator.template_fields) ) @@ -115,6 +136,8 @@ def __init__( output_format: str = "csv", csv_params: dict[str, Any] | None = None, client_parameters: dict[str, Any] | None = None, + query_tags: dict[str, str | None] | None = None, + include_airflow_query_tags: bool = True, gcp_conn_id: str = "google_cloud_default", gcs_impersonation_chain: str | Sequence[str] | None = None, **kwargs, @@ -132,6 +155,8 @@ def __init__( self.http_headers = http_headers self.catalog = catalog self.schema = schema + self.query_tags = query_tags or {} + self.include_airflow_query_tags = include_airflow_query_tags self._gcp_conn_id = gcp_conn_id self._gcs_impersonation_chain = gcs_impersonation_chain @@ -303,6 +328,20 @@ def _process_output(self, results: list[Any], descriptions: list[Sequence[Sequen return list(zip(descriptions, results)) + def _get_query_tags(self, context: Context) -> dict[str, str | None] | None: + query_tags: dict[str, str | None] = {} + + if self.include_airflow_query_tags and context is not None: + query_tags.update(_get_airflow_query_tags(context)) + + query_tags.update(self.query_tags) + + return query_tags or None + + def execute(self, context: Context) -> Any: + self.get_db_hook().query_tags = self._get_query_tags(context) + return super().execute(context) + COPY_INTO_APPROVED_FORMATS = ["CSV", "JSON", "AVRO", "ORC", "PARQUET", "TEXT", "BINARYFILE"] @@ -335,6 +374,8 @@ class DatabricksCopyIntoOperator(BaseOperator): :param catalog: An optional initial catalog to use. Requires DBR version 9.0+ :param schema: An optional initial schema to use. Requires DBR version 9.0+ :param client_parameters: Additional parameters internal to Databricks SQL Connector parameters + :param query_tags: Optional dictionary of query tags to attach to Databricks SQL queries. + :param include_airflow_query_tags: If True, add Airflow DAG/task/run metadata as query tags. :param files: optional list of files to import. Can't be specified together with ``pattern``. (templated) :param pattern: optional regex string to match file names to import. Can't be specified together with ``files``. @@ -355,6 +396,7 @@ class DatabricksCopyIntoOperator(BaseOperator): "files", "table_name", "databricks_conn_id", + "query_tags", ) def __init__( @@ -381,9 +423,11 @@ def __init__( force_copy: bool | None = None, copy_options: dict[str, str] | None = None, validate: bool | int | None = None, + query_tags: dict[str, str | None] | None = None, + include_airflow_query_tags: bool = True, **kwargs, ) -> None: - """Create a new ``DatabricksSqlOperator``.""" + """Create a new ``DatabricksCopyIntoOperator``.""" super().__init__(**kwargs) if files is not None and pattern is not None: raise AirflowException("Only one of 'pattern' or 'files' should be specified") @@ -413,6 +457,8 @@ def __init__( self._validate = validate self._http_headers = http_headers self._client_parameters = client_parameters or {} + self.query_tags = query_tags or {} + self.include_airflow_query_tags = include_airflow_query_tags if force_copy is not None: self._copy_options["force"] = "true" if force_copy else "false" self._sql: str | None = None @@ -514,10 +560,21 @@ def _create_sql_query(self) -> str: """ return sql.strip() + def _get_query_tags(self, context: Context) -> dict[str, str | None] | None: + query_tags: dict[str, str | None] = {} + + if self.include_airflow_query_tags and context is not None: + query_tags.update(_get_airflow_query_tags(context)) + + query_tags.update(self.query_tags) + + return query_tags or None + def execute(self, context: Context) -> Any: self._sql = self._create_sql_query() self.log.info("Executing: %s", self._sql) hook = self._get_hook() + hook.query_tags = self._get_query_tags(context) hook.run(self._sql) def on_kill(self) -> None: diff --git a/providers/databricks/src/airflow/providers/databricks/sensors/databricks_partition.py b/providers/databricks/src/airflow/providers/databricks/sensors/databricks_partition.py index b4501ef1d4325..2036dca97c387 100644 --- a/providers/databricks/src/airflow/providers/databricks/sensors/databricks_partition.py +++ b/providers/databricks/src/airflow/providers/databricks/sensors/databricks_partition.py @@ -130,7 +130,7 @@ def _get_hook(self) -> DatabricksSqlHook: self.http_headers, self.catalog, self.schema, - self.caller, + caller=self.caller, **self.client_parameters, **self.hook_params, ) diff --git a/providers/databricks/src/airflow/providers/databricks/sensors/databricks_sql.py b/providers/databricks/src/airflow/providers/databricks/sensors/databricks_sql.py index 68a44a20fec2a..6ab52df67b579 100644 --- a/providers/databricks/src/airflow/providers/databricks/sensors/databricks_sql.py +++ b/providers/databricks/src/airflow/providers/databricks/sensors/databricks_sql.py @@ -109,7 +109,7 @@ def hook(self) -> DatabricksSqlHook: self.http_headers, self.catalog, self.schema, - self.caller, + caller=self.caller, **self.client_parameters, **self.hook_params, ) diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py index f3f053d443c75..cd3c00e2839b1 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py @@ -32,7 +32,12 @@ from airflow.models import Connection from airflow.providers.common.compat.sdk import AirflowException, AirflowOptionalProviderFeatureException from airflow.providers.common.sql.hooks.handlers import fetch_all_handler -from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook, create_timeout_thread +from airflow.providers.databricks.hooks.databricks_sql import ( + DatabricksSqlHook, + _format_query_tag_value, + _format_query_tags, + create_timeout_thread, +) TASK_ID = "databricks-sql-operator" DEFAULT_CONN_ID = "databricks_default" @@ -792,3 +797,141 @@ def test_resolve_warehouse_name_empty_response(self, mock_requests): hook = DatabricksSqlHook(sql_endpoint_name="Test") with pytest.raises(RuntimeError, match="Can't list Databricks SQL warehouses"): hook._get_sql_endpoint_by_name("Test") + + +@mock.patch("airflow.providers.databricks.hooks.databricks_sql.sql.connect") +def test_get_conn_passes_query_tags_via_session_configuration(mock_connect, mock_get_requests): + """query_tags must be injected into session_configuration['QUERY_TAGS'], not sql.connect(query_tags=).""" + hook = DatabricksSqlHook( + databricks_conn_id=DEFAULT_CONN_ID, + http_path=HTTP_PATH, + query_tags={"airflow_dag_id": "dag_1", "airflow_task_id": "task_1"}, + ) + + hook.get_conn() + + mock_connect.assert_called_once() + session_cfg = mock_connect.call_args.kwargs["session_configuration"] + assert session_cfg is not None + assert "QUERY_TAGS" in session_cfg + query_tags_str = session_cfg["QUERY_TAGS"] + assert "airflow_dag_id:dag_1" in query_tags_str + assert "airflow_task_id:task_1" in query_tags_str + assert "query_tags" not in mock_connect.call_args.kwargs + + +@mock.patch("airflow.providers.databricks.hooks.databricks_sql.sql.connect") +def test_get_conn_merges_query_tags_with_existing_session_configuration(mock_connect, mock_get_requests): + """Existing QUERY_TAGS in session_configuration must be preserved and new tags appended.""" + hook = DatabricksSqlHook( + databricks_conn_id=DEFAULT_CONN_ID, + http_path=HTTP_PATH, + session_configuration={"QUERY_TAGS": "existing_tag:existing_value"}, + query_tags={"airflow_dag_id": "dag_1"}, + ) + + hook.get_conn() + + mock_connect.assert_called_once() + session_cfg = mock_connect.call_args.kwargs["session_configuration"] + query_tags_str = session_cfg["QUERY_TAGS"] + assert "existing_tag:existing_value" in query_tags_str + assert "airflow_dag_id:dag_1" in query_tags_str + + +@mock.patch("airflow.providers.databricks.hooks.databricks_sql.sql.connect") +def test_get_conn_no_query_tags(mock_connect, mock_get_requests): + """When no query_tags are provided, session_configuration should not gain a QUERY_TAGS key.""" + hook = DatabricksSqlHook( + databricks_conn_id=DEFAULT_CONN_ID, + http_path=HTTP_PATH, + ) + + hook.get_conn() + + mock_connect.assert_called_once() + session_cfg = mock_connect.call_args.kwargs.get("session_configuration") + assert session_cfg is None or "QUERY_TAGS" not in session_cfg + + +class TestFormatQueryTags: + def test_simple_values(self): + result = _format_query_tags({"dag_id": "my_dag", "task_id": "my_task"}) + assert "dag_id:my_dag" in result + assert "task_id:my_task" in result + + def test_none_values_omitted(self): + result = _format_query_tags({"dag_id": "my_dag", "map_index": None}) + assert "dag_id:my_dag" in result + assert "map_index" not in result + + def test_empty_dict_returns_empty_string(self): + assert _format_query_tags({}) == "" + + def test_value_escaping_comma(self): + result = _format_query_tag_value("a,b") + assert result == "a\\,b" + + def test_value_escaping_colon(self): + result = _format_query_tag_value("a:b") + assert result == "a\\:b" + + def test_value_escaping_backslash(self): + result = _format_query_tag_value("a\\b") + assert result == "a\\\\b" + + def test_value_truncated_at_128_chars(self): + long_value = "x" * 200 + result = _format_query_tag_value(long_value) + assert len(result) == 128 + + def test_format_query_tags_roundtrip(self): + tags = {"airflow_dag_id": "dag:1", "airflow_run_id": "run,2"} + result = _format_query_tags(tags) + assert "airflow_dag_id:dag\\:1" in result + assert "airflow_run_id:run\\,2" in result + + +class TestDatabricksSqlHookQueryTagsParamOrder: + """Ensure moving query_tags after caller preserves positional backward compatibility.""" + + def test_query_tags_keyword_sets_field(self): + """query_tags kwarg must be stored on the instance.""" + with patch( + "airflow.providers.databricks.hooks.databricks_sql.BaseDatabricksHook.__init__", + return_value=None, + ) as mock_base_init: + hook = DatabricksSqlHook.__new__(DatabricksSqlHook) + DatabricksSqlHook.__init__( + hook, + DEFAULT_CONN_ID, + query_tags={"key": "val"}, + ) + assert hook.query_tags == {"key": "val"} + # caller is forwarded to BaseDatabricksHook.__init__; verify the default was passed + assert mock_base_init.call_args.kwargs.get("caller") == "DatabricksSqlHook" + + def test_caller_positional_not_confused_with_query_tags(self): + """Passing caller as the 8th positional arg must not end up in query_tags.""" + with patch( + "airflow.providers.databricks.hooks.databricks_sql.BaseDatabricksHook.__init__", + return_value=None, + ) as mock_base_init: + hook = DatabricksSqlHook.__new__(DatabricksSqlHook) + # positional order: conn_id, http_path, sql_endpoint, session_cfg, + # http_headers, catalog, schema, caller + DatabricksSqlHook.__init__( + hook, + DEFAULT_CONN_ID, # databricks_conn_id + None, # http_path + None, # sql_endpoint_name + None, # session_configuration + None, # http_headers + None, # catalog + None, # schema + "CustomCaller", # caller (8th positional) + ) + # caller is forwarded to BaseDatabricksHook.__init__; verify it was not + # confused with query_tags (which comes after caller) + assert mock_base_init.call_args.kwargs.get("caller") == "CustomCaller" + assert hook.query_tags is None 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..bfd6d89437b15 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks_sql.py @@ -20,13 +20,17 @@ import json import os from collections import namedtuple +from unittest import mock from unittest.mock import patch import pytest from databricks.sql.types import Row from airflow.providers.common.sql.hooks.handlers import fetch_all_handler -from airflow.providers.databricks.operators.databricks_sql import DatabricksSqlOperator +from airflow.providers.databricks.operators.databricks_sql import ( + DatabricksSqlOperator, + _get_airflow_query_tags, +) DATE = "2017-04-20" TASK_ID = "databricks-sql-operator" @@ -453,3 +457,162 @@ 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" + + +class TestDatabricksSqlOperatorQueryTags: + """Tests for query tags support in DatabricksSqlOperator.""" + + def test_get_airflow_query_tags_returns_empty_dict_without_task_instance(self): + """_get_airflow_query_tags must return {} when context has no 'ti' key.""" + result = _get_airflow_query_tags({}) + assert result == {} + + def test_get_query_tags_with_none_context_returns_custom_tags_only(self): + """When context is None, only custom tags are returned (no Airflow tags).""" + op = DatabricksSqlOperator( + task_id=TASK_ID, + sql="SELECT 1", + query_tags={"custom_tag": "custom_value"}, + ) + result = op._get_query_tags(None) + assert result == {"custom_tag": "custom_value"} + + def test_get_query_tags_with_none_context_and_no_custom_tags_returns_none(self): + """When context is None and no custom tags, None is returned.""" + op = DatabricksSqlOperator(task_id=TASK_ID, sql="SELECT 1") + result = op._get_query_tags(None) + assert result is None + + def test_get_query_tags_with_disabled_airflow_tags(self): + """When include_airflow_query_tags=False, only custom tags are returned.""" + op = DatabricksSqlOperator( + task_id=TASK_ID, + sql="SELECT 1", + query_tags={"custom_tag": "val"}, + include_airflow_query_tags=False, + ) + mock_context = {"ti": object()} + result = op._get_query_tags(mock_context) + assert result == {"custom_tag": "val"} + + def test_get_query_tags_with_airflow_context(self): + """When context is provided and include_airflow_query_tags=True, Airflow tags are included.""" + op = DatabricksSqlOperator( + task_id=TASK_ID, + sql="SELECT 1", + query_tags={"custom_tag": "custom_value"}, + ) + mock_ti = mock.MagicMock(spec=["dag_id", "task_id", "run_id", "try_number", "map_index"]) + mock_ti.dag_id = "test_dag" + mock_ti.task_id = "test_task" + mock_ti.run_id = "test_run" + mock_ti.try_number = 1 + mock_ti.map_index = -1 + mock_context = {"ti": mock_ti} + + result = op._get_query_tags(mock_context) + + assert result is not None + assert result["airflow_dag_id"] == "test_dag" + assert result["airflow_task_id"] == "test_task" + assert result["airflow_run_id"] == "test_run" + assert result["airflow_try_number"] == "1" + assert result["airflow_map_index"] == "-1" + assert result["custom_tag"] == "custom_value" + + def test_execute_sets_query_tags_on_hook(self): + """execute() sets query_tags on the hook before delegating to SQLExecuteQueryOperator.""" + with patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as mock_cls: + mock_hook = mock_cls.return_value + mock_hook.run.return_value = [] + mock_hook.descriptions = [[]] + + op = DatabricksSqlOperator( + task_id=TASK_ID, + sql="SELECT 1", + query_tags={"env": "test"}, + include_airflow_query_tags=False, + ) + + op.execute(None) + + assert mock_hook.query_tags == {"env": "test"} + + def test_custom_tags_override_airflow_tags_on_key_collision(self): + """Custom query_tags override Airflow tags when the same key is used.""" + op = DatabricksSqlOperator( + task_id=TASK_ID, + sql="SELECT 1", + query_tags={"airflow_dag_id": "overridden"}, + ) + mock_ti = mock.MagicMock(spec=["dag_id", "task_id", "run_id", "try_number", "map_index"]) + mock_ti.dag_id = "original_dag" + mock_ti.task_id = "task" + mock_ti.run_id = "run" + mock_ti.try_number = 1 + mock_ti.map_index = -1 + mock_context = {"ti": mock_ti} + + result = op._get_query_tags(mock_context) + + assert result is not None + assert result["airflow_dag_id"] == "overridden" + + +class TestDatabricksCopyIntoOperatorQueryTags: + """Tests for query tags support in DatabricksCopyIntoOperator.""" + + def _make_op(self, **kwargs): + from airflow.providers.databricks.operators.databricks_sql import DatabricksCopyIntoOperator + + return DatabricksCopyIntoOperator( + task_id=TASK_ID, + table_name="test_table", + file_location="s3://bucket/path", + file_format="CSV", + **kwargs, + ) + + def test_get_query_tags_with_none_context_returns_custom_tags_only(self): + op = self._make_op(query_tags={"custom": "value"}) + result = op._get_query_tags(None) + assert result == {"custom": "value"} + + def test_get_query_tags_with_none_context_and_no_custom_tags_returns_none(self): + op = self._make_op() + result = op._get_query_tags(None) + assert result is None + + def test_get_query_tags_with_airflow_context(self): + op = self._make_op(query_tags={"env": "staging"}) + mock_ti = mock.MagicMock(spec=["dag_id", "task_id", "run_id", "try_number", "map_index"]) + mock_ti.dag_id = "copy_dag" + mock_ti.task_id = "copy_task" + mock_ti.run_id = "run_1" + mock_ti.try_number = 2 + mock_ti.map_index = 0 + mock_context = {"ti": mock_ti} + + result = op._get_query_tags(mock_context) + + assert result is not None + assert result["airflow_dag_id"] == "copy_dag" + assert result["airflow_task_id"] == "copy_task" + assert result["env"] == "staging" + + def test_execute_sets_query_tags_on_hook(self): + from airflow.providers.databricks.operators.databricks_sql import DatabricksCopyIntoOperator + + with patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as mock_cls: + mock_hook = mock_cls.return_value + op = DatabricksCopyIntoOperator( + task_id=TASK_ID, + table_name="test_table", + file_location="s3://bucket/path", + file_format="CSV", + query_tags={"env": "prod"}, + include_airflow_query_tags=False, + ) + op.execute(None) + + assert mock_hook.query_tags == {"env": "prod"}