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

assets.py

Blame
  • assets.py 5.99 KiB
    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