Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
eladkal commented on PR #66886: URL: https://github.com/apache/airflow/pull/66886#issuecomment-4541784501 Thanks for your work @Shaan-alpha -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
Shaan-alpha closed pull request #66886: Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator URL: https://github.com/apache/airflow/pull/66886 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
Shaan-alpha commented on PR #66886: URL: https://github.com/apache/airflow/pull/66886#issuecomment-4541761555 Thanks everyone for the reviews and feedback! Since the core objective of this PR (injecting Airflow metadata as Databricks query tags) was natively implemented in the hook and merged today via #66895, this PR is now redundant. I'm going ahead and closing this as superseded by #66895. Appreciate the support! -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
Shaan-alpha commented on PR #66886: URL: https://github.com/apache/airflow/pull/66886#issuecomment-4485792842 Thanks @potiuk for the triage — pushed `e9da47e` addressing all three categories: - **Provider tests (Compat 2.11.1 / 3.0.6 / 3.1.8 / 3.2.1, Non-DB common.compat...google, Low-dep databricks)** — these ~30 `AttributeError: 'NoneType' object has no attribute 'get'` failures all reduced to one root cause: pre-existing Databricks operator tests call `op.execute(None)`, and `_inject_query_tags` was unconditionally calling `context.get(...)`. Added an early-return for falsy context in `_inject_query_tags` so the existing tests are unaffected and the public `execute()` stays tolerant. The new behavioral test `test_query_tags_injection_with_empty_context` continues to assert the no-op path explicitly. - **MyPy providers checks** — `_escape_query_tag_value` was annotated `(value: str)` but the call site passes a value derived from `context.get(...)` / `getattr(...)`, which MyPy widens to `object`. The function already does `str(value)` internally, so broadened the parameter type to `object` — resolves the `[arg-type]` error at `databricks_sql.py:80` without changing behavior. - **CI image checks / Static checks** — two test files had multi-line parenthesized `assert` expressions that pinned `ruff==0.15.12` wanted collapsed. Ran `ruff format`. Verified locally: `ruff format --check`, `ruff check`, mypy on the file, and a focused repro that confirms `execute(None)` no longer raises while the positive injection path still produces the expected `query_tags`. Marking ready for review. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
potiuk commented on PR #66886: URL: https://github.com/apache/airflow/pull/66886#issuecomment-4483645450 @Shaan-alpha Converting to **draft** — this PR doesn't yet meet our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria). - :x: **Provider tests** — Failing: `provider distributions tests / Compat 3.0.6:P3.10:`, `provider distributions tests / Compat 3.1.8:P3.10:`, `provider distributions tests / Compat 3.2.1:P3.10:`, `provider distributions tests / Compat 2.11.1:P3.10:`, `Non-DB tests: providers / Non-DB-prov::3.10:common.compat...google` (+1 more). See [docs](https://github.com/apache/airflow/blob/main/contributing-docs/12_provider_distributions.rst). - :x: **Pre-commit / static checks** — Failing: `CI image checks / Static checks`. See [docs](https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst). - :x: **Other failing CI checks** — Failing: `MyPy providers checks`. See [docs](https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst). See the linked criteria for how to fix each item, then mark the PR "Ready for review". This is **not** a rejection — just an invitation to bring the PR up to standard. No rush. --- _Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this [two-stage triage process](https://github.com/apache/airflow/blob/main/contributing-docs/25_maintainer_pr_triage.md#why-the-first-pass-is-automated) so that our maintainers' limited time is spent where it matters most: the conversation with you._ --- Drafted-by: Claude Code (Opus 4.7); reviewed by @potiuk before posting -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
Shaan-alpha commented on PR #66886: URL: https://github.com/apache/airflow/pull/66886#issuecomment-4472658197 Thanks @SameerMesiah97 for the review! I have rebased the branch cleanly on latest main and added concise usage examples for inject_query_tags=True in both the operator documentation (sql.rst) and system test examples (example_databricks_sql.py). All tests and checks are up to date. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
Shaan-alpha commented on PR #66886: URL: https://github.com/apache/airflow/pull/66886#issuecomment-4467576711 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
SameerMesiah97 commented on PR #66886: URL: https://github.com/apache/airflow/pull/66886#issuecomment-4455831965 > @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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
Shaan-alpha commented on PR #66886: URL: https://github.com/apache/airflow/pull/66886#issuecomment-4446975652 @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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
Shaan-alpha commented on code in PR #66886:
URL: https://github.com/apache/airflow/pull/66886#discussion_r3238643742
##
providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py:
##
@@ -46,6 +46,46 @@
_DISALLOWED_SQL_TOKENS = (";", "--", "/*", "*/")
+def _escape_query_tag_value(value: str) -> str:
+"""Escape Databricks query-tag separators in a tag value."""
+return str(value).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_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'])}")
+
+return ",".join(tags)
Review Comment:
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.
##
providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py:
##
@@ -153,6 +193,10 @@ def _hook(self) -> DatabricksSqlHook:
def get_db_hook(self) -> DatabricksSqlHook:
return self._hook
+def execute(self, context: Context) -> Any:
+_inject_query_tags(self.get_db_hook(), context)
+return super().execute(context)
Review Comment:
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.
##
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"
+
+
Review Comment:
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!
##
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)}
+
+
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
SameerMesiah97 commented on code in PR #66886:
URL: https://github.com/apache/airflow/pull/66886#discussion_r3238298731
##
providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py:
##
@@ -46,6 +46,46 @@
_DISALLOWED_SQL_TOKENS = (";", "--", "/*", "*/")
+def _escape_query_tag_value(value: str) -> str:
+"""Escape Databricks query-tag separators in a tag value."""
+return str(value).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_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'])}")
+
+return ",".join(tags)
Review Comment:
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
```
##
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)}
+
+
Review Comment:
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.
##
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"
+
+
Review Comment:
I would move this test after `test_exec_write_gcs_parquet_output`. Also, the
above comment regarding incomplete test coverage applies here too.
##
providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py:
##
@@ -153,6 +193,10 @@ def _hook(self) -> DatabricksSqlHook:
def get_db_hook(self) -> DatabricksSqlHook:
return self._hook
+def execute(self, context: Context) -> Any:
+_inject_query_tags(self.get_db_hook(), context)
+return super().execute(context)
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
Shaan-alpha commented on PR #66886: URL: https://github.com/apache/airflow/pull/66886#issuecomment-821696 Refactored query-tag injection into shared helper utilities to reduce duplicated logic and improve maintainability across operators. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add Databricks query tags for DatabricksSqlOperator and DatabricksCopyIntoOperator [airflow]
boring-cyborg[bot] commented on PR #66886: URL: https://github.com/apache/airflow/pull/66886#issuecomment-618255 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](https://github.com/apache/airflow/blob/main/contributing-docs/README.rst) Here are some useful points: - Pay attention to the quality of your code (ruff, mypy and type annotations). Our [prek-hooks]( https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst#prerequisites-for-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](https://github.com/apache/airflow/blob/main/airflow-core/docs/howto/custom-operator.rst) Consider adding an example Dag that shows how users should use it. - Consider using [Breeze environment](https://github.com/apache/airflow/blob/main/dev/breeze/doc/README.rst) 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](https://www.apache.org/foundation/policies/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]( https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#coding-style-and-best-practices). - 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: [email protected] Slack: https://s.apache.org/airflow-slack -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
