From 66f94f95c2e92baad2761b5a1fa405e36c17808a Mon Sep 17 00:00:00 2001 From: Josh Fell <48934154+josh-fell@users.noreply.github.com> Date: Fri, 10 Dec 2021 05:05:37 -0500 Subject: [PATCH] Remove db call from `DatabricksHook.__init__()` (#20180) --- airflow/providers/databricks/hooks/databricks.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index ac8d9511e0676..e56b15a99b33f 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -135,11 +135,7 @@ def __init__( ) -> None: super().__init__() self.databricks_conn_id = databricks_conn_id - self.databricks_conn = self.get_connection(databricks_conn_id) - if 'host' in self.databricks_conn.extra_dejson: - self.host = self._parse_host(self.databricks_conn.extra_dejson['host']) - else: - self.host = self._parse_host(self.databricks_conn.host) + self.databricks_conn = None self.timeout_seconds = timeout_seconds if retry_limit < 1: raise ValueError('Retry limit must be greater than equal to 1') @@ -303,6 +299,15 @@ def _do_api_call(self, endpoint_info, json): :rtype: dict """ method, endpoint = endpoint_info + + if self.databricks_conn is None: + self.databricks_conn = self.get_connection(self.databricks_conn_id) + + if 'host' in self.databricks_conn.extra_dejson: + self.host = self._parse_host(self.databricks_conn.extra_dejson['host']) + else: + self.host = self._parse_host(self.databricks_conn.host) + url = f'https://{self.host}/{endpoint}' aad_headers = self._get_aad_headers()