[superset] 01/01: fix: search_path in RDS

[be...@apache.org]

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

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

commit 6322cedea38bcfa049d650d6ddd79de35f18dfa5
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Tue Jul 18 17:25:51 2023 -0700

    fix: search_path in RDS
---
 superset/db_engine_specs/base.py     | 13 +++++++++++++
 superset/db_engine_specs/postgres.py | 21 +++++----------------
 superset/models/core.py              |  2 ++
 3 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 0d778de439..ebfbfa6eb8 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -1098,6 +1098,19 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
             **cls.enforce_uri_query_params.get(uri.get_driver_name(), {}),
         }
 
+    @classmethod
+    def get_prequeries(
+        cls,
+        catalog: str | None = None,
+        schema: str | None = None,
+    ) -> [str]:
+        """
+        Return pre-session queries.
+
+        This are useful for setting the default catalog and/or schema.
+        """
+        return []
+
     @classmethod
     def patch(cls) -> None:
         """
diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py
index cdd71fdfcc..ba3a411f1a 100644
--- a/superset/db_engine_specs/postgres.py
+++ b/superset/db_engine_specs/postgres.py
@@ -252,23 +252,12 @@ class PostgresEngineSpec(PostgresBaseEngineSpec, BasicParametersMixin):
         return None
 
     @classmethod
-    def adjust_engine_params(
+    def get_prequeries(
         cls,
-        uri: URL,
-        connect_args: dict[str, Any],
-        catalog: Optional[str] = None,
-        schema: Optional[str] = None,
-    ) -> tuple[URL, dict[str, Any]]:
-        if not schema:
-            return uri, connect_args
-
-        options = parse_options(connect_args)
-        options["search_path"] = schema
-        connect_args["options"] = " ".join(
-            f"-c{key}={value}" for key, value in options.items()
-        )
-
-        return uri, connect_args
+        catalog: str | None = None,
+        schema: str | None = None,
+    ) -> [str]:
+        return [f'set search_path = "{schema}"']
 
     @classmethod
     def get_allow_cost_estimate(cls, extra: dict[str, Any]) -> bool:
diff --git a/superset/models/core.py b/superset/models/core.py
index 4ff56145e1..ad452b74df 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -521,6 +521,8 @@ class Database(
             schema=schema, nullpool=nullpool, source=source
         ) as engine:
             with closing(engine.raw_connection()) as conn:
+                for prequery in self.db_engine_spec.get_prequeries(schema=schema):
+                    conn.execute(prequery)
                 yield conn
 
     def get_default_schema_for_query(self, query: "Query") -> Optional[str]: