fix: export/import catalogs (#28408)

This commit is contained in:
Beto Dealmeida 2024-05-09 14:42:03 -04:00 committed by GitHub
parent ba2cf5dbbc
commit e6a85c5901
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 105 additions and 2 deletions

View File

@ -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",

View File

@ -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,

View File

@ -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 "")
)

View File

@ -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",

View File

@ -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",

View File

@ -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

View File

@ -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"),
)

View File

@ -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:

View File

@ -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(