From d7080bdec16eb003e5cc6f0546462dece6c2f029 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Fri, 17 Sep 2021 23:33:07 +0100 Subject: [PATCH] Move DB call out of `DatabricksHook.__init__`` Because of that using `DatabricksHook` in DAG files causes tons of issues as DB calls are made everytime files is parsed. --- airflow/providers/databricks/hooks/databricks.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index f83191ec834e8..5e7fc35c6dd0f 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -117,7 +117,7 @@ def __init__( ) -> None: super().__init__() self.databricks_conn_id = databricks_conn_id - self.databricks_conn = self.get_connection(databricks_conn_id) + self.databricks_conn = None self.timeout_seconds = timeout_seconds if retry_limit < 1: raise ValueError('Retry limit must be greater than equal to 1') @@ -166,6 +166,8 @@ def _do_api_call(self, endpoint_info, json): """ method, endpoint = endpoint_info + self.databricks_conn = self.get_connection(self.databricks_conn_id) + if 'token' in self.databricks_conn.extra_dejson: self.log.info('Using token auth. ') auth = _TokenAuth(self.databricks_conn.extra_dejson['token'])