fix: remove whitespace from http_path for databricks (#22671)

This commit is contained in:
Elizabeth Thompson 2023-01-13 14:50:18 -08:00 committed by GitHub
parent 85da86dc81
commit cd1f6d469b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 27 additions and 0 deletions

View File

@ -162,6 +162,7 @@ class DatabricksNativeEngineSpec(DatabricksODBCEngineSpec, BasicParametersMixin)
def get_extra_params(database: "Database") -> Dict[str, Any]:
"""
Add a user agent to be used in the requests.
Trim whitespace from connect_args to avoid databricks driver errors
"""
extra: Dict[str, Any] = BaseEngineSpec.get_extra_params(database)
engine_params: Dict[str, Any] = extra.setdefault("engine_params", {})
@ -170,6 +171,10 @@ class DatabricksNativeEngineSpec(DatabricksODBCEngineSpec, BasicParametersMixin)
connect_args.setdefault("http_headers", [("User-Agent", USER_AGENT)])
connect_args.setdefault("_user_agent_entry", USER_AGENT)
# trim whitespace from http_path to avoid databricks errors on connecting
if http_path := connect_args.get("http_path"):
connect_args["http_path"] = http_path.strip()
return extra
@classmethod

View File

@ -175,3 +175,25 @@ def test_get_extra_params(mocker: MockerFixture) -> None:
}
}
}
# it should also remove whitespace from http_path
database.extra = json.dumps(
{
"engine_params": {
"connect_args": {
"http_headers": [("User-Agent", "Custom user agent")],
"_user_agent_entry": "Custom user agent",
"http_path": "/some_path_here_with_whitespace ",
}
}
}
)
assert DatabricksNativeEngineSpec.get_extra_params(database) == {
"engine_params": {
"connect_args": {
"http_headers": [["User-Agent", "Custom user agent"]],
"_user_agent_entry": "Custom user agent",
"http_path": "/some_path_here_with_whitespace",
}
}
}