feat: Add Aurora Data API engine spec (#16535)

* feat: Add Aurora Data API engine spec

* Fix lint

setup.py

@@ -114,6 +114,7 @@ def get_git_sha() -> str:
     ],
     extras_require={
         "athena": ["pyathena>=1.10.8, <1.11"],
+        "aurora-data-api": ["preset-sqlalchemy-aurora-data-api>=0.2.8,<0.3"],
         "bigquery": [
             "pandas_gbq>=0.10.0",
             "pybigquery>=0.4.10",

superset/db_engine_specs/__init__.py

@@ -94,6 +94,14 @@ def get_engine_specs() -> Dict[str, Type[BaseEngineSpec]]:
     return engine_specs_map
 
 
+# there's a mismatch between the dialect name reported by the driver in these
+# libraries and the dialect name used in the URI
+backend_replacements = {
+    "drilldbapi": "drill",
+    "exasol": "exa",
+}
+
+
 def get_available_engine_specs() -> Dict[Type[BaseEngineSpec], Set[str]]:
     """
     Return available engine specs and installed drivers for them.
@@ -124,12 +132,14 @@ def get_available_engine_specs() -> Dict[Type[BaseEngineSpec], Set[str]]:
     for ep in iter_entry_points("sqlalchemy.dialects"):
         try:
             dialect = ep.load()
-        except Exception:  # pylint: disable=broad-except
-            logger.warning("Unable to load SQLAlchemy dialect: %s", dialect)
+        except Exception as ex:  # pylint: disable=broad-except
+            logger.warning("Unable to load SQLAlchemy dialect %s: %s", dialect, ex)
         else:
             backend = dialect.name
             if isinstance(backend, bytes):
                 backend = backend.decode()
+            backend = backend_replacements.get(backend, backend)
+
             driver = getattr(dialect, "driver", dialect.name)
             if isinstance(driver, bytes):
                 driver = driver.decode()

superset/db_engine_specs/aurora.py

@@ -0,0 +1,46 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from superset.db_engine_specs.mysql import MySQLEngineSpec
+from superset.db_engine_specs.postgres import PostgresEngineSpec
+
+
+class AuroraMySQLDataAPI(MySQLEngineSpec):
+
+    engine = "mysql"
+    default_driver = "auroradataapi"
+    engine_name = "Aurora MySQL (Data API)"
+    sqlalchemy_uri_placeholder = (
+        "mysql+auroradataapi://{aws_access_id}:{aws_secret_access_key}@/"
+        "{database_name}?"
+        "aurora_cluster_arn={aurora_cluster_arn}&"
+        "secret_arn={secret_arn}&"
+        "region_name={region_name}"
+    )
+
+
+class AuroraPostgresDataAPI(PostgresEngineSpec):
+
+    engine = "postgresql"
+    default_driver = "auroradataapi"
+    engine_name = "Aurora PostgreSQL (Data API)"
+    sqlalchemy_uri_placeholder = (
+        "postgresql+auroradataapi://{aws_access_id}:{aws_secret_access_key}@/"
+        "{database_name}?"
+        "aurora_cluster_arn={aurora_cluster_arn}&"
+        "secret_arn={secret_arn}&"
+        "region_name={region_name}"
+    )

superset/db_engine_specs/rockset.py

@@ -51,7 +51,6 @@ def epoch_ms_to_dttm(cls) -> str:
 
     @classmethod
     def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]:
-        print(cls)
         tt = target_type.upper()
         if tt == utils.TemporalType.DATE:
             return f"DATE '{dttm.date().isoformat()}'"

superset/models/slice.py

@@ -53,7 +53,7 @@
 logger = logging.getLogger(__name__)
 
 
-class Slice(  # pylint: disable=too-many-public-methods, too-many-instance-attributes
+class Slice(  # pylint: disable=too-many-public-methods
     Model, AuditMixinNullable, ImportExportMixin
 ):
     """A slice is essentially a report or a view on data"""