Skip to content

Commit

Permalink
fix: move dynamic schema out of base Postgres class (apache#23868)
Browse files Browse the repository at this point in the history
(cherry picked from commit ba00dfa)
  • Loading branch information
hughhhh authored and jinghua-qa committed May 1, 2023
1 parent 64a4e22 commit bf2f119
Showing 1 changed file with 52 additions and 52 deletions.
104 changes: 52 additions & 52 deletions superset/db_engine_specs/postgres.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,6 @@ class PostgresBaseEngineSpec(BaseEngineSpec):
engine = ""
engine_name = "PostgreSQL"

supports_dynamic_schema = True
supports_catalog = True

_time_grain_expressions = {
Expand Down Expand Up @@ -167,57 +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,
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 fetch_data(
cls, cursor: Any, limit: Optional[int] = None
Expand All @@ -234,6 +182,7 @@ def epoch_to_dttm(cls) -> str:
class PostgresEngineSpec(PostgresBaseEngineSpec, BasicParametersMixin):
engine = "postgresql"
engine_aliases = {"postgres"}
supports_dynamic_schema = True

default_driver = "psycopg2"
sqlalchemy_uri_placeholder = (
Expand Down Expand Up @@ -268,6 +217,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
Expand Down

0 comments on commit bf2f119

Please sign in to comment.