[superset] branch fix_rds created (now 6322cedea3)

[be...@apache.org]

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

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


      at 6322cedea3 fix: search_path in RDS

This branch includes the following new commits:

     new 6322cedea3 fix: search_path in RDS

The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.

[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]: