Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator

[Shaan-alpha]

Description

This PR addresses issue #66839 by injecting Airflow context metadata (dag_id, ask_id, and
un_id) into the session_configuration parameter as query_tags during the execution of DatabricksSqlOperator and DatabricksCopyIntoOperator.

By automatically populating query tags, it significantly enhances observability on the Databricks side, allowing administrators and developers to trace specific queries executed in Databricks SQL directly back to the Airflow task run that spawned them.

Key Changes:

• Added a _format_query_tags(context) helper to extract and safely escape the metadata strings before formatting them into a comma-separated query tag string.
• Modified the execute() methods of DatabricksSqlOperator and DatabricksCopyIntoOperator to correctly merge the generated tags into the existing session_configuration (preserving any user-defined custom query tags).
• Updated existing unit tests and added est_query_tags_injection for both operators in est_databricks_sql.py and est_databricks_copy.py to ensure injection and non-regression of user tags.

Related Issues

• Closes Improve Databricks operators with query tags #66839

Testing

• Added unit tests mimicking operator execution and asserting proper manipulation of the hook's session_config parameter.

[boring-cyborg]

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide
Here are some useful points:

• Pay attention to the quality of your code (ruff, mypy and type annotations). Our prek-hooks will help you with that.
• In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example Dag that shows how users should use it.
• Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
• Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
• Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
• Be sure to read the Airflow Coding style.
• Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
  Apache Airflow is a community-driven project and together we are making it better 🚀.
  In case of doubts contact the developers at:
  Mailing List: dev@airflow.apache.org
  Slack: https://s.apache.org/airflow-slack

[Shaan-alpha]

Refactored query-tag injection into shared helper utilities to reduce duplicated logic and improve maintainability across operators.

[SameerMesiah97]

I have left a few comments.

[SameerMesiah97]

I think this could be made more clear and explicit via a mapping driven approach. Please see the below for guidance:

_QUERY_TAG_FIELDS = {
    "airflow_dag_id": ("dag", "dag_id"),
    "airflow_task_id": ("task", "task_id"),
    "airflow_run_id": ("run_id", None),
}


def _format_query_tags(context: Context) -> str:
    tags = []

    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)

Also, you could do the same for _escape_query_tag_value:

_QUERY_TAG_ESCAPE_SEQUENCES = {
    "\\": "\\\\",
    ",": "\\,",
    ":": "\\:",
}


def _escape_query_tag_value(value: str) -> str:
    """Escape Databricks query-tag separator characters in a tag value."""
    escaped = str(value)

    for char, replacement in _QUERY_TAG_ESCAPE_SEQUENCES.items():
        escaped = escaped.replace(char, replacement)

    return escaped

[Shaan-alpha]

Thanks for the suggestion @SameerMesiah97! I agree the mapping-driven approach is cleaner and more maintainable. I've gone ahead and refactored both _format_query_tags and _escape_query_tag_value to use the _QUERY_TAG_FIELDS and _QUERY_TAG_ESCAPE_SEQUENCES mappings exactly as you suggested — this is included in commit 732df03 ("Refactor Databricks query tag helper utilities"). Please take another look and let me know if you'd like any further tweaks.

[SameerMesiah97]

Do we want this behavior configurable (operator/provider-level opt-out)? Since this mutates session_configuration automatically, some users may prefer explicit control over injected warehouse metadata.

[Shaan-alpha]

Good point — making this configurable makes sense so users retain explicit control over session_configuration. I'll add an inject_query_tags: bool = True parameter to both DatabricksSqlOperator and DatabricksCopyIntoOperator (defaulting to True to preserve the observability benefit, while allowing easy opt-out). I'll also document the new parameter in the operator docstrings. Will push the update shortly.

[SameerMesiah97]

I think this test should be moved above the openlineage tests as it is closer to core execution logic. Also, they only cover the happy-path. Could we add coverage for empty/partial context, empty existing query_tags, and values containing commas/colons/backslashes?

The escaping path is probably the most important one here since _escape_query_tag_value() is not really exercised end-to-end right now. It would also be good to verify that unrelated session_configuration values are preserved after the merge.

[Shaan-alpha]

Thanks, that ordering makes sense. I'll move test_query_tags_injection above the OpenLineage tests so it sits closer to the core execution logic. I'll also expand coverage to include: empty/partial context (missing dag/task/run_id), empty existing query_tags, values containing commas/colons/backslashes (to exercise _escape_query_tag_value end-to-end), and an assertion that unrelated keys in session_configuration are preserved after the merge. Will include in the next push.

[SameerMesiah97]

I would move this test after test_exec_write_gcs_parquet_output. Also, the above comment regarding incomplete test coverage applies here too.

[Shaan-alpha]

Will do — I'll move the new test to sit right after test_exec_write_gcs_parquet_output and mirror the same expanded coverage (empty/partial context, empty existing query_tags, escape-character values, and preservation of unrelated session_configuration keys). Thanks for the review!

[Shaan-alpha]

@SameerMesiah97 thanks again for the review — pushed 38d58f8 addressing the remaining three points:

• Opt-out control: added inject_query_tags: bool = True to both DatabricksSqlOperator and DatabricksCopyIntoOperator, documented in the operator docstrings. Defaults to True to preserve the observability benefit, but users can pass inject_query_tags=False to retain full control over session_configuration.
• Test placement: moved the query-tag tests above the OpenLineage block in test_databricks_copy.py and to right after test_exec_write_gcs_parquet_output in test_databricks_sql.py, so they sit closer to the core execution logic.
• Expanded test coverage (mirrored across both files):
  • empty / partial Airflow context (missing dag / task / run_id)
  • empty existing query_tags
  • escape-character values (,, :, \) — exercises _escape_query_tag_value end-to-end
  • preservation of unrelated session_configuration keys after the merge
  • fallback to databricks_conn.extra_dejson["session_configuration"] when hook.session_config is None
  • opt-out path: inject_query_tags=False leaves session_config untouched

Comment 1 (mapping-driven helpers) was already addressed in 732df03. Ready for another look whenever you have a moment.

[SameerMesiah97]

@SameerMesiah97 thanks again for the review — pushed 38d58f8 addressing the remaining three points:

• Opt-out control: added inject_query_tags: bool = True to both DatabricksSqlOperator and DatabricksCopyIntoOperator, documented in the operator docstrings. Defaults to True to preserve the observability benefit, but users can pass inject_query_tags=False to retain full control over session_configuration.

• Test placement: moved the query-tag tests above the OpenLineage block in test_databricks_copy.py and to right after test_exec_write_gcs_parquet_output in test_databricks_sql.py, so they sit closer to the core execution logic.

• Expanded test coverage (mirrored across both files):

  • empty / partial Airflow context (missing dag / task / run_id)
  • empty existing query_tags
  • escape-character values (,, :, \) — exercises _escape_query_tag_value end-to-end
  • preservation of unrelated session_configuration keys after the merge
  • fallback to databricks_conn.extra_dejson["session_configuration"] when hook.session_config is None
  • opt-out path: inject_query_tags=False leaves session_config untouched

Comment 1 (mapping-driven helpers) was already addressed in 732df03. Ready for another look whenever you have a moment.

Let's wait for a maintainer to trigger CI.

[Shaan-alpha]

Follow-up improvements planned based on review and maintainability considerations:

• Move duplicated test helpers (_make_context and _run_with_mocked_hook) into shared test utilities/fixtures to reduce duplication across Databricks operator tests.
• Keep query tag injection scoped to execution lifecycle while preserving existing session configuration semantics.
• Add a small usage/example snippet for inject_query_tags behavior in docs/docstrings if needed.

This PR intentionally preserves backward compatibility by:

• preserving existing query_tags
• preserving unrelated session_configuration keys
• supporting opt-out via inject_query_tags=False
• handling escaping for Databricks query tag separators.