[superset] branch master updated: feat(DB engine spec): `get_catalog_names` (#23447)

[be...@apache.org]

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

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


The following commit(s) were added to refs/heads/master by this push:
     new 8588f81855 feat(DB engine spec): `get_catalog_names` (#23447)
8588f81855 is described below

commit 8588f81855b926876b65e7481f163e0ee47d8bff
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Wed Mar 22 13:40:14 2023 -0700

    feat(DB engine spec): `get_catalog_names` (#23447)
---
 superset/db_engine_specs/base.py                   | 22 ++++++++++++++++++++++
 superset/db_engine_specs/postgres.py               | 22 ++++++++++++++++++++++
 .../db_engine_specs/postgres_tests.py              | 18 ++++++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index b8b1662057..f1a46ceb34 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -374,6 +374,14 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
     # a custom `adjust_engine_params` method.
     supports_dynamic_schema = False
 
+    # Does the DB support catalogs? A catalog here is a group of schemas, and has
+    # different names depending on the DB: BigQuery calles it a "project", Postgres calls
+    # it a "database", Trino calls it a "catalog", etc.
+    supports_catalog = False
+
+    # Can the catalog be changed on a per-query basis?
+    supports_dynamic_catalog = False
+
     @classmethod
     def supports_url(cls, url: URL) -> bool:
         """
@@ -1091,6 +1099,20 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
         TODO: Improve docstring and refactor implementation in Hive
         """
 
+    @classmethod
+    def get_catalog_names(  # pylint: disable=unused-argument
+        cls,
+        database: Database,
+        inspector: Inspector,
+    ) -> List[str]:
+        """
+        Get all catalogs from database.
+
+        This needs to be implemented per database, since SQLAlchemy doesn't offer an
+        abstraction.
+        """
+        return []
+
     @classmethod
     def get_schema_names(cls, inspector: Inspector) -> List[str]:
         """
diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py
index fac0b1b1d0..e52e69d7e8 100644
--- a/superset/db_engine_specs/postgres.py
+++ b/superset/db_engine_specs/postgres.py
@@ -23,6 +23,7 @@ from typing import Any, Dict, List, Optional, Pattern, Set, Tuple, TYPE_CHECKING
 from flask_babel import gettext as __
 from sqlalchemy.dialects.postgresql import DOUBLE_PRECISION, ENUM, JSON
 from sqlalchemy.dialects.postgresql.base import PGInspector
+from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.engine.url import URL
 from sqlalchemy.types import Date, DateTime, String
 
@@ -291,6 +292,27 @@ class PostgresEngineSpec(PostgresBaseEngineSpec, BasicParametersMixin):
     ) -> List[Dict[str, str]]:
         return [{k: str(v) for k, v in row.items()} for row in raw_cost]
 
+    @classmethod
+    def get_catalog_names(
+        cls,
+        database: "Database",
+        inspector: Inspector,
+    ) -> List[str]:
+        """
+        Return all catalogs.
+
+        In Postgres, a catalog is called a "database".
+        """
+        return sorted(
+            catalog
+            for (catalog,) in inspector.bind.execute(
+                """
+SELECT datname FROM pg_database
+WHERE datistemplate = false;
+            """
+            ).fetchall()
+        )
+
     @classmethod
     def get_table_names(
         cls, database: "Database", inspector: PGInspector, schema: Optional[str]
diff --git a/tests/integration_tests/db_engine_specs/postgres_tests.py b/tests/integration_tests/db_engine_specs/postgres_tests.py
index a6145432c2..44fb1c0474 100644
--- a/tests/integration_tests/db_engine_specs/postgres_tests.py
+++ b/tests/integration_tests/db_engine_specs/postgres_tests.py
@@ -17,6 +17,7 @@
 from textwrap import dedent
 from unittest import mock
 
+from flask.ctx import AppContext
 from sqlalchemy import column, literal_column
 from sqlalchemy.dialects import postgresql
 
@@ -24,6 +25,7 @@ from superset.db_engine_specs import load_engine_specs
 from superset.db_engine_specs.postgres import PostgresEngineSpec
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.models.sql_lab import Query
+from superset.utils.database import get_example_database
 from tests.integration_tests.db_engine_specs.base_tests import TestDbEngineSpec
 from tests.integration_tests.fixtures.certificates import ssl_certificate
 from tests.integration_tests.fixtures.database import default_db_extra
@@ -514,3 +516,19 @@ def test_base_parameters_mixin():
         },
         "required": ["database", "host", "port", "username"],
     }
+
+
+def test_get_catalog_names(app_context: AppContext) -> None:
+    """
+    Test the ``get_catalog_names`` method.
+    """
+    database = get_example_database()
+
+    if database.backend != "postgresql":
+        return
+
+    with database.get_inspector_with_context() as inspector:
+        assert PostgresEngineSpec.get_catalog_names(database, inspector) == [
+            "postgres",
+            "superset",
+        ]