Skip to content
Draft
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,26 @@ def _get_sql_endpoint_by_name(self, endpoint_name) -> dict[str, Any]:
else:
return endpoint

def _resolve_http_path(self) -> str:
"""Resolve http_path from explicit parameter, endpoint name, or connection extra."""
if self._http_path:
return 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"]
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
raise ValueError(
"http_path should be provided either explicitly, "
"or in extra parameter of Databricks connection, "
"or sql_endpoint_name should be specified"
)

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:
raise AirflowException(
"http_path should be provided either explicitly, "
"or in extra parameter of Databricks connection, "
"or sql_endpoint_name should be specified"
)
self._resolve_http_path()

prev_token = self._token
new_token = self._get_token(raise_error=True)
Expand Down Expand Up @@ -255,7 +261,7 @@ def sqlalchemy_url(self) -> URL:
)

url_query = {
"http_path": self._http_path,
"http_path": self._resolve_http_path(),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This line makes sqlalchemy_url (a property) perform a REST API call when the hook is configured via sql_endpoint_name — see the review body for the failing chain through dialect_name/insert_rows. Please resolve only from local sources (explicit param, cached value, connection extra) here.


Drafted-by: Claude Code (Fable 5); reviewed by @moomindani before posting

"catalog": self.catalog,
"schema": self.schema,
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,49 @@ def test_get_uri():
assert uri == expected_uri


CONN_ID_WITH_HTTP_PATH_EXTRA = "databricks_with_http_path_extra"


@pytest.fixture
def create_connection_with_http_path_extra(create_connection_without_db):
create_connection_without_db(
Connection(
conn_id=CONN_ID_WITH_HTTP_PATH_EXTRA,
conn_type="databricks",
host=HOST,
login=None,
password=TOKEN,
extra={"http_path": HTTP_PATH},
)
)


def test_sqlalchemy_url_with_http_path_from_connection_extra(create_connection_with_http_path_extra):
hook = DatabricksSqlHook(databricks_conn_id=CONN_ID_WITH_HTTP_PATH_EXTRA, catalog=CATALOG, schema=SCHEMA)
url = hook.sqlalchemy_url.render_as_string(hide_password=False)
expected_url = (
f"databricks://token:{TOKEN}@{HOST}?"
f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}"
)
assert url == expected_url


def test_get_uri_with_http_path_from_connection_extra(create_connection_with_http_path_extra):
hook = DatabricksSqlHook(databricks_conn_id=CONN_ID_WITH_HTTP_PATH_EXTRA, catalog=CATALOG, schema=SCHEMA)
uri = hook.get_uri()
expected_uri = (
f"databricks://token:{TOKEN}@{HOST}?"
f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}"
)
assert uri == expected_uri


def test_resolve_http_path_raises_when_not_provided():
hook = DatabricksSqlHook(databricks_conn_id=DEFAULT_CONN_ID)
with pytest.raises(ValueError, match="http_path should be provided"):
hook._resolve_http_path()


def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]:
return [(field,) for field in fields]

Expand Down
Loading