Skip to content
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(dbview): Add token request button to DuckDB and MotherDuck database modal #27908

Merged
merged 19 commits into from
Apr 15, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions superset-frontend/src/components/Form/LabeledErrorBoundInput.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -109,6 +110,8 @@ const LabeledErrorBoundInput = ({
id,
className,
visibilityToggle,
get_url,
description,
...props
}: LabeledErrorBoundInputProps) => (
<StyledFormGroup className={className}>
Expand Down Expand Up @@ -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>
);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,8 @@ export const accessTokenField = ({
validationErrors,
db,
isEditMode,
default_value,
description,
}: FieldPropTypes) => (
<ValidatedInput
id="access_token"
Expand All @@ -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}
/>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,8 @@ const DatabaseConnectionForm = ({
db,
key: field,
field,
default_value: parameters.properties[field]?.default,
description: parameters.properties[field]?.description,
isEditMode,
sslForced,
editNewDb,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
2 changes: 2 additions & 0 deletions superset-frontend/src/features/databases/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,8 @@ export interface FieldPropTypes {
db?: DatabaseObject;
dbModel?: DatabaseForm;
field: string;
default_value?: any;
description?: string;
isEditMode?: boolean;
sslForced?: boolean;
defaultDBName?: string;
Expand Down
194 changes: 189 additions & 5 deletions superset/db_engine_specs/duckdb.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,28 +19,180 @@
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
from superset.models.core import Database


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__]
Comment on lines +174 to +189
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note to self, we should not require this, it's the same for every database.



class DuckDBEngineSpec(DuckDBParametersMixin, BaseEngineSpec):
engine = "duckdb"
engine_name = "DuckDB"
default_driver = "duckdb_engine"

sqlalchemy_uri_placeholder = "duckdb:////path/to/duck.db"

Expand Down Expand Up @@ -103,9 +255,41 @@ def get_extra_params(database: Database) -> dict[str, Any]:


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)
)
53 changes: 53 additions & 0 deletions tests/unit_tests/db_engine_specs/test_duckdb.py
Original file line number Diff line number Diff line change
Expand Up @@ -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