Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -184,20 +184,38 @@ 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.

When allow_endpoint_lookup=False (used by sqlalchemy_url etc.),
sql_endpoint_name is ignored and we fall back to http_path from the
connection extra. This keeps the property fully offline. If both
sql_endpoint_name and an extra http_path are set, sqlalchemy_url may
therefore point at a different warehouse than the one get_conn() will
actually connect to. This asymmetry is intentional.
"""
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")
Comment thread
Vamsi-klu marked this conversation as resolved.

def get_conn(self) -> AirflowConnection:
"""Return a Databricks SQL connection object."""
if not self._http_path:
Comment thread
Vamsi-klu marked this conversation as resolved.
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:
raise AirflowException(
"http_path should be provided either explicitly, "
"or in extra parameter of Databricks connection, "
"or sql_endpoint_name should be specified"
)
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, "
"or sql_endpoint_name should be specified"
)

prev_token = self._token
new_token = self._get_token(raise_error=True)
Expand Down Expand Up @@ -254,8 +272,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,
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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,
},
}
]
Expand Down Expand Up @@ -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]

Expand Down