feat(dbview): Add token request button to DuckDB and MotherDuck datab…

…ase modal (apache#27908)

Co-authored-by: Beto Dealmeida <[email protected]>

superset-frontend/src/components/Form/LabeledErrorBoundInput.tsx

@@ -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>
 );

superset-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/CommonParameters.tsx

@@ -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}
   />

superset-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/index.tsx

@@ -169,6 +169,8 @@ const DatabaseConnectionForm = ({
               db,
               key: field,
               field,
+              default_value: parameters.properties[field]?.default,
+              description: parameters.properties[field]?.description,
               isEditMode,
               sslForced,
               editNewDb,

superset-frontend/src/features/databases/DatabaseModal/ModalHeader.tsx

@@ -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>

superset-frontend/src/features/databases/types.ts

@@ -291,6 +291,8 @@ export interface FieldPropTypes {
   db?: DatabaseObject;
   dbModel?: DatabaseForm;
   field: string;
+  default_value?: any;
+  description?: string;
   isEditMode?: boolean;
   sslForced?: boolean;
   defaultDBName?: string;

superset/db_engine_specs/duckdb.py

@@ -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__]
+
+
+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 @@ 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)
+        )

tests/unit_tests/db_engine_specs/test_duckdb.py

@@ -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"