fix: export/import catalogs (#28408)

superset/connectors/sqla/models.py

@@ -285,6 +285,11 @@ def connection(self) -> str | None:
         """String representing the context of the Datasource"""
         return None
 
+    @property
+    def catalog(self) -> str | None:
+        """String representing the catalog of the Datasource (if it applies)"""
+        return None
+
     @property
     def schema(self) -> str | None:
         """String representing the schema of the Datasource (if it applies)"""
@@ -330,6 +335,7 @@ def short_data(self) -> dict[str, Any]:
             "edit_url": self.url,
             "id": self.id,
             "uid": self.uid,
+            "catalog": self.catalog,
             "schema": self.schema or None,
             "name": self.name,
             "type": self.type,
@@ -384,6 +390,7 @@ def data(self) -> dict[str, Any]:
             "datasource_name": self.datasource_name,
             "table_name": self.datasource_name,
             "type": self.type,
+            "catalog": self.catalog,
             "schema": self.schema or None,
             "offset": self.offset,
             "cache_timeout": self.cache_timeout,
@@ -1135,7 +1142,9 @@ class SqlaTable(
     # The reason it does not physically exist is MySQL, PostgreSQL, etc. have a
     # different interpretation of uniqueness when it comes to NULL which is problematic
     # given the schema is optional.
-    __table_args__ = (UniqueConstraint("database_id", "schema", "table_name"),)
+    __table_args__ = (
+        UniqueConstraint("database_id", "catalog", "schema", "table_name"),
+    )
 
     table_name = Column(String(250), nullable=False)
     main_dttm_col = Column(String(250))
@@ -1166,6 +1175,7 @@ class SqlaTable(
         "database_id",
         "offset",
         "cache_timeout",
+        "catalog",
         "schema",
         "sql",
         "params",

superset/models/sql_lab.py

@@ -169,6 +169,7 @@ def to_dict(self) -> dict[str, Any]:
             "limitingFactor": self.limiting_factor,
             "progress": self.progress,
             "rows": self.rows,
+            "catalog": self.catalog,
             "schema": self.schema,
             "ctas": self.select_as_cta,
             "serverId": self.id,
@@ -251,6 +252,7 @@ def data(self) -> dict[str, Any]:
             "owners": self.owners_data,
             "database": {"id": self.database_id, "backend": self.database.backend},
             "order_by_choices": order_by_choices,
+            "catalog": self.catalog,
             "schema": self.schema,
             "verbose_map": {},
         }
@@ -415,6 +417,7 @@ class SavedQuery(
 
     export_parent = "database"
     export_fields = [
+        "catalog",
         "schema",
         "label",
         "description",
@@ -557,6 +560,7 @@ def to_dict(self) -> dict[str, Any]:
             "id": self.id,
             "tab_state_id": self.tab_state_id,
             "database_id": self.database_id,
+            "catalog": self.catalog,
             "schema": self.schema,
             "table": self.table,
             "description": description,

superset/security/manager.py

@@ -486,7 +486,10 @@ def can_access_schema(self, datasource: "BaseDatasource") -> bool:
         return (
             self.can_access_all_datasources()
             or self.can_access_database(datasource.database)
-            or self.can_access_catalog(datasource.database, datasource.catalog)
+            or (
+                datasource.catalog
+                and self.can_access_catalog(datasource.database, datasource.catalog)
+            )
             or self.can_access("schema_access", datasource.schema_perm or "")
         )
 

tests/integration_tests/datasets/commands_tests.py

@@ -95,6 +95,7 @@ def test_export_dataset_command(self, mock_g):
 
         assert metadata == {
             "cache_timeout": None,
+            "catalog": None,
             "columns": [
                 {
                     "column_name": "source",
@@ -224,6 +225,7 @@ def test_export_dataset_command_key_order(self, mock_g):
             "default_endpoint",
             "offset",
             "cache_timeout",
+            "catalog",
             "schema",
             "sql",
             "params",

tests/integration_tests/queries/saved_queries/commands_tests.py

@@ -75,6 +75,7 @@ def test_export_query_command(self, mock_g):
             contents["queries/examples/schema1/The_answer.yaml"]()
         )
         assert metadata == {
+            "catalog": None,
             "schema": "schema1",
             "label": "The answer",
             "description": "Answer to the Ultimate Question of Life, the Universe, and Everything",
@@ -134,6 +135,7 @@ def test_export_query_command_key_order(self, mock_g):
             contents["queries/examples/schema1/The_answer.yaml"]()
         )
         assert list(metadata.keys()) == [
+            "catalog",
             "schema",
             "label",
             "description",

tests/unit_tests/commands/report/base_test.py

@@ -15,6 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from __future__ import annotations
+
 import logging
 from datetime import timedelta
 from functools import wraps

tests/unit_tests/connectors/sqla/models_test.py

@@ -18,10 +18,14 @@
 import pytest
 from pytest_mock import MockerFixture
 from sqlalchemy import create_engine
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm.session import Session
 
 from superset.connectors.sqla.models import SqlaTable
+from superset.daos.dataset import DatasetDAO
 from superset.exceptions import OAuth2RedirectError
 from superset.models.core import Database
+from superset.sql_parse import Table
 from superset.superset_typing import QueryObjectDict
 
 
@@ -187,3 +191,75 @@ def test_query_datasources_by_permissions_with_catalog_schema(
         "tables.schema_perm IN ('[my_db].[db1].[schema1]', '[my_other_db].[schema]') OR "
         "tables.catalog_perm IN ('[my_db].[db1]')"
     )
+
+
+def test_dataset_uniqueness(session: Session) -> None:
+    """
+    Test dataset uniqueness constraints.
+    """
+    Database.metadata.create_all(session.bind)
+
+    database = Database(database_name="my_db", sqlalchemy_uri="sqlite://")
+
+    # add prod.schema.table
+    dataset = SqlaTable(
+        database=database,
+        catalog="prod",
+        schema="schema",
+        table_name="table",
+    )
+    session.add(dataset)
+    session.commit()
+
+    # add dev.schema.table
+    dataset = SqlaTable(
+        database=database,
+        catalog="dev",
+        schema="schema",
+        table_name="table",
+    )
+    session.add(dataset)
+    session.commit()
+
+    # try to add dev.schema.table again, fails
+    dataset = SqlaTable(
+        database=database,
+        catalog="dev",
+        schema="schema",
+        table_name="table",
+    )
+    session.add(dataset)
+    with pytest.raises(IntegrityError):
+        session.commit()
+    session.rollback()
+
+    # add schema.table
+    dataset = SqlaTable(
+        database=database,
+        catalog=None,
+        schema="schema",
+        table_name="table",
+    )
+    session.add(dataset)
+    session.commit()
+
+    # add schema.table again, works because in SQL `NULlL != NULL`
+    dataset = SqlaTable(
+        database=database,
+        catalog=None,
+        schema="schema",
+        table_name="table",
+    )
+    session.add(dataset)
+    session.commit()
+
+    # but the DAO enforces application logic for uniqueness
+    assert not DatasetDAO.validate_uniqueness(
+        database.id,
+        Table("table", "schema", None),
+    )
+
+    assert DatasetDAO.validate_uniqueness(
+        database.id,
+        Table("table", "schema", "some_catalog"),
+    )

tests/unit_tests/datasets/commands/export_test.py

@@ -68,6 +68,7 @@ def test_export(session: Session) -> None:
         description="This is the description",
         is_featured=1,
         cache_timeout=3600,
+        catalog="public",
         schema="my_schema",
         sql=None,
         params=json.dumps(
@@ -111,6 +112,7 @@ def test_export(session: Session) -> None:
 default_endpoint: null
 offset: -8
 cache_timeout: 3600
+catalog: public
 schema: my_schema
 sql: null
 params:

tests/unit_tests/datasets/commands/importers/v1/import_test.py

@@ -61,6 +61,7 @@ def test_import_dataset(mocker: MockFixture, session: Session) -> None:
         "default_endpoint": None,
         "offset": -8,
         "cache_timeout": 3600,
+        "catalog": "public",
         "schema": "my_schema",
         "sql": None,
         "params": {
@@ -115,6 +116,7 @@ def test_import_dataset(mocker: MockFixture, session: Session) -> None:
     assert sqla_table.default_endpoint is None
     assert sqla_table.offset == -8
     assert sqla_table.cache_timeout == 3600
+    assert sqla_table.catalog == "public"
     assert sqla_table.schema == "my_schema"
     assert sqla_table.sql is None
     assert sqla_table.params == json.dumps(