feat(db): add dynamic schema support for athena (#36003)

tests/unit_tests/db_engine_specs/test_athena.py

@@ -20,6 +20,7 @@ from datetime import datetime
 from typing import Optional
 
 import pytest
+from sqlalchemy.engine.url import make_url
 
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
@@ -120,3 +121,72 @@ def test_handle_boolean_filter() -> None:
         str(result_false.compile(compile_kwargs={"literal_binds": True}))
         == "test_col = false"
     )
+
+
+def test_adjust_engine_params() -> None:
+    """
+    Test `adjust_engine_params`.
+
+    The method can be used to adjust the schema dynamically.
+    """
+    from superset.db_engine_specs.athena import AthenaEngineSpec
+
+    url = make_url("awsathena+rest://athena.us-east-1.amazonaws.com:443/default")
+
+    uri = AthenaEngineSpec.adjust_engine_params(url, {})[0]
+    assert str(uri) == "awsathena+rest://athena.us-east-1.amazonaws.com:443/default"
+
+    uri = AthenaEngineSpec.adjust_engine_params(
+        url,
+        {},
+        schema="new_schema",
+    )[0]
+    assert str(uri) == "awsathena+rest://athena.us-east-1.amazonaws.com:443/new_schema"
+
+    uri = AthenaEngineSpec.adjust_engine_params(
+        url,
+        {},
+        catalog="new_catalog",
+    )[0]
+    assert (
+        str(uri)
+        == "awsathena+rest://athena.us-east-1.amazonaws.com:443/default?catalog_name=new_catalog"
+    )
+
+    uri = AthenaEngineSpec.adjust_engine_params(
+        url,
+        {},
+        catalog="new_catalog",
+        schema="new_schema",
+    )[0]
+    assert (
+        str(uri)
+        == "awsathena+rest://athena.us-east-1.amazonaws.com:443/new_schema?catalog_name=new_catalog"
+    )
+
+
+def test_get_schema_from_engine_params() -> None:
+    """
+    Test the ``get_schema_from_engine_params`` method.
+    """
+    from superset.db_engine_specs.athena import AthenaEngineSpec
+
+    assert (
+        AthenaEngineSpec.get_schema_from_engine_params(
+            make_url(
+                "awsathena+rest://athena.us-east-1.amazonaws.com:443/default?s3_staging_dir=s3%3A%2F%2Fathena-staging"
+            ),
+            {},
+        )
+        == "default"
+    )
+
+    assert (
+        AthenaEngineSpec.get_schema_from_engine_params(
+            make_url(
+                "awsathena+rest://athena.us-east-1.amazonaws.com:443?s3_staging_dir=s3%3A%2F%2Fathena-staging"
+            ),
+            {},
+        )
+        is None
+    )