You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by el...@apache.org on 2026/09/15 05:17:54 UTC
(superset) 01/01: feat(athena): add dynamic connection form parameters schema
This is an automated email from the ASF dual-hosted git repository.
eschutho pushed a commit to branch sc-120994-athena-new-form
in repository https://gitbox.apache.org/repos/asf/superset.git
commit caba0fc9069edc9ce5a661027959bd88527c41a7
Author: Elizabeth Thompson <es...@gmail.com>
AuthorDate: Tue Sep 15 05:16:35 2026 +0000
feat(athena): add dynamic connection form parameters schema
Add a custom `parameters_schema` to `AthenaEngineSpec` so Athena connections
render the new dynamic connection form (individual fields) instead of the raw
SQLAlchemy URI box.
Athena's PyAthena URI
(`awsathena+rest://{key}:{secret}@athena.{region}.amazonaws.com/{schema}
?s3_staging_dir=...&work_group=...`) is not a simple host/port/user/pass/db,
so this follows the custom-schema pattern used by bigquery/gsheets/snowflake
rather than `BasicParametersMixin`:
- `AthenaParametersSchema` / `AthenaParametersType` with aws_access_key_id,
aws_secret_access_key, region_name, s3_staging_dir, schema_name, work_group
- `build_sqlalchemy_uri` / `get_parameters_from_uri` round-trip (secret masked
via the URI password machinery)
- `parameters_json_schema`, `validate_parameters`, `sqlalchemy_uri_placeholder`
and `default_driver = "rest"`
Wire the new fields into the frontend dynamic connection form
(FormFieldOrder / FORM_FIELD_MAP / FIELD_TEXT_MAP and DatabaseParameters type)
so they render as labeled inputs.
Co-Authored-By: Claude Opus 4.8 <no...@anthropic.com>
---
.../DatabaseConnectionForm/ValidatedInputField.tsx | 27 ++++
.../DatabaseConnectionForm/constants.ts | 12 ++
superset-frontend/src/features/databases/types.ts | 6 +
superset/db_engine_specs/athena.py | 168 ++++++++++++++++++++-
tests/unit_tests/db_engine_specs/test_athena.py | 153 +++++++++++++++++++
5 files changed, 363 insertions(+), 3 deletions(-)
diff --git a/superset-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/ValidatedInputField.tsx b/superset-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/ValidatedInputField.tsx
index 74bcce3f795..2a942e6c742 100644
--- a/superset-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/ValidatedInputField.tsx
+++ b/superset-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/ValidatedInputField.tsx
@@ -38,6 +38,33 @@ const FIELD_TEXT_MAP = {
placeholder: t('e.g. AccountAdmin'),
className: 'form-group-w-50',
},
+ aws_access_key_id: {
+ label: t('AWS Access Key ID'),
+ placeholder: t('e.g. AKIAIOSFODNN7EXAMPLE'),
+ },
+ aws_secret_access_key: {
+ label: t('AWS Secret Access Key'),
+ placeholder: t('e.g. ********'),
+ },
+ region_name: {
+ label: t('Region'),
+ placeholder: t('e.g. us-east-1'),
+ className: 'form-group-w-50',
+ },
+ s3_staging_dir: {
+ label: t('S3 Staging Directory'),
+ placeholder: t('e.g. s3://my-bucket/staging/'),
+ },
+ schema_name: {
+ label: t('Schema'),
+ placeholder: t('e.g. default'),
+ className: 'form-group-w-50',
+ },
+ work_group: {
+ label: t('Work Group'),
+ placeholder: t('e.g. primary'),
+ className: 'form-group-w-50',
+ },
};
type FieldTextMapKey = keyof typeof FIELD_TEXT_MAP;
diff --git a/superset-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/constants.ts b/superset-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/constants.ts
index 4512d1f707f..058454c9d95 100644
--- a/superset-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/constants.ts
+++ b/superset-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/constants.ts
@@ -60,6 +60,12 @@ export const FormFieldOrder = [
'account',
'warehouse',
'role',
+ 'aws_access_key_id',
+ 'aws_secret_access_key',
+ 'region_name',
+ 's3_staging_dir',
+ 'schema_name',
+ 'work_group',
'ssh',
'oauth2_client_info',
];
@@ -90,6 +96,12 @@ export const FORM_FIELD_MAP = {
warehouse: validatedInputField,
role: validatedInputField,
account: validatedInputField,
+ aws_access_key_id: validatedInputField,
+ aws_secret_access_key: validatedInputField,
+ region_name: validatedInputField,
+ s3_staging_dir: validatedInputField,
+ schema_name: validatedInputField,
+ work_group: validatedInputField,
ssh: SSHTunnelSwitchComponent,
project_id: projectIdfield,
};
diff --git a/superset-frontend/src/features/databases/types.ts b/superset-frontend/src/features/databases/types.ts
index 162980f7f81..c0ba6f5ddb6 100644
--- a/superset-frontend/src/features/databases/types.ts
+++ b/superset-frontend/src/features/databases/types.ts
@@ -60,6 +60,12 @@ export type DatabaseParameters = {
warehouse?: string;
role?: string;
account?: string;
+ aws_access_key_id?: string;
+ aws_secret_access_key?: string;
+ region_name?: string;
+ s3_staging_dir?: string;
+ schema_name?: string;
+ work_group?: string;
ssh?: boolean;
project_id?: string;
};
diff --git a/superset/db_engine_specs/athena.py b/superset/db_engine_specs/athena.py
index 02deb7212b4..9b2c9ca387e 100644
--- a/superset/db_engine_specs/athena.py
+++ b/superset/db_engine_specs/athena.py
@@ -17,20 +17,77 @@
import re
from datetime import datetime
from re import Pattern
-from typing import Any, Optional
+from typing import Any, Optional, TypedDict
+from apispec import APISpec
+from apispec.ext.marshmallow import MarshmallowPlugin
from flask_babel import gettext as __
+from marshmallow import fields, Schema
from sqlalchemy import types
from sqlalchemy.engine.url import URL
from superset.constants import TimeGrain
-from superset.db_engine_specs.base import BaseEngineSpec, DatabaseCategory
-from superset.errors import SupersetErrorType
+from superset.databases.utils import make_url_safe
+from superset.db_engine_specs.base import (
+ BaseEngineSpec,
+ BasicPropertiesType,
+ DatabaseCategory,
+)
+from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
SYNTAX_ERROR_REGEX = re.compile(
": mismatched input '(?P<syntax_error>.*?)'. Expecting: "
)
+# Host portion of the PyAthena SQLAlchemy URI, e.g.
+# ``athena.us-east-1.amazonaws.com`` — the region is the only variable part.
+ATHENA_HOST_REGEX = re.compile(r"^athena\.(?P<region_name>.+)\.amazonaws\.com$")
+
+
+class AthenaParametersSchema(Schema):
+ aws_access_key_id = fields.Str(
+ required=False,
+ allow_none=True,
+ metadata={"description": __("AWS access key ID")},
+ )
+ aws_secret_access_key = fields.Str(
+ required=False,
+ allow_none=True,
+ metadata={"description": __("AWS secret access key")},
+ )
+ region_name = fields.Str(
+ required=True,
+ metadata={"description": __("AWS region, e.g. us-east-1")},
+ )
+ s3_staging_dir = fields.Str(
+ required=True,
+ metadata={
+ "description": __(
+ "S3 location where Athena query results are stored, "
+ "e.g. s3://my-bucket/staging/"
+ )
+ },
+ )
+ schema_name = fields.Str(
+ required=False,
+ allow_none=True,
+ metadata={"description": __("Default schema (database) to connect to")},
+ )
+ work_group = fields.Str(
+ required=False,
+ allow_none=True,
+ metadata={"description": __("Athena work group to run queries against")},
+ )
+
+
+class AthenaParametersType(TypedDict, total=False):
+ aws_access_key_id: Optional[str]
+ aws_secret_access_key: Optional[str]
+ region_name: str
+ s3_staging_dir: str
+ schema_name: Optional[str]
+ work_group: Optional[str]
+
class AthenaEngineSpec(BaseEngineSpec):
engine = "awsathena"
@@ -41,6 +98,17 @@ class AthenaEngineSpec(BaseEngineSpec):
use_equality_for_boolean_filters = True
supports_dynamic_schema = True
+ # Render the dynamic connection form (individual parameters) instead of the
+ # raw SQLAlchemy URI box. Athena's URI is not a simple host/port/user/pass/db,
+ # so a custom schema is used rather than ``BasicParametersMixin``.
+ parameters_schema = AthenaParametersSchema()
+ default_driver = "rest"
+ sqlalchemy_uri_placeholder = (
+ "awsathena+rest://{aws_access_key_id}:{aws_secret_access_key}@"
+ "athena.{region_name}.amazonaws.com/{schema_name}"
+ "?s3_staging_dir={s3_staging_dir}"
+ )
+
metadata = {
"description": (
"Amazon Athena is an interactive query service for "
@@ -190,3 +258,97 @@ class AthenaEngineSpec(BaseEngineSpec):
awsathena+rest://athena.{region_name}.amazonaws.com:443/{schema_name}?catalog_name={catalog_name}&s3_staging_dir={s3_staging_dir}
"""
return sqlalchemy_uri.database
+
+ @classmethod
+ def build_sqlalchemy_uri(
+ cls,
+ parameters: AthenaParametersType,
+ encrypted_extra: Optional[ # pylint: disable=unused-argument
+ dict[str, Any]
+ ] = None,
+ ) -> str:
+ query: dict[str, str] = {"s3_staging_dir": parameters["s3_staging_dir"]}
+ if work_group := parameters.get("work_group"):
+ query["work_group"] = work_group
+
+ # SQLAlchemy 2.0 made URL.__str__() hide the password by default
+ # (it rendered in full under 1.4); render_as_string(hide_password=
+ # False) is required here since this URI is stored/used to actually
+ # connect, not just displayed.
+ return URL.create(
+ f"{cls.engine}+{cls.default_driver}",
+ username=parameters.get("aws_access_key_id"),
+ password=parameters.get("aws_secret_access_key"),
+ host=f"athena.{parameters['region_name']}.amazonaws.com",
+ database=parameters.get("schema_name"),
+ query=query,
+ ).render_as_string(hide_password=False)
+
+ @classmethod
+ def get_parameters_from_uri(
+ cls,
+ uri: str,
+ encrypted_extra: Optional[ # pylint: disable=unused-argument
+ dict[str, Any]
+ ] = None,
+ ) -> AthenaParametersType:
+ url = make_url_safe(uri)
+ query = dict(url.query.items())
+
+ host = url.host or ""
+ if match := ATHENA_HOST_REGEX.match(host):
+ region_name = match["region_name"]
+ else:
+ region_name = host
+
+ parameters: AthenaParametersType = {
+ "aws_access_key_id": url.username,
+ "aws_secret_access_key": url.password,
+ "region_name": region_name,
+ "s3_staging_dir": query.get("s3_staging_dir", ""),
+ "schema_name": url.database,
+ }
+ if "work_group" in query:
+ parameters["work_group"] = query["work_group"]
+ return parameters
+
+ @classmethod
+ def validate_parameters(
+ cls,
+ properties: BasicPropertiesType,
+ ) -> list[SupersetError]:
+ errors: list[SupersetError] = []
+ required = {"region_name", "s3_staging_dir"}
+ parameters = properties.get("parameters", {})
+ 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.0",
+ plugins=[MarshmallowPlugin()],
+ )
+
+ spec.components.schema(cls.__name__, schema=cls.parameters_schema)
+ return spec.to_dict()["components"]["schemas"][cls.__name__]
diff --git a/tests/unit_tests/db_engine_specs/test_athena.py b/tests/unit_tests/db_engine_specs/test_athena.py
index d205572d2cb..cd58f3c13f0 100644
--- a/tests/unit_tests/db_engine_specs/test_athena.py
+++ b/tests/unit_tests/db_engine_specs/test_athena.py
@@ -190,3 +190,156 @@ def test_get_schema_from_engine_params() -> None:
)
is None
)
+
+
+def test_build_sqlalchemy_uri_full() -> None:
+ """
+ Test ``build_sqlalchemy_uri`` with every parameter supplied.
+ """
+ from superset.db_engine_specs.athena import AthenaEngineSpec
+
+ uri = AthenaEngineSpec.build_sqlalchemy_uri(
+ {
+ "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
+ "aws_secret_access_key": "secret/with+special=chars",
+ "region_name": "us-east-1",
+ "s3_staging_dir": "s3://my-bucket/staging/",
+ "schema_name": "default",
+ "work_group": "primary",
+ }
+ )
+ assert uri == (
+ "awsathena+rest://AKIAIOSFODNN7EXAMPLE:secret%2Fwith+special%3Dchars@"
+ "athena.us-east-1.amazonaws.com/default"
+ "?s3_staging_dir=s3%3A%2F%2Fmy-bucket%2Fstaging%2F&work_group=primary"
+ )
+
+
+def test_build_sqlalchemy_uri_minimal() -> None:
+ """
+ Test ``build_sqlalchemy_uri`` with only the required parameters.
+
+ Without explicit credentials (e.g. IAM role based auth) the URI has no
+ user/password, no schema, and no work group.
+ """
+ from superset.db_engine_specs.athena import AthenaEngineSpec
+
+ uri = AthenaEngineSpec.build_sqlalchemy_uri(
+ {
+ "region_name": "eu-west-1",
+ "s3_staging_dir": "s3://another-bucket/",
+ }
+ )
+ assert uri == (
+ "awsathena+rest://athena.eu-west-1.amazonaws.com"
+ "?s3_staging_dir=s3%3A%2F%2Fanother-bucket%2F"
+ )
+
+
+def test_get_parameters_from_uri() -> None:
+ """
+ Test ``get_parameters_from_uri`` extracts each parameter back out.
+ """
+ from superset.db_engine_specs.athena import AthenaEngineSpec
+
+ parameters = AthenaEngineSpec.get_parameters_from_uri(
+ "awsathena+rest://AKIAIOSFODNN7EXAMPLE:secret%2Fkey@"
+ "athena.us-east-1.amazonaws.com/default"
+ "?s3_staging_dir=s3%3A%2F%2Fmy-bucket%2Fstaging%2F&work_group=primary"
+ )
+ assert parameters == {
+ "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
+ "aws_secret_access_key": "secret/key",
+ "region_name": "us-east-1",
+ "s3_staging_dir": "s3://my-bucket/staging/",
+ "schema_name": "default",
+ "work_group": "primary",
+ }
+
+
+def test_parameters_round_trip() -> None:
+ """
+ Building a URI and parsing it back should yield the original parameters.
+ """
+ from superset.db_engine_specs.athena import (
+ AthenaEngineSpec,
+ AthenaParametersType,
+ )
+
+ parameters: AthenaParametersType = {
+ "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
+ "aws_secret_access_key": "secret/with+special=chars",
+ "region_name": "us-east-1",
+ "s3_staging_dir": "s3://my-bucket/staging/",
+ "schema_name": "default",
+ "work_group": "primary",
+ }
+ uri = AthenaEngineSpec.build_sqlalchemy_uri(parameters)
+ assert AthenaEngineSpec.get_parameters_from_uri(uri) == parameters
+
+
+def test_get_parameters_from_uri_non_standard_host() -> None:
+ """
+ A host that doesn't match the ``athena.<region>.amazonaws.com`` shape is
+ returned verbatim as ``region_name``.
+ """
+ from superset.db_engine_specs.athena import AthenaEngineSpec
+
+ parameters = AthenaEngineSpec.get_parameters_from_uri(
+ "awsathena+rest://vpce-1234.athena.us-east-1.vpce.amazonaws.com/default"
+ "?s3_staging_dir=s3%3A%2F%2Fbucket%2F"
+ )
+ assert parameters["region_name"] == "vpce-1234.athena.us-east-1.vpce.amazonaws.com"
+ assert parameters["s3_staging_dir"] == "s3://bucket/"
+
+
+def test_parameters_json_schema() -> None:
+ """
+ Test that the OpenAPI schema exposes the expected fields.
+ """
+ from superset.db_engine_specs.athena import AthenaEngineSpec
+
+ schema = AthenaEngineSpec.parameters_json_schema()
+ assert set(schema["properties"]) == {
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "region_name",
+ "s3_staging_dir",
+ "schema_name",
+ "work_group",
+ }
+ assert sorted(schema["required"]) == ["region_name", "s3_staging_dir"]
+
+
+def test_validate_parameters_missing_required() -> None:
+ """
+ Test that ``validate_parameters`` flags missing required parameters.
+ """
+ from superset.db_engine_specs.athena import AthenaEngineSpec
+
+ errors = AthenaEngineSpec.validate_parameters(
+ {"parameters": {"region_name": "us-east-1"}} # type: ignore
+ )
+ assert len(errors) == 1
+ assert errors[0].error_type == (
+ SupersetErrorType.CONNECTION_MISSING_PARAMETERS_ERROR
+ )
+ assert errors[0].extra is not None
+ assert errors[0].extra["missing"] == ["s3_staging_dir"]
+
+
+def test_validate_parameters_valid() -> None:
+ """
+ Test that ``validate_parameters`` passes when required parameters exist.
+ """
+ from superset.db_engine_specs.athena import AthenaEngineSpec
+
+ errors = AthenaEngineSpec.validate_parameters(
+ {
+ "parameters": { # type: ignore
+ "region_name": "us-east-1",
+ "s3_staging_dir": "s3://my-bucket/staging/",
+ }
+ }
+ )
+ assert errors == []