From 257f698ba03d5c57d3432ee7af2c7b22793b41cf Mon Sep 17 00:00:00 2001
From: Leah Tacke genannt Unterberg <leah.tgu@pads.rwth-aachen.de>
Date: Thu, 20 Mar 2025 10:28:33 +0100
Subject: [PATCH] fixed some superset definition generation stuff
---
.../superset/definition_bundles.py | 31 ++++++++++++++++++-
.../superset/definitions/core.py | 5 +--
.../superset/definitions/post_processing.py | 13 +++-----
.../superset/factories/dataset.py | 2 +-
.../superset/factories/mitm_dataset.py | 30 ++++++++++++++++--
.../superset/from_intermediate.py | 9 +++---
.../transformation/superset/from_sql.py | 12 ++++---
pyproject.toml | 2 +-
8 files changed, 77 insertions(+), 27 deletions(-)
diff --git a/mitm_tooling/transformation/superset/definition_bundles.py b/mitm_tooling/transformation/superset/definition_bundles.py
index b87ab43..bbea4c0 100644
--- a/mitm_tooling/transformation/superset/definition_bundles.py
+++ b/mitm_tooling/transformation/superset/definition_bundles.py
@@ -1,5 +1,6 @@
from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, Self
+from uuid import UUID
import pydantic
@@ -24,6 +25,14 @@ class SupersetDatasourceBundle(SupersetAssetBundle):
database: SupersetDatabaseDef
datasets: list[SupersetDatasetDef] = pydantic.Field(default_factory=list)
+ @property
+ def database_uuid(self) -> UUID:
+ return self.database.uuid
+
+ @property
+ def dataset_uuids(self) -> list[UUID]:
+ return [ds.uuid for ds in self.datasets]
+
@property
def placeholder_dataset_identifiers(self) -> dict[TableName, DatasourceIdentifier]:
return {ds.table_name: DatasourceIdentifier(uuid=ds.uuid) for ds in self.datasets}
@@ -36,6 +45,14 @@ class SupersetVisualizationBundle(SupersetAssetBundle):
charts: list[SupersetChartDef] = pydantic.Field(default_factory=list)
dashboards: list[SupersetDashboardDef] = pydantic.Field(default_factory=list)
+ @property
+ def chart_uuids(self) -> list[UUID]:
+ return [ch.uuid for ch in self.charts]
+
+ @property
+ def dashboard_uuids(self) -> list[UUID]:
+ return [da.uuid for da in self.dashboards]
+
def to_import(self) -> SupersetAssetsImport:
return mk_assets_import(charts=self.charts, dashboards=self.dashboards)
@@ -45,6 +62,18 @@ class SupersetMitMDatasetBundle(SupersetAssetBundle):
datasource_bundle: SupersetDatasourceBundle
visualization_bundle: SupersetVisualizationBundle = pydantic.Field(default_factory=SupersetVisualizationBundle)
+ def with_visualization_bundle(self, visualization_bundle: SupersetVisualizationBundle) -> Self:
+ mitm_ds = self.mitm_dataset
+ from mitm_tooling.transformation.superset.factories.mitm_dataset import mk_mitm_dataset
+ return self.__class__(
+ mitm_dataset=mk_mitm_dataset(name=mitm_ds.dataset_name, mitm=mitm_ds.mitm, uuid=mitm_ds.uuid,
+ database_uuid=self.datasource_bundle.database_uuid,
+ table_uuids=self.datasource_bundle.dataset_uuids,
+ slice_uuids=visualization_bundle.chart_uuids,
+ dashboard_uuids=visualization_bundle.dashboard_uuids),
+ datasource_bundle=self.datasource_bundle,
+ visualization_bundle=visualization_bundle)
+
def to_import(self) -> SupersetMitMDatasetImport:
base_assets = mk_assets_import(databases=[self.datasource_bundle.database],
datasets=self.datasource_bundle.datasets,
diff --git a/mitm_tooling/transformation/superset/definitions/core.py b/mitm_tooling/transformation/superset/definitions/core.py
index 3a31bda..7fad78f 100644
--- a/mitm_tooling/transformation/superset/definitions/core.py
+++ b/mitm_tooling/transformation/superset/definitions/core.py
@@ -16,10 +16,7 @@ class DatasourceIdentifier(FrozenSupersetDefinition):
class SupersetPostProcessing(pydantic.BaseModel, ABC):
- @pydantic.computed_field()
- @property
- def operation(self) -> str:
- raise NotImplementedError()
+ operation: str
class SupersetColumn(FrozenSupersetDefinition):
diff --git a/mitm_tooling/transformation/superset/definitions/post_processing.py b/mitm_tooling/transformation/superset/definitions/post_processing.py
index e7d7d41..6c00e40 100644
--- a/mitm_tooling/transformation/superset/definitions/post_processing.py
+++ b/mitm_tooling/transformation/superset/definitions/post_processing.py
@@ -13,9 +13,7 @@ class PivotOptions(FrozenSupersetDefinition):
class Pivot(SupersetPostProcessing):
- @property
- def operation(self) -> str:
- return 'pivot'
+ operation: Literal['pivot'] = 'pivot'
options: PivotOptions
@@ -27,14 +25,11 @@ class RenameOptions(FrozenSupersetDefinition):
class Rename(SupersetPostProcessing):
- @property
- def operation(self) -> str:
- return 'flatten'
+ operation: Literal['rename'] = 'rename'
options: RenameOptions
class Flatten(SupersetPostProcessing):
- @property
- def operation(self) -> str:
- return 'flatten'
+ operation: Literal['flatten'] = 'flatten'
+
diff --git a/mitm_tooling/transformation/superset/factories/dataset.py b/mitm_tooling/transformation/superset/factories/dataset.py
index 24183a0..477fe20 100644
--- a/mitm_tooling/transformation/superset/factories/dataset.py
+++ b/mitm_tooling/transformation/superset/factories/dataset.py
@@ -19,7 +19,7 @@ def mk_dataset(tm: TableMetaInfo, database_uuid: UUID, dialect: sa.Dialect | Non
cols.append(
mk_column(c, dt, dialect=dialect),
)
- if dt in {MITMDataType.Numeric, MITMDataType.Integer}:
+ if dt in {MITMDataType.Numeric}:
metrics.extend((
mk_metric(c, SupersetAggregate.AVG),
mk_metric(c, SupersetAggregate.SUM),
diff --git a/mitm_tooling/transformation/superset/factories/mitm_dataset.py b/mitm_tooling/transformation/superset/factories/mitm_dataset.py
index a55f4b7..fae7f79 100644
--- a/mitm_tooling/transformation/superset/factories/mitm_dataset.py
+++ b/mitm_tooling/transformation/superset/factories/mitm_dataset.py
@@ -1,18 +1,44 @@
+from typing import Sequence, Literal, Iterable
from uuid import UUID
from mitm_tooling.definition import MITM
from ..common import name_plus_uuid
from ..definitions import SupersetMitMDatasetDef
+from ..definitions.mitm_dataset import RelatedTable, RelatedSlice, RelatedDashboard
from ..factories.utils import mk_uuid
-def mk_mitm_dataset(name: str, mitm: MITM, database_uuid: UUID, uuid: UUID | None = None, uniquify_name: bool = False) -> SupersetMitMDatasetDef:
+def mk_related_obj(kind: Literal['table', 'slice', 'dashboard'],
+ uuid: UUID) -> RelatedTable | RelatedSlice | RelatedDashboard | None:
+ match kind:
+ case 'table':
+ return RelatedTable(table_uuid=uuid)
+ case 'slice':
+ return RelatedSlice(slice_uuid=uuid)
+ case 'dashboard':
+ return RelatedDashboard(dashboard_uuid=uuid)
+
+
+def mk_related_objs(kind: Literal['table', 'slice', 'dashboard'], uuids: Iterable[UUID]) -> Iterable[RelatedTable] | \
+ Iterable[RelatedSlice] | \
+ Iterable[
+ RelatedDashboard] | None:
+ if uuids:
+ return [mk_related_obj(kind, uuid) for uuid in uuids]
+
+
+def mk_mitm_dataset(name: str, mitm: MITM, database_uuid: UUID, table_uuids: list[UUID] | None = None,
+ slice_uuids: Sequence[UUID] | None = None, dashboard_uuids: Sequence[UUID] | None = None,
+ uuid: UUID | None = None, uniquify_name: bool = False) -> SupersetMitMDatasetDef:
uuid = uuid or mk_uuid()
if uniquify_name:
name = name_plus_uuid(name, uuid)
return SupersetMitMDatasetDef(
dataset_name=name,
mitm=mitm,
+ uuid=uuid or mk_uuid(),
database_uuid=database_uuid,
- uuid=uuid or mk_uuid()
+ tables=mk_related_objs('table', table_uuids),
+ slices=mk_related_objs('slice', slice_uuids),
+ dashboards=mk_related_objs('dashboard', dashboard_uuids),
)
diff --git a/mitm_tooling/transformation/superset/from_intermediate.py b/mitm_tooling/transformation/superset/from_intermediate.py
index d3b8615..83bb728 100644
--- a/mitm_tooling/transformation/superset/from_intermediate.py
+++ b/mitm_tooling/transformation/superset/from_intermediate.py
@@ -1,8 +1,7 @@
from mitm_tooling.representation import Header
-
from .common import SupersetDBConnectionInfo
-from .definitions.mitm_dataset import MitMDatasetIdentifier
from .definition_bundles import SupersetDatasourceBundle, SupersetMitMDatasetBundle, SupersetVisualizationBundle
+from .definitions.mitm_dataset import MitMDatasetIdentifier
from .mitm_specific import get_mitm_visualization_factory
@@ -28,6 +27,8 @@ def header_into_superset_mitm_dataset_bundle(header: Header,
db_meta = header_into_db_meta(header)
mitm_dataset_bundle = db_meta_into_mitm_dataset_bundle(db_meta, db_conn_info, dataset_identifier, header.mitm)
if include_visualizations:
- mitm_dataset_bundle.visualization_bundle = header_into_superset_visualization_bundle(header,
- mitm_dataset_bundle.datasource_bundle)
+ mitm_dataset_bundle = mitm_dataset_bundle.with_visualization_bundle(
+ header_into_superset_visualization_bundle(header,
+ mitm_dataset_bundle.datasource_bundle))
+
return mitm_dataset_bundle
diff --git a/mitm_tooling/transformation/superset/from_sql.py b/mitm_tooling/transformation/superset/from_sql.py
index 1751e7d..359b16c 100644
--- a/mitm_tooling/transformation/superset/from_sql.py
+++ b/mitm_tooling/transformation/superset/from_sql.py
@@ -2,9 +2,9 @@ from mitm_tooling.extraction.sql.data_models import DBMetaInfo
from mitm_tooling.extraction.sql.data_models.db_meta import DBMetaInfoBase
from mitm_tooling.extraction.sql.db import connect_and_reflect
from .common import SupersetDBConnectionInfo
-from .definitions.mitm_dataset import MitMDatasetIdentifier
from .common import _mk_engine, SQLiteFileOrEngine
from .definition_bundles import SupersetDatasourceBundle, SupersetMitMDatasetBundle
+from .definitions.mitm_dataset import MitMDatasetIdentifier
from .factories.database import mk_database
from .factories.dataset import mk_dataset
from .factories.mitm_dataset import mk_mitm_dataset
@@ -33,8 +33,9 @@ def db_meta_into_mitm_dataset_bundle(db_meta: DBMetaInfoBase,
dataset_identifier: MitMDatasetIdentifier,
mitm: MITM) -> SupersetMitMDatasetBundle:
datasource_bundle = db_meta_into_superset_datasource_bundle(db_meta, db_conn_info)
- db_uuid = datasource_bundle.database.uuid
- mitm_dataset = mk_mitm_dataset(dataset_identifier.dataset_name, mitm, db_uuid, uuid=dataset_identifier.uuid)
+ mitm_dataset = mk_mitm_dataset(dataset_identifier.dataset_name, mitm, uuid=dataset_identifier.uuid,
+ database_uuid=datasource_bundle.database_uuid,
+ table_uuids=datasource_bundle.dataset_uuids)
return SupersetMitMDatasetBundle(mitm_dataset=mitm_dataset, datasource_bundle=datasource_bundle)
@@ -53,6 +54,7 @@ def db_into_mitm_dataset_bundle(arg: SQLiteFileOrEngine,
dataset_identifier: MitMDatasetIdentifier,
mitm: MITM) -> SupersetMitMDatasetBundle:
datasource_bundle = db_into_superset_datasource_bundle(arg, db_conn_info)
- db_uuid = datasource_bundle.database.uuid
- mitm_dataset = mk_mitm_dataset(dataset_identifier.dataset_name, mitm, db_uuid, uuid=dataset_identifier.uuid)
+ mitm_dataset = mk_mitm_dataset(dataset_identifier.dataset_name, mitm, uuid=dataset_identifier.uuid,
+ database_uuid=datasource_bundle.database_uuid,
+ table_uuids=datasource_bundle.dataset_uuids)
return SupersetMitMDatasetBundle(mitm_dataset=mitm_dataset, datasource_bundle=datasource_bundle)
diff --git a/pyproject.toml b/pyproject.toml
index c5ad1fa..118d222 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[tool.poetry]
name = "mitm-tooling"
-version = "0.4.6"
+version = "0.4.7"
description = ""
authors = ["Leah Tacke genannt Unterberg <leah.tgu@pads.rwth-aachen.de>"]
readme = "README.md"
--
GitLab