This is an automated email from the ASF dual-hosted git repository.
xddeng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/master by this push:
new 966a06d Fetching databricks host from connection if not supplied in
extras. (#10762)
966a06d is described below
commit 966a06d96bbfe330f1d2825f7b7eaa16d43b7a00
Author: Shubham Joshi <[email protected]>
AuthorDate: Fri Sep 18 16:45:11 2020 +0530
Fetching databricks host from connection if not supplied in extras. (#10762)
* Fetching databricks host from connection if not supplied in extras.
* Fixing formatting issue in databricks test
Co-authored-by: joshi95 <[email protected]>
---
airflow/providers/databricks/hooks/databricks.py | 5 ++++-
tests/providers/databricks/hooks/test_databricks.py | 11 +++++++++++
2 files changed, 15 insertions(+), 1 deletion(-)
diff --git a/airflow/providers/databricks/hooks/databricks.py
b/airflow/providers/databricks/hooks/databricks.py
index 608d63a..22a76cf 100644
--- a/airflow/providers/databricks/hooks/databricks.py
+++ b/airflow/providers/databricks/hooks/databricks.py
@@ -157,7 +157,10 @@ class DatabricksHook(BaseHook): # noqa
if 'token' in self.databricks_conn.extra_dejson:
self.log.info('Using token auth. ')
auth = _TokenAuth(self.databricks_conn.extra_dejson['token'])
- host = self._parse_host(self.databricks_conn.extra_dejson['host'])
+ if 'host' in self.databricks_conn.extra_dejson:
+ host =
self._parse_host(self.databricks_conn.extra_dejson['host'])
+ else:
+ host = self.databricks_conn.host
else:
self.log.info('Using basic auth. ')
auth = (self.databricks_conn.login, self.databricks_conn.password)
diff --git a/tests/providers/databricks/hooks/test_databricks.py
b/tests/providers/databricks/hooks/test_databricks.py
index 212be39..318643a 100644
--- a/tests/providers/databricks/hooks/test_databricks.py
+++ b/tests/providers/databricks/hooks/test_databricks.py
@@ -454,6 +454,17 @@ class TestDatabricksHookToken(unittest.TestCase):
self.assertEqual(kwargs['auth'].token, TOKEN)
+class
TestDatabricksHookTokenWhenNoHostIsProvidedInExtra(TestDatabricksHookToken):
+ @provide_session
+ def setUp(self, session=None):
+ conn = session.query(Connection).filter(Connection.conn_id ==
DEFAULT_CONN_ID).first()
+ conn.extra = json.dumps({'token': TOKEN})
+
+ session.commit()
+
+ self.hook = DatabricksHook()
+
+
class TestRunState(unittest.TestCase):
def test_is_terminal_true(self):
terminal_states = ['TERMINATED', 'SKIPPED', 'INTERNAL_ERROR']
