[superset] 01/01: move dynamic schema to Postgres class

[hu...@apache.org]

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

hugh pushed a commit to branch fix-connect-eng
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 3d612d438d0ee6f46775b2f97a05a5e83e97ef2c
Author: hughhhh <hu...@gmail.com>
AuthorDate: Fri Apr 28 12:22:52 2023 -0500

    move dynamic schema to Postgres class
---
 superset/db_engine_specs/postgres.py | 95 ++++++++++++++++++++++++------------
 1 file changed, 64 insertions(+), 31 deletions(-)

diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py
index 99a927541e..7364965cba 100644
--- a/superset/db_engine_specs/postgres.py
+++ b/superset/db_engine_specs/postgres.py
@@ -95,7 +95,6 @@ class PostgresBaseEngineSpec(BaseEngineSpec):
     engine = ""
     engine_name = "PostgreSQL"
 
-    supports_dynamic_schema = True
     supports_catalog = True
 
     _time_grain_expressions = {
@@ -167,25 +166,6 @@ class PostgresBaseEngineSpec(BaseEngineSpec):
         ),
     }
 
-    @classmethod
-    def adjust_engine_params(
-        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
-
     @classmethod
     def get_schema_from_engine_params(
         cls,
@@ -205,17 +185,19 @@ class PostgresBaseEngineSpec(BaseEngineSpec):
         to determine the schema for a non-qualified table in a query. In cases like
         that we raise an exception.
         """
-        options = parse_options(connect_args)
-        if search_path := options.get("search_path"):
-            schemas = search_path.split(",")
-            if len(schemas) > 1:
-                raise Exception(
-                    "Multiple schemas are configured in the search path, which means "
-                    "Superset is unable to determine the schema of unqualified table "
-                    "names and enforce permissions."
-                )
-            return schemas[0]
-
+        options = re.split(r"-c\s?", connect_args.get("options", ""))
+        for option in options:
+            if "=" not in option:
+                continue
+            key, value = option.strip().split("=", 1)
+            if key.strip() == "search_path":
+                if "," in value:
+                    raise Exception(
+                        "Multiple schemas are configured in the search path, which means "
+                        "Superset is unable to determine the schema of unqualified table "
+                        "names and enforce permissions."
+                    )
+                return value.strip()
         return None
 
     @classmethod
@@ -268,6 +250,57 @@ class PostgresEngineSpec(PostgresBaseEngineSpec, BasicParametersMixin):
         ),
     )
 
+    @classmethod
+    def get_schema_from_engine_params(
+        cls,
+        sqlalchemy_uri: URL,
+        connect_args: Dict[str, Any],
+    ) -> Optional[str]:
+        """
+        Return the configured schema.
+
+        While Postgres doesn't support connecting directly to a given schema, it allows
+        users to specify a "search path" that is used to resolve non-qualified table
+        names; this can be specified in the database ``connect_args``.
+
+        One important detail is that the search path can be a comma separated list of
+        schemas. While this is supported by the SQLAlchemy dialect, it shouldn't be used
+        in Superset because it breaks schema-level permissions, since it's impossible
+        to determine the schema for a non-qualified table in a query. In cases like
+        that we raise an exception.
+        """
+        options = parse_options(connect_args)
+        if search_path := options.get("search_path"):
+            schemas = search_path.split(",")
+            if len(schemas) > 1:
+                raise Exception(
+                    "Multiple schemas are configured in the search path, which means "
+                    "Superset is unable to determine the schema of unqualified table "
+                    "names and enforce permissions."
+                )
+            return schemas[0]
+
+        return None
+
+    @classmethod
+    def adjust_engine_params(
+        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
+
     @classmethod
     def get_allow_cost_estimate(cls, extra: Dict[str, Any]) -> bool:
         return True