feat(dbview): Add token request button to DuckDB and MotherDuck database modal (#27908)
Co-authored-by: Beto Dealmeida <roberto@dealmeida.net>
This commit is contained in:
Guen Prawiroatmodjo
GitHub
parent
caad29b5b3
commit
08aaebbf7c
|
|
@ -21,6 +21,7 @@ import { Input, Tooltip } from 'antd';
|
|||
import { styled, css, SupersetTheme, t } from '@superset-ui/core';
|
||||
import InfoTooltip from 'src/components/InfoTooltip';
|
||||
import Icons from 'src/components/Icons';
|
||||
import Button from 'src/components/Button';
|
||||
import errorIcon from 'src/assets/images/icons/error.svg';
|
||||
import FormItem from './FormItem';
|
||||
import FormLabel from './FormLabel';
|
||||
|
|
@ -109,6 +110,8 @@ const LabeledErrorBoundInput = ({
|
|||
id,
|
||||
className,
|
||||
visibilityToggle,
|
||||
get_url,
|
||||
description,
|
||||
...props
|
||||
}: LabeledErrorBoundInputProps) => (
|
||||
<StyledFormGroup className={className}>
|
||||
|
|
@ -149,6 +152,21 @@ const LabeledErrorBoundInput = ({
|
|||
) : (
|
||||
<StyledInput {...props} {...validationMethods} />
|
||||
)}
|
||||
{get_url && description ? (
|
||||
<Button
|
||||
type="link"
|
||||
htmlType="button"
|
||||
buttonStyle="default"
|
||||
onClick={() => {
|
||||
window.open(get_url);
|
||||
return true;
|
||||
}}
|
||||
>
|
||||
Get {description}
|
||||
</Button>
|
||||
) : (
|
||||
<br />
|
||||
)}
|
||||
</FormItem>
|
||||
</StyledFormGroup>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -163,6 +163,8 @@ export const accessTokenField = ({
|
|||
validationErrors,
|
||||
db,
|
||||
isEditMode,
|
||||
default_value,
|
||||
description,
|
||||
}: FieldPropTypes) => (
|
||||
<ValidatedInput
|
||||
id="access_token"
|
||||
|
|
@ -172,7 +174,13 @@ export const accessTokenField = ({
|
|||
value={db?.parameters?.access_token}
|
||||
validationMethods={{ onBlur: getValidation }}
|
||||
errorMessage={validationErrors?.access_token}
|
||||
placeholder={t('e.g. ********')}
|
||||
placeholder={t('Paste your access token here')}
|
||||
get_url={
|
||||
typeof default_value === 'string' && default_value.includes('https://')
|
||||
? default_value
|
||||
: null
|
||||
}
|
||||
description={description}
|
||||
label={t('Access token')}
|
||||
onChange={changeMethods.onParametersChange}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -169,6 +169,8 @@ const DatabaseConnectionForm = ({
|
|||
db,
|
||||
key: field,
|
||||
field,
|
||||
default_value: parameters.properties[field]?.default,
|
||||
description: parameters.properties[field]?.description,
|
||||
isEditMode,
|
||||
sslForced,
|
||||
editNewDb,
|
||||
|
|
|
|||
|
|
@ -149,8 +149,7 @@ const ModalHeader = ({
|
|||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
>
|
||||
{t('connecting to %(dbModelName)s.', { dbModelName: dbModel.name })}
|
||||
.
|
||||
{t('connecting to %(dbModelName)s', { dbModelName: dbModel.name })}.
|
||||
</a>
|
||||
</p>
|
||||
</StyledFormHeader>
|
||||
|
|
|
|||
|
|
@ -291,6 +291,8 @@ export interface FieldPropTypes {
|
|||
db?: DatabaseObject;
|
||||
dbModel?: DatabaseForm;
|
||||
field: string;
|
||||
default_value?: any;
|
||||
description?: string;
|
||||
isEditMode?: boolean;
|
||||
sslForced?: boolean;
|
||||
defaultDBName?: string;
|
||||
|
|
|
|||
|
|
@ -19,16 +19,21 @@ from __future__ import annotations
|
|||
import re
|
||||
from datetime import datetime
|
||||
from re import Pattern
|
||||
from typing import Any, TYPE_CHECKING
|
||||
from typing import Any, TYPE_CHECKING, TypedDict
|
||||
|
||||
from flask_babel import gettext as __
|
||||
from apispec import APISpec
|
||||
from apispec.ext.marshmallow import MarshmallowPlugin
|
||||
from flask_babel import gettext as __, lazy_gettext as _
|
||||
from marshmallow import fields, Schema
|
||||
from sqlalchemy import types
|
||||
from sqlalchemy.engine.reflection import Inspector
|
||||
from sqlalchemy.engine.url import URL
|
||||
|
||||
from superset.config import VERSION_STRING
|
||||
from superset.constants import TimeGrain, USER_AGENT
|
||||
from superset.databases.utils import make_url_safe
|
||||
from superset.db_engine_specs.base import BaseEngineSpec
|
||||
from superset.errors import SupersetErrorType
|
||||
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# prevent circular imports
|
||||
|
|
@ -36,11 +41,158 @@ if TYPE_CHECKING:
|
|||
|
||||
|
||||
COLUMN_DOES_NOT_EXIST_REGEX = re.compile("no such column: (?P<column_name>.+)")
|
||||
DEFAULT_ACCESS_TOKEN_URL = (
|
||||
"https://app.motherduck.com/token-request?appName=Superset&close=y"
|
||||
)
|
||||
|
||||
|
||||
class DuckDBEngineSpec(BaseEngineSpec):
|
||||
# schema for adding a database by providing parameters instead of the
|
||||
# full SQLAlchemy URI
|
||||
class DuckDBParametersSchema(Schema):
|
||||
access_token = fields.String(
|
||||
allow_none=True,
|
||||
metadata={"description": __("MotherDuck token")},
|
||||
load_default=DEFAULT_ACCESS_TOKEN_URL,
|
||||
)
|
||||
database = fields.String(
|
||||
required=False, metadata={"description": __("Database name")}
|
||||
)
|
||||
query = fields.Dict(
|
||||
keys=fields.Str(),
|
||||
values=fields.Raw(),
|
||||
metadata={"description": __("Additional parameters")},
|
||||
)
|
||||
|
||||
|
||||
class DuckDBParametersType(TypedDict, total=False):
|
||||
access_token: str | None
|
||||
database: str
|
||||
query: dict[str, Any]
|
||||
|
||||
|
||||
class DuckDBPropertiesType(TypedDict):
|
||||
parameters: DuckDBParametersType
|
||||
|
||||
|
||||
class DuckDBParametersMixin:
|
||||
"""
|
||||
Mixin for configuring DB engine specs via a dictionary.
|
||||
|
||||
With this mixin the SQLAlchemy engine can be configured through
|
||||
individual parameters, instead of the full SQLAlchemy URI. This
|
||||
mixin is for DuckDB:
|
||||
|
||||
duckdb:///file_path[?key=value&key=value...]
|
||||
duckdb:///md:database[?key=value&key=value...]
|
||||
|
||||
"""
|
||||
|
||||
engine = "duckdb"
|
||||
|
||||
# schema describing the parameters used to configure the DB
|
||||
parameters_schema = DuckDBParametersSchema()
|
||||
|
||||
# recommended driver name for the DB engine spec
|
||||
default_driver = ""
|
||||
|
||||
# query parameter to enable encryption in the database connection
|
||||
# for Postgres this would be `{"sslmode": "verify-ca"}`, eg.
|
||||
encryption_parameters: dict[str, str] = {}
|
||||
|
||||
@staticmethod
|
||||
def _is_motherduck(database: str) -> bool:
|
||||
return "md:" in database
|
||||
|
||||
@classmethod
|
||||
def build_sqlalchemy_uri( # pylint: disable=unused-argument
|
||||
cls,
|
||||
parameters: DuckDBParametersType,
|
||||
encrypted_extra: dict[str, str] | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build SQLAlchemy URI for connecting to a DuckDB database.
|
||||
If an access token is specified, return a URI to connect to a MotherDuck database.
|
||||
"""
|
||||
if parameters is None:
|
||||
parameters = {}
|
||||
query = parameters.get("query", {})
|
||||
database = parameters.get("database", ":memory:")
|
||||
token = parameters.get("access_token")
|
||||
|
||||
if cls._is_motherduck(database) or (
|
||||
token and token != DEFAULT_ACCESS_TOKEN_URL
|
||||
):
|
||||
return MotherDuckEngineSpec.build_sqlalchemy_uri(parameters)
|
||||
|
||||
return str(URL(drivername=cls.engine, database=database, query=query))
|
||||
|
||||
@classmethod
|
||||
def get_parameters_from_uri( # pylint: disable=unused-argument
|
||||
cls, uri: str, encrypted_extra: dict[str, Any] | None = None
|
||||
) -> DuckDBParametersType:
|
||||
url = make_url_safe(uri)
|
||||
query = {
|
||||
key: value
|
||||
for (key, value) in url.query.items()
|
||||
if (key, value) not in cls.encryption_parameters.items()
|
||||
}
|
||||
access_token = query.pop("motherduck_token", "")
|
||||
return {
|
||||
"access_token": access_token,
|
||||
"database": url.database,
|
||||
"query": query,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def validate_parameters(
|
||||
cls, properties: DuckDBPropertiesType
|
||||
) -> list[SupersetError]:
|
||||
"""
|
||||
Validates any number of parameters, for progressive validation.
|
||||
"""
|
||||
errors: list[SupersetError] = []
|
||||
|
||||
parameters = properties.get("parameters", {})
|
||||
if cls._is_motherduck(parameters.get("database", "")):
|
||||
required = {"access_token"}
|
||||
else:
|
||||
required = set()
|
||||
present = {key for key in parameters if parameters.get(key, ())}
|
||||
|
||||
if missing := sorted(required - present):
|
||||
errors.append(
|
||||
SupersetError(
|
||||
message=f'One or more parameters are missing: {", ".join(missing)}',
|
||||
error_type=SupersetErrorType.CONNECTION_MISSING_PARAMETERS_ERROR,
|
||||
level=ErrorLevel.WARNING,
|
||||
extra={"missing": missing},
|
||||
),
|
||||
)
|
||||
|
||||
return errors
|
||||
|
||||
@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.2",
|
||||
plugins=[MarshmallowPlugin()],
|
||||
)
|
||||
spec.components.schema(cls.__name__, schema=cls.parameters_schema)
|
||||
return spec.to_dict()["components"]["schemas"][cls.__name__]
|
||||
|
||||
|
||||
class DuckDBEngineSpec(DuckDBParametersMixin, BaseEngineSpec):
|
||||
engine = "duckdb"
|
||||
engine_name = "DuckDB"
|
||||
default_driver = "duckdb_engine"
|
||||
|
||||
sqlalchemy_uri_placeholder = "duckdb:////path/to/duck.db"
|
||||
|
||||
|
|
@ -103,9 +255,41 @@ class DuckDBEngineSpec(BaseEngineSpec):
|
|||
|
||||
|
||||
class MotherDuckEngineSpec(DuckDBEngineSpec):
|
||||
engine = "duckdb"
|
||||
engine = "motherduck"
|
||||
engine_name = "MotherDuck"
|
||||
engine_aliases: set[str] = {"duckdb"}
|
||||
|
||||
sqlalchemy_uri_placeholder = (
|
||||
"duckdb:///md:{database_name}?motherduck_token={SERVICE_TOKEN}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_motherduck(database: str) -> bool:
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def build_sqlalchemy_uri(
|
||||
cls,
|
||||
parameters: DuckDBParametersType,
|
||||
encrypted_extra: dict[str, str] | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build SQLAlchemy URI for connecting to a MotherDuck database
|
||||
"""
|
||||
# make a copy so that we don't update the original
|
||||
query = parameters.get("query", {}).copy()
|
||||
database = parameters.get("database", "")
|
||||
token = parameters.get("access_token", "")
|
||||
|
||||
if not database.startswith("md:"):
|
||||
database = f"md:{database}"
|
||||
if token and token != DEFAULT_ACCESS_TOKEN_URL:
|
||||
query["motherduck_token"] = token
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Need MotherDuck token to connect to database '{database}'."
|
||||
)
|
||||
|
||||
return str(
|
||||
URL(drivername=DuckDBEngineSpec.engine, database=database, query=query)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -72,3 +72,56 @@ def test_get_extra_params(mocker: MockerFixture) -> None:
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_build_sqlalchemy_uri() -> None:
|
||||
"""Test DuckDBEngineSpec.build_sqlalchemy_uri"""
|
||||
from superset.db_engine_specs.duckdb import DuckDBEngineSpec, DuckDBParametersType
|
||||
|
||||
# No database provided, default to :memory:
|
||||
parameters = DuckDBParametersType()
|
||||
uri = DuckDBEngineSpec.build_sqlalchemy_uri(parameters)
|
||||
assert "duckdb:///:memory:" == uri
|
||||
|
||||
# Database provided
|
||||
parameters = DuckDBParametersType(database="/path/to/duck.db")
|
||||
uri = DuckDBEngineSpec.build_sqlalchemy_uri(parameters)
|
||||
assert "duckdb:////path/to/duck.db" == uri
|
||||
|
||||
|
||||
def test_md_build_sqlalchemy_uri() -> None:
|
||||
"""Test MotherDuckEngineSpec.build_sqlalchemy_uri"""
|
||||
from superset.db_engine_specs.duckdb import (
|
||||
DuckDBParametersType,
|
||||
MotherDuckEngineSpec,
|
||||
)
|
||||
|
||||
# No access token provided, throw ValueError
|
||||
parameters = DuckDBParametersType(database="my_db")
|
||||
with pytest.raises(ValueError):
|
||||
MotherDuckEngineSpec.build_sqlalchemy_uri(parameters)
|
||||
|
||||
# No database provided, default to "md:"
|
||||
parameters = DuckDBParametersType(access_token="token")
|
||||
uri = MotherDuckEngineSpec.build_sqlalchemy_uri(parameters)
|
||||
assert "duckdb:///md:?motherduck_token=token"
|
||||
|
||||
# Database and access_token provided
|
||||
parameters = DuckDBParametersType(database="my_db", access_token="token")
|
||||
uri = MotherDuckEngineSpec.build_sqlalchemy_uri(parameters)
|
||||
assert "duckdb:///md:my_db?motherduck_token=token" == uri
|
||||
|
||||
|
||||
def test_get_parameters_from_uri() -> None:
|
||||
from superset.db_engine_specs.duckdb import DuckDBEngineSpec
|
||||
|
||||
uri = "duckdb:////path/to/duck.db"
|
||||
parameters = DuckDBEngineSpec.get_parameters_from_uri(uri)
|
||||
|
||||
assert parameters["database"] == "/path/to/duck.db"
|
||||
|
||||
uri = "duckdb:///md:my_db?motherduck_token=token"
|
||||
parameters = DuckDBEngineSpec.get_parameters_from_uri(uri)
|
||||
|
||||
assert parameters["database"] == "md:my_db"
|
||||
assert parameters["access_token"] == "token"
|
||||
|
|
|
|||
Loading…
Reference in New Issue