Remove db call from DatabricksHook.__init__() (#20180)

apache · Dec 10, 2021 · 66f94f9 · 66f94f9
1 parent 214b62d
commit 66f94f9
Showing 1 changed file with 10 additions and 5 deletions.
diff --git 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()