From 6322cedea38bcfa049d650d6ddd79de35f18dfa5 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Tue, 18 Jul 2023 17:25:51 -0700
Subject: [PATCH] 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 0d778de43987a..ebfbfa6eb8c92 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -1098,6 +1098,19 @@ def adjust_engine_params(  # pylint: disable=unused-argument
             **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 cdd71fdfccbcc..ba3a411f1a387 100644
--- a/superset/db_engine_specs/postgres.py
+++ b/superset/db_engine_specs/postgres.py
@@ -252,23 +252,12 @@ def get_schema_from_engine_params(
         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 4ff56145e1bd5..ad452b74dff85 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -521,6 +521,8 @@ def get_raw_connection(
             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]: