diff --git a/airflow/providers/databricks/.latest-doc-only-change.txt b/airflow/providers/databricks/.latest-doc-only-change.txt index 28124098645cf..76d3da63c0788 100644 --- a/airflow/providers/databricks/.latest-doc-only-change.txt +++ b/airflow/providers/databricks/.latest-doc-only-change.txt @@ -1 +1 @@ -6c3a67d4fccafe4ab6cd9ec8c7bacf2677f17038 +c3867781e09b7e0e0d19c0991865a2453194d9a8 diff --git a/airflow/providers/databricks/provider.yaml b/airflow/providers/databricks/provider.yaml index 133899996c3f9..47fec6a0df4d3 100644 --- a/airflow/providers/databricks/provider.yaml +++ b/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: diff --git a/airflow/providers/databricks/sensors/__init__.py b/airflow/providers/databricks/sensors/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/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. diff --git a/airflow/providers/databricks/sensors/sql.py b/airflow/providers/databricks/sensors/sql.py new file mode 100644 index 0000000000000..033f0a26b02b1 --- /dev/null +++ b/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` (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() diff --git a/docs/apache-airflow-providers-databricks/operators/sql.rst b/docs/apache-airflow-providers-databricks/operators/sql.rst index 7e80a6b7cffd0..01f106bf69ba1 100644 --- a/docs/apache-airflow-providers-databricks/operators/sql.rst +++ b/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] diff --git a/tests/providers/databricks/sensors/__init__.py b/tests/providers/databricks/sensors/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/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. diff --git a/tests/providers/databricks/sensors/test_sql.py b/tests/providers/databricks/sensors/test_sql.py new file mode 100644 index 0000000000000..fae35bd4a01ca --- /dev/null +++ b/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) diff --git a/tests/system/providers/databricks/example_databricks_sensor.py b/tests/system/providers/databricks/example_databricks_sensor.py new file mode 100644 index 0000000000000..b38e2140f2f69 --- /dev/null +++ b/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)