Skip to content

export

Base store export class.

ExportSpec

Specification of MLTE store objects to be exported.

Source code in mlte/store/export/export.py 
 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
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
class ExportSpec:
    """Specification of MLTE store objects to be exported."""

    models: dict[str, list[str]] | None = None
    """
    Dict of models to be exported. Key is model ID, value is list of versions.

    An empty dict will export all models, all versions. An empty list of versions will export all versions.
    """

    custom_lists: list[CustomListName] | None = None
    """
    List of custom lists to be exported.

    An empty list will export all custom lists.
    """

    users: list[str] | None = None
    """
    List of user IDs for users to be exported.

    An empty list will export all users.
    """

    catalogs: list[str] | None = None
    """
    List of catalogs to be exported.

    An empty list will export all catalogs.
    """

    def __init__(
        self,
        artifact_store: ArtifactStore,
        user_store: UserStore,
        catalog_stores: CatalogStoreGroup,
        models: dict[str, list[str]] | None = None,
        custom_lists: list[CustomListName] | None = None,
        users: list[str] | None = None,
        catalogs: list[str] | None = None,
    ) -> None:
        self._setup_artifacts(artifact_store, models)
        self._setup_custom_lists(custom_lists)
        self._setup_users(user_store, users)
        self._setup_catalogs(catalog_stores, catalogs)

    def _setup_artifacts(
        self,
        artifact_store: ArtifactStore,
        models: dict[str, list[str]] | None = None,
    ) -> None:
        """Setup artifact export, accounts for the all option of an empty dict for models, lists for versions."""
        self.models = models

        if self.models is None:
            return

        with ManagedArtifactSession(
            artifact_store.session()
        ) as artifact_store_session:
            if self.models == {}:
                for model_id in artifact_store_session.model_mapper.list():
                    self.models[model_id] = []

            for model_id in self.models:
                if self.models[model_id] == []:
                    self.models[model_id] = (
                        artifact_store_session.version_mapper.list(model_id)
                    )

    def _setup_custom_lists(
        self, custom_lists: list[CustomListName] | None = None
    ) -> None:
        """Setup custom list export, accounts for the all option of an empty list."""
        self.custom_lists = custom_lists

        if self.custom_lists == []:
            for custom_list_id in CustomListName:
                self.custom_lists.append(custom_list_id)

    def _setup_users(
        self, user_store: UserStore, users: list[str] | None = None
    ) -> None:
        """Setup user export, accounts for the all option of an empty list."""
        self.users = users

        if self.users == []:
            with ManagedUserSession(user_store.session()) as user_store_session:
                self.users = user_store_session.user_mapper.list()

    def _setup_catalogs(
        self,
        catalog_stores: CatalogStoreGroup,
        catalogs: list[str] | None = None,
    ) -> None:
        """Setup catalog export, accounts for the all option of an empty list."""
        self.catalogs = catalogs

        if self.catalogs == []:
            self.catalogs = list(catalog_stores.catalogs.keys())

catalogs = None class-attribute instance-attribute

List of catalogs to be exported.

An empty list will export all catalogs.

custom_lists = None class-attribute instance-attribute

List of custom lists to be exported.

An empty list will export all custom lists.

models = None class-attribute instance-attribute

Dict of models to be exported. Key is model ID, value is list of versions.

An empty dict will export all models, all versions. An empty list of versions will export all versions.

users = None class-attribute instance-attribute

List of user IDs for users to be exported.

An empty list will export all users.

export_to_file(export_spec, output_path, artifact_store, custom_list_store, user_store, catalog_stores)

Export store data, writes the exported JSON to output_path as zip file.

Parameters:

Name Type Description Default
export_spec ExportSpec

Selection of MLTE store objects to be exported

 required
output_path Path

Path to write zipped JSON to

 required
Source code in mlte/store/export/export.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def export_to_file(
    export_spec: ExportSpec,
    output_path: Path,
    artifact_store: ArtifactStore,
    custom_list_store: CustomListStore,
    user_store: UserStore,
    catalog_stores: CatalogStoreGroup,
) -> None:
    """
    Export store data, writes the exported JSON to output_path as zip file.

    :param export_spec: Selection of MLTE store objects to be exported
    :param output_path: Path to write zipped JSON to
    """
    export_json = _export(
        export_spec,
        artifact_store,
        custom_list_store,
        user_store,
        catalog_stores,
    )

    # Ensure output directory exists
    os.makedirs(output_path, exist_ok=True)

    with zipfile.ZipFile(
        os.path.join(output_path, EXPORT_ZIP_FILE), "w", zipfile.ZIP_DEFLATED
    ) as zip_export_file:
        zip_export_file.writestr(
            EXPORT_JSON_FILE, json.dumps(export_json, indent=4)
        )