fix: handle empty catalog when DB supports them (#29840)

betodealmeida · sadpandajoe · commit 9677fa97ff30 · 2024-08-13T13:23:17.000-07:00
(cherry picked from commit 39209c2)
diff --git a/scripts/change_detector.py b/scripts/change_detector.py
@@ -52,7 +52,7 @@
 def fetch_files_github_api(url: str):  # type: ignore
     """Fetches data using GitHub API."""
     req = Request(url)
-    req.add_header("Authorization", f"token {GITHUB_TOKEN}")
+    req.add_header("Authorization", f"Bearer {GITHUB_TOKEN}")
     req.add_header("Accept", "application/vnd.github.v3+json")
 
     print(f"Fetching from {url}")
diff --git a/superset/cachekeys/schemas.py b/superset/cachekeys/schemas.py
@@ -32,6 +32,10 @@ class Datasource(Schema):
     datasource_name = fields.String(
         metadata={"description": datasource_name_description},
     )
+    catalog = fields.String(
+        allow_none=True,
+        metadata={"description": "Datasource catalog"},
+    )
     schema = fields.String(
         metadata={"description": "Datasource schema"},
     )
diff --git a/superset/commands/dataset/create.py b/superset/commands/dataset/create.py
@@ -54,23 +54,28 @@ def run(self) -> Model:
     def validate(self) -> None:
         exceptions: list[ValidationError] = []
         database_id = self._properties["database"]
-        schema = self._properties.get("schema")
         catalog = self._properties.get("catalog")
+        schema = self._properties.get("schema")
+        table_name = self._properties["table_name"]
         sql = self._properties.get("sql")
         owner_ids: Optional[list[int]] = self._properties.get("owners")
 
-        table = Table(self._properties["table_name"], schema, catalog)
-
-        # Validate uniqueness
-        if not DatasetDAO.validate_uniqueness(database_id, table):
-            exceptions.append(DatasetExistsValidationError(table))
-
         # Validate/Populate database
         database = DatasetDAO.get_database_by_id(database_id)
         if not database:
             exceptions.append(DatabaseNotFoundValidationError())
         self._properties["database"] = database
 
+        # Validate uniqueness
+        if database:
+            if not catalog:
+                catalog = self._properties["catalog"] = database.get_default_catalog()
+
+            table = Table(table_name, schema, catalog)
+
+            if not DatasetDAO.validate_uniqueness(database, table):
+                exceptions.append(DatasetExistsValidationError(table))
+
         # Validate table exists on dataset if sql is not provided
         # This should be validated when the dataset is physical
         if (
diff --git a/superset/commands/dataset/importers/v1/utils.py b/superset/commands/dataset/importers/v1/utils.py
@@ -166,7 +166,7 @@ def import_dataset(
 
     try:
         table_exists = dataset.database.has_table(
-            Table(dataset.table_name, dataset.schema),
+            Table(dataset.table_name, dataset.schema, dataset.catalog),
         )
     except Exception:  # pylint: disable=broad-except
         # MySQL doesn't play nice with GSheets table names
diff --git a/superset/commands/dataset/update.py b/superset/commands/dataset/update.py
@@ -79,10 +79,12 @@ def run(self) -> Model:
     def validate(self) -> None:
         exceptions: list[ValidationError] = []
         owner_ids: Optional[list[int]] = self._properties.get("owners")
+
         # Validate/populate model exists
         self._model = DatasetDAO.find_by_id(self._model_id)
         if not self._model:
             raise DatasetNotFoundError()
+
         # Check ownership
         try:
             security_manager.raise_for_ownership(self._model)
@@ -91,22 +93,30 @@ def validate(self) -> None:
 
         database_id = self._properties.get("database")
 
+        catalog = self._properties.get("catalog")
+        if not catalog:
+            catalog = self._properties["catalog"] = (
+                self._model.database.get_default_catalog()
+            )
+
         table = Table(
             self._properties.get("table_name"),  # type: ignore
             self._properties.get("schema"),
-            self._properties.get("catalog"),
+            catalog,
         )
 
         # Validate uniqueness
         if not DatasetDAO.validate_update_uniqueness(
-            self._model.database_id,
+            self._model.database,
             table,
             self._model_id,
         ):
             exceptions.append(DatasetExistsValidationError(table))
+
         # Validate/Populate database not allowed to change
         if database_id and database_id != self._model:
             exceptions.append(DatabaseChangeValidationError())
+
         # Validate/Populate owner
         try:
             owners = self.compute_owners(
@@ -116,6 +126,7 @@ def validate(self) -> None:
             self._properties["owners"] = owners
         except ValidationError as ex:
             exceptions.append(ex)
+
         # Validate columns
         if columns := self._properties.get("columns"):
             self._validate_columns(columns, exceptions)
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
@@ -461,9 +461,11 @@ def data_for_slices(  # pylint: disable=too-many-locals
                 )
             else:
                 _columns = [
-                    utils.get_column_name(column_)
-                    if utils.is_adhoc_column(column_)
-                    else column_
+                    (
+                        utils.get_column_name(column_)
+                        if utils.is_adhoc_column(column_)
+                        else column_
+                    )
                     for column_param in COLUMN_FORM_DATA_PARAMS
                     for column_ in utils.as_list(form_data.get(column_param) or [])
                 ]
@@ -1963,7 +1965,7 @@ class and any keys added via `ExtraCache`.
         if self.has_extra_cache_key_calls(query_obj):
             sqla_query = self.get_sqla_query(**query_obj)
             extra_cache_keys += sqla_query.extra_cache_keys
-        return extra_cache_keys
+        return list(set(extra_cache_keys))
 
     @property
     def quote_identifier(self) -> Callable[[str], str]:
diff --git a/superset/daos/dataset.py b/superset/daos/dataset.py
@@ -84,15 +84,19 @@ def validate_table_exists(
 
     @staticmethod
     def validate_uniqueness(
-        database_id: int,
+        database: Database,
         table: Table,
         dataset_id: int | None = None,
     ) -> bool:
+        # The catalog might not be set even if the database supports catalogs, in case
+        # multi-catalog is disabled.
+        catalog = table.catalog or database.get_default_catalog()
+
         dataset_query = db.session.query(SqlaTable).filter(
             SqlaTable.table_name == table.table,
             SqlaTable.schema == table.schema,
-            SqlaTable.catalog == table.catalog,
-            SqlaTable.database_id == database_id,
+            SqlaTable.catalog == catalog,
+            SqlaTable.database_id == database.id,
         )
 
         if dataset_id:
@@ -103,15 +107,19 @@ def validate_uniqueness(
 
     @staticmethod
     def validate_update_uniqueness(
-        database_id: int,
+        database: Database,
         table: Table,
         dataset_id: int,
     ) -> bool:
+        # The catalog might not be set even if the database supports catalogs, in case
+        # multi-catalog is disabled.
+        catalog = table.catalog or database.get_default_catalog()
+
         dataset_query = db.session.query(SqlaTable).filter(
             SqlaTable.table_name == table.table,
-            SqlaTable.database_id == database_id,
+            SqlaTable.database_id == database.id,
             SqlaTable.schema == table.schema,
-            SqlaTable.catalog == table.catalog,
+            SqlaTable.catalog == catalog,
             SqlaTable.id != dataset_id,
         )
         return not db.session.query(dataset_query.exists()).scalar()
diff --git a/superset/databases/api.py b/superset/databases/api.py
@@ -1159,7 +1159,7 @@ def select_star(
         self.incr_stats("init", self.select_star.__name__)
         try:
             result = database.select_star(
-                Table(table_name, schema_name),
+                Table(table_name, schema_name, database.get_default_catalog()),
                 latest_partition=True,
             )
         except NoSuchTableError:
diff --git a/superset/jinja_context.py b/superset/jinja_context.py
@@ -565,7 +565,7 @@ def process_template(self, sql: str, **kwargs: Any) -> str:
         """
         Makes processing a template a noop
         """
-        return sql
+        return str(sql)
 
 
 class PrestoTemplateProcessor(JinjaTemplateProcessor):
diff --git a/superset/models/core.py b/superset/models/core.py
@@ -490,7 +490,7 @@ def _get_sqla_engine(  # pylint: disable=too-many-locals
                 g.user.id,
                 self.db_engine_spec,
             )
-            if hasattr(g, "user") and hasattr(g.user, "id") and oauth2_config
+            if oauth2_config and hasattr(g, "user") and hasattr(g.user, "id")
             else None
         )
         # If using MySQL or Presto for example, will set url.username
diff --git a/superset/models/slice.py b/superset/models/slice.py
@@ -369,6 +369,7 @@ def set_related_perm(_mapper: Mapper, _connection: Connection, target: Slice) ->
         ds = db.session.query(src_class).filter_by(id=int(id_)).first()
         if ds:
             target.perm = ds.perm
+            target.catalog_perm = ds.catalog_perm
             target.schema_perm = ds.schema_perm
 
 
diff --git a/superset/security/manager.py b/superset/security/manager.py
@@ -774,6 +774,9 @@ def get_schemas_accessible_by_user(
         # pylint: disable=import-outside-toplevel
         from superset.connectors.sqla.models import SqlaTable
 
+        default_catalog = database.get_default_catalog()
+        catalog = catalog or default_catalog
+
         if hierarchical and (
             self.can_access_database(database)
             or (catalog and self.can_access_catalog(database, catalog))
@@ -783,7 +786,6 @@ def get_schemas_accessible_by_user(
         # schema_access
         accessible_schemas: set[str] = set()
         schema_access = self.user_view_menu_names("schema_access")
-        default_catalog = database.get_default_catalog()
         default_schema = database.get_default_schema(default_catalog)
 
         for perm in schema_access:
@@ -800,7 +802,7 @@ def get_schemas_accessible_by_user(
             # [database].[catalog].[schema] matches when the catalog is equal to the
             # requested catalog or, when no catalog specified, it's equal to the default
             # catalog.
-            elif len(parts) == 3 and parts[1] == (catalog or default_catalog):
+            elif len(parts) == 3 and parts[1] == catalog:
                 accessible_schemas.add(parts[2])
 
         # datasource_access
@@ -906,16 +908,16 @@ def get_datasources_accessible_by_user(  # pylint: disable=invalid-name
         if self.can_access_database(database):
             return datasource_names
 
+        catalog = catalog or database.get_default_catalog()
         if catalog:
             catalog_perm = self.get_catalog_perm(database.database_name, catalog)
             if catalog_perm and self.can_access("catalog_access", catalog_perm):
                 return datasource_names
 
         if schema:
-            default_catalog = database.get_default_catalog()
             schema_perm = self.get_schema_perm(
                 database.database_name,
-                catalog or default_catalog,
+                catalog,
                 schema,
             )
             if schema_perm and self.can_access("schema_access", schema_perm):
@@ -2183,6 +2185,7 @@ def raise_for_access(
                 database = query.database
 
             database = cast("Database", database)
+            default_catalog = database.get_default_catalog()
 
             if self.can_access_database(database):
                 return
@@ -2196,19 +2199,19 @@ def raise_for_access(
                 # from the SQLAlchemy URI if possible; if not, we use the SQLAlchemy
                 # inspector to read it.
                 default_schema = database.get_default_schema_for_query(query)
-                # Determining the default catalog is much easier, because DB engine
-                # specs need explicit support for catalogs.
-                default_catalog = database.get_default_catalog()
                 tables = {
                     Table(
                         table_.table,
                         table_.schema or default_schema,
-                        table_.catalog or default_catalog,
+                        table_.catalog or query.catalog or default_catalog,
                     )
                     for table_ in extract_tables_from_jinja_sql(query.sql, database)
                 }
             elif table:
-                tables = {table}
+                # Make sure table has the default catalog, if not specified.
+                tables = {
+                    Table(table.table, table.schema, table.catalog or default_catalog)
+                }
 
             denied = set()
 
diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py
@@ -284,6 +284,7 @@ def export_csv(self, client_id: str) -> CsvResponse:
             "client_id": client_id,
             "row_count": row_count,
             "database": query.database.name,
+            "catalog": query.catalog,
             "schema": query.schema,
             "sql": query.sql,
             "exported_format": "csv",
diff --git a/superset/sqllab/sqllab_execution_context.py b/superset/sqllab/sqllab_execution_context.py
@@ -125,6 +125,8 @@ def select_as_cta(self) -> bool:
     def set_database(self, database: Database) -> None:
         self._validate_db(database)
         self.database = database
+        if self.catalog is None:
+            self.catalog = database.get_default_catalog()
         if self.select_as_cta:
             schema_name = self._get_ctas_target_schema_name(database)
             self.create_table_as_select.target_schema_name = schema_name  # type: ignore
diff --git a/superset/views/sql_lab/views.py b/superset/views/sql_lab/views.py
@@ -239,13 +239,15 @@ def post(self) -> FlaskResponse:
         db.session.query(TableSchema).filter(
             TableSchema.tab_state_id == table["queryEditorId"],
             TableSchema.database_id == table["dbId"],
+            TableSchema.catalog == table["catalog"],
             TableSchema.schema == table["schema"],
             TableSchema.table == table["name"],
         ).delete(synchronize_session=False)
 
         table_schema = TableSchema(
             tab_state_id=table["queryEditorId"],
             database_id=table["dbId"],
+            catalog=table["catalog"],
             schema=table["schema"],
             table=table["name"],
             description=json.dumps(table),
diff --git a/tests/integration_tests/databases/api_tests.py b/tests/integration_tests/databases/api_tests.py
@@ -1563,34 +1563,6 @@ def test_get_select_star_not_allowed(self):
         rv = self.client.get(uri)
         self.assertEqual(rv.status_code, 404)
 
-    def test_get_select_star_datasource_access(self):
-        """
-        Database API: Test get select star with datasource access
-        """
-        table = SqlaTable(
-            schema="main", table_name="ab_permission", database=get_main_database()
-        )
-        db.session.add(table)
-        db.session.commit()
-
-        tmp_table_perm = security_manager.find_permission_view_menu(
-            "datasource_access", table.get_perm()
-        )
-        gamma_role = security_manager.find_role("Gamma")
-        security_manager.add_permission_role(gamma_role, tmp_table_perm)
-
-        self.login(GAMMA_USERNAME)
-        main_db = get_main_database()
-        uri = f"api/v1/database/{main_db.id}/select_star/ab_permission/"
-        rv = self.client.get(uri)
-        self.assertEqual(rv.status_code, 200)
-
-        # rollback changes
-        security_manager.del_permission_role(gamma_role, tmp_table_perm)
-        db.session.delete(table)
-        db.session.delete(main_db)
-        db.session.commit()
-
     def test_get_select_star_not_found_database(self):
         """
         Database API: Test get select star not found database
diff --git a/tests/integration_tests/datasets/api_tests.py b/tests/integration_tests/datasets/api_tests.py
diff --git a/tests/integration_tests/db_engine_specs/postgres_tests.py b/tests/integration_tests/db_engine_specs/postgres_tests.py
diff --git a/tests/integration_tests/security_tests.py b/tests/integration_tests/security_tests.py
diff --git a/tests/integration_tests/sqla_models_tests.py b/tests/integration_tests/sqla_models_tests.py
diff --git a/tests/unit_tests/connectors/sqla/models_test.py b/tests/unit_tests/connectors/sqla/models_test.py
diff --git a/tests/unit_tests/dao/dataset_test.py b/tests/unit_tests/dao/dataset_test.py
diff --git a/tests/unit_tests/security/manager_test.py b/tests/unit_tests/security/manager_test.py

Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,10 @@ class Datasource(Schema):`
`32`	`32`	`datasource_name = fields.String(`
`33`	`33`	`metadata={"description": datasource_name_description},`
`34`	`34`	`)`
	`35`	`+ catalog = fields.String(`
	`36`	`+ allow_none=True,`
	`37`	`+ metadata={"description": "Datasource catalog"},`
	`38`	`+ )`
`35`	`39`	`schema = fields.String(`
`36`	`40`	`metadata={"description": "Datasource schema"},`
`37`	`41`	`)`
Original file line number	Diff line number	Diff line change
`@@ -166,7 +166,7 @@ def import_dataset(`
`166`	`166`
`167`	`167`	`try:`
`168`	`168`	`table_exists = dataset.database.has_table(`
`169`		`- Table(dataset.table_name, dataset.schema),`
	`169`	`+ Table(dataset.table_name, dataset.schema, dataset.catalog),`
`170`	`170`	`)`
`171`	`171`	`except Exception: # pylint: disable=broad-except`
`172`	`172`	`# MySQL doesn't play nice with GSheets table names`
Original file line number	Diff line number	Diff line change
`@@ -1159,7 +1159,7 @@ def select_star(`
`1159`	`1159`	`self.incr_stats("init", self.select_star.__name__)`
`1160`	`1160`	`try:`
`1161`	`1161`	`result = database.select_star(`
`1162`		`- Table(table_name, schema_name),`
	`1162`	`+ Table(table_name, schema_name, database.get_default_catalog()),`
`1163`	`1163`	`latest_partition=True,`
`1164`	`1164`	`)`
`1165`	`1165`	`except NoSuchTableError:`
Original file line number	Diff line number	Diff line change
`@@ -490,7 +490,7 @@ def _get_sqla_engine( # pylint: disable=too-many-locals`
`490`	`490`	`g.user.id,`
`491`	`491`	`self.db_engine_spec,`
`492`	`492`	`)`
`493`		`- if hasattr(g, "user") and hasattr(g.user, "id") and oauth2_config`
	`493`	`+ if oauth2_config and hasattr(g, "user") and hasattr(g.user, "id")`
`494`	`494`	`else None`
`495`	`495`	`)`
`496`	`496`	`# If using MySQL or Presto for example, will set url.username`