(superset) 01/01: feat: catalog support for Databricks native

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch databricks-catalogs
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 771f1245bb24e3e52bd605b136da3d42df0dc467
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Wed May 8 17:06:20 2024 -0400

    feat: catalog support for Databricks native
---
 superset/db_engine_specs/databricks.py             | 32 +++++++++++++++++++++-
 .../unit_tests/db_engine_specs/test_databricks.py  | 19 +++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py
index 4b2f93ca5d..f2f5b81dcd 100644
--- a/superset/db_engine_specs/databricks.py
+++ b/superset/db_engine_specs/databricks.py
@@ -39,7 +39,6 @@ if TYPE_CHECKING:
     from superset.models.core import Database
 
 
-#
 class DatabricksParametersSchema(Schema):
     """
     This is the list of fields that are expected
@@ -160,6 +159,8 @@ class DatabricksNativeEngineSpec(BasicParametersMixin, DatabricksODBCEngineSpec)
     )
     encryption_parameters = {"ssl": "1"}
 
+    supports_dynamic_schema = supports_catalog = supports_dynamic_catalog = True
+
     @staticmethod
     def get_extra_params(database: Database) -> dict[str, Any]:
         """
@@ -367,3 +368,32 @@ class DatabricksNativeEngineSpec(BasicParametersMixin, DatabricksODBCEngineSpec)
         )
         spec.components.schema(cls.__name__, schema=cls.properties_schema)
         return spec.to_dict()["components"]["schemas"][cls.__name__]
+
+    @classmethod
+    def get_default_catalog(
+        cls,
+        database: Database,
+    ) -> str | None:
+        with database.get_inspector() as inspector:
+            return inspector.bind.execute("SELECT current_catalog()").scalar()
+
+    @classmethod
+    def get_prequeries(
+        cls,
+        catalog: str | None = None,
+        schema: str | None = None,
+    ) -> list[str]:
+        prequeries = []
+        if catalog:
+            prequeries.append(f"USE CATALOG {catalog}")
+        if schema:
+            prequeries.append(f"USE SCHEMA {schema}")
+        return prequeries
+
+    @classmethod
+    def get_catalog_names(
+        cls,
+        database: Database,
+        inspector: Inspector,
+    ) -> set[str]:
+        return {catalog for (catalog,) in inspector.bind.execute("SHOW CATALOGS")}
diff --git a/tests/unit_tests/db_engine_specs/test_databricks.py b/tests/unit_tests/db_engine_specs/test_databricks.py
index de06f919be..03379f5c2e 100644
--- a/tests/unit_tests/db_engine_specs/test_databricks.py
+++ b/tests/unit_tests/db_engine_specs/test_databricks.py
@@ -245,3 +245,22 @@ def test_convert_dttm(
     from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec as spec
 
     assert_convert_dttm(spec, target_type, expected_result, dttm)
+
+
+def test_get_prequeries() -> None:
+    """
+    Test the ``get_prequeries`` method.
+    """
+    from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+
+    assert DatabricksNativeEngineSpec.get_prequeries() == []
+    assert DatabricksNativeEngineSpec.get_prequeries(schema="test") == [
+        "USE SCHEMA test",
+    ]
+    assert DatabricksNativeEngineSpec.get_prequeries(catalog="test") == [
+        "USE CATALOG test",
+    ]
+    assert DatabricksNativeEngineSpec.get_prequeries(catalog="foo", schema="bar") == [
+        "USE CATALOG foo",
+        "USE SCHEMA bar",
+    ]