diff --git a/airflow/providers/databricks/provider.yaml b/airflow/providers/databricks/provider.yaml index c48a47f2da3fd..e542bfd66b91e 100644 --- a/airflow/providers/databricks/provider.yaml +++ b/airflow/providers/databricks/provider.yaml @@ -99,6 +99,10 @@ triggers: python-modules: - airflow.providers.databricks.triggers.databricks +sensors: + - integration-name: Databricks + python-modules: + - airflow.providers.databricks.sensors.databricks_sql connection-types: - hook-class-name: airflow.providers.databricks.hooks.databricks.DatabricksHook 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/databricks_sql.py b/airflow/providers/databricks/sensors/databricks_sql.py new file mode 100644 index 0000000000000..0cb6f2a88d431 --- /dev/null +++ b/airflow/providers/databricks/sensors/databricks_sql.py @@ -0,0 +1,134 @@ +# +# 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.exceptions import AirflowException +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 that runs a SQL query on Databricks. + + :param databricks_conn_id: Reference to :ref:`Databricks + connection id` (templated), defaults to + DatabricksSqlHook.default_conn_name. + :param sql_warehouse_name: Optional name of Databricks SQL warehouse. If not specified, ``http_path`` + must be provided as described below, defaults to None + :param http_path: Optional string specifying HTTP path of Databricks SQL warehouse or All Purpose cluster. + If not specified, it should be either specified in the Databricks connection's + extra parameters, or ``sql_warehouse_name`` must be specified. + :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 Databricks Runtime version 9.0+ (templated), defaults to "" + :param schema: An optional initial schema to use. + Requires Databricks Runtime 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_warehouse_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: + """Creates DatabricksSqlSensor object using the specified input arguments.""" + self.databricks_conn_id = databricks_conn_id + self._http_path = http_path + self._sql_warehouse_name = sql_warehouse_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 hook(self) -> DatabricksSqlHook: + """Creates and returns a DatabricksSqlHook object.""" + return DatabricksSqlHook( + self.databricks_conn_id, + self._http_path, + self._sql_warehouse_name, + self.session_config, + self.http_headers, + self.catalog, + self.schema, + self.caller, + **self.client_parameters, + **self.hook_params, + ) + + def _get_results(self) -> bool: + """Uses the Databricks SQL hook and runs the specified SQL query.""" + if not (self._http_path or self._sql_warehouse_name): + raise AirflowException( + "Databricks SQL warehouse/cluster configuration missing. Please specify either http_path or " + "sql_warehouse_name." + ) + hook = self.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: + """Sensor poke function to get and return results from the SQL sensor.""" + 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..fd0535283df17 100644 --- a/docs/apache-airflow-providers-databricks/operators/sql.rst +++ b/docs/apache-airflow-providers-databricks/operators/sql.rst @@ -22,14 +22,14 @@ DatabricksSqlOperator ===================== Use the :class:`~airflow.providers.databricks.operators.databricks_sql.DatabricksSqlOperator` to execute SQL -on a `Databricks SQL endpoint `_ or a +on a `Databricks SQL warehouse `_ or a `Databricks cluster `_. Using the Operator ------------------ -Operator executes given SQL queries against configured endpoint. The only required parameters are: +Operator executes given SQL queries against configured warehouse. The only required parameters are: * ``sql`` - SQL queries to execute. There are 3 ways of specifying SQL queries: @@ -37,7 +37,7 @@ Operator executes given SQL queries against configured endpoint. The only requir 2. List of strings representing SQL statements. 3. Name of the file with SQL queries. File must have ``.sql`` extension. Each query should finish with ``;`` -* One of ``sql_endpoint_name`` (name of Databricks SQL endpoint to use) or ``http_path`` (HTTP path for Databricks SQL endpoint or Databricks cluster). +* One of ``sql_warehouse_name`` (name of Databricks SQL warehouse to use) or ``http_path`` (HTTP path for Databricks SQL warehouse or Databricks cluster). Other parameters are optional and could be found in the class documentation. @@ -84,3 +84,39 @@ 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 table accessible via a Databricks SQL warehouse or interactive cluster. + +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_warehouse_name`` (name of Databricks SQL warehouse to use) or ``http_path`` (HTTP path for Databricks SQL warehouse or Databricks cluster). + +Other parameters are optional and could be found in the class documentation. + +Examples +-------- +Configuring Databricks connection to be used with the Sensor. + +.. exampleinclude:: /../../tests/system/providers/databricks/example_databricks_sensors.py + :language: python + :dedent: 4 + :start-after: [START howto_sensor_databricks_connection_setup] + :end-before: [END howto_sensor_databricks_connection_setup] + +Poking the specific table for existence of data/partition: + +.. exampleinclude:: /../../tests/system/providers/databricks/example_databricks_sensors.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_databricks_sql.py b/tests/providers/databricks/sensors/test_databricks_sql.py new file mode 100644 index 0000000000000..eb77b66ac605a --- /dev/null +++ b/tests/providers/databricks/sensors/test_databricks_sql.py @@ -0,0 +1,97 @@ +# +# 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.databricks_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" +PERSONAL_ACCESS_TOKEN = "token" + +DEFAULT_SCHEMA = "schema1" +DEFAULT_CATALOG = "catalog1" +DEFAULT_TABLE = "table1" +DEFAULT_HTTP_PATH = "/sql/1.0/warehouses/xxxxx" +DEFAULT_SQL_WAREHOUSE = "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_warehouse_name=DEFAULT_SQL_WAREHOUSE, + 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_warehouse_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) + + def test_sql_warehouse_http_path(self): + """Neither SQL warehouse name not HTTP path has been specified.""" + _sensor_without_sql_warehouse_http = DatabricksSqlSensor( + task_id="task2", + databricks_conn_id=DEFAULT_CONN_ID, + dag=self.dag, + sql=DEFAULT_SQL, + schema=DEFAULT_SCHEMA, + catalog=DEFAULT_CATALOG, + timeout=30, + poke_interval=15, + ) + with pytest.raises(AirflowException): + _sensor_without_sql_warehouse_http._get_results() diff --git a/tests/system/providers/databricks/example_databricks_sensors.py b/tests/system/providers/databricks/example_databricks_sensors.py new file mode 100644 index 0000000000000..d2507118e2f27 --- /dev/null +++ b/tests/system/providers/databricks/example_databricks_sensors.py @@ -0,0 +1,86 @@ +# +# 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.databricks_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_warehouse_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 in a table. + sql_sensor = DatabricksSqlSensor( + databricks_conn_id=connection_id, + sql_warehouse_name=sql_warehouse_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 and does not + # affect the execution of the single DAG task which would run regardless of its presence. + # 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)