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 <[email protected]> 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", + ]
