From e774b4b3e2a5d43acf5c9412011329d7274a445c Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Sat, 11 Jul 2026 06:46:18 +0000 Subject: [PATCH 1/2] Fix DatabricksSqlHook sqlalchemy_url missing http_path from connection extra get_conn correctly resolved http_path from explicit arg, sql_endpoint_name API, and extra_dejson, but sqlalchemy_url only used explicit arg. This caused get_uri() to return URL without http_path when connection defined via UI extra, breaking SQLAlchemy usage. Centralize resolution in _resolve_http_path helper, allow extra lookup in property without triggering API call. Add tests for extra path. Fixes: #69031 --- .../databricks/hooks/databricks_sql.py | 31 ++++++++++++++----- 1 file changed, 24 insertions(+), 7 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 ce2bc28b45966..068bcfee7308c 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py @@ -184,15 +184,30 @@ def _get_sql_endpoint_by_name(self, endpoint_name) -> dict[str, Any]: else: return endpoint + def _resolve_http_path(self, allow_endpoint_lookup: bool = True) -> str | None: + """ + Resolve http_path from explicit arg, connection extra, or endpoint name. + + :param allow_endpoint_lookup: If True, may call API to resolve sql_endpoint_name. + Set False for offline-safe paths like sqlalchemy_url. + :return: resolved http_path or None if not found. + """ + if self._http_path: + return self._http_path + if "http_path" in self.databricks_conn.extra_dejson: + self._http_path = self.databricks_conn.extra_dejson["http_path"] + return self._http_path + if allow_endpoint_lookup and self._sql_endpoint_name: + endpoint = self._get_sql_endpoint_by_name(self._sql_endpoint_name) + self._http_path = endpoint["odbc_params"]["path"] + return self._http_path + return self._http_path + def get_conn(self) -> AirflowConnection: """Return a Databricks SQL connection object.""" if not self._http_path: - if self._sql_endpoint_name: - endpoint = self._get_sql_endpoint_by_name(self._sql_endpoint_name) - self._http_path = endpoint["odbc_params"]["path"] - elif "http_path" in self.databricks_conn.extra_dejson: - self._http_path = self.databricks_conn.extra_dejson["http_path"] - else: + self._http_path = self._resolve_http_path(allow_endpoint_lookup=True) + if not self._http_path: raise AirflowException( "http_path should be provided either explicitly, " "or in extra parameter of Databricks connection, " @@ -254,8 +269,10 @@ def sqlalchemy_url(self) -> URL: "Install it with: pip install 'apache-airflow-providers-databricks[sqlalchemy]'" ) + http_path = self._resolve_http_path(allow_endpoint_lookup=False) + url_query = { - "http_path": self._http_path, + "http_path": http_path, "catalog": self.catalog, "schema": self.schema, } From fde316dc6092446b0ef867b143207b52043de6ff Mon Sep 17 00:00:00 2001 From: probe Date: Sun, 12 Jul 2026 19:52:17 +0000 Subject: [PATCH 2/2] Preserve Databricks SQL endpoint precedence The SQLAlchemy URL path should be available from connection extras without changing the existing get_conn precedence for named SQL endpoints, so the resolution path needs explicit coverage for both behaviors. --- .../databricks/hooks/databricks_sql.py | 10 ++--- .../databricks/hooks/test_databricks_sql.py | 38 ++++++++++++++++++- 2 files changed, 40 insertions(+), 8 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 068bcfee7308c..d6bec663f7ff6 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py @@ -186,7 +186,7 @@ def _get_sql_endpoint_by_name(self, endpoint_name) -> dict[str, Any]: def _resolve_http_path(self, allow_endpoint_lookup: bool = True) -> str | None: """ - Resolve http_path from explicit arg, connection extra, or endpoint name. + Resolve http_path from explicit arg, endpoint name, or connection extra. :param allow_endpoint_lookup: If True, may call API to resolve sql_endpoint_name. Set False for offline-safe paths like sqlalchemy_url. @@ -194,14 +194,10 @@ def _resolve_http_path(self, allow_endpoint_lookup: bool = True) -> str | None: """ if self._http_path: return self._http_path - if "http_path" in self.databricks_conn.extra_dejson: - self._http_path = self.databricks_conn.extra_dejson["http_path"] - return self._http_path if allow_endpoint_lookup and self._sql_endpoint_name: endpoint = self._get_sql_endpoint_by_name(self._sql_endpoint_name) - self._http_path = endpoint["odbc_params"]["path"] - return self._http_path - return self._http_path + return endpoint["odbc_params"]["path"] + return self.databricks_conn.extra_dejson.get("http_path") def get_conn(self) -> AirflowConnection: """Return a Databricks SQL connection object.""" 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 ec0718025ead4..4d736014efd36 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py @@ -45,6 +45,7 @@ HOST_WITH_SCHEME = "https://xx.cloud.databricks.com" TOKEN = "token" HTTP_PATH = "sql/protocolv1/o/1234567890123456/0123-456789-abcd123" +ENDPOINT_HTTP_PATH = "/sql/1.0/endpoints/1264e5078741679a" SCHEMA = "test_schema" CATALOG = "test_catalog" @@ -94,7 +95,7 @@ def mock_get_requests(): "name": "Test", "odbc_params": { "hostname": "xx.cloud.databricks.com", - "path": "/sql/1.0/endpoints/1264e5078741679a", + "path": ENDPOINT_HTTP_PATH, }, } ] @@ -150,6 +151,41 @@ def test_get_uri(): assert uri == expected_uri +@mock.patch("airflow.providers.databricks.hooks.databricks_sql.sql.connect") +def test_get_conn_prefers_sql_endpoint_name_over_connection_extra_http_path(mock_connect, mock_get_requests): + hook = DatabricksSqlHook(databricks_conn_id=DEFAULT_CONN_ID, sql_endpoint_name="Test") + hook.databricks_conn = Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + password=TOKEN, + extra={"http_path": "sql/protocolv1/o/extra/path"}, + ) + + hook.get_conn() + + mock_connect.assert_called_once() + assert mock_connect.call_args.args[1] == ENDPOINT_HTTP_PATH + + +def test_sqlalchemy_url_uses_connection_extra_http_path_without_endpoint_lookup(): + hook = DatabricksSqlHook(databricks_conn_id=DEFAULT_CONN_ID, sql_endpoint_name="Test") + hook.databricks_conn = Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + password=TOKEN, + extra={"http_path": "sql/protocolv1/o/extra/path"}, + ) + + with mock.patch.object(DatabricksSqlHook, "_get_sql_endpoint_by_name", autospec=True) as mock_endpoint: + url = hook.sqlalchemy_url.render_as_string(hide_password=False) + + assert "http_path=sql%2Fprotocolv1%2Fo%2Fextra%2Fpath" in url + assert hook._http_path is None + mock_endpoint.assert_not_called() + + def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: return [(field,) for field in fields]