This is an automated email from the ASF dual-hosted git repository. beto pushed a commit to branch sip-95 in repository https://gitbox.apache.org/repos/asf/superset.git
commit 02b54deece8981264843852b5a943ebbf2ab9767 Author: Beto Dealmeida <[email protected]> AuthorDate: Fri Apr 12 11:53:46 2024 -0400 WIP --- .../databases/DatabaseModal/ExtraOptions.tsx | 23 +++++- superset-frontend/src/features/databases/types.ts | 3 + superset/commands/database/test_connection.py | 2 +- superset/commands/dataset/importers/v1/utils.py | 5 +- superset/connectors/sqla/models.py | 1 + superset/databases/api.py | 60 ++++++++++++++ superset/databases/schemas.py | 22 +++++- superset/db_engine_specs/README.md | 6 +- superset/db_engine_specs/base.py | 56 ++++++++++--- superset/db_engine_specs/bigquery.py | 16 +++- superset/db_engine_specs/postgres.py | 40 ++++++++-- superset/extensions/metadb.py | 6 +- superset/models/core.py | 92 ++++++++++++++++++---- superset/models/dashboard.py | 7 -- superset/models/sql_lab.py | 4 + superset/sql_lab.py | 10 ++- superset/sql_validators/presto_db.py | 10 ++- 17 files changed, 305 insertions(+), 58 deletions(-) diff --git a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx index 28e02a02f0..61e465def3 100644 --- a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx +++ b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx @@ -56,6 +56,8 @@ const ExtraOptions = ({ const createAsOpen = !!(db?.allow_ctas || db?.allow_cvas); const isFileUploadSupportedByEngine = db?.engine_information?.supports_file_upload; + const supportsDynamicCatalog = + db?.engine_information?.supports_dynamic_catalog; // JSON.parse will deep parse engine_params // if it's an object, and we want to keep it a string @@ -191,7 +193,8 @@ const ExtraOptions = ({ <IndeterminateCheckbox id="allows_virtual_table_explore" indeterminate={false} - checked={!!extraJson?.allows_virtual_table_explore} + // when `allows_virtual_table_explore` is not present in `extra` it defaults to true + checked={extraJson?.allows_virtual_table_explore !== false} onChange={onExtraInputChange} labelText={t('Allow this database to be explored')} /> @@ -587,6 +590,24 @@ const ExtraOptions = ({ /> </div> </StyledInputContainer> + {supportsDynamicCatalog && ( + <StyledInputContainer css={no_margin_bottom}> + <div className="input-container"> + <IndeterminateCheckbox + id="allow_multi_catalog" + indeterminate={false} + checked={!!extraJson?.allow_multi_catalog} + onChange={onExtraInputChange} + labelText={t('Allow multiple catalogs')} + /> + <InfoTooltip + tooltip={t( + 'Allow access to multiple catalogs in a single database connection', + )} + /> + </div> + </StyledInputContainer> + )} </Collapse.Panel> </Collapse> ); diff --git a/superset-frontend/src/features/databases/types.ts b/superset-frontend/src/features/databases/types.ts index 2dff61d5e8..7e5b0d67fb 100644 --- a/superset-frontend/src/features/databases/types.ts +++ b/superset-frontend/src/features/databases/types.ts @@ -109,6 +109,7 @@ export type DatabaseObject = { engine_information?: { supports_file_upload?: boolean; disable_ssh_tunneling?: boolean; + supports_dynamic_catalog?: boolean; }; // SSH Tunnel information @@ -202,6 +203,7 @@ export type DatabaseForm = { engine_information: { supports_file_upload: boolean; disable_ssh_tunneling: boolean; + supports_dynamic_catalog: boolean; }; }; @@ -223,6 +225,7 @@ export interface ExtraJson { cost_estimate_enabled?: boolean; // in SQL Lab disable_data_preview?: boolean; // in SQL Lab disable_drill_to_detail?: boolean; + allow_multi_catalog?: boolean; engine_params?: { catalog?: Record<string, string>; connect_args?: { diff --git a/superset/commands/database/test_connection.py b/superset/commands/database/test_connection.py index 431918c6bc..9cb00c10d8 100644 --- a/superset/commands/database/test_connection.py +++ b/superset/commands/database/test_connection.py @@ -139,7 +139,7 @@ class TestConnectionDatabaseCommand(BaseCommand): return engine.dialect.do_ping(conn) with database.get_sqla_engine_with_context( - override_ssh_tunnel=ssh_tunnel + override_ssh_tunnel=ssh_tunnel, ) as engine: try: alive = func_timeout( diff --git a/superset/commands/dataset/importers/v1/utils.py b/superset/commands/dataset/importers/v1/utils.py index 04fc81e241..ed188f1758 100644 --- a/superset/commands/dataset/importers/v1/utils.py +++ b/superset/commands/dataset/importers/v1/utils.py @@ -217,7 +217,10 @@ def load_data(data_uri: str, dataset: SqlaTable, database: Database) -> None: ) else: logger.warning("Loading data outside the import transaction") - with database.get_sqla_engine_with_context() as engine: + with database.get_sqla_engine_with_context( + catalog=dataset.catalog, + schema=dataset.schema, + ) as engine: df.to_sql( dataset.table_name, con=engine, diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index afd2791c9c..b920e2a8d1 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -1146,6 +1146,7 @@ class SqlaTable( foreign_keys=[database_id], ) schema = Column(String(255)) + catalog = Column(String(256), nullable=True, default=None) sql = Column(MediumText()) is_sqllab_view = Column(Boolean, default=False) template_params = Column(Text) diff --git a/superset/databases/api.py b/superset/databases/api.py index 5eb9de90d4..b4b34fafa7 100644 --- a/superset/databases/api.py +++ b/superset/databases/api.py @@ -565,6 +565,66 @@ class DatabaseRestApi(BaseSupersetModelRestApi): ) return self.response_422(message=str(ex)) + @expose("/<int:pk>/catalogs/") + @protect() + @safe + @rison(database_schemas_query_schema) + @statsd_metrics + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}" f".schemas", + log_to_statsd=False, + ) + def catalogs(self, pk: int, **kwargs: Any) -> FlaskResponse: + """Get all catalogs from a database. + --- + get: + summary: Get all catalogs from a database + parameters: + - in: path + schema: + type: integer + name: pk + description: The database id + - in: query + name: q + content: + application/json: + schema: + $ref: '#/components/schemas/database_schemas_query_schema' + responses: + 200: + description: A List of all catalogs from the database + content: + application/json: + schema: + $ref: "#/components/schemas/CatalogsResponseSchema" + 400: + $ref: '#/components/responses/400' + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + 500: + $ref: '#/components/responses/500' + """ + database = self.datamodel.get(pk, self._base_filters) + if not database: + return self.response_404() + try: + schemas = database.get_all_schema_names( + cache=database.schema_cache_enabled, + cache_timeout=database.schema_cache_timeout or None, + force=kwargs["rison"].get("force", False), + ) + schemas = security_manager.get_schemas_accessible_by_user(database, schemas) + return self.response(200, result=schemas) + except OperationalError: + return self.response( + 500, message="There was an error connecting to the database" + ) + except SupersetException as ex: + return self.response(ex.status, message=ex.message) + @expose("/<int:pk>/schemas/") @protect() @safe diff --git a/superset/databases/schemas.py b/superset/databases/schemas.py index 791cfd360e..52f8405be5 100644 --- a/superset/databases/schemas.py +++ b/superset/databases/schemas.py @@ -650,6 +650,12 @@ class SelectStarResponseSchema(Schema): result = fields.String(metadata={"description": "SQL select star"}) +class CatalogsResponseSchema(Schema): + result = fields.List( + fields.String(metadata={"description": "A database catalog name"}) + ) + + class SchemasResponseSchema(Schema): result = fields.List( fields.String(metadata={"description": "A database schema name"}) @@ -897,6 +903,20 @@ class DatabaseSchemaAccessForFileUploadResponse(Schema): ) +class EngineInformationSchema(Schema): + supports_file_upload = fields.Boolean( + metadata={"description": "Users can upload files to the database"} + ) + disable_ssh_tunneling = fields.Boolean( + metadata={"description": "SSH tunnel is not available to the database"} + ) + supports_dynamic_catalog = fields.Boolean( + metadata={ + "description": "The database supports multiple catalogs in a single connection" + } + ) + + class DatabaseConnectionSchema(Schema): """ Schema with database connection information. @@ -930,7 +950,7 @@ class DatabaseConnectionSchema(Schema): driver = fields.String( allow_none=True, metadata={"description": "SQLAlchemy driver to use"} ) - engine_information = fields.Dict(keys=fields.String(), values=fields.Raw()) + engine_information = fields.Nested(EngineInformationSchema) expose_in_sqllab = fields.Boolean( metadata={"description": expose_in_sqllab_description} ) diff --git a/superset/db_engine_specs/README.md b/superset/db_engine_specs/README.md index f158a3a41b..4b17ad144a 100644 --- a/superset/db_engine_specs/README.md +++ b/superset/db_engine_specs/README.md @@ -660,7 +660,7 @@ This way, when a user selects a column that doesn't exist Superset can return a ### Dynamic schema -In SQL Lab it's possible to select a database, and then a schema in that database. Ideally, when running a query in SQL Lab, any unqualified table names (eg, `table`, instead of `schema.table`) should be in the selected schema. For example, if the user select `dev` as the schema and then runs the following query: +In SQL Lab it's possible to select a database, and then a schema in that database. Ideally, when running a query in SQL Lab, any unqualified table names (eg, `table`, instead of `schema.table`) should be in the selected schema. For example, if the user selects `dev` as the schema and then runs the following query: ```sql SELECT * FROM my_table @@ -674,7 +674,7 @@ Implementing this method is also important for usability. When the method is not ### Catalog -In general, databases support a hierarchy of concepts of one-to-many concepts: +In general, databases support a hierarchy of one-to-many concepts: 1. Database 2. Catalog @@ -692,7 +692,7 @@ These concepts have different names depending on the database. For example, Post BigQuery, on the other hand: -1. Bigquery (database) +1. BigQuery (database) 2. Project (catalog) 3. Schema (namespace) 4. Table diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index bcb4035c9c..36ca064f2a 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -760,18 +760,28 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods def get_engine( cls, database: Database, + catalog: str | None = None, schema: str | None = None, source: utils.QuerySource | None = None, ) -> ContextManager[Engine]: """ Return an engine context manager. - >>> with DBEngineSpec.get_engine(database, schema, source) as engine: + >>> with DBEngineSpec.get_engine( + ... database, + ... catalog, + ... schema, + ... source + ... ) as engine: ... connection = engine.connect() ... connection.execute(sql) """ - return database.get_sqla_engine_with_context(schema=schema, source=source) + return database.get_sqla_engine_with_context( + catalog=catalog, + schema=schema, + source=source, + ) @classmethod def get_timestamp_expr( @@ -1224,7 +1234,11 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods # Only add schema when it is preset and non-empty. to_sql_kwargs["schema"] = table.schema - with cls.get_engine(database) as engine: + with cls.get_engine( + database=database, + catalog=table.catalog, + schema=table.schema, + ) as engine: if engine.dialect.supports_multivalues_insert: to_sql_kwargs["method"] = "multi" @@ -1632,12 +1646,22 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods return sql @classmethod - def estimate_statement_cost(cls, statement: str, cursor: Any) -> dict[str, Any]: + def estimate_statement_cost( + cls, + statement: str, + database: Database, + catalog: str | None = None, + schema: str | None = None, + source: utils.QuerySource | None = None, + ) -> dict[str, Any]: """ Generate a SQL query that estimates the cost of a given statement. :param statement: A single SQL statement - :param cursor: Cursor instance + :param database: A database instance + :param catalog: The database catalog + :param schema: The database schema + :param source: Source of the query (eg, "sql_lab") :return: Dictionary with different costs """ raise Exception( # pylint: disable=broad-exception-raised @@ -1684,6 +1708,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods def estimate_query_cost( cls, database: Database, + catalog: str | None, schema: str, sql: str, source: utils.QuerySource | None = None, @@ -1692,9 +1717,11 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods Estimate the cost of a multiple statement SQL query. :param database: Database instance + :param catalog: Database catalog :param schema: Database schema :param sql: SQL query with possibly multiple statements :param source: Source of the query (eg, "sql_lab") + :return: List of dictionaries with different costs """ extra = database.get_extra() or {} if not cls.get_allow_cost_estimate(extra): @@ -1706,11 +1733,17 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods statements = parsed_query.get_statements() costs = [] - with database.get_raw_connection(schema=schema, source=source) as conn: - cursor = conn.cursor() - for statement in statements: - processed_statement = cls.process_statement(statement, database) - costs.append(cls.estimate_statement_cost(processed_statement, cursor)) + for statement in statements: + processed_statement = cls.process_statement(statement, database) + costs.append( + cls.estimate_statement_cost( + processed_statement, + database, + catalog, + schema, + source, + ) + ) return costs @@ -2131,7 +2164,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods @classmethod def get_public_information(cls) -> dict[str, Any]: """ - Construct a Dict with properties we want to expose. + Construct database properties we want to expose to the frontend. :returns: Dict with properties of our class like supports_file_upload and disable_ssh_tunneling @@ -2139,6 +2172,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods return { "supports_file_upload": cls.supports_file_upload, "disable_ssh_tunneling": cls.disable_ssh_tunneling, + "supports_dynamic_catalog": cls.supports_dynamic_catalog, } @classmethod diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py index a8d834276e..c4a3c8d5f1 100644 --- a/superset/db_engine_specs/bigquery.py +++ b/superset/db_engine_specs/bigquery.py @@ -378,7 +378,11 @@ class BigQueryEngineSpec(BaseEngineSpec): # pylint: disable=too-many-public-met raise SupersetException("The table schema must be defined") to_gbq_kwargs = {} - with cls.get_engine(database) as engine: + with cls.get_engine( + database=database, + catalog=table.catalog, + schema=table.schema, + ) as engine: to_gbq_kwargs = { "destination_table": str(table), "project_id": engine.url.host, @@ -419,6 +423,7 @@ class BigQueryEngineSpec(BaseEngineSpec): # pylint: disable=too-many-public-met def estimate_query_cost( cls, database: "Database", + catalog: Optional[str], schema: str, sql: str, source: Optional[utils.QuerySource] = None, @@ -427,6 +432,7 @@ class BigQueryEngineSpec(BaseEngineSpec): # pylint: disable=too-many-public-met Estimate the cost of a multiple statement SQL query. :param database: Database instance + :param catalog: Database project :param schema: Database schema :param sql: SQL query with possibly multiple statements :param source: Source of the query (eg, "sql_lab") @@ -467,8 +473,12 @@ class BigQueryEngineSpec(BaseEngineSpec): # pylint: disable=too-many-public-met return True @classmethod - def estimate_statement_cost(cls, statement: str, cursor: Any) -> dict[str, Any]: - with cls.get_engine(cursor) as engine: + def estimate_statement_cost( + cls, + statement: str, + database: "Database", + ) -> dict[str, Any]: + with cls.get_engine(database) as engine: client = cls._get_client(engine) job_config = bigquery.QueryJobConfig(dry_run=True) query_job = client.query( diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index ce87aa1f9b..962e8ce093 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -101,8 +101,6 @@ class PostgresBaseEngineSpec(BaseEngineSpec): engine = "" engine_name = "PostgreSQL" - supports_catalog = True - _time_grain_expressions = { None: "{col}", TimeGrain.SECOND: "DATE_TRUNC('second', {col})", @@ -199,7 +197,10 @@ class PostgresBaseEngineSpec(BaseEngineSpec): class PostgresEngineSpec(BasicParametersMixin, PostgresBaseEngineSpec): engine = "postgresql" engine_aliases = {"postgres"} + supports_dynamic_schema = True + supports_catalog = True + supports_dynamic_catalog = True default_driver = "psycopg2" sqlalchemy_uri_placeholder = ( @@ -296,6 +297,22 @@ class PostgresEngineSpec(BasicParametersMixin, PostgresBaseEngineSpec): return super().get_default_schema_for_query(database, query) + @classmethod + def adjust_engine_params( + cls, + uri: URL, + connect_args: dict[str, Any], + catalog: str | None = None, + schema: str | None = None, + ) -> tuple[URL, dict[str, Any]]: + """ + Set the catalog (database). + """ + if catalog: + uri = uri.set(database=catalog) + + return uri, connect_args + @classmethod def get_prequeries( cls, @@ -321,11 +338,24 @@ class PostgresEngineSpec(BasicParametersMixin, PostgresBaseEngineSpec): return True @classmethod - def estimate_statement_cost(cls, statement: str, cursor: Any) -> dict[str, Any]: + def estimate_statement_cost( + cls, + statement: str, + database: Database, + catalog: str | None = None, + schema: str | None = None, + source: utils.QuerySource | None = None, + ) -> dict[str, Any]: sql = f"EXPLAIN {statement}" - cursor.execute(sql) + with database.get_raw_connection( + catalog=catalog, + schema=schema, + source=source, + ) as conn: + cursor = conn.cursor() + cursor.execute(sql) + result = cursor.fetchone()[0] - result = cursor.fetchone()[0] match = re.search(r"cost=([\d\.]+)\.\.([\d\.]+)", result) if match: return { diff --git a/superset/extensions/metadb.py b/superset/extensions/metadb.py index bdfe1ae1e7..9667d79ec6 100644 --- a/superset/extensions/metadb.py +++ b/superset/extensions/metadb.py @@ -272,9 +272,6 @@ class SupersetShillelaghAdapter(Adapter): self.schema = parts.pop(-1) if parts else None self.catalog = parts.pop(-1) if parts else None - if self.catalog: - raise NotImplementedError("Catalogs are not currently supported") - # If the table has a single integer primary key we use that as the row ID in order # to perform updates and deletes. Otherwise we can only do inserts and selects. self._rowid: str | None = None @@ -316,7 +313,8 @@ class SupersetShillelaghAdapter(Adapter): # store this callable for later whenever we need an engine self.engine_context = partial( database.get_sqla_engine_with_context, - self.schema, + catalog=self.catalog, + schema=self.schema, ) # fetch column names and types diff --git a/superset/models/core.py b/superset/models/core.py index 92f6946f1e..dd7da0684b 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -382,8 +382,9 @@ class Database( ) @contextmanager - def get_sqla_engine_with_context( + def get_sqla_engine_with_context( # pylint: disable=too-many-arguments self, + catalog: str | None = None, schema: str | None = None, nullpool: bool = True, source: utils.QuerySource | None = None, @@ -419,6 +420,7 @@ class Database( ) yield self._get_sqla_engine( schema=schema, + catalog=catalog, nullpool=nullpool, source=source, sqlalchemy_uri=sqlalchemy_uri, @@ -426,6 +428,7 @@ class Database( def _get_sqla_engine( self, + catalog: str | None = None, schema: str | None = None, nullpool: bool = True, source: utils.QuerySource | None = None, @@ -461,7 +464,7 @@ class Database( sqlalchemy_url, connect_args = self.db_engine_spec.adjust_engine_params( uri=sqlalchemy_url, connect_args=connect_args, - catalog=None, + catalog=catalog, schema=schema, ) @@ -527,12 +530,16 @@ class Database( @contextmanager def get_raw_connection( self, + catalog: str | None = None, schema: str | None = None, nullpool: bool = True, source: utils.QuerySource | None = None, ) -> Connection: with self.get_sqla_engine_with_context( - schema=schema, nullpool=nullpool, source=source + catalog=catalog, + schema=schema, + nullpool=nullpool, + source=source, ) as engine: with closing(engine.raw_connection()) as conn: # pre-session queries are used to set the selected schema and, in the @@ -570,11 +577,15 @@ class Database( def get_df( # pylint: disable=too-many-locals self, sql: str, + catalog: str | None = None, schema: str | None = None, mutator: Callable[[pd.DataFrame], None] | None = None, ) -> pd.DataFrame: sqls = self.db_engine_spec.parse_sql(sql) - with self.get_sqla_engine_with_context(schema) as engine: + with self.get_sqla_engine_with_context( + catalog=catalog, + schema=schema, + ) as engine: engine_url = engine.url mutate_after_split = config["MUTATE_AFTER_SPLIT"] sql_query_mutator = config["SQL_QUERY_MUTATOR"] @@ -635,8 +646,16 @@ class Database( return df - def compile_sqla_query(self, qry: Select, schema: str | None = None) -> str: - with self.get_sqla_engine_with_context(schema) as engine: + def compile_sqla_query( + self, + qry: Select, + catalog: str | None = None, + schema: str | None = None, + ) -> str: + with self.get_sqla_engine_with_context( + catalog=catalog, + schema=schema, + ) as engine: sql = str(qry.compile(engine, compile_kwargs={"literal_binds": True})) # pylint: disable=protected-access @@ -648,6 +667,7 @@ class Database( def select_star( # pylint: disable=too-many-arguments self, table_name: str, + catalog: str | None = None, schema: str | None = None, limit: int = 100, show_cols: bool = False, @@ -656,7 +676,10 @@ class Database( cols: list[ResultSetColumnType] | None = None, ) -> str: """Generates a ``select *`` statement in the proper dialect""" - with self.get_sqla_engine_with_context(schema) as engine: + with self.get_sqla_engine_with_context( + catalog=catalog, + schema=schema, + ) as engine: return self.db_engine_spec.select_star( self, table_name, @@ -751,10 +774,15 @@ class Database( @contextmanager def get_inspector_with_context( - self, ssh_tunnel: SSHTunnel | None = None + self, + catalog: str | None = None, + schema: str | None = None, + ssh_tunnel: SSHTunnel | None = None, ) -> Inspector: with self.get_sqla_engine_with_context( - override_ssh_tunnel=ssh_tunnel + catalog=catalog, + schema=schema, + override_ssh_tunnel=ssh_tunnel, ) as engine: yield sqla.inspect(engine) @@ -832,10 +860,18 @@ class Database( def update_params_from_encrypted_extra(self, params: dict[str, Any]) -> None: self.db_engine_spec.update_params_from_encrypted_extra(self, params) - def get_table(self, table_name: str, schema: str | None = None) -> Table: + def get_table( + self, + table_name: str, + catalog: str | None, + schema: str | None = None, + ) -> Table: extra = self.get_extra() meta = MetaData(**extra.get("metadata_params", {})) - with self.get_sqla_engine_with_context() as engine: + with self.get_sqla_engine_with_context( + catalog=catalog, + schema=schema, + ) as engine: return Table( table_name, meta, @@ -939,11 +975,22 @@ class Database( return self.perm # type: ignore def has_table(self, table: Table) -> bool: - with self.get_sqla_engine_with_context() as engine: + with self.get_sqla_engine_with_context( + catalog=table.catalog, + schema=table.schema, + ) as engine: return engine.has_table(table.table_name, table.schema or None) - def has_table_by_name(self, table_name: str, schema: str | None = None) -> bool: - with self.get_sqla_engine_with_context() as engine: + def has_table_by_name( + self, + table_name: str, + catalog: str | None, + schema: str | None = None, + ) -> bool: + with self.get_sqla_engine_with_context( + catalog=catalog, + schema=schema, + ) as engine: return engine.has_table(table_name, schema) @classmethod @@ -961,10 +1008,21 @@ class Database( logger.warning("Has view failed", exc_info=True) return view_name in view_names - def has_view(self, view_name: str, schema: str | None = None) -> bool: - with self.get_sqla_engine_with_context(schema) as engine: + def has_view( + self, + view_name: str, + catalog: str | None, + schema: str | None = None, + ) -> bool: + with self.get_sqla_engine_with_context( + catalog=catalog, + schema=schema, + ) as engine: return engine.run_callable( - self._has_view, engine.dialect, view_name, schema + self._has_view, + engine.dialect, + view_name, + schema, ) def has_view_by_name(self, view_name: str, schema: str | None = None) -> bool: diff --git a/superset/models/dashboard.py b/superset/models/dashboard.py index 0a0d789c7a..9d29664dcc 100644 --- a/superset/models/dashboard.py +++ b/superset/models/dashboard.py @@ -214,13 +214,6 @@ class Dashboard(AuditMixinNullable, ImportExportMixin, Model): def charts(self) -> list[str]: return [slc.chart for slc in self.slices] - @property - def sqla_metadata(self) -> None: - # pylint: disable=no-member - with self.get_sqla_engine_with_context() as engine: - meta = MetaData(bind=engine) - meta.reflect() - @property def status(self) -> utils.DashboardStatus: if self.published: diff --git a/superset/models/sql_lab.py b/superset/models/sql_lab.py index 455dfd1aff..94e82217ac 100644 --- a/superset/models/sql_lab.py +++ b/superset/models/sql_lab.py @@ -107,6 +107,7 @@ class Query( status = Column(String(16), default=QueryStatus.PENDING) tab_name = Column(String(256)) sql_editor_id = Column(String(256)) + catalog = Column(String(256), nullable=True, default=None) schema = Column(String(256)) sql = Column(MediumText()) # Query to retrieve the results, @@ -384,6 +385,7 @@ class SavedQuery( id = Column(Integer, primary_key=True) user_id = Column(Integer, ForeignKey("ab_user.id"), nullable=True) db_id = Column(Integer, ForeignKey("dbs.id"), nullable=True) + catalog = Column(String(256), nullable=True, default=None) schema = Column(String(128)) label = Column(String(256)) description = Column(Text) @@ -472,6 +474,7 @@ class TabState(AuditMixinNullable, ExtraJSONMixin, Model): # selected DB and schema database_id = Column(Integer, ForeignKey("dbs.id", ondelete="CASCADE")) database = relationship("Database", foreign_keys=[database_id]) + catalog = Column(String(256), nullable=True, default=None) schema = Column(String(256)) # tables that are open in the schema browser and their data previews @@ -533,6 +536,7 @@ class TableSchema(AuditMixinNullable, ExtraJSONMixin, Model): Integer, ForeignKey("dbs.id", ondelete="CASCADE"), nullable=False ) database = relationship("Database", foreign_keys=[database_id]) + catalog = Column(String(256), nullable=True, default=None) schema = Column(String(256)) table = Column(String(256)) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index e34f7e2fde..b38a20cf4f 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -470,7 +470,11 @@ def execute_sql_statements( ) ) - with database.get_raw_connection(query.schema, source=QuerySource.SQL_LAB) as conn: + with database.get_raw_connection( + catalog=query.catalog, + schema=query.schema, + source=QuerySource.SQL_LAB, + ) as conn: # Sharing a single connection and cursor across the # execution of all statements (if many) cursor = conn.cursor() @@ -645,7 +649,9 @@ def cancel_query(query: Query) -> bool: return False with query.database.get_sqla_engine_with_context( - query.schema, source=QuerySource.SQL_LAB + catalog=query.catalog, + schema=query.schema, + source=QuerySource.SQL_LAB, ) as engine: with closing(engine.raw_connection()) as conn: with closing(conn.cursor()) as cursor: diff --git a/superset/sql_validators/presto_db.py b/superset/sql_validators/presto_db.py index 4852f70ee4..849cd688fc 100644 --- a/superset/sql_validators/presto_db.py +++ b/superset/sql_validators/presto_db.py @@ -145,7 +145,11 @@ class PrestoDBSQLValidator(BaseSQLValidator): @classmethod def validate( - cls, sql: str, schema: Optional[str], database: Database + cls, + sql: str, + catalog: Optional[str], + schema: Optional[str], + database: Database, ) -> list[SQLValidationAnnotation]: """ Presto supports query-validation queries by running them with a @@ -161,7 +165,9 @@ class PrestoDBSQLValidator(BaseSQLValidator): # todo(hughhh): update this to use new database.get_raw_connection() # this function keeps stalling CI with database.get_sqla_engine_with_context( - schema, source=QuerySource.SQL_LAB + catalog=catalog, + schema=schema, + source=QuerySource.SQL_LAB, ) as engine: # Sharing a single connection and cursor across the # execution of all statements (if many)
