From 345ff4b412fe2ad572b42d8aa1dada06f57a6e18 Mon Sep 17 00:00:00 2001 From: Alex Ott Date: Wed, 9 Mar 2022 20:46:26 +0100 Subject: [PATCH 1/2] initial version of ReposUpdate operator --- .../providers/databricks/hooks/databricks.py | 24 ++++ .../databricks/operators/databricks_repos.py | 106 ++++++++++++++++++ 2 files changed, 130 insertions(+) create mode 100644 airflow/providers/databricks/operators/databricks_repos.py diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index cdab4d0e159f1..c4f7edc32d95e 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -45,6 +45,8 @@ LIST_JOBS_ENDPOINT = ('GET', 'api/2.1/jobs/list') +WORKSPACE_GET_STATUS_ENDPOINT = ('GET', 'api/2.0/workspace/get-status') + RUN_LIFE_CYCLE_STATES = ['PENDING', 'RUNNING', 'TERMINATING', 'TERMINATED', 'SKIPPED', 'INTERNAL_ERROR'] @@ -328,3 +330,25 @@ def uninstall(self, json: dict) -> None: :param json: json dictionary containing cluster_id and an array of library """ self._do_api_call(UNINSTALL_LIBS_ENDPOINT, json) + + def update_repo(self, repo_id: str, json: dict) -> dict: + """ + + :param repo_id: + :param json: + :return: + """ + repos_endpoint = ('PATCH', f'api/2.0/repos/{repo_id}') + return self._do_api_call(repos_endpoint, json) + + def get_repo_by_path(self, path: str) -> Optional[str]: + """ + + :param path: + :return: + """ + result = self._do_api_call(WORKSPACE_GET_STATUS_ENDPOINT, {'path': path}) + if result.get('object_type', '') == 'REPO': + return str(result['object_id']) + + return None diff --git a/airflow/providers/databricks/operators/databricks_repos.py b/airflow/providers/databricks/operators/databricks_repos.py new file mode 100644 index 0000000000000..30cff02992920 --- /dev/null +++ b/airflow/providers/databricks/operators/databricks_repos.py @@ -0,0 +1,106 @@ +# +# 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 operators.""" + +import time +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union + +from airflow.exceptions import AirflowException +from airflow.models import BaseOperator +from airflow.providers.databricks.hooks.databricks import DatabricksHook + +if TYPE_CHECKING: + from airflow.utils.context import Context + + +class DatabricksReposUpdateOperator(BaseOperator): + """ + Updates specified repository to a given branch or tag using + `api/2.0/repos/ + `_ + API endpoint. + + :param branch: optional name of branch to update to. Should be specified if tag is omitted + :param tag: optional name of tag to update to. Should be specified if branch is omitted + :param repo_id: optional ID of existing repository. Should be specified if repo_path is omitted + :param repo_path: optional path of existing repository. Should be specified if repo_id is omitted + :param databricks_conn_id: Reference to the :ref:`Databricks connection `. + By default and in the common case this will be ``databricks_default``. To use + token based authentication, provide the key ``token`` in the extra field for the + connection and create the key ``host`` and leave the ``host`` field empty. + :param databricks_retry_limit: Amount of times retry if the Databricks backend is + unreachable. Its value must be greater than or equal to 1. + :param databricks_retry_delay: Number of seconds to wait between retries (it + might be a floating point number). + """ + + # Used in airflow.models.BaseOperator + template_fields: Sequence[str] = ('repo_path', 'tag', 'branch') + + def __init__( + self, + *, + branch: Optional[str] = None, + tag: Optional[str] = None, + repo_id: Optional[str] = None, + repo_path: Optional[str] = None, + databricks_conn_id: str = 'databricks_default', + databricks_retry_limit: int = 3, + databricks_retry_delay: int = 1, + **kwargs, + ) -> None: + """Creates a new ``DatabricksSubmitRunOperator``.""" + super().__init__(**kwargs) + self.databricks_conn_id = databricks_conn_id + self.databricks_retry_limit = databricks_retry_limit + self.databricks_retry_delay = databricks_retry_delay + if branch is not None and tag is not None: + raise AirflowException("Only one of branch or tag should be provided, but not both") + if branch is None and tag is None: + raise AirflowException("One of branch or tag should be provided") + if repo_id is not None and repo_path is not None: + raise AirflowException( + "Only one of repo_id or repo_path should be provided, but not both") + if repo_id is None and repo_path is None: + raise AirflowException("One of repo_id repo_path tag should be provided") + self.repo_path = repo_path + self.repo_id = repo_id + self.branch = branch + self.tag = tag + + def _get_hook(self) -> DatabricksHook: + return DatabricksHook( + self.databricks_conn_id, + retry_limit=self.databricks_retry_limit, + retry_delay=self.databricks_retry_delay, + ) + + def execute(self, context: 'Context'): + hook = self._get_hook() + if self.repo_path is not None: + self.repo_id = hook.get_repo_by_path(self.repo_path) + if self.repo_id is None: + raise AirflowException(f"Can't find Repo ID for path '{self.repo_path}'") + if self.branch is not None: + payload = {'branch': self.branch} + else: + payload = {'tag': self.tag} + + result = hook.update_repo(self.repo_id, payload) + return result['head_commit_id'] From 9109ccc1016af959d2649390437d467b8b7e59d9 Mon Sep 17 00:00:00 2001 From: Alex Ott Date: Tue, 15 Mar 2022 11:07:01 +0100 Subject: [PATCH 2/2] Add documentation, example & tests for DatabricksReposUpdateOperator --- .../example_dags/example_databricks_repos.py | 56 ++++++++++++++++ .../providers/databricks/hooks/databricks.py | 2 +- .../databricks/operators/databricks_repos.py | 20 +++--- airflow/providers/databricks/provider.yaml | 9 +++ .../operators/repos_update.rst | 66 +++++++++++++++++++ .../operators/run_now.rst | 2 +- .../operators/test_databricks_repos.py | 64 ++++++++++++++++++ 7 files changed, 206 insertions(+), 13 deletions(-) create mode 100644 airflow/providers/databricks/example_dags/example_databricks_repos.py create mode 100644 docs/apache-airflow-providers-databricks/operators/repos_update.rst create mode 100644 tests/providers/databricks/operators/test_databricks_repos.py diff --git a/airflow/providers/databricks/example_dags/example_databricks_repos.py b/airflow/providers/databricks/example_dags/example_databricks_repos.py new file mode 100644 index 0000000000000..458f7cb8ce72b --- /dev/null +++ b/airflow/providers/databricks/example_dags/example_databricks_repos.py @@ -0,0 +1,56 @@ +# 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 datetime import datetime + +from airflow import DAG +from airflow.providers.databricks.operators.databricks import DatabricksSubmitRunOperator +from airflow.providers.databricks.operators.databricks_repos import DatabricksReposUpdateOperator + +default_args = { + 'owner': 'airflow', + 'databricks_conn_id': 'my-shard-pat', +} + +with DAG( + dag_id='example_databricks_operator', + schedule_interval='@daily', + start_date=datetime(2021, 1, 1), + tags=['example'], + catchup=False, +) as dag: + # [START howto_operator_databricks_repo_update] + # Example of updating a Databricks Repo to the latest code + repo_path = "/Repos/user@domain.com/demo-repo" + update_repo = DatabricksReposUpdateOperator(task_id='update_repo', repo_path=repo_path, branch="releases") + # [END howto_operator_databricks_repo_update] + + notebook_task_params = { + 'new_cluster': { + 'spark_version': '9.1.x-scala2.12', + 'node_type_id': 'r3.xlarge', + 'aws_attributes': {'availability': 'ON_DEMAND'}, + 'num_workers': 8, + }, + 'notebook_task': { + 'notebook_path': f'{repo_path}/PrepareData', + }, + } + + notebook_task = DatabricksSubmitRunOperator(task_id='notebook_task', json=notebook_task_params) + + (update_repo >> notebook_task) diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index c4f7edc32d95e..977800edb77bc 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -331,7 +331,7 @@ def uninstall(self, json: dict) -> None: """ self._do_api_call(UNINSTALL_LIBS_ENDPOINT, json) - def update_repo(self, repo_id: str, json: dict) -> dict: + def update_repo(self, repo_id: str, json: Dict[str, Any]) -> dict: """ :param repo_id: diff --git a/airflow/providers/databricks/operators/databricks_repos.py b/airflow/providers/databricks/operators/databricks_repos.py index 30cff02992920..fc50730d03d06 100644 --- a/airflow/providers/databricks/operators/databricks_repos.py +++ b/airflow/providers/databricks/operators/databricks_repos.py @@ -18,8 +18,7 @@ # """This module contains Databricks operators.""" -import time -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union +from typing import TYPE_CHECKING, Optional, Sequence from airflow.exceptions import AirflowException from airflow.models import BaseOperator @@ -36,10 +35,10 @@ class DatabricksReposUpdateOperator(BaseOperator): `_ API endpoint. - :param branch: optional name of branch to update to. Should be specified if tag is omitted - :param tag: optional name of tag to update to. Should be specified if branch is omitted - :param repo_id: optional ID of existing repository. Should be specified if repo_path is omitted - :param repo_path: optional path of existing repository. Should be specified if repo_id is omitted + :param branch: optional name of branch to update to. Should be specified if ``tag`` is omitted + :param tag: optional name of tag to update to. Should be specified if ``branch`` is omitted + :param repo_id: optional ID of existing repository. Should be specified if ``repo_path`` is omitted + :param repo_path: optional path of existing repository. Should be specified if ``repo_id`` is omitted :param databricks_conn_id: Reference to the :ref:`Databricks connection `. By default and in the common case this will be ``databricks_default``. To use token based authentication, provide the key ``token`` in the extra field for the @@ -75,8 +74,7 @@ def __init__( if branch is None and tag is None: raise AirflowException("One of branch or tag should be provided") if repo_id is not None and repo_path is not None: - raise AirflowException( - "Only one of repo_id or repo_path should be provided, but not both") + raise AirflowException("Only one of repo_id or repo_path should be provided, but not both") if repo_id is None and repo_path is None: raise AirflowException("One of repo_id repo_path tag should be provided") self.repo_path = repo_path @@ -98,9 +96,9 @@ def execute(self, context: 'Context'): if self.repo_id is None: raise AirflowException(f"Can't find Repo ID for path '{self.repo_path}'") if self.branch is not None: - payload = {'branch': self.branch} + payload = {'branch': str(self.branch)} else: - payload = {'tag': self.tag} + payload = {'tag': str(self.tag)} - result = hook.update_repo(self.repo_id, payload) + result = hook.update_repo(str(self.repo_id), payload) return result['head_commit_id'] diff --git a/airflow/providers/databricks/provider.yaml b/airflow/providers/databricks/provider.yaml index 3d08e20abef4e..ae307c7b79cff 100644 --- a/airflow/providers/databricks/provider.yaml +++ b/airflow/providers/databricks/provider.yaml @@ -53,6 +53,12 @@ integrations: - /docs/apache-airflow-providers-databricks/operators/copy_into.rst logo: /integration-logos/databricks/Databricks.png tags: [service] + - integration-name: Databricks Repos + external-doc-url: https://docs.databricks.com/repos/index.html + how-to-guide: + - /docs/apache-airflow-providers-databricks/operators/repos_update.rst + logo: /integration-logos/databricks/Databricks.png + tags: [service] operators: - integration-name: Databricks @@ -61,6 +67,9 @@ operators: - integration-name: Databricks SQL python-modules: - airflow.providers.databricks.operators.databricks_sql + - integration-name: Databricks Repos + python-modules: + - airflow.providers.databricks.operators.databricks_repos hooks: - integration-name: Databricks diff --git a/docs/apache-airflow-providers-databricks/operators/repos_update.rst b/docs/apache-airflow-providers-databricks/operators/repos_update.rst new file mode 100644 index 0000000000000..0f63c2468551e --- /dev/null +++ b/docs/apache-airflow-providers-databricks/operators/repos_update.rst @@ -0,0 +1,66 @@ + .. 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. + + + +DatabricksReposUpdateOperator +============================= + +Use the :class:`~airflow.providers.databricks.operators.DatabricksReposUpdateOperator` to update code in an existing +`Databricks Repos `_ to a given Git branch or tag +via `api/2.0/repos/ `_ API endpoint. + + +Using the Operator +^^^^^^^^^^^^^^^^^^ + +Usually this operator is used to update a source code of the Databricks job before its execution. +To use this operator you need to provide either ``branch`` or ``tag`` and either ``repo_path`` or ``repo_id``. + +.. list-table:: + :widths: 15 25 + :header-rows: 1 + + * - Parameter + - Input + * - branch: str + - Name of the existing Git branch to update to (required if ``tag`` isn't provided). + * - tag: str + - Name of the existing Git tag to update to (required if ``branch`` isn't provided). + * - repo_path: str + - Path to existing Databricks Repos, like, ``/Repos//repo_name`` (required if ``repo_id`` isn't provided). + * - repo_id: str + - ID of existing Databricks Repos (required if ``repo_path`` isn't provided). + * - databricks_conn_id: string + - the name of the Airflow connection to use. + * - databricks_retry_limit: integer + - amount of times retry if the Databricks backend is unreachable. + * - databricks_retry_delay: decimal + - number of seconds to wait between retries. + +Examples +-------- + +Updating Databricks Repo by specifying path +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +An example usage of the DatabricksReposUpdateOperator is as follows: + +.. exampleinclude:: /../../airflow/providers/databricks/example_dags/example_databricks_repos.py + :language: python + :start-after: [START howto_operator_databricks_repo_update] + :end-before: [END howto_operator_databricks_repo_update] diff --git a/docs/apache-airflow-providers-databricks/operators/run_now.rst b/docs/apache-airflow-providers-databricks/operators/run_now.rst index 62fb4fd3f2673..f77cfb2f53139 100644 --- a/docs/apache-airflow-providers-databricks/operators/run_now.rst +++ b/docs/apache-airflow-providers-databricks/operators/run_now.rst @@ -18,7 +18,7 @@ DatabricksRunNowOperator -=========================== +======================== Use the :class:`~airflow.providers.databricks.operators.DatabricksRunNowOperator` to trigger a run of an existing Databricks job via `api/2.1/jobs/run-now `_ API endpoint. diff --git a/tests/providers/databricks/operators/test_databricks_repos.py b/tests/providers/databricks/operators/test_databricks_repos.py new file mode 100644 index 0000000000000..ad8ccdc82ef8b --- /dev/null +++ b/tests/providers/databricks/operators/test_databricks_repos.py @@ -0,0 +1,64 @@ +# +# 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. +# +import unittest +from unittest import mock + +from airflow.providers.databricks.operators.databricks_repos import DatabricksReposUpdateOperator + +TASK_ID = 'databricks-operator' +DEFAULT_CONN_ID = 'databricks_default' + + +class TestDatabricksReposUpdateOperator(unittest.TestCase): + @mock.patch('airflow.providers.databricks.operators.databricks_repos.DatabricksHook') + def test_update_with_id(self, db_mock_class): + """ + Test the execute function in case where the run is successful. + """ + op = DatabricksReposUpdateOperator(task_id=TASK_ID, branch="releases", repo_id="123") + db_mock = db_mock_class.return_value + db_mock.update_repo.return_value = {'head_commit_id': '123456'} + + op.execute(None) + + db_mock_class.assert_called_once_with( + DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit, retry_delay=op.databricks_retry_delay + ) + + db_mock.update_repo.assert_called_once_with('123', {'branch': 'releases'}) + + @mock.patch('airflow.providers.databricks.operators.databricks_repos.DatabricksHook') + def test_update_with_path(self, db_mock_class): + """ + Test the execute function in case where the run is successful. + """ + op = DatabricksReposUpdateOperator( + task_id=TASK_ID, tag="v1.0.0", repo_path="/Repos/user@domain.com/test-repo" + ) + db_mock = db_mock_class.return_value + db_mock.get_repo_by_path.return_value = '123' + db_mock.update_repo.return_value = {'head_commit_id': '123456'} + + op.execute(None) + + db_mock_class.assert_called_once_with( + DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit, retry_delay=op.databricks_retry_delay + ) + + db_mock.update_repo.assert_called_once_with('123', {'tag': 'v1.0.0'})