QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-05-29 11:45:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: Create BigQuery Parameters for DatabaseModal (#14721)

Browse Source
This commit is contained in:
Hugh A. Miles II
2021-05-23 12:45:48 -04:00
committed by GitHub
parent 9e13476714
commit 6d33432b58
8 changed files with 132 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
superset/databases/api.py
View File

@@ -66,7 +66,6 @@ from superset.databases.schemas import (
)
from superset.databases.utils import get_table_metadata
from superset.db_engine_specs import get_available_engine_specs
from superset.db_engine_specs.base import BasicParametersMixin
from superset.exceptions import InvalidPayloadFormatError, InvalidPayloadSchemaError
from superset.extensions import security_manager
from superset.models.core import Database
@@ -909,11 +908,15 @@ class DatabaseRestApi(BaseSupersetModelRestApi):
"preferred": engine_spec.engine in preferred_databases,
}
if issubclass(engine_spec, BasicParametersMixin):
payload["parameters"] = engine_spec.parameters_json_schema()
if hasattr(engine_spec, "parameters_json_schema") and hasattr(
engine_spec, "sqlalchemy_uri_placeholder"
):
payload[
"parameters"
] = engine_spec.parameters_json_schema() # type: ignore
payload[
"sqlalchemy_uri_placeholder"
] = engine_spec.sqlalchemy_uri_placeholder
] = engine_spec.sqlalchemy_uri_placeholder # type: ignore
available_databases.append(payload)

12
superset/databases/commands/validate.py
View File

@@ -14,6 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
import json
from contextlib import closing
from typing import Any, Dict, Optional
@@ -81,9 +82,16 @@ class ValidateDatabaseParametersCommand(BaseCommand):
if errors:
raise InvalidParametersError(errors)
serialized_encrypted_extra = self._properties.get("encrypted_extra", "{}")
try:
encrypted_extra = json.loads(serialized_encrypted_extra)
except json.decoder.JSONDecodeError:
encrypted_extra = {}
# try to connect
sqlalchemy_uri = engine_spec.build_sqlalchemy_uri(
self._properties["parameters"] # type: ignore
self._properties["parameters"], # type: ignore
encrypted_extra,
)
if self._model and sqlalchemy_uri == self._model.safe_sqlalchemy_uri():
sqlalchemy_uri = self._model.sqlalchemy_uri_decrypted
@@ -91,7 +99,7 @@ class ValidateDatabaseParametersCommand(BaseCommand):
server_cert=self._properties.get("server_cert", ""),
extra=self._properties.get("extra", "{}"),
impersonate_user=self._properties.get("impersonate_user", False),
encrypted_extra=self._properties.get("encrypted_extra", "{}"),
encrypted_extra=serialized_encrypted_extra,
)
database.set_sqlalchemy_uri(sqlalchemy_uri)
database.db_engine_spec.mutate_db_for_connection_test(database)

22
superset/databases/schemas.py
View File

@@ -246,6 +246,12 @@ class DatabaseParametersSchemaMixin:
the constructed SQLAlchemy URI to be passed.
"""
parameters = data.pop("parameters", None)
serialized_encrypted_extra = data.get("encrypted_extra", "{}")
try:
encrypted_extra = json.loads(serialized_encrypted_extra)
except json.decoder.JSONDecodeError:
encrypted_extra = {}
if parameters:
if "engine" not in parameters:
raise ValidationError(
@@ -264,12 +270,14 @@ class DatabaseParametersSchemaMixin:
[_('Engine "%(engine)s" is not a valid engine.', engine=engine,)]
)
engine_spec = engine_specs[engine]
if hasattr(engine_spec, "build_sqlalchemy_uri"):
data[
"sqlalchemy_uri"
] = engine_spec.build_sqlalchemy_uri( # type: ignore
parameters
parameters, encrypted_extra
)
return data
@@ -546,3 +554,15 @@ class ImportV1DatabaseSchema(Schema):
password = make_url(uri).password
if password == PASSWORD_MASK and data.get("password") is None:
raise ValidationError("Must provide a password for the database")
class EncryptedField(fields.String):
pass
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

6
superset/db_engine_specs/base.py
View File

@@ -1348,7 +1348,11 @@ class BasicParametersMixin:
encryption_parameters: Dict[str, str] = {}
@classmethod
def build_sqlalchemy_uri(cls, parameters: BasicParametersType) -> str:
def build_sqlalchemy_uri(
cls,
parameters: BasicParametersType,
encryted_extra: Optional[Dict[str, str]] = None,
) -> str:
query = parameters.get("query", {})
if parameters.get("encryption"):
if not cls.encryption_parameters:

66
superset/db_engine_specs/bigquery.py
View File

@@ -19,12 +19,18 @@ from datetime import datetime
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 Schema
from sqlalchemy import literal_column
from sqlalchemy.sql.expression import ColumnClause
from typing_extensions import TypedDict
from superset.databases.schemas import encrypted_field_properties, EncryptedField
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 +44,18 @@ CONNECTION_DATABASE_PERMISSIONS_REGEX = re.compile(
+ "permission in project (?P<project>.+?)"
)
ma_plugin = MarshmallowPlugin()
class BigQueryParametersSchema(Schema):
credentials_info = EncryptedField(
description="Contents of BigQuery JSON credentials.",
)
class BigQueryParametersType(TypedDict):
credentials_info: Dict[str, Any]
class BigQueryEngineSpec(BaseEngineSpec):
"""Engine spec for Google's BigQuery
@@ -48,6 +66,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 +304,47 @@ class BigQueryEngineSpec(BaseEngineSpec):
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.",
)
@classmethod
def get_parameters_from_uri(
cls, _: str, encrypted_extra: Optional[Dict[str, str]] = None
) -> Any:
# BigQuery doesn't have parameters
if encrypted_extra:
return encrypted_extra
raise SupersetGenericDBErrorException(
message="Big Query encrypted_extra is not available.",
)
@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__]

1
superset/models/core.py
View File

@@ -241,6 +241,7 @@ class Database(
def parameters(self) -> Dict[str, Any]:
# Build parameters if db_engine_spec is a subclass of BasicParametersMixin
parameters = {"engine": self.backend}
if hasattr(self.db_engine_spec, "parameters_schema") and hasattr(
self.db_engine_spec, "get_parameters_from_uri"
):

18
tests/databases/api_tests.py
View File

@@ -36,6 +36,7 @@ from superset import db, security_manager
from superset.connectors.sqla.models import SqlaTable
from superset.db_engine_specs.mysql import MySQLEngineSpec
from superset.db_engine_specs.postgres import PostgresEngineSpec
from superset.db_engine_specs.bigquery import BigQueryEngineSpec
from superset.db_engine_specs.hana import HanaEngineSpec
from superset.errors import SupersetError
from superset.models.core import Database, ConfigurationMethod
@@ -1373,6 +1374,7 @@ class TestDatabaseApi(SupersetTestCase):
app.config = {"PREFERRED_DATABASES": ["postgresql"]}
get_available_engine_specs.return_value = [
PostgresEngineSpec,
BigQueryEngineSpec,
HanaEngineSpec,
]
@@ -1428,6 +1430,22 @@ class TestDatabaseApi(SupersetTestCase):
"preferred": True,
"sqlalchemy_uri_placeholder": "postgresql+psycopg2://user:password@host:port/dbname[?key=value&key=value...]",
},
{
"engine": "bigquery",
"name": "Google BigQuery",
"parameters": {
"properties": {
"credentials_info": {
"description": "Contents of BigQuery JSON credentials.",
"type": "string",
"x-encrypted-extra": True,
}
},
"type": "object",
},
"preferred": False,
"sqlalchemy_uri_placeholder": "bigquery://{project_id}",
},
{"engine": "hana", "name": "SAP HANA", "preferred": False},
]
}

5
tests/db_engine_specs/postgres_tests.py
View File

@@ -432,7 +432,10 @@ def test_base_parameters_mixin():
"query": {"foo": "bar"},
"encryption": True,
}
sqlalchemy_uri = PostgresEngineSpec.build_sqlalchemy_uri(parameters)
encrypted_extra = None
sqlalchemy_uri = PostgresEngineSpec.build_sqlalchemy_uri(
parameters, encrypted_extra
)
assert sqlalchemy_uri == (
"postgresql+psycopg2://username:password@localhost:5432/dbname?"
"foo=bar&sslmode=verify-ca"