Add query tags: DatabricksSqlSensor, DatabricksPartitionSensor #68704
Draft
cruseakshay wants to merge 2 commits into
Draft
Add query tags: DatabricksSqlSensor, DatabricksPartitionSensor #68704cruseakshay wants to merge 2 commits into
cruseakshay wants to merge 2 commits into
Conversation
…ricksSQLStatementsOperator, DatabricksSQLStatementsSensor
03f70ea to
af92c2e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Extends the session-level query-tags instrumentation added in #66895 to the remaining Databricks operators and sensors that send queries to Databricks:
DatabricksSqlSensorDatabricksPartitionSensorDatabricksSQLStatementsOperatorDatabricksSQLStatementsSensorCloses: #68582
What changed
Two mechanisms are used because these components talk to Databricks in two different ways:
DatabricksSqlSensorandDatabricksPartitionSensorgo throughDatabricksSqlHook, so they reuse the existingQUERY_TAGSsession-parameter plumbing from Add session-level query tags to Databricks SQL operators #66895. The merged tags are set on the hook before the query runs.DatabricksSQLStatementsOperatorandDatabricksSQLStatementsSensoruse the REST Statement Execution API (/api/2.0/sql/statements/), which does not acceptsession_configuration. The API exposes a nativequery_tagsfield ([{"key": ..., "value": ...}]), so tags are injected directly into the request body. For the sensor, tags are only applied when it submits a new statement; if astatement_idis passed in, nothing is submitted and no tags are attached.Each component gains two parameters mirroring the existing SQL operators:
query_tags: dict[str, str | None] | None— user-supplied tags, templatedinclude_airflow_query_tags: bool = True— merge in Airflow context tags:dag_idtask_idrun_idtry_numbermap_indexUser-supplied tags win on key collision.
The Airflow-context tag logic that previously lived in
operators/databricks_sql.pyis extracted to a sharedutils/query_tags.pymodule:get_airflow_query_tagsbuild_query_tagsdict_to_query_tag_listThis lets all five operators and sensors share one implementation.
This is a pure relocation:
DatabricksSqlOperatorandDatabricksCopyIntoOperatorbehavior is unchanged.Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code, following the [guidelines]