Databricks SQL Sensor

airflow/providers/databricks/.latest-doc-only-change.txt

@@ -1 +1 @@
-6c3a67d4fccafe4ab6cd9ec8c7bacf2677f17038
+c3867781e09b7e0e0d19c0991865a2453194d9a8

airflow/providers/databricks/provider.yaml

@@ -92,6 +92,11 @@ hooks:
     python-modules:
       - airflow.providers.databricks.hooks.databricks_sql
 
+sensors:
+  - integration-name: Databricks
+    python-modules:
+      - airflow.providers.databricks.sensors.sql
+
 triggers:
   - integration-name: Databricks
     python-modules:

airflow/providers/databricks/sensors/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/databricks/sensors/sql.py

@@ -0,0 +1,124 @@
+#
+# 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.
+#
+"""This module contains Databricks sensors."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence
+
+from airflow.compat.functools import cached_property
+from airflow.providers.common.sql.hooks.sql import fetch_all_handler
+from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook
+from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class DatabricksSqlSensor(BaseSensorOperator):
+    """
+    Sensor to execute SQL statements on a Delta table via Databricks.
+
+    :param databricks_conn_id: Reference to :ref:`Databricks
+        connection id<howto/connection:databricks>` (templated), defaults to
+        DatabricksSqlHook.default_conn_name
+    :param http_path: Optional string specifying HTTP path of Databricks SQL Endpoint or cluster.
+        If not specified, it should be either specified in the Databricks connection's
+        extra parameters, or ``sql_endpoint_name`` must be specified.
+    :param sql_endpoint_name: Optional name of Databricks SQL Endpoint. If not specified, ``http_path``
+        must be provided as described above, defaults to None
+    :param session_configuration: An optional dictionary of Spark session parameters. If not specified,
+        it could be specified in the Databricks connection's extra parameters., defaults to None
+    :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), defaults to ""
+    :param schema: An optional initial schema to use.
+        Requires DBR version 9.0+ (templated), defaults to "default"
+    :param sql: SQL statement to be executed.
+    :param handler: Handler for DbApiHook.run() to return results, defaults to fetch_all_handler
+    :param client_parameters: Additional parameters internal to Databricks SQL Connector parameters.
+    """
+
+    template_fields: Sequence[str] = (
+        "databricks_conn_id",
+        "sql",
+        "catalog",
+        "schema",
+        "http_headers",
+    )
+
+    template_ext: Sequence[str] = (".sql",)
+    template_fields_renderers = {"sql": "sql"}
+
+    def __init__(
+        self,
+        *,
+        databricks_conn_id: str = DatabricksSqlHook.default_conn_name,
+        http_path: str | None = None,
+        sql_endpoint_name: str | None = None,
+        session_configuration=None,
+        http_headers: list[tuple[str, str]] | None = None,
+        catalog: str = "",
+        schema: str = "default",
+        sql: str | Iterable[str],
+        handler: Callable[[Any], Any] = fetch_all_handler,
+        client_parameters: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> None:
+        self.databricks_conn_id = databricks_conn_id
+        self._http_path = http_path
+        self._sql_endpoint_name = sql_endpoint_name
+        self.session_config = session_configuration
+        self.http_headers = http_headers
+        self.catalog = catalog
+        self.schema = schema
+        self.sql = sql
+        self.caller = "DatabricksSqlSensor"
+        self.client_parameters = client_parameters or {}
+        self.hook_params = kwargs.pop("hook_params", {})
+        self.handler = handler
+        super().__init__(**kwargs)
+
+    @cached_property
+    def _get_hook(self) -> DatabricksSqlHook:
+        return DatabricksSqlHook(
+            self.databricks_conn_id,
+            self._http_path,
+            self._sql_endpoint_name,
+            self.session_config,
+            self.http_headers,
+            self.catalog,
+            self.schema,
+            self.caller,
+            **self.client_parameters,
+            **self.hook_params,
+        )
+
+    def _get_results(self) -> bool:
+        hook = self._get_hook
+        sql_result = hook.run(
+            self.sql,
+            handler=self.handler if self.do_xcom_push else None,
+        )
+        self.log.debug("SQL result: %s", sql_result)
+        return bool(sql_result)
+
+    def poke(self, context: Context) -> bool:
+        return self._get_results()

docs/apache-airflow-providers-databricks/operators/sql.rst

@@ -84,3 +84,32 @@ An example usage of the DatabricksSqlOperator to perform statements from a file
     :language: python
     :start-after: [START howto_operator_databricks_sql_multiple_file]
     :end-before: [END howto_operator_databricks_sql_multiple_file]
+
+
+DatabricksSqlSensor
+===================
+
+Use the :class:`~airflow.providers.databricks.sensors.sql.DatabricksSqlSensor` to run the sensor
+for a specific Delta table on a Databricks SQL workspace.
+
+Using the Sensor
+----------------
+
+The sensor executes the SQL statement supplied by the user. The only required parameters are:
+
+* ``sql`` - SQL query to execute for the sensor.
+
+* One of ``sql_endpoint_name`` (name of Databricks SQL endpoint to use) or ``http_path`` (HTTP path for Databricks SQL endpoint or Databricks cluster).
+
+Other parameters are optional and could be found in the class documentation.
+
+Examples
+--------
+
+Poking the specific table for existence of data/partition:
+
+.. exampleinclude:: /../../tests/system/providers/databricks/example_databricks_sensor.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_sensor_databricks_sql]
+    :end-before: [END howto_sensor_databricks_sql]

tests/providers/databricks/sensors/__init__.py

@@ -0,0 +1,16 @@
+# 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.

tests/providers/databricks/sensors/test_sql.py

@@ -0,0 +1,81 @@
+#
+# 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 datetime import datetime, timedelta
+from unittest.mock import patch
+
+import pytest
+
+from airflow.exceptions import AirflowException
+from airflow.models import DAG
+from airflow.providers.databricks.sensors.sql import DatabricksSqlSensor
+from airflow.utils import timezone
+
+TASK_ID = "db-sensor"
+DEFAULT_CONN_ID = "databricks_default"
+HOST = "xx.cloud.databricks.com"
+HOST_WITH_SCHEME = "https://xx.cloud.databricks.com"
+TOKEN = "token"
+
+DEFAULT_SCHEMA = "schema1"
+DEFAULT_CATALOG = "catalog1"
+DEFAULT_TABLE = "table1"
+DEFAULT_SQL_ENDPOINT = "sql_warehouse_default"
+DEFAULT_CALLER = "TestDatabricksSqlSensor"
+DEFAULT_SQL = f"select 1 from {DEFAULT_CATALOG}.{DEFAULT_SCHEMA}.{DEFAULT_TABLE} LIMIT 1"
+DEFAULT_DATE = timezone.datetime(2017, 1, 1)
+
+TIMESTAMP_TEST = datetime.now() - timedelta(days=30)
+
+
+class TestDatabricksSqlSensor:
+    def setup_method(self):
+        args = {"owner": "airflow", "start_date": DEFAULT_DATE}
+        self.dag = DAG("test_dag_id", default_args=args)
+
+        self.sensor = DatabricksSqlSensor(
+            task_id=TASK_ID,
+            databricks_conn_id=DEFAULT_CONN_ID,
+            sql_endpoint_name=DEFAULT_SQL_ENDPOINT,
+            dag=self.dag,
+            sql=DEFAULT_SQL,
+            schema=DEFAULT_SCHEMA,
+            catalog=DEFAULT_CATALOG,
+            timeout=30,
+            poke_interval=15,
+        )
+
+    def test_init(self):
+        assert self.sensor.databricks_conn_id == "databricks_default"
+        assert self.sensor.task_id == "db-sensor"
+        assert self.sensor._sql_endpoint_name == "sql_warehouse_default"
+        assert self.sensor.poke_interval == 15
+
+    @pytest.mark.parametrize(
+        argnames=("sensor_poke_result", "expected_poke_result"), argvalues=[(True, True), (False, False)]
+    )
+    @patch.object(DatabricksSqlSensor, "poke")
+    def test_poke(self, mock_poke, sensor_poke_result, expected_poke_result):
+        mock_poke.return_value = sensor_poke_result
+        assert self.sensor.poke({}) == expected_poke_result
+
+    def test_unsupported_conn_type(self):
+        with pytest.raises(AirflowException):
+            self.sensor.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)

tests/system/providers/databricks/example_databricks_sensor.py

@@ -0,0 +1,85 @@
+#
+# 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
+
+import os
+import textwrap
+from datetime import datetime
+
+from airflow import DAG
+from airflow.providers.databricks.sensors.sql import DatabricksSqlSensor
+
+# [Env variable to be used from the OS]
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+# [DAG name to be shown on Airflow UI]
+DAG_ID = "example_databricks_sensor"
+
+with DAG(
+    dag_id=DAG_ID,
+    schedule="@daily",
+    start_date=datetime(2021, 1, 1),
+    tags=["example"],
+    catchup=False,
+) as dag:
+    dag.doc_md = textwrap.dedent(
+        """
+
+        This is an example DAG which uses the DatabricksSqlSensor
+        sensor. The example task in the DAG executes the provided
+        SQL query against the Databricks SQL warehouse and if a
+        result is returned, the sensor returns True/succeeds.
+        If no results are returned, the sensor returns False/
+        fails.
+
+        """
+    )
+    # [START howto_sensor_databricks_connection_setup]
+    # Connection string setup for Databricks workspace.
+    connection_id = "databricks_default"
+    sql_endpoint_name = "Starter Warehouse"
+    # [END howto_sensor_databricks_connection_setup]
+
+    # [START howto_sensor_databricks_sql]
+    # Example of using the Databricks SQL Sensor to check existence of data/partitions for a Delta table.
+    sql_sensor = DatabricksSqlSensor(
+        databricks_conn_id=connection_id,
+        sql_endpoint_name=sql_endpoint_name,
+        catalog="hive_metastore",
+        task_id="sql_sensor_task",
+        sql="select * from hive_metastore.temp.sample_table_3 limit 1",
+        timeout=60 * 2,
+    )
+    # [END howto_sensor_databricks_sql]
+
+    # This DAG contains only one task, so the below pattern (task1) is not necessary,
+    # but it is present here as a pattern to be expanded for users.
+    # For example, (task1 >> task 2 >> task3)
+    (sql_sensor)
+
+    from tests.system.utils.watcher import watcher
+
+    # This example does not need a watcher in order to properly mark success/failure
+    # since it is a single task, but it is given here as an example for users to
+    # extend it to their use cases.
+    # when "tearDown" task with trigger rule is part of the DAG
+    list(dag.tasks) >> watcher()
+
+from tests.system.utils import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)