diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index ffa77570d5ab9..79116604124ba 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -27,6 +27,8 @@ """ from typing import Any, Dict, List, Optional +from requests import exceptions as requests_exceptions + from airflow.exceptions import AirflowException from airflow.providers.databricks.hooks.databricks_base import BaseDatabricksHook @@ -364,12 +366,16 @@ def create_repo(self, json: Dict[str, Any]) -> dict: 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']) + Obtains Repos ID by path + :param path: path to a repository + :return: Repos ID if it exists, None if doesn't. + """ + try: + result = self._do_api_call(WORKSPACE_GET_STATUS_ENDPOINT, {'path': path}, wrap_http_errors=False) + if result.get('object_type', '') == 'REPO': + return str(result['object_id']) + except requests_exceptions.HTTPError as e: + if e.response.status_code != 404: + raise e return None diff --git a/airflow/providers/databricks/hooks/databricks_base.py b/airflow/providers/databricks/hooks/databricks_base.py index 1a418fd04e625..6e0f1b44d8ae7 100644 --- a/airflow/providers/databricks/hooks/databricks_base.py +++ b/airflow/providers/databricks/hooks/databricks_base.py @@ -307,7 +307,12 @@ def _get_token(self, raise_error: bool = False) -> Optional[str]: def _log_request_error(self, attempt_num: int, error: str) -> None: self.log.error('Attempt %s API Request to Databricks failed with reason: %s', attempt_num, error) - def _do_api_call(self, endpoint_info: Tuple[str, str], json: Optional[Dict[str, Any]] = None): + def _do_api_call( + self, + endpoint_info: Tuple[str, str], + json: Optional[Dict[str, Any]] = None, + wrap_http_errors: bool = True, + ): """ Utility function to perform an API call with retries @@ -362,7 +367,12 @@ def _do_api_call(self, endpoint_info: Tuple[str, str], json: Optional[Dict[str, except RetryError: raise AirflowException(f'API requests to Databricks failed {self.retry_limit} times. Giving up.') except requests_exceptions.HTTPError as e: - raise AirflowException(f'Response: {e.response.content}, Status Code: {e.response.status_code}') + if wrap_http_errors: + raise AirflowException( + f'Response: {e.response.content}, Status Code: {e.response.status_code}' + ) + else: + raise e @staticmethod def _get_error_code(exception: BaseException) -> str: