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..d6bec663f7ff6 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,26 @@ 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, 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. + :return: resolved http_path or None if not found. + """ + if self._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) + 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.""" 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 +265,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, } 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]