Skip to content

Commit

Permalink
Update databend engine spec
Browse files Browse the repository at this point in the history
 * For use with databend dialect version > 0.3.2
  • Loading branch information
rad-pat committed May 17, 2024
1 parent 81cce06 commit 8f032a7
Showing 1 changed file with 11 additions and 5 deletions.
16 changes: 11 additions & 5 deletions superset/db_engine_specs/databend.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,10 +207,11 @@ class DatabendConnectEngineSpec(DatabendEngineSpec, BasicParametersMixin):
_function_names: list[str] = []

sqlalchemy_uri_placeholder = (
"databend://user:password@host[:port][/dbname][?secure=value&=value...]"
"databend://user:password@host[:port][/dbname][?sslmode=value&=value...]"
)
parameters_schema = DatabendParametersSchema()
encryption_parameters = {"secure": "true"}
encryption_parameters = {"sslmode": "require"}
encryption_disable_parameters = {"sslmode": "disable"}

@classmethod
def get_dbapi_exception_mapping(cls) -> dict[type[Exception], type[Exception]]:
Expand Down Expand Up @@ -248,10 +249,12 @@ def build_sqlalchemy_uri(
cls, parameters: BasicParametersType, *_args: dict[str, str] | None
) -> str:
url_params = parameters.copy()
query = parameters.get("query", {}).copy()
if url_params.get("encryption"):
query = parameters.get("query", {}).copy()
query.update(cls.encryption_parameters)
url_params["query"] = query
else:
query.update(cls.encryption_disable_parameters)
url_params["query"] = query
if not url_params.get("database"):
url_params["database"] = "__default__"
url_params.pop("encryption", None)
Expand All @@ -263,7 +266,10 @@ def get_parameters_from_uri(
) -> BasicParametersType:
url = make_url_safe(uri)
query = url.query
if "secure" in query:
if "sslmode" in query:
encryption = url.query.get("sslmode") == "require"
query.pop("sslmode")
elif "secure" in query:
encryption = url.query.get("secure") == "true"
query.pop("secure")
else:
Expand Down

0 comments on commit 8f032a7

Please sign in to comment.