diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index efe7d68da2182..1e57996998183 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -472,7 +472,7 @@ def get_default_schema_for_query( Determining the correct schema is crucial for managing access to data, so please make sure you understand this logic when working on a new DB engine spec. """ - # default schema varies on a per-query basis + # dynamic schema varies on a per-query basis if cls.supports_dynamic_schema: return query.schema @@ -1061,25 +1061,27 @@ def adjust_engine_params( # pylint: disable=unused-argument cls, uri: URL, connect_args: Dict[str, Any], + catalog: Optional[str], schema: Optional[str], ) -> Tuple[URL, Dict[str, Any]]: """ - Return a modified URL with a new database component. + Return a new URL and ``connect_args`` for a specific catalog/schema. - The URI here represents the URI as entered when saving the database, - ``selected_schema`` is the schema currently active presumably in - the SQL Lab dropdown. Based on that, for some database engine, - we can return a new altered URI that connects straight to the - active schema, meaning the users won't have to prefix the object - names by the schema name. + This is used in SQL Lab, allowing users to select a schema from the list of + schemas available in a given database, and have the query run with that schema as + the default one. - Some databases engines have 2 level of namespacing: database and - schema (postgres, oracle, mssql, ...) - For those it's probably better to not alter the database - component of the URI with the schema name, it won't work. + For some databases (like MySQL, Presto, Snowflake) this requires modifying the + SQLAlchemy URI before creating the connection. For others (like Postgres), it + requires additional parameters in ``connect_args``. - Some database drivers like Presto accept '{catalog}/{schema}' in - the database component of the URL, that can be handled here. + When a DB engine spec implements this method it should also have the attribute + ``supports_dynamic_schema`` set to true, so that Superset knows in which schema a + given query is running in order to enforce permissions (see #23385 and #23401). + + Currently, changing the catalog is not supported. The method acceps a catalog so + that when catalog support is added to Superse the interface remains the same. This + is important because DB engine specs can be installed from 3rd party packages. """ return uri, connect_args diff --git a/superset/db_engine_specs/drill.py b/superset/db_engine_specs/drill.py index 5f7e9baf51170..ec7788ead2d07 100644 --- a/superset/db_engine_specs/drill.py +++ b/superset/db_engine_specs/drill.py @@ -75,6 +75,7 @@ def adjust_engine_params( cls, uri: URL, connect_args: Dict[str, Any], + catalog: Optional[str], schema: Optional[str], ) -> Tuple[URL, Dict[str, Any]]: if schema: diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py index f70ff2db03dfd..792ef947350a1 100644 --- a/superset/db_engine_specs/hive.py +++ b/superset/db_engine_specs/hive.py @@ -264,6 +264,7 @@ 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 schema: diff --git a/superset/db_engine_specs/mysql.py b/superset/db_engine_specs/mysql.py index 572db3aab14a4..e5ff964f868da 100644 --- a/superset/db_engine_specs/mysql.py +++ b/superset/db_engine_specs/mysql.py @@ -195,6 +195,7 @@ 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 schema: diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index 3999c98e64967..7555b660116b5 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -154,6 +154,7 @@ 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: diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py index 4004483abf52f..c0b4f2c6dd881 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -305,6 +305,7 @@ 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]]: database = uri.database diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py index f3bff7021cf6b..033b637e48d1d 100644 --- a/superset/db_engine_specs/snowflake.py +++ b/superset/db_engine_specs/snowflake.py @@ -139,6 +139,7 @@ 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]]: database = uri.database diff --git a/superset/models/core.py b/superset/models/core.py index 45973e44d9061..d7a38cdc033a5 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -432,11 +432,29 @@ def _get_sqla_engine( params["poolclass"] = NullPool connect_args = params.get("connect_args", {}) + # The ``adjust_database_uri`` method was renamed to ``adjust_engine_params`` and + # had its signature changed in order to support more DB engine specs. Since DB + # engine specs can be released as 3rd party modules we want to make sure the old + # method is still supported so we don't introduce a breaking change. + if hasattr(self.db_engine_spec, "adjust_database_uri"): + sqlalchemy_url = self.db_engine_spec.adjust_database_uri( + sqlalchemy_url, + schema, + ) + logger.warning( + "DB engine spec %s implements the method `adjust_database_uri`, which is " + "deprecated and will be removed in version 3.0. Please update it to " + "implement `adjust_engine_params` instead.", + self.db_engine_spec, + ) + sqlalchemy_url, connect_args = self.db_engine_spec.adjust_engine_params( - sqlalchemy_url, - connect_args, - schema, + uri=sqlalchemy_url, + connect_args=connect_args, + catalog=None, + schema=schema, ) + effective_username = self.get_effective_user(sqlalchemy_url) # If using MySQL or Presto for example, will set url.username # If using Hive, will not do anything yet since that relies on a diff --git a/tests/unit_tests/db_engine_specs/test_postgres.py b/tests/unit_tests/db_engine_specs/test_postgres.py index 828f155cc9208..fef864795962e 100644 --- a/tests/unit_tests/db_engine_specs/test_postgres.py +++ b/tests/unit_tests/db_engine_specs/test_postgres.py @@ -141,7 +141,7 @@ def test_adjust_engine_params() -> None: uri = make_url("postgres://user:password@host/catalog") - assert PostgresEngineSpec.adjust_engine_params(uri, {}, "secret") == ( + assert PostgresEngineSpec.adjust_engine_params(uri, {}, None, "secret") == ( uri, {"options": "-csearch_path=secret"}, ) @@ -149,6 +149,7 @@ def test_adjust_engine_params() -> None: assert PostgresEngineSpec.adjust_engine_params( uri, {"foo": "bar", "options": "-csearch_path=default -c debug=1"}, + None, "secret", ) == ( uri,