This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/master by this push:
     new e6a85c5901 fix: export/import catalogs (#28408)
e6a85c5901 is described below

commit e6a85c5901b05f86fb871977a97ab93eaf4bfbf3
Author: Beto Dealmeida <[email protected]>
AuthorDate: Thu May 9 14:42:03 2024 -0400

    fix: export/import catalogs (#28408)
---
 superset/connectors/sqla/models.py                 | 12 +++-
 superset/models/sql_lab.py                         |  4 ++
 superset/security/manager.py                       |  5 +-
 tests/integration_tests/datasets/commands_tests.py |  2 +
 .../queries/saved_queries/commands_tests.py        |  2 +
 tests/unit_tests/commands/report/base_test.py      |  2 +
 tests/unit_tests/connectors/sqla/models_test.py    | 76 ++++++++++++++++++++++
 tests/unit_tests/datasets/commands/export_test.py  |  2 +
 .../datasets/commands/importers/v1/import_test.py  |  2 +
 9 files changed, 105 insertions(+), 2 deletions(-)

diff --git a/superset/connectors/sqla/models.py 
b/superset/connectors/sqla/models.py
index ca1adba0ac..8c74dfd589 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -285,6 +285,11 @@ class BaseDatasource(AuditMixinNullable, 
ImportExportMixin):  # pylint: disable=
         """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 @@ class BaseDatasource(AuditMixinNullable, 
ImportExportMixin):  # pylint: disable=
             "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 @@ class BaseDatasource(AuditMixinNullable, 
ImportExportMixin):  # pylint: disable=
             "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",
diff --git a/superset/models/sql_lab.py b/superset/models/sql_lab.py
index 4c06867501..41647ea43b 100644
--- a/superset/models/sql_lab.py
+++ b/superset/models/sql_lab.py
@@ -169,6 +169,7 @@ class Query(
             "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 @@ class Query(
             "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 @@ class TableSchema(AuditMixinNullable, ExtraJSONMixin, 
Model):
             "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,
diff --git a/superset/security/manager.py b/superset/security/manager.py
index d28ed17893..4ee5e7a1db 100644
--- a/superset/security/manager.py
+++ b/superset/security/manager.py
@@ -486,7 +486,10 @@ class SupersetSecurityManager(  # pylint: 
disable=too-many-public-methods
         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 "")
         )
 
diff --git a/tests/integration_tests/datasets/commands_tests.py 
b/tests/integration_tests/datasets/commands_tests.py
index 405ea52c44..53bd7fa051 100644
--- a/tests/integration_tests/datasets/commands_tests.py
+++ b/tests/integration_tests/datasets/commands_tests.py
@@ -95,6 +95,7 @@ class TestExportDatasetsCommand(SupersetTestCase):
 
         assert metadata == {
             "cache_timeout": None,
+            "catalog": None,
             "columns": [
                 {
                     "column_name": "source",
@@ -224,6 +225,7 @@ class TestExportDatasetsCommand(SupersetTestCase):
             "default_endpoint",
             "offset",
             "cache_timeout",
+            "catalog",
             "schema",
             "sql",
             "params",
diff --git a/tests/integration_tests/queries/saved_queries/commands_tests.py 
b/tests/integration_tests/queries/saved_queries/commands_tests.py
index 8dac4e77f5..8babd7efb9 100644
--- a/tests/integration_tests/queries/saved_queries/commands_tests.py
+++ b/tests/integration_tests/queries/saved_queries/commands_tests.py
@@ -75,6 +75,7 @@ class TestExportSavedQueriesCommand(SupersetTestCase):
             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 @@ class TestExportSavedQueriesCommand(SupersetTestCase):
             contents["queries/examples/schema1/The_answer.yaml"]()
         )
         assert list(metadata.keys()) == [
+            "catalog",
             "schema",
             "label",
             "description",
diff --git a/tests/unit_tests/commands/report/base_test.py 
b/tests/unit_tests/commands/report/base_test.py
index 499682a1e6..871f3a511b 100644
--- a/tests/unit_tests/commands/report/base_test.py
+++ b/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
diff --git a/tests/unit_tests/connectors/sqla/models_test.py 
b/tests/unit_tests/connectors/sqla/models_test.py
index 687295baf8..c1e06f3755 100644
--- a/tests/unit_tests/connectors/sqla/models_test.py
+++ b/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"),
+    )
diff --git a/tests/unit_tests/datasets/commands/export_test.py 
b/tests/unit_tests/datasets/commands/export_test.py
index fbfa8d346c..9104e5b76e 100644
--- a/tests/unit_tests/datasets/commands/export_test.py
+++ b/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 @@ description: This is the description
 default_endpoint: null
 offset: -8
 cache_timeout: 3600
+catalog: public
 schema: my_schema
 sql: null
 params:
diff --git a/tests/unit_tests/datasets/commands/importers/v1/import_test.py 
b/tests/unit_tests/datasets/commands/importers/v1/import_test.py
index 511b60188a..6c306007d3 100644
--- a/tests/unit_tests/datasets/commands/importers/v1/import_test.py
+++ b/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(

Reply via email to