-
Notifications
You must be signed in to change notification settings - Fork 14.1k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat: Create BigQuery Parameters for DatabaseModal #14721
Changes from 13 commits
139cd49
f252d72
b77a3e5
bb24664
4a1683e
1fc08d4
b8640a1
8ca9b80
8adfe4c
7d2b7d9
9bddea7
cdc6580
0506752
401873d
77ce531
5a604b4
bb66ffb
dffb040
b4e9a02
d9d7379
eaecbb6
a5cd2b7
361e755
65822ed
42e7247
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -83,7 +83,8 @@ def run(self) -> None: | |
|
||
# try to connect | ||
sqlalchemy_uri = engine_spec.build_sqlalchemy_uri( | ||
self._properties["parameters"] # type: ignore | ||
self._properties["parameters"], # type: ignore | ||
self._properties.get("encrypted_extra", "{}"), | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We need to deserialize this: serialized_encrypted_extra = self._properties.get("encrypted_extra", "{}")
try:
encrypted_extra = json.loads(serialized_encrypted_extra)
except json.decoder.JSONDecodeError:
encrypted_extra = {} Then: sqlalchemy_uri = engine_spec.build_sqlalchemy_uri(
self._properties["parameters"],
encrypted_extra,
) You can then pass |
||
) | ||
if self._model and sqlalchemy_uri == self._model.safe_sqlalchemy_uri(): | ||
sqlalchemy_uri = self._model.sqlalchemy_uri_decrypted | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -246,6 +246,8 @@ def build_sqlalchemy_uri( | |
the constructed SQLAlchemy URI to be passed. | ||
""" | ||
parameters = data.pop("parameters", None) | ||
encrypted_extra = data.get("encrypted_extra", None) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Same here, we need to try to deserialize from JSON like above. |
||
|
||
if parameters: | ||
if "engine" not in parameters: | ||
raise ValidationError( | ||
|
@@ -275,7 +277,9 @@ def build_sqlalchemy_uri( | |
] | ||
) | ||
|
||
data["sqlalchemy_uri"] = engine_spec.build_sqlalchemy_uri(parameters) | ||
data["sqlalchemy_uri"] = engine_spec.build_sqlalchemy_uri( | ||
parameters, encrypted_extra | ||
) | ||
return data | ||
|
||
|
||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -19,12 +19,17 @@ | |
from typing import Any, Dict, List, Optional, Pattern, Tuple, TYPE_CHECKING | ||
|
||
import pandas as pd | ||
from apispec import APISpec | ||
from apispec.ext.marshmallow import MarshmallowPlugin | ||
from flask_babel import gettext as __ | ||
from marshmallow import fields, Schema | ||
from sqlalchemy import literal_column | ||
from sqlalchemy.sql.expression import ColumnClause | ||
from typing_extensions import TypedDict | ||
|
||
from superset.db_engine_specs.base import BaseEngineSpec | ||
from superset.errors import SupersetErrorType | ||
from superset.exceptions import SupersetGenericDBErrorException | ||
from superset.sql_parse import Table | ||
from superset.utils import core as utils | ||
from superset.utils.hashing import md5_sha_from_str | ||
|
@@ -38,6 +43,28 @@ | |
+ "permission in project (?P<project>.+?)" | ||
) | ||
|
||
ma_plugin = MarshmallowPlugin() | ||
|
||
|
||
class EncryptedField(fields.String): | ||
pass | ||
|
||
|
||
class BigQueryParametersSchema(Schema): | ||
credentials_info = EncryptedField(description="credentials.json file for BigQuery") | ||
hughhhh marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
|
||
def encrypted_field_properties(self, field: Any, **_) -> Dict[str, Any]: # type: ignore | ||
ret = {} | ||
if isinstance(field, EncryptedField): | ||
if self.openapi_version.major > 2: | ||
ret["x-encrypted-extra"] = True | ||
return ret | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Let's move this and |
||
|
||
|
||
class BigQueryParametersType(TypedDict): | ||
pass | ||
hughhhh marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
|
||
class BigQueryEngineSpec(BaseEngineSpec): | ||
"""Engine spec for Google's BigQuery | ||
|
@@ -48,6 +75,10 @@ class BigQueryEngineSpec(BaseEngineSpec): | |
engine_name = "Google BigQuery" | ||
max_column_name_length = 128 | ||
|
||
parameters_schema = BigQueryParametersSchema() | ||
drivername = engine | ||
sqlalchemy_uri_placeholder = "bigquery://{project_id}" | ||
|
||
# BigQuery doesn't maintain context when running multiple statements in the | ||
# same cursor, so we need to run all statements at once | ||
run_multiple_statements_as_one = True | ||
|
@@ -282,3 +313,40 @@ def df_to_sql( | |
to_gbq_kwargs[key] = to_sql_kwargs[key] | ||
|
||
pandas_gbq.to_gbq(df, **to_gbq_kwargs) | ||
|
||
@classmethod | ||
def build_sqlalchemy_uri( | ||
cls, _: BigQueryParametersType, encrypted_extra: Optional[Dict[str, str]] = None | ||
) -> str: | ||
if encrypted_extra: | ||
project_id = encrypted_extra.get("project_id") | ||
return f"{cls.drivername}://{project_id}" | ||
|
||
raise SupersetGenericDBErrorException( | ||
message="Big Query encrypted_extra is not available", | ||
hughhhh marked this conversation as resolved.
Show resolved
Hide resolved
|
||
) | ||
|
||
@classmethod | ||
def get_parameters_from_uri(cls, _: str) -> Any: | ||
# BigQuery doesn't have parameters | ||
return None | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Ideally:
Ie, this should be the opposite of We need to return |
||
|
||
@classmethod | ||
def parameters_json_schema(cls) -> Any: | ||
""" | ||
Return configuration parameters as OpenAPI. | ||
""" | ||
if not cls.parameters_schema: | ||
return None | ||
|
||
spec = APISpec( | ||
title="Database Parameters", | ||
version="1.0.0", | ||
openapi_version="3.0.0", | ||
plugins=[ma_plugin], | ||
) | ||
|
||
ma_plugin.init_spec(spec) | ||
ma_plugin.converter.add_attribute_function(encrypted_field_properties) | ||
spec.components.schema(cls.__name__, schema=cls.parameters_schema) | ||
return spec.to_dict()["components"]["schemas"][cls.__name__] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@betodealmeida let me know if this is good enough of a guard statement moving forward to make sure we don't call parameters_json_schema for engines we have implemented yet
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks great!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Though I think ideally we'd have all these methods in
BaseEngineSpec
and we'd check ifparameters_schema
is notNone
.