Skip to content
Snippets Groups Projects
Select Git revision
  • eecb58b36ec408bbabc13b54f08f8f216f32069b
  • master default protected
  • develop
3 results

importable.py

Blame
    • importable.py 2.80 KiB
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84 
    

    

    

    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