improved superset visualization definition generation architecture

mitm_tooling/definition/definition_representation.py

@@ -8,6 +8,7 @@ import pydantic
 from pydantic import Field
 
 from mitm_tooling.data_types.data_types import MITMDataType
+from mitm_tooling.utilities.identifiers import naive_pluralize
 from mitm_tooling.utilities.python_utils import unpack_singleton, map_vals, normalize_list_of_mixed, combine_dicts
 
 COLUMN_GROUPS = Literal['kind', 'type', 'identity-relations', 'inline-relations', 'foreign-relations', 'attributes']
@@ -52,6 +53,9 @@ class OwnedRelations(pydantic.BaseModel):
         return {n: fk_rel_info.fk_relations for n, fk_rel_info in self.foreign.items() if
                 fk_rel_info.target_concept == concept}
 
+    @property
+    def relation_names(self) -> tuple[RelationName, ...]:
+        return tuple(itertools.chain(self.identity.keys(), self.inline.keys(), self.foreign.keys()))
 
 class ConceptProperties(pydantic.BaseModel):
     nature: tuple[ConceptLevel, ConceptKind]
@@ -208,7 +212,7 @@ class ConceptPropertiesFile(pydantic.BaseModel):
 
         kwargs = combine_dicts(inherited_kwargs,
                                {'typing_concept': self.typing_concept, 'permit_attributes': self.permit_attributes,
-                                'plural': self.plural}, {'plural': concept + 's'})
+                                'plural': self.plural}, {'plural': naive_pluralize(concept)})
 
         opt_kwargs = {}
         if self.override_column_group_ordering is not None:

mitm_tooling/transformation/superset/__init__.py

-from . import definitions, factories, mitm_specific
+from . import definitions, factories, asset_bundles, visualizations
 from . import exporting, from_sql, from_intermediate
 from . import interface
 from .exporting import write_superset_import_as_zip
 from .interface import mk_superset_datasource_bundle, mk_superset_visualization_bundle, mk_superset_mitm_dataset_bundle
+from .visualizations.registry import VisualizationType, MAEDVisualizationType
+from .visualizations.registry import mk_visualization, get_mitm_visualization_creator

mitm_tooling/transformation/superset/asset_bundles/__init__.py

+from .asset_bundles import *

mitm_tooling/transformation/superset/definition_bundles.py → mitm_tooling/transformation/superset/asset_bundles/asset_bundles.py

+import itertools
 from abc import ABC, abstractmethod
 from typing import Any, Self
 from uuid import UUID
 
 import pydantic
 
-from .definitions import SupersetDatabaseDef, SupersetMitMDatasetDef, \
+from mitm_tooling.representation import TableName
+from ..definitions import SupersetDatabaseDef, SupersetMitMDatasetDef, \
     SupersetChartDef, SupersetDashboardDef, SupersetAssetsImport, SupersetDatasetDef, \
     SupersetMitMDatasetImport, SupersetDefFolder, DatasourceIdentifier
-from .factories.importable import mk_assets_import, mk_mitm_dataset_import
-from ...representation import TableName
+from ..factories.importable import mk_assets_import, mk_mitm_dataset_import
 
 
 class SupersetAssetBundle(SupersetDefFolder, ABC):
@@ -53,6 +54,15 @@ class SupersetVisualizationBundle(SupersetAssetBundle):
     def dashboard_uuids(self) -> list[UUID]:
         return [da.uuid for da in self.dashboards]
 
+    @classmethod
+    def combine(cls, *bundles: Self) -> Self:
+        if not bundles or len(bundles) == 0:
+            return cls()
+
+        charts, dashboards = itertools.chain(*(b.charts for b in bundles)), itertools.chain(*(b.dashboards for b in
+                                                                                              bundles))
+        return cls(charts=list(charts), dashboards=list(dashboards))
+
     def to_import(self) -> SupersetAssetsImport:
         return mk_assets_import(charts=self.charts, dashboards=self.dashboards)
 

mitm_tooling/transformation/superset/common.py

-from uuid import UUID
+from typing import Type
 
 import pydantic
 import sqlalchemy as sa
+from pydantic import AnyUrl, ConfigDict
+
 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
 from mitm_tooling.utilities.io_utils import FilePath
 from mitm_tooling.utilities.sql_utils import create_sa_engine, dialect_cls_from_url, any_url_into_sa_url
-from pydantic import AnyUrl, ConfigDict
-from typing import Type
+from .definitions import StrUrl
 
 SQLiteFileOrEngine = FilePath | sa.Engine
 
 
-class SupersetDBConnectionInfo(pydantic.BaseModel):
+class DBConnectionInfo(pydantic.BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     sql_alchemy_uri: StrUrl
     explicit_db_name: str | None = None
     schema_name: SchemaName = SQL_REPRESENTATION_DEFAULT_SCHEMA
+    catalog: str | None = None  # for future use
 
     @property
     def sa_url(self) -> sa.URL:
@@ -38,10 +38,6 @@ class SupersetDBConnectionInfo(pydantic.BaseModel):
         return dialect_cls_from_url(self.sql_alchemy_uri)
 
 
-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:
     if isinstance(arg, sa.Engine):
         return arg

mitm_tooling/transformation/superset/definitions/__init__.py

+from .chart import *
 from .constants import *
 from .core import *
+from .dashboard import *
 from .database import *
 from .dataset import *
-from .chart import *
-from .dashboard import *
 from .importable import *
 from .post_processing import *

mitm_tooling/transformation/superset/definitions/chart.py

-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

mitm_tooling/transformation/superset/definitions/constants.py

 from abc import ABC, abstractmethod
 from datetime import datetime
-from datetime import datetime
 from enum import StrEnum, IntEnum
 from typing import Annotated, Literal, Self, Union
 from uuid import UUID
@@ -81,6 +80,7 @@ class SupersetVizType(StrEnum):
     PIE = 'pie'
     TIMESERIES_BAR = 'echarts_timeseries_bar'
     TIMESERIES_LINE = 'echarts_timeseries_line'
+    MAED_CUSTOM = 'maed_custom'
 
 
 class ExpressionType(StrEnum):

mitm_tooling/transformation/superset/definitions/core.py

@@ -178,5 +178,3 @@ class QueryContext(BaseSupersetDefinition):
     result_format: ChartDataResultFormat = ChartDataResultFormat.JSON
     force: bool = False
     custom_cache_timeout: int | None = None
-
-

mitm_tooling/transformation/superset/definitions/dashboard.py

@@ -3,7 +3,7 @@ from typing import Literal
 
 import pydantic
 
-from . import SupersetDefFile, StrUUID, StrUrl
+from . import SupersetDefFile, StrUrl
 from .constants import StrUUID
 
 DashboardInternalID = str

mitm_tooling/transformation/superset/definitions/mitm_dataset.py

@@ -16,15 +16,15 @@ class RelatedObjectIdentifier(BaseSupersetDefinition):
 
 
 class RelatedTable(RelatedObjectIdentifier):
-    pass
+    table_name: str | None = None
 
 
 class RelatedSlice(RelatedObjectIdentifier):
-    pass
+    slice_name: str | None = None
 
 
 class RelatedDashboard(RelatedObjectIdentifier):
-    pass
+    dashboard_title: str | None = None
 
 
 class SupersetMitMDatasetDef(SupersetDefFile):
@@ -38,6 +38,10 @@ class SupersetMitMDatasetDef(SupersetDefFile):
     dashboards: list[RelatedDashboard] | None = None
     version: str = '1.0.0'
 
+    @property
+    def identifier(self) -> MitMDatasetIdentifier:
+        return MitMDatasetIdentifier(dataset_name=self.dataset_name, uuid=self.uuid)
+
     @property
     def filename(self) -> str:
         return self.dataset_name

mitm_tooling/transformation/superset/definitions/post_processing.py

@@ -32,4 +32,3 @@ class Rename(SupersetPostProcessing):
 
 class Flatten(SupersetPostProcessing):
     operation: Literal['flatten'] = 'flatten'
-

mitm_tooling/transformation/superset/exporting.py

@@ -3,11 +3,11 @@ import zipfile
 
 import yaml
 
-from mitm_tooling.utilities.io_utils import FilePath, ByteSink, use_bytes_io
+from mitm_tooling.utilities.io_utils import ByteSink, use_bytes_io
 from .definitions import SupersetDefFile, SupersetDefFolder
 
 
-def write_superset_import_as_zip(target: ByteSink, superset_def: SupersetDefFolder):
+def write_superset_import_as_zip(target: ByteSink, superset_def: SupersetDefFolder) -> None:
     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:

mitm_tooling/transformation/superset/factories/chart.py

 from uuid import UUID
 
-from .utils import mk_uuid
+from mitm_tooling.utilities.identifiers import mk_uuid
 from ..definitions import SupersetChartDef, SupersetVizType, ChartParams, QueryContext
 
 
@@ -15,5 +15,3 @@ def mk_chart_def(name: str, dataset_uuid: UUID, viz_type: SupersetVizType, param
         uuid=uuid or mk_uuid(),
         **kwargs
     )
-
-

mitm_tooling/transformation/superset/factories/custom_charts.py

+from typing import Literal
+
+from mitm_tooling.transformation.superset.definitions import SupersetChartDef, SupersetVizType, ChartParams, \
+    DatasourceIdentifier
+from mitm_tooling.transformation.superset.definitions.mitm_dataset import MitMDatasetIdentifier
+from mitm_tooling.transformation.superset.factories.chart import mk_chart_def
+from mitm_tooling.transformation.superset.factories.query import mk_query_object, mk_query_context
+
+
+class MAEDCustomChartParams(ChartParams):
+    viz_type: Literal[SupersetVizType.MAED_CUSTOM] = SupersetVizType.MAED_CUSTOM
+    mitm_dataset: MitMDatasetIdentifier
+
+
+def mk_maed_custom_chart(name: str,
+                         mitm_dataset_identifier: MitMDatasetIdentifier,
+                         datasource_identifier: DatasourceIdentifier,
+                         **kwargs) -> SupersetChartDef:
+    params = MAEDCustomChartParams(datasource=datasource_identifier, mitm_dataset=mitm_dataset_identifier)
+    qo = mk_query_object()
+    qc = mk_query_context(datasource=datasource_identifier, queries=[qo], form_data=params)
+
+    return mk_chart_def(name,
+                        dataset_uuid=datasource_identifier.uuid,
+                        viz_type=SupersetVizType.MAED_CUSTOM,
+                        params=params,
+                        qc=qc,
+                        **kwargs)

mitm_tooling/transformation/superset/factories/dashboard.py

@@ -2,9 +2,11 @@ from uuid import UUID
 
 from mitm_tooling.representation import ColumnName
 from mitm_tooling.transformation.superset.definitions import SupersetChartDef
-from mitm_tooling.transformation.superset.factories.utils import mk_uuid, mk_short_uuid_str
+from mitm_tooling.utilities.identifiers import mk_uuid, mk_short_uuid_str
 from ..definitions.dashboard import *
 
+ChartGrid = list[list[DashboardChart]]
+
 
 def mk_filter_config(name: str, target_cols: list[tuple[ColumnName, UUID]], ft: FilterType = FilterType.FILTER_SELECT,
                      control_values: ControlValues = ControlValues()) -> NativeFilterConfig:
@@ -42,7 +44,7 @@ def chart_def_into_dashboard_chart(chart_def: SupersetChartDef, width: int, heig
     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:
+def mk_dashboard_position_data(header_text: str, chart_grid: ChartGrid) -> DashboardPositionData:
     row_ids = []
     elements = {}
     for row in chart_grid:
@@ -60,12 +62,15 @@ def mk_dashboard_position_data(header_text: str, chart_grid: list[list[Dashboard
     return res
 
 
-def mk_dashboard_metadata(native_filters: list[NativeFilterConfig]) -> DashboardMetadata:
-    return DashboardMetadata(native_filter_configuration=native_filters)
+def mk_dashboard_metadata(native_filters: list[NativeFilterConfig] | None = None) -> DashboardMetadata:
+    return DashboardMetadata(native_filter_configuration=native_filters or [])
 
 
-def mk_dashboard_def(dashboard_title, chart_grid: list[list[DashboardChart]], native_filters: list[NativeFilterConfig],
-                     description: str | None = None, uuid: UUID | None = None) -> SupersetDashboardDef:
+def mk_dashboard_def(dashboard_title: str,
+                     chart_grid: ChartGrid,
+                     native_filters: list[NativeFilterConfig] | None = None,
+                     description: str | None = None,
+                     uuid: UUID | None = None) -> SupersetDashboardDef:
     position_data = mk_dashboard_position_data(dashboard_title, chart_grid)
     dashboard_metadata = mk_dashboard_metadata(native_filters=native_filters)
     return SupersetDashboardDef(

mitm_tooling/transformation/superset/factories/database.py

@@ -2,8 +2,7 @@ from uuid import UUID
 
 from pydantic import AnyUrl
 
-from .utils import mk_uuid
-from ..common import name_plus_uuid
+from mitm_tooling.utilities.identifiers import mk_uuid, name_plus_uuid
 from ..definitions import SupersetDatabaseDef
 
 

mitm_tooling/transformation/superset/factories/dataset.py

 from uuid import UUID
 
-import pydantic
 import sqlalchemy as sa
 
 from mitm_tooling.data_types import MITMDataType
 from mitm_tooling.extraction.sql.data_models import TableMetaInfo
+from mitm_tooling.utilities.identifiers import mk_uuid
 from .core import mk_column, mk_metric
-from .utils import mk_uuid
 from ..definitions import SupersetAggregate, SupersetDatasetDef
 
 

mitm_tooling/transformation/superset/factories/mitm_dataset.py

@@ -2,10 +2,9 @@ from typing import Sequence, Literal, Iterable
 from uuid import UUID
 
 from mitm_tooling.definition import MITM
-from ..common import name_plus_uuid
+from mitm_tooling.utilities.identifiers import mk_uuid, name_plus_uuid
 from ..definitions import SupersetMitMDatasetDef
 from ..definitions.mitm_dataset import RelatedTable, RelatedSlice, RelatedDashboard
-from ..factories.utils import mk_uuid
 
 
 def mk_related_obj(kind: Literal['table', 'slice', 'dashboard'],