Operator for updating Databricks Repos

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)

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[str, Any]) -> 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

airflow/providers/databricks/operators/databricks_repos.py

@@ -0,0 +1,104 @@
+#
+# 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."""
+
+from typing import TYPE_CHECKING, Optional, Sequence
+
+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/
+    <https://docs.databricks.com/dev-tools/api/latest/repos.html#operation/update-repo>`_
+    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 <howto/connection:databricks>`.
+        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': str(self.branch)}
+        else:
+            payload = {'tag': str(self.tag)}
+
+        result = hook.update_repo(str(self.repo_id), payload)
+        return result['head_commit_id']

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

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 <https://docs.databricks.com/repos/index.html>`_ to a given Git branch or tag
+via `api/2.0/repos/ <https://docs.databricks.com/dev-tools/api/latest/repos.html#operation/update-repo>`_ 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/<user_email>/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]

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 <https://docs.databricks.com/dev-tools/api/latest/jobs.html#operation/JobsRunNow>`_ API endpoint.

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'})