From 77b0218902e1499c7be6fb5c988e7295dd887cca Mon Sep 17 00:00:00 2001 From: Brian Li Date: Thu, 14 Aug 2025 01:02:00 -0700 Subject: [PATCH 1/2] fix: use _get_token in url construction --- .../databricks/hooks/databricks_sql.py | 6 ++-- .../databricks/hooks/test_databricks_sql.py | 33 +++++++++---------- 2 files changed, 17 insertions(+), 22 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py index c689324f97dca..d86d0453fa592 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py @@ -179,7 +179,6 @@ def sqlalchemy_url(self) -> URL: :return: the extracted sqlalchemy.engine.URL object. """ - conn = self.get_conn() url_query = { "http_path": self._http_path, "catalog": self.catalog, @@ -189,9 +188,8 @@ def sqlalchemy_url(self) -> URL: return URL.create( drivername="databricks", username="token", - password=conn.password, - host=conn.host, - port=conn.port, + password=self._get_token(raise_error=True), + host=self.host, query=url_query, ) diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py index 9e2ee58d60069..acb961f3add3b 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py @@ -39,7 +39,6 @@ DEFAULT_CONN_ID = "databricks_default" HOST = "xx.cloud.databricks.com" HOST_WITH_SCHEME = "https://xx.cloud.databricks.com" -PORT = 443 TOKEN = "token" HTTP_PATH = "sql/protocolv1/o/1234567890123456/0123-456789-abcd123" SCHEMA = "test_schema" @@ -112,38 +111,36 @@ def mock_timer(): yield mock_timer -def make_mock_connection(): - return Connection( - conn_id=DEFAULT_CONN_ID, - conn_type="databricks", - host=HOST, - port=PORT, - login="token", - password=TOKEN, - ) - - -def test_sqlachemy_url_property(mock_get_conn): - mock_get_conn.return_value = make_mock_connection() +def test_sqlachemy_url_property(): hook = DatabricksSqlHook( databricks_conn_id=DEFAULT_CONN_ID, http_path=HTTP_PATH, catalog=CATALOG, schema=SCHEMA ) url = hook.sqlalchemy_url.render_as_string(hide_password=False) expected_url = ( - f"databricks://token:{TOKEN}@{HOST}:{PORT}?" + f"databricks://token:{TOKEN}@{HOST}?" f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}" ) assert url == expected_url -def test_get_uri(mock_get_conn): - mock_get_conn.return_value = make_mock_connection() +def test_get_sqlalchemy_engine(): + hook = DatabricksSqlHook( + databricks_conn_id=DEFAULT_CONN_ID, http_path=HTTP_PATH, catalog=CATALOG, schema=SCHEMA + ) + hook.get_sqlalchemy_engine() + assert hook.sqlalchemy_url.render_as_string(hide_password=False) == ( + f"databricks://token:{TOKEN}@{HOST}?" + f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}" + ) + + +def test_get_uri(): hook = DatabricksSqlHook( databricks_conn_id=DEFAULT_CONN_ID, http_path=HTTP_PATH, catalog=CATALOG, schema=SCHEMA ) uri = hook.get_uri() expected_uri = ( - f"databricks://token:{TOKEN}@{HOST}:{PORT}?" + f"databricks://token:{TOKEN}@{HOST}?" f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}" ) assert uri == expected_uri From 09209850eeee9687a83fb069c782dd942d04a911 Mon Sep 17 00:00:00 2001 From: Brian Li Date: Thu, 14 Aug 2025 01:34:36 -0700 Subject: [PATCH 2/2] fix: assert the correct url --- .../tests/unit/databricks/hooks/test_databricks_sql.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py index acb961f3add3b..38fd5e26bdd4e 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py @@ -127,8 +127,8 @@ def test_get_sqlalchemy_engine(): hook = DatabricksSqlHook( databricks_conn_id=DEFAULT_CONN_ID, http_path=HTTP_PATH, catalog=CATALOG, schema=SCHEMA ) - hook.get_sqlalchemy_engine() - assert hook.sqlalchemy_url.render_as_string(hide_password=False) == ( + engine = hook.get_sqlalchemy_engine() + assert engine.url.render_as_string(hide_password=False) == ( f"databricks://token:{TOKEN}@{HOST}?" f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}" )