BW-2722: add param to indicate metadata fetch query

CoverGenius · Aug 8, 2024 · f166544 · f166544
1 parent eb089ed
commit f166544
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py
@@ -456,7 +456,7 @@ def get_catalog_names(
         In BigQuery, a catalog is called a "project".
         """
         engine: Engine
-        with database.get_sqla_engine_with_context() as engine:
+        with database.get_sqla_engine_with_context(fetch_metadata=True) as engine:
             client = cls._get_client(engine)
             projects = client.list_projects()
 

diff --git a/superset/models/core.py b/superset/models/core.py
@@ -383,6 +383,7 @@ def get_sqla_engine_with_context(
         nullpool: bool = True,
         source: utils.QuerySource | None = None,
         override_ssh_tunnel: SSHTunnel | None = None,
+        fetch_metadata: bool = False
     ) -> Engine:
         from superset.daos.database import (  # pylint: disable=import-outside-toplevel
             DatabaseDAO,
@@ -417,6 +418,7 @@ def get_sqla_engine_with_context(
                 nullpool=nullpool,
                 source=source,
                 sqlalchemy_uri=sqlalchemy_uri,
+                fetch_metadata=fetch_metadata
             )
 
     def _get_sqla_engine(
@@ -425,6 +427,7 @@ def _get_sqla_engine(
         nullpool: bool = True,
         source: utils.QuerySource | None = None,
         sqlalchemy_uri: str | None = None,
+        fetch_metadata: bool = False
     ) -> Engine:
         sqlalchemy_url = make_url_safe(
             sqlalchemy_uri if sqlalchemy_uri else self.sqlalchemy_uri_decrypted
@@ -500,6 +503,7 @@ def _get_sqla_engine(
                 effective_username,
                 security_manager,
                 source,
+                fetch_metadata
             )
         try:
             return create_engine(sqlalchemy_url, **params)
@@ -736,7 +740,8 @@ def get_inspector_with_context(
         self, ssh_tunnel: SSHTunnel | None = None
     ) -> Inspector:
         with self.get_sqla_engine_with_context(
-            override_ssh_tunnel=ssh_tunnel
+            override_ssh_tunnel=ssh_tunnel,
+            fetch_metadata = True
         ) as engine:
             yield sqla.inspect(engine)
 
@@ -817,7 +822,7 @@ def update_params_from_encrypted_extra(self, params: dict[str, Any]) -> None:
     def get_table(self, table_name: str, schema: str | None = None) -> Table:
         extra = self.get_extra()
         meta = MetaData(**extra.get("metadata_params", {}))
-        with self.get_sqla_engine_with_context() as engine:
+        with self.get_sqla_engine_with_context(fetch_metadata=True) as engine:
             return Table(
                 table_name,
                 meta,
@@ -921,11 +926,11 @@ def get_perm(self) -> str:
         return self.perm  # type: ignore
 
     def has_table(self, table: Table) -> bool:
-        with self.get_sqla_engine_with_context() as engine:
+        with self.get_sqla_engine_with_context(fetch_metadata=True) as engine:
             return engine.has_table(table.table_name, table.schema or None)
 
     def has_table_by_name(self, table_name: str, schema: str | None = None) -> bool:
-        with self.get_sqla_engine_with_context() as engine:
+        with self.get_sqla_engine_with_context(fetch_metadata=True) as engine:
             return engine.has_table(table_name, schema)
 
     @classmethod
@@ -944,7 +949,7 @@ def _has_view(
         return view_name in view_names
 
     def has_view(self, view_name: str, schema: str | None = None) -> bool:
-        with self.get_sqla_engine_with_context(schema) as engine:
+        with self.get_sqla_engine_with_context(schema, fetch_metadata=True) as engine:
             return engine.run_callable(
                 self._has_view, engine.dialect, view_name, schema
             )