diff --git a/mitm_tooling/extraction/sql/db/db_reflection.py b/mitm_tooling/extraction/sql/db/db_reflection.py
index 205b40124359bd6e269d3a62eea9048f3d56ba8b..a26bf30a20cf9502aeec425f43a71624a8833151 100644
--- a/mitm_tooling/extraction/sql/db/db_reflection.py
+++ b/mitm_tooling/extraction/sql/db/db_reflection.py
@@ -1,6 +1,7 @@
from collections.abc import Collection
from typing import TypedDict, Any
+import pydantic
from sqlalchemy import MetaData, Engine, inspect
from ..data_models import TableMetaInfo
@@ -12,7 +13,7 @@ class TableDoesNotExist(Exception):
pass
-class AdditionalMeta(TypedDict):
+class AdditionalMeta(pydantic.BaseModel):
default_schema: SchemaName
diff --git a/mitm_tooling/representation/__init__.py b/mitm_tooling/representation/__init__.py
index dc87f676aef137a43e05cf61b89a03c03989ffcd..bb63420b1cf468fc045a0dafa66c27d6430d8a8f 100644
--- a/mitm_tooling/representation/__init__.py
+++ b/mitm_tooling/representation/__init__.py
@@ -8,3 +8,4 @@ from .df_representation import MITMDataFrames
from .file_representation import write_header_file, write_data_file, read_data_file, read_header_file
from .intermediate_representation import HeaderEntry, Header, MITMData, StreamingMITMData, StreamingConceptData
from .sql_representation import mk_sql_rep_schema, insert_mitm_data, mk_sqlite, SQLRepresentationSchema
+from .sql_representation import TableName, SchemaName, QualifiedTableName, Queryable
diff --git a/mitm_tooling/representation/sql_representation.py b/mitm_tooling/representation/sql_representation.py
index 462b9a30d953257b308ad70ed6628b006cc7f760..f0495beb1e8c2606b6cb78a7ac6cbef4d53adb0d 100644
--- a/mitm_tooling/representation/sql_representation.py
+++ b/mitm_tooling/representation/sql_representation.py
@@ -188,23 +188,40 @@ def _gen_denormalized_views(mitm: MITM, concept_tables: ConceptTablesDict, type_
TableName, Queryable], None, None]:
mitm_def = get_mitm_def(mitm)
- for concept in mitm_def.main_concepts:
- view_name = mk_concept_table_name(mitm, concept) + '_view'
- q = None
- if (concept_t := concept_tables.get(concept)) is not None:
+ for main_concept in mitm_def.main_concepts:
+ for concept in mitm_def.get_leafs(main_concept):
+ view_name = mk_concept_table_name(mitm, concept) + '_denormalized_view'
+ q = None
if has_type_tables(mitm, concept):
- shared_cols = {c.name for c in concept_t.columns}
- if concept_type_tables := type_tables.get(concept):
- selections = []
- for type_name, type_t in concept_type_tables.items():
- selection = (c if c.name in shared_cols else sa.label(_prefix_col_name(type_name, c.name), c)
- for c in type_t.columns)
- selections.append(sa.select(*selection))
- q = sa.union(*selections).subquery()
+ selections = []
+
+ for leaf_concept in mitm_def.get_leafs(concept):
+ if concept_type_tables := type_tables.get(leaf_concept):
+ col_sets = [{(c.name, str(c.type)) for c in t.columns} for t in concept_type_tables.values()]
+ shared_cols = set.intersection(*col_sets)
+ all_cols = set.union(*col_sets)
+
+ for type_name, type_t in concept_type_tables.items():
+ selection = []
+ for (col_name, col_type) in all_cols:
+ if (c := type_t.columns.get(col_name)) is not None and str(c.type) == col_type:
+ selection.append(c)
+ else:
+ selection.append(sa.null().label(col_name))
+
+ # selection = (c if (c.name, str(c.type)) in shared_cols else sa.label(_prefix_col_name(type_name, c.name), c)
+ # for c in type_t.columns)
+ selections.append(sa.select(*selection))
+
+ if selections:
+ q = sa.union_all(*selections).subquery()
else:
- q = sa.select(concept_t)
- if q is not None:
- yield view_name, q
+ if (concept_t := concept_tables.get(concept)) is not None:
+ # base_cols = {(c.name, str(c.type)) for c in concept_t.columns}
+ q = sa.select(concept_t)
+
+ if q is not None:
+ yield view_name, q
for parent_concept, subs in mitm_def.sub_concept_map.items():
if (concept_t := concept_tables.get(parent_concept)) is not None:
@@ -221,8 +238,9 @@ _view_generators: tuple[MitMDBViewsGenerator, ...] = (_gen_denormalized_views,)
def mk_sql_rep_schema(header: Header,
view_generators: Iterable[MitMDBViewsGenerator] | None = (_gen_denormalized_views,),
override_schema: SchemaName | None = None) -> SQLRepresentationSchema:
+ schema_name = override_schema if override_schema else SQL_REPRESENTATION_DEFAULT_SCHEMA
mitm_def = get_mitm_def(header.mitm)
- meta = sa.MetaData(schema=override_schema if override_schema else SQL_REPRESENTATION_DEFAULT_SCHEMA)
+ meta = sa.MetaData(schema=schema_name)
concept_tables: ConceptTablesDict = {}
type_tables: ConceptTypeTablesDict = {}
@@ -247,7 +265,7 @@ def mk_sql_rep_schema(header: Header,
mitm_def.resolve_foreign_types(concept).items() for name, dt in
resolved_fk.items()]
}, additional_column_generators=(_gen_within_concept_id_col,), schema_item_generators=(
- _gen_unique_constraint, _gen_pk_constraint, _gen_index,), override_schema=override_schema)
+ _gen_unique_constraint, _gen_pk_constraint, _gen_index,), override_schema=schema_name)
concept_tables[concept] = t
for he in header.header_entries:
@@ -276,7 +294,7 @@ def mk_sql_rep_schema(header: Header,
schema_item_generators=(
_gen_unique_constraint, _gen_pk_constraint, _gen_index,
_gen_foreign_key_constraints),
- override_schema=override_schema)
+ override_schema=schema_name)
if he_concept not in type_tables:
type_tables[he_concept] = {}
@@ -285,7 +303,7 @@ def mk_sql_rep_schema(header: Header,
if view_generators is not None:
for generator in view_generators:
for name, queryable in generator(header.mitm, concept_tables, type_tables):
- views[name] = create_view(name, queryable, meta)
+ views[name] = create_view(name, queryable, meta) # TODO make `create_view` schema-aware and add `schema=schema_name`
return SQLRepresentationSchema(meta=meta, concept_tables=concept_tables, type_tables=type_tables, views=views)
diff --git a/mitm_tooling/transformation/superset/__init__.py b/mitm_tooling/transformation/superset/__init__.py
index c5cb6390f1406c92392c78e769fda18d64a85a2c..03c0f2a80d338d27df50dd97e9e2070005eee786 100644
--- a/mitm_tooling/transformation/superset/__init__.py
+++ b/mitm_tooling/transformation/superset/__init__.py
@@ -1,5 +1,5 @@
from . import definitions, factories, mitm_specific
from . import exporting, from_sql, from_intermediate
from . import interface
-from .exporting import write_superset_assets_def
-from .interface import mk_superset_datasource_import, mk_superset_visualization_import, mk_superset_mitm_dataset_import
+from .exporting import write_superset_import_as_zip
+from .interface import mk_superset_datasource_bundle, mk_superset_visualization_bundle, mk_superset_mitm_dataset_bundle
diff --git a/mitm_tooling/transformation/superset/common.py b/mitm_tooling/transformation/superset/common.py
index 1d5f717d8c8a0b83d9a152884a4687b0cf7022c6..00af82b3c1ca1c041077f797675a655ec8e327e0 100644
--- a/mitm_tooling/transformation/superset/common.py
+++ b/mitm_tooling/transformation/superset/common.py
@@ -2,7 +2,7 @@ from uuid import UUID
import pydantic
import sqlalchemy as sa
-from mitm_tooling.representation.sql.common import SchemaName
+from mitm_tooling.representation import SchemaName
from mitm_tooling.representation.sql_representation import SQL_REPRESENTATION_DEFAULT_SCHEMA
from .definitions import StrUrl
from .factories.utils import mk_short_uuid_str
@@ -22,9 +22,12 @@ class SupersetDBConnectionInfo(pydantic.BaseModel):
schema_name: SchemaName = SQL_REPRESENTATION_DEFAULT_SCHEMA
@property
- def db_name_in_uri(self) -> AnyUrl | None:
- if self.sql_alchemy_uri is not None:
- return any_url_into_sa_url(self.sql_alchemy_uri).database
+ def sa_url(self) -> sa.URL:
+ return any_url_into_sa_url(self.sql_alchemy_uri)
+
+ @property
+ def db_name_in_uri(self) -> str:
+ return self.sa_url.database
@property
def db_name(self) -> str:
@@ -35,8 +38,8 @@ class SupersetDBConnectionInfo(pydantic.BaseModel):
return dialect_cls_from_url(self.sql_alchemy_uri)
-def name_plus_uuid(name: str, uuid: UUID | None = None) -> str:
- return f'{name}-{mk_short_uuid_str(uuid)}'
+def name_plus_uuid(name: str, uuid: UUID | None = None, sep: str = '-') -> str:
+ return f'{name}{sep}{mk_short_uuid_str(uuid)}'
def _mk_engine(arg: SQLiteFileOrEngine) -> sa.Engine:
diff --git a/mitm_tooling/transformation/superset/definition_bundles.py b/mitm_tooling/transformation/superset/definition_bundles.py
index 3869657b39608f5ea9060531107c10f38205e760..8e42b248d5111d4add3e5ed6311f43fa0a725f27 100644
--- a/mitm_tooling/transformation/superset/definition_bundles.py
+++ b/mitm_tooling/transformation/superset/definition_bundles.py
@@ -4,20 +4,20 @@ from typing import Any
import pydantic
from .definitions import SupersetDatabaseDef, SupersetMitMDatasetDef, \
- SupersetChartDef, SupersetDashboardDef, BaseSupersetDefinition, SupersetAssetsDef, SupersetDatasetDef, \
- ExtendedSupersetAssetsDef, SupersetDefFolder, DatasourceIdentifier
-from .factories.assets import mk_assets, mk_extended_assets
-from ...representation.sql.common import TableName
+ SupersetChartDef, SupersetDashboardDef, SupersetAssetsImport, SupersetDatasetDef, \
+ SupersetMitMDatasetImport, SupersetDefFolder, DatasourceIdentifier
+from .factories.importable import mk_assets_import, mk_mitm_dataset_import
+from ...representation import TableName
class SupersetAssetBundle(SupersetDefFolder, ABC):
@abstractmethod
- def to_assets(self) -> SupersetAssetsDef | ExtendedSupersetAssetsDef:
+ def to_import(self) -> SupersetAssetsImport | SupersetMitMDatasetImport:
pass
@property
def folder_dict(self) -> dict[str, Any]:
- return self.to_assets().folder_dict
+ return self.to_import().folder_dict
class SupersetDatasourceBundle(SupersetAssetBundle):
@@ -28,16 +28,16 @@ class SupersetDatasourceBundle(SupersetAssetBundle):
def placeholder_dataset_identifiers(self) -> dict[TableName, DatasourceIdentifier]:
return {ds.table_name: DatasourceIdentifier(dataset_uuid=ds.uuid) for ds in self.datasets}
- def to_assets(self) -> SupersetAssetsDef:
- return mk_assets(databases=[self.database], datasets=self.datasets)
+ def to_import(self) -> SupersetAssetsImport:
+ return mk_assets_import(databases=[self.database], datasets=self.datasets)
class SupersetVisualizationBundle(SupersetAssetBundle):
charts: list[SupersetChartDef] = pydantic.Field(default_factory=list)
dashboards: list[SupersetDashboardDef] = pydantic.Field(default_factory=list)
- def to_assets(self) -> SupersetAssetsDef:
- return mk_assets(charts=self.charts, dashboards=self.dashboards)
+ def to_import(self) -> SupersetAssetsImport:
+ return mk_assets_import(charts=self.charts, dashboards=self.dashboards)
class SupersetMitMDatasetBundle(SupersetAssetBundle):
@@ -45,9 +45,9 @@ class SupersetMitMDatasetBundle(SupersetAssetBundle):
datasource_bundle: SupersetDatasourceBundle
visualization_bundle: SupersetVisualizationBundle = pydantic.Field(default_factory=SupersetVisualizationBundle)
- def to_assets(self) -> ExtendedSupersetAssetsDef:
- base_assets = mk_assets(databases=[self.datasource_bundle.database],
- datasets=self.datasource_bundle.datasets,
- charts=self.visualization_bundle.charts,
- dashboards=self.visualization_bundle.dashboards)
- return mk_extended_assets(mitm_datasets=[self.mitm_dataset], base_assets=base_assets)
+ def to_import(self) -> SupersetMitMDatasetImport:
+ base_assets = mk_assets_import(databases=[self.datasource_bundle.database],
+ datasets=self.datasource_bundle.datasets,
+ charts=self.visualization_bundle.charts,
+ dashboards=self.visualization_bundle.dashboards)
+ return mk_mitm_dataset_import(mitm_datasets=[self.mitm_dataset], base_assets=base_assets)
diff --git a/mitm_tooling/transformation/superset/definitions/__init__.py b/mitm_tooling/transformation/superset/definitions/__init__.py
index f9ac7e553f6e90ca8463efc22a910fda1caa7a62..77e9f60906d7b569ae6bc8268f5ad714c6ffbcdc 100644
--- a/mitm_tooling/transformation/superset/definitions/__init__.py
+++ b/mitm_tooling/transformation/superset/definitions/__init__.py
@@ -1,5 +1,8 @@
-from .assets import *
-from .charts import *
from .constants import *
from .core import *
+from .database import *
+from .dataset import *
+from .chart import *
+from .dashboard import *
+from .importable import *
from .post_processing import *
diff --git a/mitm_tooling/transformation/superset/definitions/assets.py b/mitm_tooling/transformation/superset/definitions/assets.py
deleted file mode 100644
index b9afac1fd925c2f8fae571e421c957eff4ef110c..0000000000000000000000000000000000000000
--- a/mitm_tooling/transformation/superset/definitions/assets.py
+++ /dev/null
@@ -1,195 +0,0 @@
-from abc import abstractmethod
-from collections import defaultdict
-from datetime import UTC
-
-from mitm_tooling.definition import MITM
-from .charts import *
-from .dashboard import DashboardPositionData, DashboardMetadata
-from .post_processing import *
-
-
-class MetadataType(StrEnum):
- Database = 'Database'
- SqlaTable = 'SqlaTable'
- Slice = 'Slice'
- Chart = 'Chart'
- Dashboard = 'Dashboard'
- Asset = 'Asset'
- MitMDataset = 'MitMDataset'
-
-
-class SupersetDefFile(BaseSupersetDefinition, ABC):
-
- @property
- @abstractmethod
- def filename(self) -> str:
- pass
-
-
-class SupersetDefFolder(BaseSupersetDefinition, ABC):
-
- @property
- @abstractmethod
- def folder_dict(self) -> dict[str, Any]:
- pass
-
-
-class SupersetMetadataDef(SupersetDefFile):
- version: str = '1.0.0'
- type: MetadataType = MetadataType.SqlaTable
- timestamp: StrDatetime = pydantic.Field(default_factory=lambda: datetime.now(UTC))
-
- @property
- def filename(self) -> str:
- return 'metadata'
-
-
-class SupersetDatabaseDef(SupersetDefFile):
- database_name: str
- sqlalchemy_uri: StrUrl
- uuid: StrUUID
- # verbose_name : str | None = None
- cache_timeout: str | None = None
- expose_in_sqllab: bool = True
- allow_run_async: bool = False
- allow_ctas: bool = False
- allow_cvas: bool = False
- allow_dml: bool = False
- allow_file_upload: bool = False
- extra: dict[str, Any] = pydantic.Field(default_factory=lambda: {
- 'allows_virtual_table_explore': True
- })
- impersonate_user: bool = False
- version: str = '1.0.0'
- ssh_tunnel: None = None
-
- @property
- def filename(self):
- return self.database_name
-
-
-class SupersetDatasetDef(SupersetDefFile):
- table_name: str
- schema_name: str = pydantic.Field(alias='schema')
- uuid: StrUUID
- database_uuid: StrUUID
- main_dttm_col: str | None = None
- description: str | None = None
- default_endpoint: str | None = None
- offset: int = 0
- cache_timeout: str | None = None
- catalog: str | None = None
- sql: str | None = None
- params: Any = None
- template_params: Any = None
- filter_select_enabled: bool = True
- fetch_values_predicate: str | None = None
- extra: dict[str, Any] = pydantic.Field(default_factory=dict)
- normalize_columns: bool = False
- always_filter_main_dttm: bool = False
- metrics: list[SupersetMetric] = pydantic.Field(default_factory=list)
- columns: list[SupersetColumn] = pydantic.Field(default_factory=list)
- version: str = '1.0.0'
-
- @property
- def filename(self):
- return self.table_name
-
-
-class SupersetChartDef(SupersetDefFile):
- uuid: StrUUID
- slice_name: str
- viz_type: SupersetVizType
- dataset_uuid: StrUUID
- description: str | None = None
- certified_by: str | None = None
- certification_details: str | None = None
- params: ChartParams | None = None
- query_context: Annotated[pydantic.Json | QueryContext | None, pydantic.PlainSerializer(
- lambda x: x.model_dump_json(by_alias=True, exclude_none=True, serialize_as_any=True) if isinstance(x,
- pydantic.BaseModel) else x,
- return_type=pydantic.Json), pydantic.Field(default=None)]
- cache_timeout: int | None = None
- version: str = '1.0.0'
- is_managed_externally: bool = False
- external_url: StrUrl | None = None
-
- @property
- def filename(self) -> str:
- return f'{self.slice_name}_{self.dataset_uuid}'
-
-
-class SupersetDashboardDef(SupersetDefFile):
- uuid: StrUUID
- dashboard_title: str
- position: DashboardPositionData
- metadata: DashboardMetadata
- description: str | None = None
- css: str | None = None
- slug: str | None = None
- is_managed_externally: bool | None = False
- external_url: StrUrl | None = None
- certified_by: str | None = None
- certification_details: str | None = None
- published: bool | None = False
- version: str = '1.0.0'
-
- @property
- def filename(self) -> str:
- return f'{self.dashboard_title}_{self.uuid}'
-
-
-class SupersetMitMDatasetDef(SupersetDefFile):
- uuid: StrUUID
- dataset_name: str
- mitm: MITM
- database_uuid: StrUUID
- version: str = '1.0.0'
-
- @property
- def filename(self) -> str:
- return self.dataset_name
-
-
-class SupersetAssetsDef(SupersetDefFolder):
- databases: list[SupersetDatabaseDef] | None = None
- datasets: list[SupersetDatasetDef] | None = None
- charts: list[SupersetChartDef] | None = None
- dashboards: list[SupersetDashboardDef] | None = None
- metadata: SupersetMetadataDef = pydantic.Field(default_factory=SupersetMetadataDef)
-
- @property
- def folder_dict(self) -> dict[str, Any]:
- folder_dict = {'.': self.metadata}
- dbs = {}
- if self.databases:
- dbs |= {db.uuid: db.database_name for db in self.databases}
- folder_dict['databases'] = [db for db in self.databases]
- if self.datasets:
- db_dss = defaultdict(list)
- for ds in self.datasets:
- db_dss[dbs[ds.database_uuid]].append(ds)
- folder_dict['datasets'] = db_dss
- if self.charts:
- folder_dict['charts'] = self.charts
- if self.dashboards:
- folder_dict['dashboards'] = self.dashboards
- return {'my_import': folder_dict}
-
-
-class ExtendedSupersetAssetsDef(SupersetDefFolder):
- mitm_datasets: list[SupersetMitMDatasetDef] | None
- base_assets: SupersetAssetsDef | None
-
- @property
- def folder_dict(self) -> dict[str, Any]:
- asset_folder_dict = self.base_assets.folder_dict if self.base_assets else {'my_import': {}}
- dbs = {}
- if self.base_assets.databases:
- dbs = {db.uuid: db.database_name for db in self.base_assets.databases}
- if self.mitm_datasets:
- mitm_dss = defaultdict(list)
- for mitm_ds in self.mitm_datasets:
- mitm_dss[dbs[mitm_ds.database_uuid]].append(mitm_ds)
- asset_folder_dict['my_import']['mitm_datasets'] = mitm_dss
- return asset_folder_dict
diff --git a/mitm_tooling/transformation/superset/definitions/charts.py b/mitm_tooling/transformation/superset/definitions/chart.py
similarity index 69%
rename from mitm_tooling/transformation/superset/definitions/charts.py
rename to mitm_tooling/transformation/superset/definitions/chart.py
index ce8af2bc838f77181608420a3ab30c0f7945ef02..d8bc00d8737f4fd81dbf014a429aa614902c2709 100644
--- a/mitm_tooling/transformation/superset/definitions/charts.py
+++ b/mitm_tooling/transformation/superset/definitions/chart.py
@@ -1,6 +1,13 @@
+from typing import Annotated
+
+import pydantic
+
+from . import SupersetDefFile, StrUUID, SupersetVizType, QueryContext, StrUrl
from .core import *
+
+
class ChartParams(FormData):
datasource: str | DatasourceIdentifier
viz_type: SupersetVizType
@@ -72,3 +79,29 @@ class TimeSeriesLineParams(TimeSeriesChartParams):
opacity: float = 0.2
markerSize: int = 6
seriesType: str = 'line'
+
+
+class SupersetChartDef(SupersetDefFile):
+ uuid: StrUUID
+ slice_name: str
+ viz_type: SupersetVizType
+ dataset_uuid: StrUUID
+ description: str | None = None
+ certified_by: str | None = None
+ certification_details: str | None = None
+ params: ChartParams | None = None
+ query_context: QueryContext | None = None
+
+ # query_context: Annotated[pydantic.Json | QueryContext | None, pydantic.PlainSerializer(
+ # lambda x: x.model_dump_json(by_alias=True, serialize_as_any=True, exclude_none=True) if isinstance(x,
+ # pydantic.BaseModel) else x,
+ # return_type=pydantic.Json), pydantic.Field(default=None)]
+
+ cache_timeout: int | None = None
+ version: str = '1.0.0'
+ is_managed_externally: bool = False
+ external_url: StrUrl | None = None
+
+ @property
+ def filename(self) -> str:
+ return f'{self.slice_name}_{self.dataset_uuid}'
diff --git a/mitm_tooling/transformation/superset/definitions/constants.py b/mitm_tooling/transformation/superset/definitions/constants.py
index dbe2a73d64bcaa279e1d72bfc79573f041d75db6..786d525a0885a0a0e2b6030072db164687ec7324 100644
--- a/mitm_tooling/transformation/superset/definitions/constants.py
+++ b/mitm_tooling/transformation/superset/definitions/constants.py
@@ -1,3 +1,4 @@
+from abc import ABC, abstractmethod
from datetime import datetime
from datetime import datetime
from enum import StrEnum, IntEnum
@@ -190,3 +191,11 @@ class BaseSupersetDefinition(pydantic.BaseModel):
class FrozenSupersetDefinition(BaseSupersetDefinition):
model_config = ConfigDict(arbitrary_types_allowed=True, use_enum_values=True, populate_by_name=True, frozen=True)
+
+
+class SupersetDefFile(BaseSupersetDefinition, ABC):
+
+ @property
+ @abstractmethod
+ def filename(self) -> str:
+ pass
diff --git a/mitm_tooling/transformation/superset/definitions/core.py b/mitm_tooling/transformation/superset/definitions/core.py
index 0122835721f73e97b10386945ab25f477562c6a0..c53eeda10142a8b076f2a832bc43a37b5023e601 100644
--- a/mitm_tooling/transformation/superset/definitions/core.py
+++ b/mitm_tooling/transformation/superset/definitions/core.py
@@ -1,4 +1,3 @@
-from abc import ABC
from typing import Any
from mitm_tooling.representation import ColumnName
@@ -13,7 +12,7 @@ class SupersetPostProcessing(pydantic.BaseModel, ABC):
class DatasourceIdentifier(FrozenSupersetDefinition):
- id: SupersetId = '-1' # -1 as a placeholder
+ id: SupersetId = -1 # -1 as a placeholder
type: Literal['table', 'annotation'] = 'table'
dataset_uuid: StrUUID = pydantic.Field(exclude=True)
@@ -175,12 +174,12 @@ class FormData(BaseSupersetDefinition):
class QueryContext(BaseSupersetDefinition):
- model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
-
datasource: DatasourceIdentifier
queries: list[QueryObject] = pydantic.Field(default_factory=list)
- form_data: FormData | None = pydantic.Field(default=None)
+ form_data: FormData | None = pydantic.Field(default=None, strict=False)
result_type: ChartDataResultType = ChartDataResultType.FULL
result_format: ChartDataResultFormat = ChartDataResultFormat.JSON
force: bool = False
custom_cache_timeout: int | None = None
+
+
diff --git a/mitm_tooling/transformation/superset/definitions/dashboard.py b/mitm_tooling/transformation/superset/definitions/dashboard.py
index a2263ed2caf00b93050e462e0cb188747d21c745..a2c6b18286c16030bb84de78667a21be61c15c0b 100644
--- a/mitm_tooling/transformation/superset/definitions/dashboard.py
+++ b/mitm_tooling/transformation/superset/definitions/dashboard.py
@@ -3,6 +3,7 @@ from typing import Literal
import pydantic
+from . import SupersetDefFile, StrUUID, StrUrl
from .constants import StrUUID
DashboardInternalID = str
@@ -111,3 +112,23 @@ class DashboardMetadata(pydantic.BaseModel):
color_scheme: str = 'blueToGreen'
cross_filters_enabled: bool = True
native_filter_configuration: list[NativeFilterConfig] = pydantic.Field(default_factory=list)
+
+
+class SupersetDashboardDef(SupersetDefFile):
+ uuid: StrUUID
+ dashboard_title: str
+ position: DashboardPositionData
+ metadata: DashboardMetadata
+ description: str | None = None
+ css: str | None = None
+ slug: str | None = None
+ is_managed_externally: bool | None = False
+ external_url: StrUrl | None = None
+ certified_by: str | None = None
+ certification_details: str | None = None
+ published: bool | None = False
+ version: str = '1.0.0'
+
+ @property
+ def filename(self) -> str:
+ return f'{self.dashboard_title}_{self.uuid}'
diff --git a/mitm_tooling/transformation/superset/definitions/database.py b/mitm_tooling/transformation/superset/definitions/database.py
new file mode 100644
index 0000000000000000000000000000000000000000..32bc282eefd9a25625aa9c52d505389310585e65
--- /dev/null
+++ b/mitm_tooling/transformation/superset/definitions/database.py
@@ -0,0 +1,29 @@
+from typing import Any
+
+import pydantic
+
+from mitm_tooling.transformation.superset.definitions import SupersetDefFile, StrUrl, StrUUID
+
+
+class SupersetDatabaseDef(SupersetDefFile):
+ database_name: str
+ sqlalchemy_uri: StrUrl
+ uuid: StrUUID
+ # verbose_name : str | None = None
+ cache_timeout: str | None = None
+ expose_in_sqllab: bool = True
+ allow_run_async: bool = False
+ allow_ctas: bool = False
+ allow_cvas: bool = False
+ allow_dml: bool = False
+ allow_file_upload: bool = False
+ extra: dict[str, Any] = pydantic.Field(default_factory=lambda: {
+ 'allows_virtual_table_explore': True
+ })
+ impersonate_user: bool = False
+ version: str = '1.0.0'
+ ssh_tunnel: None = None
+
+ @property
+ def filename(self):
+ return self.database_name
diff --git a/mitm_tooling/transformation/superset/definitions/dataset.py b/mitm_tooling/transformation/superset/definitions/dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..150ea5953a112063be3a3911b504fdd5a9997609
--- /dev/null
+++ b/mitm_tooling/transformation/superset/definitions/dataset.py
@@ -0,0 +1,33 @@
+from typing import Any
+
+import pydantic
+
+from mitm_tooling.transformation.superset.definitions import SupersetDefFile, StrUUID, SupersetMetric, SupersetColumn
+
+
+class SupersetDatasetDef(SupersetDefFile):
+ table_name: str
+ schema_name: str = pydantic.Field(alias='schema')
+ uuid: StrUUID
+ database_uuid: StrUUID
+ main_dttm_col: str | None = None
+ description: str | None = None
+ default_endpoint: str | None = None
+ offset: int = 0
+ cache_timeout: str | None = None
+ catalog: str | None = None
+ sql: str | None = None
+ params: Any = None
+ template_params: Any = None
+ filter_select_enabled: bool = True
+ fetch_values_predicate: str | None = None
+ extra: dict[str, Any] = pydantic.Field(default_factory=dict)
+ normalize_columns: bool = False
+ always_filter_main_dttm: bool = False
+ metrics: list[SupersetMetric] = pydantic.Field(default_factory=list)
+ columns: list[SupersetColumn] = pydantic.Field(default_factory=list)
+ version: str = '1.0.0'
+
+ @property
+ def filename(self):
+ return self.table_name
diff --git a/mitm_tooling/transformation/superset/definitions/importable.py b/mitm_tooling/transformation/superset/definitions/importable.py
new file mode 100644
index 0000000000000000000000000000000000000000..c52f5a2ba659d17bc51cc55b894a81e652c97417
--- /dev/null
+++ b/mitm_tooling/transformation/superset/definitions/importable.py
@@ -0,0 +1,83 @@
+from collections import defaultdict
+from datetime import UTC
+
+from .chart import *
+from .dashboard import SupersetDashboardDef
+from .database import SupersetDatabaseDef
+from .dataset import SupersetDatasetDef
+from .mitm_dataset import SupersetMitMDatasetDef
+from .post_processing import *
+
+
+class MetadataType(StrEnum):
+ Database = 'Database'
+ SqlaTable = 'SqlaTable'
+ Slice = 'Slice'
+ Chart = 'Chart'
+ Dashboard = 'Dashboard'
+ Asset = 'Asset'
+ MitMDataset = 'MitMDataset'
+
+
+class SupersetDefFolder(BaseSupersetDefinition, ABC):
+
+ @property
+ @abstractmethod
+ def folder_dict(self) -> dict[str, Any]:
+ pass
+
+
+class SupersetMetadataDef(SupersetDefFile):
+ type: MetadataType
+ version: str = '1.0.0'
+ timestamp: StrDatetime = pydantic.Field(default_factory=lambda: datetime.now(UTC))
+
+ @property
+ def filename(self) -> str:
+ return 'metadata'
+
+
+class SupersetAssetsImport(SupersetDefFolder):
+ databases: list[SupersetDatabaseDef] | None = None
+ datasets: list[SupersetDatasetDef] | None = None
+ charts: list[SupersetChartDef] | None = None
+ dashboards: list[SupersetDashboardDef] | None = None
+ metadata: SupersetMetadataDef = pydantic.Field(default_factory=SupersetMetadataDef)
+
+ @property
+ def folder_dict(self) -> dict[str, Any]:
+ folder_dict = {'.': self.metadata}
+ dbs = {}
+ if self.databases:
+ dbs |= {db.uuid: db.database_name for db in self.databases}
+ folder_dict['databases'] = [db for db in self.databases]
+ if self.datasets:
+ db_dss = defaultdict(list)
+ for ds in self.datasets:
+ db_dss[dbs[ds.database_uuid]].append(ds)
+ folder_dict['datasets'] = db_dss
+ if self.charts:
+ folder_dict['charts'] = self.charts
+ if self.dashboards:
+ folder_dict['dashboards'] = self.dashboards
+ return {'my_import': folder_dict}
+
+
+class SupersetMitMDatasetImport(SupersetDefFolder):
+ mitm_datasets: list[SupersetMitMDatasetDef] | None
+ base_assets: SupersetAssetsImport | None
+ metadata: SupersetMetadataDef = pydantic.Field(default_factory=lambda : SupersetMetadataDef(type=MetadataType.MitMDataset))
+
+ @property
+ def folder_dict(self) -> dict[str, Any]:
+ asset_folder_dict = self.base_assets.folder_dict if self.base_assets else {'my_import': {}}
+ asset_folder_dict['my_import']['.'] = self.metadata
+ dbs = {}
+ if self.base_assets.databases:
+ dbs = {db.uuid: db.database_name for db in self.base_assets.databases}
+ if self.mitm_datasets:
+ mitm_dss = defaultdict(list)
+ for mitm_ds in self.mitm_datasets:
+ mitm_dss[dbs[mitm_ds.database_uuid]].append(mitm_ds)
+ asset_folder_dict['my_import']['mitm_datasets'] = mitm_dss
+ return asset_folder_dict
diff --git a/mitm_tooling/transformation/superset/definitions/mitm_dataset.py b/mitm_tooling/transformation/superset/definitions/mitm_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a259eb649e341838924b337f5d8777bae5fc42b
--- /dev/null
+++ b/mitm_tooling/transformation/superset/definitions/mitm_dataset.py
@@ -0,0 +1,32 @@
+from mitm_tooling.definition import MITM
+from mitm_tooling.representation import Header
+from mitm_tooling.transformation.superset.definitions import SupersetDefFile, StrUUID, BaseSupersetDefinition, \
+ SupersetId
+
+
+class RelatedTable(BaseSupersetDefinition):
+ table_id : SupersetId | None = None
+ table_uuid : StrUUID
+
+class RelatedSlice(BaseSupersetDefinition):
+ slice_id : SupersetId | None = None
+ slice_uuid : StrUUID
+
+class RelatedDashboard(BaseSupersetDefinition):
+ dashboard_id : SupersetId | None = None
+ dashboard_uuid : StrUUID
+
+class SupersetMitMDatasetDef(SupersetDefFile):
+ uuid: StrUUID
+ dataset_name: str
+ mitm: MITM
+ mitm_header: Header | None = None
+ database_uuid: StrUUID
+ tables: list[RelatedTable] | None = None
+ slices: list[RelatedSlice] | None = None
+ dashboards: list[RelatedDashboard] | None = None
+ version: str = '1.0.0'
+
+ @property
+ def filename(self) -> str:
+ return self.dataset_name
diff --git a/mitm_tooling/transformation/superset/exporting.py b/mitm_tooling/transformation/superset/exporting.py
index 52dca2901982890ee8f18c31eb560830775acbbc..98e1b554b72ea4e20f9f9f10f4aa1c7fda284da6 100644
--- a/mitm_tooling/transformation/superset/exporting.py
+++ b/mitm_tooling/transformation/superset/exporting.py
@@ -7,7 +7,7 @@ from mitm_tooling.utilities.io_utils import FilePath, ByteSink, use_bytes_io
from .definitions import SupersetDefFile, SupersetDefFolder
-def write_superset_def_as_zip(target: ByteSink, superset_def: SupersetDefFolder):
+def write_superset_import_as_zip(target: ByteSink, superset_def: SupersetDefFolder):
folder_structure = superset_def.folder_dict
with use_bytes_io(target, expected_file_ext='.zip', mode='wb', create_file_if_necessary=True) as f:
with zipfile.ZipFile(f, 'w', zipfile.ZIP_DEFLATED) as zf:
@@ -16,8 +16,8 @@ def write_superset_def_as_zip(target: ByteSink, superset_def: SupersetDefFolder)
fn = f'{arg.filename}.yaml'
if prefix:
fn = os.path.join(prefix, fn)
- dump = arg.model_dump(by_alias=True, mode='python', exclude_none=True, serialize_as_any=True)
- s = yaml.dump(dump, default_flow_style=False)
+ dump = arg.model_dump(by_alias=True, mode='json', exclude_none=True, serialize_as_any=True)
+ s = yaml.safe_dump(dump, default_flow_style=False)
zf.writestr(fn, s)
# with zf.open(fn, 'w') as df:
@@ -38,8 +38,4 @@ def write_superset_def_as_zip(target: ByteSink, superset_def: SupersetDefFolder)
zf.mkdir(path)
mk_node(folder_content, prefix=path)
- mk_node(folder_structure)
-
-
-def write_superset_assets_def(output_path: FilePath, superset_def: SupersetDefFolder):
- write_superset_def_as_zip(output_path, superset_def)
+ mk_node(folder_structure)
\ No newline at end of file
diff --git a/mitm_tooling/transformation/superset/factories/assets.py b/mitm_tooling/transformation/superset/factories/assets.py
deleted file mode 100644
index d9fb2a7a4bbb014dec4cea5251333b4134ec59ca..0000000000000000000000000000000000000000
--- a/mitm_tooling/transformation/superset/factories/assets.py
+++ /dev/null
@@ -1,21 +0,0 @@
-from ..definitions import SupersetMetadataDef, SupersetDatabaseDef, SupersetDashboardDef
-from ..definitions.assets import MetadataType, SupersetAssetsDef, SupersetDatasetDef, SupersetChartDef, \
- ExtendedSupersetAssetsDef, SupersetMitMDatasetDef
-
-
-def mk_metadata(metadata_type: MetadataType) -> SupersetMetadataDef:
- return SupersetMetadataDef(type=metadata_type)
-
-
-def mk_assets(databases: list[SupersetDatabaseDef] = None,
- datasets: list[SupersetDatasetDef] = None,
- charts: list[SupersetChartDef] = None,
- dashboards: list[SupersetDashboardDef] = None,
- metadata_type: MetadataType | None = None) -> SupersetAssetsDef:
- return SupersetAssetsDef(databases=databases, datasets=datasets, charts=charts, dashboards=dashboards,
- metadata=SupersetMetadataDef(type=metadata_type or MetadataType.Asset))
-
-
-def mk_extended_assets(mitm_datasets: list[SupersetMitMDatasetDef],
- base_assets: SupersetAssetsDef) -> ExtendedSupersetAssetsDef:
- return ExtendedSupersetAssetsDef(mitm_datasets=mitm_datasets, base_assets=base_assets)
diff --git a/mitm_tooling/transformation/superset/factories/chart.py b/mitm_tooling/transformation/superset/factories/chart.py
new file mode 100644
index 0000000000000000000000000000000000000000..eef34774f84d1046860e32bf53285fca744b0cac
--- /dev/null
+++ b/mitm_tooling/transformation/superset/factories/chart.py
@@ -0,0 +1,19 @@
+from uuid import UUID
+
+from .utils import mk_uuid
+from ..definitions import SupersetChartDef, SupersetVizType, ChartParams, QueryContext
+
+
+def mk_chart_def(name: str, dataset_uuid: UUID, viz_type: SupersetVizType, params: ChartParams,
+ query_context: QueryContext, uuid: UUID | None = None, **kwargs) -> SupersetChartDef:
+ return SupersetChartDef(
+ slice_name=name,
+ viz_type=viz_type,
+ dataset_uuid=dataset_uuid,
+ params=params,
+ query_context=query_context,
+ uuid=uuid or mk_uuid(),
+ **kwargs
+ )
+
+
diff --git a/mitm_tooling/transformation/superset/factories/dashboard.py b/mitm_tooling/transformation/superset/factories/dashboard.py
index f72fe18e53847399d3d55ae59139c6e34589528e..3bb6542993323af9d653de906626ea29c8263da8 100644
--- a/mitm_tooling/transformation/superset/factories/dashboard.py
+++ b/mitm_tooling/transformation/superset/factories/dashboard.py
@@ -1,8 +1,7 @@
-from typing import Any
from uuid import UUID
from mitm_tooling.representation import ColumnName
-from mitm_tooling.transformation.superset.definitions import SupersetDashboardDef, SupersetChartDef
+from mitm_tooling.transformation.superset.definitions import SupersetChartDef
from mitm_tooling.transformation.superset.factories.utils import mk_uuid, mk_short_uuid_str
from ..definitions.dashboard import *
@@ -38,9 +37,11 @@ def mk_dashboard_chart(chart_uuid: UUID, width: int, height: int, slice_name: st
return DashboardChart(id=id or f'CHART-{mk_short_uuid_str()}',
meta=ChartMeta(uuid=chart_uuid, width=width, height=height, sliceName=slice_name))
+
def chart_def_into_dashboard_chart(chart_def: SupersetChartDef, width: int, height: int) -> DashboardChart:
return mk_dashboard_chart(chart_uuid=chart_def.uuid, width=width, height=height)
+
def mk_dashboard_position_data(header_text: str, chart_grid: list[list[DashboardChart]]) -> DashboardPositionData:
row_ids = []
elements = {}
diff --git a/mitm_tooling/transformation/superset/factories/dataset.py b/mitm_tooling/transformation/superset/factories/dataset.py
index 6f172977159061e233a0e38b5625dd6fdc1b2e4d..24183a0dfecb7b9200564cd9b5296c299f6b0112 100644
--- a/mitm_tooling/transformation/superset/factories/dataset.py
+++ b/mitm_tooling/transformation/superset/factories/dataset.py
@@ -5,10 +5,9 @@ import sqlalchemy as sa
from mitm_tooling.data_types import MITMDataType
from mitm_tooling.extraction.sql.data_models import TableMetaInfo
-from mitm_tooling.transformation.superset.definitions import SupersetDatasetDef
from .core import mk_column, mk_metric
from .utils import mk_uuid
-from ..definitions import SupersetAggregate
+from ..definitions import SupersetAggregate, SupersetDatasetDef
def mk_dataset(tm: TableMetaInfo, database_uuid: UUID, dialect: sa.Dialect | None = None,
diff --git a/mitm_tooling/transformation/superset/factories/charts.py b/mitm_tooling/transformation/superset/factories/generic_charts.py
similarity index 72%
rename from mitm_tooling/transformation/superset/factories/charts.py
rename to mitm_tooling/transformation/superset/factories/generic_charts.py
index cb7939551c409591735155aca2685369e88d27a9..3f8973dc71828f3234aa44557519823766296a3b 100644
--- a/mitm_tooling/transformation/superset/factories/charts.py
+++ b/mitm_tooling/transformation/superset/factories/generic_charts.py
@@ -1,18 +1,18 @@
-from uuid import UUID
from mitm_tooling.data_types import MITMDataType
+from mitm_tooling.representation import ColumnName
from mitm_tooling.utilities.python_utils import unique
-from .core import mk_empty_adhoc_time_filter, mk_adhoc_metric, mk_pivot_post_processing, mk_adhoc_column
-from .query import mk_query_object, mk_query_context, \
- mk_empty_query_object_time_filter_clause
-from .utils import mk_uuid
-from ..definitions import SupersetChartDef, PieChartParams, DatasourceIdentifier, SupersetAggregate, \
- SupersetVizType, TimeSeriesBarParams, TimeGrain, QueryObjectFilterClause, SupersetAdhocFilter, \
+from .chart import mk_chart_def
+from .core import mk_adhoc_metric, mk_empty_adhoc_time_filter, \
+ mk_pivot_post_processing, mk_adhoc_column
+from .query import mk_query_object, \
+ mk_empty_query_object_time_filter_clause, mk_query_context
+from ..definitions import DatasourceIdentifier, SupersetChartDef, SupersetAggregate, \
+ PieChartParams, SupersetVizType, SupersetAdhocFilter, TimeGrain, TimeSeriesBarParams, QueryObjectFilterClause, \
TimeSeriesLineParams, QueryObjectExtras
-from mitm_tooling.representation import ColumnName
def mk_pie_chart(name: str, datasource_identifier: DatasourceIdentifier, col: ColumnName, dt: MITMDataType,
- groupby_cols: list[ColumnName] | None = None, uuid: UUID | None = None) -> SupersetChartDef:
+ groupby_cols: list[ColumnName] | None = None) -> SupersetChartDef:
groupby_cols = groupby_cols or []
metric = mk_adhoc_metric(col, agg=SupersetAggregate.COUNT, dt=dt)
params = PieChartParams(datasource=datasource_identifier,
@@ -24,12 +24,11 @@ def mk_pie_chart(name: str, datasource_identifier: DatasourceIdentifier, col: Co
filters=[mk_empty_query_object_time_filter_clause()])
qc = mk_query_context(datasource=datasource_identifier, queries=[qo], form_data=params)
- return SupersetChartDef(slice_name=name,
- viz_type=SupersetVizType.PIE,
- dataset_uuid=datasource_identifier.dataset_uuid,
- params=params,
- query_context=qc,
- uuid=uuid or mk_uuid())
+ return mk_chart_def(name=name,
+ viz_type=SupersetVizType.PIE,
+ dataset_uuid=datasource_identifier.dataset_uuid,
+ params=params,
+ query_context=qc)
def mk_time_series_bar_chart(name: str,
@@ -39,7 +38,6 @@ def mk_time_series_bar_chart(name: str,
x_col: ColumnName,
groupby_cols: list[ColumnName] | None = None,
filters: list[SupersetAdhocFilter] | None = None,
- uuid: UUID | None = None,
time_grain: TimeGrain | None = None) -> SupersetChartDef:
groupby_cols = groupby_cols or []
metric = mk_adhoc_metric(y_col, agg=SupersetAggregate.COUNT, dt=y_dt)
@@ -64,12 +62,11 @@ def mk_time_series_bar_chart(name: str,
series_columns=[y_col])
qc = mk_query_context(datasource=datasource_identifier, queries=[qo], form_data=params)
- return SupersetChartDef(slice_name=name,
- viz_type=SupersetVizType.TIMESERIES_BAR,
- dataset_uuid=datasource_identifier.dataset_uuid,
- params=params,
- query_context=qc,
- uuid=uuid or mk_uuid())
+ return mk_chart_def(name=name,
+ viz_type=SupersetVizType.TIMESERIES_BAR,
+ dataset_uuid=datasource_identifier.dataset_uuid,
+ params=params,
+ query_context=qc)
def mk_avg_count_time_series_chart(name: str,
@@ -77,7 +74,6 @@ def mk_avg_count_time_series_chart(name: str,
groupby_cols: list[ColumnName],
time_col: ColumnName = 'time',
filters: list[SupersetAdhocFilter] | None = None,
- uuid: UUID | None = None,
time_grain: TimeGrain | None = None):
groupby_cols = groupby_cols or []
metric = mk_adhoc_metric(time_col, agg=SupersetAggregate.COUNT, dt=MITMDataType.Datetime)
@@ -103,9 +99,8 @@ def mk_avg_count_time_series_chart(name: str,
extras=QueryObjectExtras(time_grain_sqla=time_grain))
qc = mk_query_context(datasource=datasource_identifier, queries=[qo], form_data=params)
- return SupersetChartDef(slice_name=name,
- viz_type=SupersetVizType.TIMESERIES_LINE,
- dataset_uuid=datasource_identifier.dataset_uuid,
- params=params,
- query_context=qc,
- uuid=uuid or mk_uuid())
+ return mk_chart_def(name=name,
+ viz_type=SupersetVizType.TIMESERIES_LINE,
+ dataset_uuid=datasource_identifier.dataset_uuid,
+ params=params,
+ query_context=qc)
diff --git a/mitm_tooling/transformation/superset/factories/importable.py b/mitm_tooling/transformation/superset/factories/importable.py
new file mode 100644
index 0000000000000000000000000000000000000000..c601702a4c294ec77a2ca8e79e828016f112da7b
--- /dev/null
+++ b/mitm_tooling/transformation/superset/factories/importable.py
@@ -0,0 +1,22 @@
+from ..definitions import SupersetMetadataDef, SupersetDatabaseDef, SupersetDashboardDef, SupersetDatasetDef, \
+ SupersetChartDef, SupersetMitMDatasetDef
+from ..definitions.importable import MetadataType, SupersetAssetsImport, SupersetMitMDatasetImport
+
+
+def mk_metadata(metadata_type: MetadataType) -> SupersetMetadataDef:
+ return SupersetMetadataDef(type=metadata_type)
+
+
+def mk_assets_import(databases: list[SupersetDatabaseDef] = None,
+ datasets: list[SupersetDatasetDef] = None,
+ charts: list[SupersetChartDef] = None,
+ dashboards: list[SupersetDashboardDef] = None,
+ metadata_type: MetadataType = MetadataType.Asset) -> SupersetAssetsImport:
+ return SupersetAssetsImport(databases=databases, datasets=datasets, charts=charts, dashboards=dashboards,
+ metadata=mk_metadata(metadata_type))
+
+
+def mk_mitm_dataset_import(mitm_datasets: list[SupersetMitMDatasetDef],
+ base_assets: SupersetAssetsImport) -> SupersetMitMDatasetImport:
+ return SupersetMitMDatasetImport(mitm_datasets=mitm_datasets, base_assets=base_assets,
+ metadata=mk_metadata(MetadataType.MitMDataset))
diff --git a/mitm_tooling/transformation/superset/factories/mitm_specific/__init__.py b/mitm_tooling/transformation/superset/factories/mitm_specific/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/mitm_tooling/transformation/superset/from_intermediate.py b/mitm_tooling/transformation/superset/from_intermediate.py
index 717ec5cb6579f69c042ea538d7464e61508c8bb1..29123bf6c5a86fe00247d4cd41357ad6a52d6257 100644
--- a/mitm_tooling/transformation/superset/from_intermediate.py
+++ b/mitm_tooling/transformation/superset/from_intermediate.py
@@ -1,8 +1,8 @@
from mitm_tooling.representation import Header
from .common import SupersetDBConnectionInfo
-from .definition_bundles import SupersetDatasourceBundle, SupersetMitMDatasetBundle
-from ...definition import MITM
+from .definition_bundles import SupersetDatasourceBundle, SupersetMitMDatasetBundle, SupersetVisualizationBundle
+from .mitm_specific import get_mitm_visualization_factory
def header_into_superset_datasource_bundle(header: Header,
@@ -13,10 +13,20 @@ def header_into_superset_datasource_bundle(header: Header,
return db_meta_into_superset_datasource_bundle(db_meta, db_conn_info)
-def header_into_superset_mitm_dataset(header: Header,
- db_conn_info: SupersetDBConnectionInfo,
- dataset_name: str) -> SupersetMitMDatasetBundle:
+def header_into_superset_visualization_bundle(header: Header,
+ datasource_bundle: SupersetDatasourceBundle) -> SupersetVisualizationBundle:
+ return get_mitm_visualization_factory(header.mitm)(header, datasource_bundle)
+
+
+def header_into_superset_mitm_dataset_bundle(header: Header,
+ db_conn_info: SupersetDBConnectionInfo,
+ dataset_name: str,
+ include_visualizations: bool = False) -> SupersetMitMDatasetBundle:
from ..sql.from_intermediate import header_into_db_meta
from .from_sql import db_meta_into_mitm_dataset_bundle
db_meta = header_into_db_meta(header)
- return db_meta_into_mitm_dataset_bundle(db_meta, db_conn_info, dataset_name, header.mitm)
+ mitm_dataset_bundle = db_meta_into_mitm_dataset_bundle(db_meta, db_conn_info, dataset_name, header.mitm)
+ if include_visualizations:
+ mitm_dataset_bundle.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 de25df12660122b9b34047da67f17e7d081829ca..93b6bf05712c0e0cb917b254b4dcad7d462ffd67 100644
--- a/mitm_tooling/transformation/superset/from_sql.py
+++ b/mitm_tooling/transformation/superset/from_sql.py
@@ -40,8 +40,6 @@ def db_meta_into_mitm_dataset_bundle(db_meta: DBMetaInfoBase,
def db_into_superset_datasource_bundle(arg: SQLiteFileOrEngine,
db_conn_info: SupersetDBConnectionInfo) -> SupersetDatasourceBundle:
engine = _mk_engine(arg)
- # db_name = db_conn_info.db_name or db_conn_info.db_name_in_uri
- # db_name = os.path.splitext(os.path.basename(sqlite_file_path))[0]
meta, _ = connect_and_reflect(engine, allowed_schemas=[db_conn_info.schema_name])
db_meta = DBMetaInfo.from_sa_meta(meta, default_schema=db_conn_info.schema_name)
diff --git a/mitm_tooling/transformation/superset/interface.py b/mitm_tooling/transformation/superset/interface.py
index d9c40282f1d15bd6daa2d07222819e837d8736c6..048e6169977a5d29c547a7da9f09268def448b61 100644
--- a/mitm_tooling/transformation/superset/interface.py
+++ b/mitm_tooling/transformation/superset/interface.py
@@ -1,37 +1,21 @@
-from mitm_tooling.definition import MITM
from mitm_tooling.representation import Header
from mitm_tooling.transformation.superset.definition_bundles import SupersetDatasourceBundle, \
SupersetVisualizationBundle, SupersetMitMDatasetBundle
-from mitm_tooling.transformation.superset.from_intermediate import header_into_superset_mitm_dataset
-from pydantic import AnyUrl
-from typing import Callable
-
-from . import mitm_specific
+from mitm_tooling.transformation.superset.from_intermediate import header_into_superset_mitm_dataset_bundle
+from mitm_tooling.transformation.superset.from_intermediate import header_into_superset_visualization_bundle
from .common import SupersetDBConnectionInfo
from .from_intermediate import header_into_superset_datasource_bundle
-from ...representation.sql_representation import SQL_REPRESENTATION_DEFAULT_SCHEMA, SchemaName
-
-mitm_specific_visualization_factories: dict[
- MITM, Callable[[Header, SupersetDatasourceBundle], SupersetVisualizationBundle]] = {
- MITM.MAED: mitm_specific.mk_maed_visualization,
-}
-def mk_superset_datasource_import(header: Header, sql_alchemy_uri: AnyUrl, explicit_db_name: str | None = None,
- schema_name: SchemaName = SQL_REPRESENTATION_DEFAULT_SCHEMA) -> SupersetDatasourceBundle:
- db_conn_info = SupersetDBConnectionInfo(sql_alchemy_uri=sql_alchemy_uri, explicit_db_name=explicit_db_name,
- schema_name=schema_name)
+def mk_superset_datasource_bundle(header: Header, db_conn_info: SupersetDBConnectionInfo) -> SupersetDatasourceBundle:
return header_into_superset_datasource_bundle(header, db_conn_info)
-def mk_superset_mitm_dataset_import(header: Header, sql_alchemy_uri: AnyUrl, dataset_name: str,
- explicit_db_name: str | None = None,
- schema_name: SchemaName = SQL_REPRESENTATION_DEFAULT_SCHEMA) -> SupersetMitMDatasetBundle:
- db_conn_info = SupersetDBConnectionInfo(sql_alchemy_uri=sql_alchemy_uri, explicit_db_name=explicit_db_name,
- schema_name=schema_name)
- return header_into_superset_mitm_dataset(header, db_conn_info, dataset_name=dataset_name)
+def mk_superset_visualization_bundle(header: Header,
+ superset_datasource_bundle: SupersetDatasourceBundle) -> SupersetVisualizationBundle:
+ return header_into_superset_visualization_bundle(header, superset_datasource_bundle)
-def mk_superset_visualization_import(header: Header,
- superset_datasource_bundle: SupersetDatasourceBundle) -> SupersetVisualizationBundle:
- return mitm_specific_visualization_factories[header.mitm](header, superset_datasource_bundle)
+def mk_superset_mitm_dataset_bundle(dataset_name: str, header: Header, db_conn_info: SupersetDBConnectionInfo,
+ include_visualizations: bool = False) -> SupersetMitMDatasetBundle:
+ return header_into_superset_mitm_dataset_bundle(header, db_conn_info, dataset_name=dataset_name)
diff --git a/mitm_tooling/transformation/superset/mitm_specific/__init__.py b/mitm_tooling/transformation/superset/mitm_specific/__init__.py
index 77f4b30d91836d7c2859a115dd4f9141e2866a76..dbd365983a6f8f0ca9607639fca4777e42502377 100644
--- a/mitm_tooling/transformation/superset/mitm_specific/__init__.py
+++ b/mitm_tooling/transformation/superset/mitm_specific/__init__.py
@@ -1 +1,2 @@
from .maed_visualization import mk_maed_visualization
+from .registry import get_mitm_visualization_factory
\ No newline at end of file
diff --git a/mitm_tooling/transformation/superset/mitm_specific/fallback.py b/mitm_tooling/transformation/superset/mitm_specific/fallback.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e076ab1b3c4669232ca5f953343fd5013c3d8c1
--- /dev/null
+++ b/mitm_tooling/transformation/superset/mitm_specific/fallback.py
@@ -0,0 +1,7 @@
+from mitm_tooling.representation import Header
+from mitm_tooling.transformation.superset.definition_bundles import SupersetDatasourceBundle, \
+ SupersetVisualizationBundle
+
+
+def mk_empty_visualization(header: Header, datasource_bundle: SupersetDatasourceBundle) -> SupersetVisualizationBundle:
+ return SupersetVisualizationBundle()
\ No newline at end of file
diff --git a/mitm_tooling/transformation/superset/factories/mitm_specific/maed_charts.py b/mitm_tooling/transformation/superset/mitm_specific/maed_charts.py
similarity index 89%
rename from mitm_tooling/transformation/superset/factories/mitm_specific/maed_charts.py
rename to mitm_tooling/transformation/superset/mitm_specific/maed_charts.py
index c1f37ca897f9802764197b25f83d2925cdea63ab..523e3e175332309ea8f3a78225ce9d5e2ef276b3 100644
--- a/mitm_tooling/transformation/superset/factories/mitm_specific/maed_charts.py
+++ b/mitm_tooling/transformation/superset/mitm_specific/maed_charts.py
@@ -1,10 +1,10 @@
from mitm_tooling.data_types import MITMDataType
from mitm_tooling.representation import SQLRepresentationSchema, Header, mk_sql_rep_schema
from mitm_tooling.utilities.python_utils import take_first
-from ..charts import mk_avg_count_time_series_chart, mk_pie_chart, mk_time_series_bar_chart
-from ..core import mk_adhoc_filter
-from ...definition_bundles import SupersetDatasourceBundle
-from ...definitions import SupersetChartDef, FilterOperator
+from ..factories.generic_charts import mk_pie_chart, mk_time_series_bar_chart, mk_avg_count_time_series_chart
+from ..factories.core import mk_adhoc_filter
+from ..definition_bundles import SupersetDatasourceBundle
+from ..definitions import SupersetChartDef, FilterOperator
def mk_maed_charts(header: Header, superset_datasource_bundle: SupersetDatasourceBundle,
diff --git a/mitm_tooling/transformation/superset/factories/mitm_specific/maed_dashboards.py b/mitm_tooling/transformation/superset/mitm_specific/maed_dashboards.py
similarity index 67%
rename from mitm_tooling/transformation/superset/factories/mitm_specific/maed_dashboards.py
rename to mitm_tooling/transformation/superset/mitm_specific/maed_dashboards.py
index 94233a54200509c85a6a2f928bf45424b002801b..e83b5c46c7f4c578a2e98cec5c38057b6ba5f3e6 100644
--- a/mitm_tooling/transformation/superset/factories/mitm_specific/maed_dashboards.py
+++ b/mitm_tooling/transformation/superset/mitm_specific/maed_dashboards.py
@@ -1,8 +1,8 @@
from mitm_tooling.representation import Header
-from mitm_tooling.transformation.superset.definition_bundles import SupersetDatasourceBundle
-from mitm_tooling.transformation.superset.definitions import SupersetDashboardDef, SupersetChartDef
-from mitm_tooling.transformation.superset.factories.dashboard import mk_dashboard_def, mk_dashboard_chart
-from mitm_tooling.transformation.superset.factories.mitm_specific.maed_charts import mk_maed_charts
+from ..definition_bundles import SupersetDatasourceBundle
+from ..definitions import SupersetDashboardDef, SupersetChartDef
+from ..factories.dashboard import mk_dashboard_def, mk_dashboard_chart
+from .maed_charts import mk_maed_charts
def mk_maed_dashboard(header: Header, datasource_bundle: SupersetDatasourceBundle) -> tuple[SupersetDashboardDef, list[SupersetChartDef]]:
diff --git a/mitm_tooling/transformation/superset/mitm_specific/maed_visualization.py b/mitm_tooling/transformation/superset/mitm_specific/maed_visualization.py
index ff48e1a184a4e0487453c73c5854cfa4b86e8f60..2ab8b7ebb51d30501ab25a6b125295787d000333 100644
--- a/mitm_tooling/transformation/superset/mitm_specific/maed_visualization.py
+++ b/mitm_tooling/transformation/superset/mitm_specific/maed_visualization.py
@@ -1,13 +1,10 @@
from mitm_tooling.representation import Header
-from mitm_tooling.transformation.superset.definition_bundles import SupersetVisualizationBundle, \
+from ..definition_bundles import SupersetVisualizationBundle, \
SupersetDatasourceBundle
-from mitm_tooling.transformation.superset.factories.mitm_specific.maed_charts import mk_maed_charts
-from mitm_tooling.transformation.superset.factories.mitm_specific.maed_dashboards import mk_maed_dashboard
+from .maed_dashboards import mk_maed_dashboard
def mk_maed_visualization(header: Header,
superset_datasource_bundle: SupersetDatasourceBundle) -> SupersetVisualizationBundle:
- ds_ids = superset_datasource_bundle.placeholder_dataset_identifiers
-
dashboard, charts = mk_maed_dashboard(header, superset_datasource_bundle)
return SupersetVisualizationBundle(charts=charts, dashboards=[dashboard])
diff --git a/mitm_tooling/transformation/superset/mitm_specific/registry.py b/mitm_tooling/transformation/superset/mitm_specific/registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..746062e8684048109424e975f12b211ad4cc3b88
--- /dev/null
+++ b/mitm_tooling/transformation/superset/mitm_specific/registry.py
@@ -0,0 +1,23 @@
+from typing import Callable
+
+from mitm_tooling.definition import MITM
+from mitm_tooling.representation import Header
+from .fallback import mk_empty_visualization
+from .maed_visualization import mk_maed_visualization
+from ..definition_bundles import SupersetDatasourceBundle, SupersetVisualizationBundle
+
+VisualizationBundleFactory = Callable[[Header, SupersetDatasourceBundle], SupersetVisualizationBundle]
+
+
+
+mitm_specific_visualization_factories: dict[
+ MITM, VisualizationBundleFactory] = {
+ MITM.MAED: mk_maed_visualization,
+}
+
+
+def get_mitm_visualization_factory(mitm: MITM) -> VisualizationBundleFactory:
+ if factory := mitm_specific_visualization_factories.get(mitm):
+ return factory
+ else:
+ return mk_empty_visualization
diff --git a/pyproject.toml b/pyproject.toml
index 417693ee7d236926ca05e55e46990209d32ff1fb..f39c1ff51932550d3e663d0b6ea8ddc8649bd19a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[tool.poetry]
name = "mitm-tooling"
-version = "0.4.2"
+version = "0.4.3"
description = ""
authors = ["Leah Tacke genannt Unterberg <leah.tgu@pads.rwth-aachen.de>"]
readme = "README.md"
diff --git a/test/something.py b/test/something.py
index 8a9850e7a347ce196dcaa0955ccf5bde25166e01..26100fd75ec554929fc560646fd4f98b48766f86 100644
--- a/test/something.py
+++ b/test/something.py
@@ -3,7 +3,7 @@ import unittest
from pydantic import AnyUrl
-from mitm_tooling.transformation.superset.common import name_plus_uuid
+from mitm_tooling.transformation.superset.common import name_plus_uuid, SupersetDBConnectionInfo
class MyTestCase(unittest.TestCase):
@@ -64,27 +64,29 @@ class MyTestCase(unittest.TestCase):
def test_superset_dataset_assets(self):
from mitm_tooling.io import importing, MITM
syn = importing.read_zip('synthetic.maed', MITM.MAED, header_only=True)
- from mitm_tooling.transformation.superset import mk_superset_mitm_dataset_import, write_superset_assets_def
- dataset_import = mk_superset_mitm_dataset_import(syn.header,
- sql_alchemy_uri=AnyUrl('sqlite://synthetic-variation.sqlite'),
- dataset_name=name_plus_uuid('SyntheticExampleDataset'),
- explicit_db_name='SyntheticExampleDB',
- schema_name='main')
- write_superset_assets_def('superset_dataset_import', dataset_import.to_assets())
+ from mitm_tooling.transformation.superset import mk_superset_mitm_dataset_bundle, write_superset_import_as_zip
+ dataset_import = mk_superset_mitm_dataset_bundle(name_plus_uuid('SyntheticExampleDataset'), syn.header,
+ SupersetDBConnectionInfo(
+ sql_alchemy_uri=AnyUrl(
+ 'sqlite://synthetic-variation.sqlite'),
+ explicit_db_name='SyntheticExampleDB',
+ schema_name='main'))
+ write_superset_import_as_zip('superset_dataset_import', dataset_import.to_import())
def test_superset_viz_assets(self):
from mitm_tooling.io import importing, MITM
syn = importing.read_zip('synthetic.maed', MITM.MAED, header_only=True)
- from mitm_tooling.transformation.superset import mk_superset_mitm_dataset_import, write_superset_assets_def
- dataset_import = mk_superset_mitm_dataset_import(syn.header,
- sql_alchemy_uri=AnyUrl('sqlite://synthetic-variation.sqlite'),
- dataset_name=name_plus_uuid('SyntheticExampleDataset'),
- explicit_db_name='SyntheticExampleDB',
- schema_name='main')
+ from mitm_tooling.transformation.superset import mk_superset_mitm_dataset_bundle, write_superset_import_as_zip
+ dataset_import = mk_superset_mitm_dataset_bundle(name_plus_uuid('SyntheticExampleDataset'), syn.header,
+ SupersetDBConnectionInfo(
+ sql_alchemy_uri=AnyUrl(
+ 'sqlite://synthetic-variation.sqlite'),
+ explicit_db_name='SyntheticExampleDB',
+ schema_name='main'))
- from mitm_tooling.transformation.superset import mk_superset_visualization_import
- visualization_import = mk_superset_visualization_import(syn.header, dataset_import.datasource_bundle)
- write_superset_assets_def('superset_viz_import', visualization_import.to_assets())
+ from mitm_tooling.transformation.superset import mk_superset_visualization_bundle
+ visualization_import = mk_superset_visualization_bundle(syn.header, dataset_import.datasource_bundle)
+ write_superset_import_as_zip('superset_viz_import', visualization_import.to_import())
if __name__ == '__main__':