This is an automated email from the ASF dual-hosted git repository.
mobuchowski pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 7f1fe8167a7 Fix Databricks sqlalchemy URL construction (#54478)
7f1fe8167a7 is described below
commit 7f1fe8167a78d9044035f2616b4f21390b155c4b
Author: Chang-Yen (Brian) Li <[email protected]>
AuthorDate: Thu Aug 14 17:15:30 2025 +0800
Fix Databricks sqlalchemy URL construction (#54478)
* fix: use _get_token in url construction
* fix: assert the correct url
---
.../providers/databricks/hooks/databricks_sql.py | 6 ++--
.../unit/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 c689324f97d..d86d0453fa5 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 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
: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 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
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 9e2ee58d600..38fd5e26bdd 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 @@ TASK_ID = "databricks-sql-operator"
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
+ )
+ 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}"
+ )
+
+
+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
