You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "jroachgolf84 (via GitHub)" <gi...@apache.org> on 2025/04/21 18:04:54 UTC
[PR] Adding `DatabricksSQLStatementSensor` Operator with Deferrability [airflow]
jroachgolf84 opened a new pull request, #49516:
URL: https://github.com/apache/airflow/pull/49516
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2070196244
##########
providers/databricks/tests/system/databricks/example_databricks_sensors.py:
##########
@@ -67,6 +70,18 @@
)
# [END howto_sensor_databricks_sql]
+ # [START howto_sensor_databricks_sql_statement]
+ # Example of using the DatabricksSQLStatementSensor to wait for a query
+ # to successfully run.
+ sql_statement_sensor = DatabricksSQLStatementsSensor(
+ task_id="sql_statement_sensor_task",
+ databricks_conn_id=connection_id,
+ warehouse_id=WAREHOUSE_ID,
+ statement_id="statement_id",
Review Comment:
Changed.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2070197819
##########
providers/databricks/src/airflow/providers/databricks/sensors/databricks.py:
##########
@@ -0,0 +1,152 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.databricks.hooks.databricks import DatabricksHook, SQLStatementState
+from airflow.providers.databricks.operators.databricks import DEFER_METHOD_NAME
+from airflow.providers.databricks.utils.mixins import DatabricksSQLStatementsMixin
+from airflow.providers.databricks.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseSensorOperator
+else:
+ from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+ from airflow.utils.context import Context
+
+XCOM_STATEMENT_ID_KEY = "statement_id"
+
+
+class DatabricksSQLStatementsSensor(DatabricksSQLStatementsMixin, BaseSensorOperator):
+ """DatabricksSQLStatementsSensor."""
+
+ template_fields: Sequence[str] = (
+ "databricks_conn_id",
+ "statement",
+ "statement_id",
+ )
+ template_ext: Sequence[str] = (".json-tpl",)
+ ui_color = "#1CB1C2"
+ ui_fgcolor = "#fff"
+
+ def __init__(
+ self,
+ warehouse_id: str,
+ *,
+ statement: str | None = None,
+ statement_id: str | None = None,
+ catalog: str | None = None,
+ schema: str | None = None,
+ parameters: list[dict[str, Any]] | None = None,
+ databricks_conn_id: str = "databricks_default",
+ polling_period_seconds: int = 30,
+ databricks_retry_limit: int = 3,
+ databricks_retry_delay: int = 1,
+ databricks_retry_args: dict[Any, Any] | None = None,
+ do_xcom_push: bool = True,
+ wait_for_termination: bool = True,
+ timeout: float = 3600,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.statement = statement
+ self.warehouse_id = warehouse_id
+ self.catalog = catalog
+ self.schema = schema
+ self.statement_id = statement_id # Optional
Review Comment:
Changed - good callout.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2072906746
##########
airflow-core/tests/unit/always/test_project_structure.py:
##########
@@ -113,6 +113,7 @@ def test_providers_modules_should_have_tests(self):
"providers/common/messaging/tests/unit/common/messaging/providers/test_base_provider.py",
"providers/common/messaging/tests/unit/common/messaging/providers/test_sqs.py",
"providers/databricks/tests/unit/databricks/test_version_compat.py",
+ "providers/databricks/tests/unit/databricks/utils/test_mixins.py",
Review Comment:
See the comment in line 60.
> The test below had a but for quite a while and we missed a lot of modules to have tess
> We should make sure that one goes to 0
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2072116369
##########
providers/databricks/docs/operators/sql_statements.rst:
##########
@@ -55,3 +55,49 @@ An example usage of the ``DatabricksSQLStatementsOperator`` is as follows:
:language: python
:start-after: [START howto_operator_sql_statements]
:end-before: [END howto_operator_sql_statements]
+
+
+.. _howto/sensor:DatabricksSQLStatementsSensor:
+
+DatabricksSQLStatementsOperator
+===============================
+
+Use the :class:`~airflow.providers.databricks.sensor.databricks.DatabricksSQLStatementsSensor` to either submit a
+Databricks SQL Statement to Databricks using the
+`Databricks SQL Statement Execution API <https://docs.databricks.com/api/workspace/statementexecution>`_, or pass
+a Statement ID to the Sensor and await for the query to terminate execution.
+
+
+Using the Operator
+------------------
+
+The ``DatabricksSQLStatementsSensor`` does one of two things. The Sensor can submit SQL statements to Databricks using
+the `/api/2.0/sql/statements/ <https://docs.databricks.com/api/workspace/statementexecution/executestatement>`_
+endpoint. However, the Sensor can also take the Statement ID of an already-submitted SQL Statement and handle the
+response to that execution.
+
+It supports configurable execution parameters such as warehouse selection, catalog, schema, and parameterized queries.
+The operator can either synchronously poll for query completion or run in a deferrable mode for improved efficiency.
+
+The only required parameters for using the Sensor are:
+
+* One of ``statement`` or ``statement_id`` - The SQL statement to execute. The statement can optionally be
+ parameterized, see parameters.
+* ``warehouse_id`` - Warehouse upon which to execute a statement.
+
+All other parameters are optional and described in the documentation for ``DatabricksSQLStatementsSensor`` including
+but not limited to:
+
+* ``catalog``
+* ``schema``
+* ``parameters``
+
+Examples
+--------
+
+An example usage of the ``DatabricksSQLStatementsOperator`` is as follows:
Review Comment:
Changes added locally, thanks!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2072847690
##########
airflow-core/tests/unit/always/test_project_structure.py:
##########
@@ -113,6 +113,7 @@ def test_providers_modules_should_have_tests(self):
"providers/common/messaging/tests/unit/common/messaging/providers/test_base_provider.py",
"providers/common/messaging/tests/unit/common/messaging/providers/test_sqs.py",
"providers/databricks/tests/unit/databricks/test_version_compat.py",
+ "providers/databricks/tests/unit/databricks/utils/test_mixins.py",
Review Comment:
Why?
The list here should go down not up.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on PR #49516:
URL: https://github.com/apache/airflow/pull/49516#issuecomment-2864386052
Are we good to go ahead and merge this?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "RNHTTR (via GitHub)" <gi...@apache.org>.
RNHTTR commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2087191219
##########
airflow-core/tests/unit/always/test_project_structure.py:
##########
@@ -113,6 +113,7 @@ def test_providers_modules_should_have_tests(self):
"providers/common/messaging/tests/unit/common/messaging/providers/test_base_provider.py",
"providers/common/messaging/tests/unit/common/messaging/providers/test_sqs.py",
"providers/databricks/tests/unit/databricks/test_version_compat.py",
+ "providers/databricks/tests/unit/databricks/utils/test_mixins.py",
Review Comment:
@eladkal correct me if I'm wrong but we shouldn't be adding tests to `OVERLOOKED_TESTS`
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2072115830
##########
providers/databricks/docs/operators/sql_statements.rst:
##########
@@ -55,3 +55,49 @@ An example usage of the ``DatabricksSQLStatementsOperator`` is as follows:
:language: python
:start-after: [START howto_operator_sql_statements]
:end-before: [END howto_operator_sql_statements]
+
+
+.. _howto/sensor:DatabricksSQLStatementsSensor:
+
+DatabricksSQLStatementsOperator
Review Comment:
Changes added locally, thanks!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2072113153
##########
providers/databricks/tests/system/databricks/example_databricks_sensors.py:
##########
@@ -30,6 +31,8 @@
# [DAG name to be shown on Airflow UI]
DAG_ID = "example_databricks_sensor"
+WAREHOUSE_ID = os.environ.get("WAREHOUSE_ID", "") # Defaults to an empty string to satisfy typing
Review Comment:
Logic added.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2071321959
##########
providers/databricks/tests/system/databricks/example_databricks_sensors.py:
##########
@@ -30,6 +31,8 @@
# [DAG name to be shown on Airflow UI]
DAG_ID = "example_databricks_sensor"
+WAREHOUSE_ID = os.environ.get("WAREHOUSE_ID", "") # Defaults to an empty string to satisfy typing
Review Comment:
warehouse_id is a required param for the APIs to work. I guess we have mentioned those as positional params but if it receives empty values I would expect the APIs to fail and report an error.
But yes, like mentioned by @alexott might make sense to raise an error from the operator/sensor if we get an empty warehouse_id
##########
providers/databricks/docs/operators/sql_statements.rst:
##########
@@ -55,3 +55,49 @@ An example usage of the ``DatabricksSQLStatementsOperator`` is as follows:
:language: python
:start-after: [START howto_operator_sql_statements]
:end-before: [END howto_operator_sql_statements]
+
+
+.. _howto/sensor:DatabricksSQLStatementsSensor:
+
+DatabricksSQLStatementsOperator
+===============================
+
+Use the :class:`~airflow.providers.databricks.sensor.databricks.DatabricksSQLStatementsSensor` to either submit a
+Databricks SQL Statement to Databricks using the
+`Databricks SQL Statement Execution API <https://docs.databricks.com/api/workspace/statementexecution>`_, or pass
+a Statement ID to the Sensor and await for the query to terminate execution.
+
+
+Using the Operator
Review Comment:
```suggestion
Using the Sensor
```
##########
providers/databricks/docs/operators/sql_statements.rst:
##########
@@ -55,3 +55,49 @@ An example usage of the ``DatabricksSQLStatementsOperator`` is as follows:
:language: python
:start-after: [START howto_operator_sql_statements]
:end-before: [END howto_operator_sql_statements]
+
+
+.. _howto/sensor:DatabricksSQLStatementsSensor:
+
+DatabricksSQLStatementsOperator
Review Comment:
```suggestion
DatabricksSQLStatementsSensor
```
##########
providers/databricks/docs/operators/sql_statements.rst:
##########
@@ -55,3 +55,49 @@ An example usage of the ``DatabricksSQLStatementsOperator`` is as follows:
:language: python
:start-after: [START howto_operator_sql_statements]
:end-before: [END howto_operator_sql_statements]
+
+
+.. _howto/sensor:DatabricksSQLStatementsSensor:
Review Comment:
Since this is a sensor, might make sense to move the documentation to a separate `sensors.rst` or something. Just an opinion.
##########
providers/databricks/docs/operators/sql_statements.rst:
##########
@@ -55,3 +55,49 @@ An example usage of the ``DatabricksSQLStatementsOperator`` is as follows:
:language: python
:start-after: [START howto_operator_sql_statements]
:end-before: [END howto_operator_sql_statements]
+
+
+.. _howto/sensor:DatabricksSQLStatementsSensor:
+
+DatabricksSQLStatementsOperator
+===============================
+
+Use the :class:`~airflow.providers.databricks.sensor.databricks.DatabricksSQLStatementsSensor` to either submit a
+Databricks SQL Statement to Databricks using the
+`Databricks SQL Statement Execution API <https://docs.databricks.com/api/workspace/statementexecution>`_, or pass
+a Statement ID to the Sensor and await for the query to terminate execution.
+
+
+Using the Operator
+------------------
+
+The ``DatabricksSQLStatementsSensor`` does one of two things. The Sensor can submit SQL statements to Databricks using
+the `/api/2.0/sql/statements/ <https://docs.databricks.com/api/workspace/statementexecution/executestatement>`_
+endpoint. However, the Sensor can also take the Statement ID of an already-submitted SQL Statement and handle the
+response to that execution.
+
+It supports configurable execution parameters such as warehouse selection, catalog, schema, and parameterized queries.
+The operator can either synchronously poll for query completion or run in a deferrable mode for improved efficiency.
+
+The only required parameters for using the Sensor are:
+
+* One of ``statement`` or ``statement_id`` - The SQL statement to execute. The statement can optionally be
+ parameterized, see parameters.
+* ``warehouse_id`` - Warehouse upon which to execute a statement.
+
+All other parameters are optional and described in the documentation for ``DatabricksSQLStatementsSensor`` including
+but not limited to:
+
+* ``catalog``
+* ``schema``
+* ``parameters``
+
+Examples
+--------
+
+An example usage of the ``DatabricksSQLStatementsOperator`` is as follows:
Review Comment:
```suggestion
An example usage of the ``DatabricksSQLStatementsSensor`` is as follows:
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "alexott (via GitHub)" <gi...@apache.org>.
alexott commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2070168781
##########
providers/databricks/src/airflow/providers/databricks/sensors/databricks.py:
##########
@@ -0,0 +1,152 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.databricks.hooks.databricks import DatabricksHook, SQLStatementState
+from airflow.providers.databricks.operators.databricks import DEFER_METHOD_NAME
+from airflow.providers.databricks.utils.mixins import DatabricksSQLStatementsMixin
+from airflow.providers.databricks.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseSensorOperator
+else:
+ from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+ from airflow.utils.context import Context
+
+XCOM_STATEMENT_ID_KEY = "statement_id"
+
+
+class DatabricksSQLStatementsSensor(DatabricksSQLStatementsMixin, BaseSensorOperator):
+ """DatabricksSQLStatementsSensor."""
+
+ template_fields: Sequence[str] = (
+ "databricks_conn_id",
+ "statement",
+ "statement_id",
+ )
+ template_ext: Sequence[str] = (".json-tpl",)
+ ui_color = "#1CB1C2"
+ ui_fgcolor = "#fff"
+
+ def __init__(
+ self,
+ warehouse_id: str,
+ *,
+ statement: str | None = None,
+ statement_id: str | None = None,
+ catalog: str | None = None,
+ schema: str | None = None,
+ parameters: list[dict[str, Any]] | None = None,
+ databricks_conn_id: str = "databricks_default",
+ polling_period_seconds: int = 30,
+ databricks_retry_limit: int = 3,
+ databricks_retry_delay: int = 1,
+ databricks_retry_args: dict[Any, Any] | None = None,
+ do_xcom_push: bool = True,
+ wait_for_termination: bool = True,
+ timeout: float = 3600,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.statement = statement
+ self.warehouse_id = warehouse_id
+ self.catalog = catalog
+ self.schema = schema
+ self.statement_id = statement_id # Optional
Review Comment:
We need to fail if both statement and statement id are None, or both are defined
##########
providers/databricks/tests/system/databricks/example_databricks_sensors.py:
##########
@@ -67,6 +70,18 @@
)
# [END howto_sensor_databricks_sql]
+ # [START howto_sensor_databricks_sql_statement]
+ # Example of using the DatabricksSQLStatementSensor to wait for a query
+ # to successfully run.
+ sql_statement_sensor = DatabricksSQLStatementsSensor(
+ task_id="sql_statement_sensor_task",
+ databricks_conn_id=connection_id,
+ warehouse_id=WAREHOUSE_ID,
+ statement_id="statement_id",
Review Comment:
Statement ID exists only when it's created by a corresponding api. SQL could be more appropriate here
##########
providers/databricks/tests/system/databricks/example_databricks_sensors.py:
##########
@@ -30,6 +31,8 @@
# [DAG name to be shown on Airflow UI]
DAG_ID = "example_databricks_sensor"
+WAREHOUSE_ID = os.environ.get("WAREHOUSE_ID", "cf414a2206dfb397")
Review Comment:
I would prefer to avoid having a default value, or at least put it into a separate constant and add a note to change it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Operator with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on PR #49516:
URL: https://github.com/apache/airflow/pull/49516#issuecomment-2819150833
cc: @pankajkoti, I've created a draft PR. Do you mind take a look at this?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on PR #49516:
URL: https://github.com/apache/airflow/pull/49516#issuecomment-2844694987
I'll go ahead and make those changes. Thanks!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2072116116
##########
providers/databricks/docs/operators/sql_statements.rst:
##########
@@ -55,3 +55,49 @@ An example usage of the ``DatabricksSQLStatementsOperator`` is as follows:
:language: python
:start-after: [START howto_operator_sql_statements]
:end-before: [END howto_operator_sql_statements]
+
+
+.. _howto/sensor:DatabricksSQLStatementsSensor:
+
+DatabricksSQLStatementsOperator
+===============================
+
+Use the :class:`~airflow.providers.databricks.sensor.databricks.DatabricksSQLStatementsSensor` to either submit a
+Databricks SQL Statement to Databricks using the
+`Databricks SQL Statement Execution API <https://docs.databricks.com/api/workspace/statementexecution>`_, or pass
+a Statement ID to the Sensor and await for the query to terminate execution.
+
+
+Using the Operator
Review Comment:
Changes added locally, thanks!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "alexott (via GitHub)" <gi...@apache.org>.
alexott commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2071108545
##########
providers/databricks/tests/system/databricks/example_databricks_sensors.py:
##########
@@ -30,6 +31,8 @@
# [DAG name to be shown on Airflow UI]
DAG_ID = "example_databricks_sensor"
+WAREHOUSE_ID = os.environ.get("WAREHOUSE_ID", "") # Defaults to an empty string to satisfy typing
Review Comment:
It's best to throw an exception when it's empty...
##########
providers/databricks/src/airflow/providers/databricks/sensors/databricks.py:
##########
@@ -0,0 +1,156 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.databricks.hooks.databricks import DatabricksHook, SQLStatementState
+from airflow.providers.databricks.operators.databricks import DEFER_METHOD_NAME
+from airflow.providers.databricks.utils.mixins import DatabricksSQLStatementsMixin
+from airflow.providers.databricks.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseSensorOperator
+else:
+ from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+ from airflow.utils.context import Context
+
+XCOM_STATEMENT_ID_KEY = "statement_id"
+
+
+class DatabricksSQLStatementsSensor(DatabricksSQLStatementsMixin, BaseSensorOperator):
+ """DatabricksSQLStatementsSensor."""
+
+ template_fields: Sequence[str] = (
+ "databricks_conn_id",
+ "statement",
+ "statement_id",
+ )
+ template_ext: Sequence[str] = (".json-tpl",)
+ ui_color = "#1CB1C2"
+ ui_fgcolor = "#fff"
+
+ def __init__(
+ self,
+ warehouse_id: str,
+ *,
+ statement: str | None = None,
+ statement_id: str | None = None,
+ catalog: str | None = None,
+ schema: str | None = None,
+ parameters: list[dict[str, Any]] | None = None,
+ databricks_conn_id: str = "databricks_default",
+ polling_period_seconds: int = 30,
+ databricks_retry_limit: int = 3,
+ databricks_retry_delay: int = 1,
+ databricks_retry_args: dict[Any, Any] | None = None,
+ do_xcom_push: bool = True,
+ wait_for_termination: bool = True,
+ timeout: float = 3600,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ if statement and statement_id:
Review Comment:
We still need a check that both aren't set...
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #49516:
URL: https://github.com/apache/airflow/pull/49516#issuecomment-2844661975
LGTM. The code looks good but I'd appreciate another review @alexott maybe ?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on PR #49516:
URL: https://github.com/apache/airflow/pull/49516#issuecomment-2844727055
Thanks for the feedback - I've implemented those changes.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2072114888
##########
providers/databricks/docs/operators/sql_statements.rst:
##########
@@ -55,3 +55,49 @@ An example usage of the ``DatabricksSQLStatementsOperator`` is as follows:
:language: python
:start-after: [START howto_operator_sql_statements]
:end-before: [END howto_operator_sql_statements]
+
+
+.. _howto/sensor:DatabricksSQLStatementsSensor:
Review Comment:
For other providers (such as AWS), such do not have a `sensors/` directory, or a `sensors.rst` file. I think it might make sense to adhere to that same logic?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on PR #49516:
URL: https://github.com/apache/airflow/pull/49516#issuecomment-2843908783
Team, can I please request a review on this issue?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2072112776
##########
providers/databricks/src/airflow/providers/databricks/sensors/databricks.py:
##########
@@ -0,0 +1,156 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.databricks.hooks.databricks import DatabricksHook, SQLStatementState
+from airflow.providers.databricks.operators.databricks import DEFER_METHOD_NAME
+from airflow.providers.databricks.utils.mixins import DatabricksSQLStatementsMixin
+from airflow.providers.databricks.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseSensorOperator
+else:
+ from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+ from airflow.utils.context import Context
+
+XCOM_STATEMENT_ID_KEY = "statement_id"
+
+
+class DatabricksSQLStatementsSensor(DatabricksSQLStatementsMixin, BaseSensorOperator):
+ """DatabricksSQLStatementsSensor."""
+
+ template_fields: Sequence[str] = (
+ "databricks_conn_id",
+ "statement",
+ "statement_id",
+ )
+ template_ext: Sequence[str] = (".json-tpl",)
+ ui_color = "#1CB1C2"
+ ui_fgcolor = "#fff"
+
+ def __init__(
+ self,
+ warehouse_id: str,
+ *,
+ statement: str | None = None,
+ statement_id: str | None = None,
+ catalog: str | None = None,
+ schema: str | None = None,
+ parameters: list[dict[str, Any]] | None = None,
+ databricks_conn_id: str = "databricks_default",
+ polling_period_seconds: int = 30,
+ databricks_retry_limit: int = 3,
+ databricks_retry_delay: int = 1,
+ databricks_retry_args: dict[Any, Any] | None = None,
+ do_xcom_push: bool = True,
+ wait_for_termination: bool = True,
+ timeout: float = 3600,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+
+ if statement and statement_id:
Review Comment:
Added.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on PR #49516:
URL: https://github.com/apache/airflow/pull/49516#issuecomment-2848082970
Changes made and implemented!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2070195540
##########
providers/databricks/tests/system/databricks/example_databricks_sensors.py:
##########
@@ -30,6 +31,8 @@
# [DAG name to be shown on Airflow UI]
DAG_ID = "example_databricks_sensor"
+WAREHOUSE_ID = os.environ.get("WAREHOUSE_ID", "cf414a2206dfb397")
Review Comment:
Changed.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adding `DatabricksSQLStatementSensor` Sensor with Deferrability [airflow]
Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #49516:
URL: https://github.com/apache/airflow/pull/49516#discussion_r2072893670
##########
airflow-core/tests/unit/always/test_project_structure.py:
##########
@@ -113,6 +113,7 @@ def test_providers_modules_should_have_tests(self):
"providers/common/messaging/tests/unit/common/messaging/providers/test_base_provider.py",
"providers/common/messaging/tests/unit/common/messaging/providers/test_sqs.py",
"providers/databricks/tests/unit/databricks/test_version_compat.py",
+ "providers/databricks/tests/unit/databricks/utils/test_mixins.py",
Review Comment:
Can you elaborate a bit more? I'm not sure I follow
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org