Improve import time of aiida.orm and aiida.storage (#6382)

1. Don't import `aiida.storage` from `aiida.orm`
2. Use `defer_build=True` for all pydantic models
3. Improve import of `storage.sqlite_zip`

src/aiida/brokers/rabbitmq/broker.py

@@ -5,8 +5,6 @@
 import functools
 import typing as t
 
-from packaging.version import parse
-
 from aiida.brokers.broker import Broker
 from aiida.common.log import AIIDA_LOGGER
 from aiida.manage.configuration import get_config_option
@@ -110,11 +108,15 @@ def is_rabbitmq_version_supported(self) -> bool:
 
         :return: boolean whether the current RabbitMQ version is supported.
         """
+        from packaging.version import parse
+
         return parse('3.6.0') <= self.get_rabbitmq_version() < parse('3.8.15')
 
     def get_rabbitmq_version(self):
         """Return the version of the RabbitMQ server that the current profile connects to.
 
         :return: :class:`packaging.version.Version`
         """
+        from packaging.version import parse
+
         return parse(self.get_communicator().server_properties['version'].decode('utf-8'))

src/aiida/orm/nodes/data/code/abstract.py

@@ -45,7 +45,7 @@ class AbstractCode(Data, metaclass=abc.ABCMeta):
     _KEY_ATTRIBUTE_WRAP_CMDLINE_PARAMS: str = 'wrap_cmdline_params'
     _KEY_EXTRA_IS_HIDDEN: str = 'hidden'  # Should become ``is_hidden`` once ``Code`` is dropped
 
-    class Model(BaseModel):
+    class Model(BaseModel, defer_build=True):
         """Model describing required information to create an instance."""
 
         label: str = MetadataField(

src/aiida/orm/nodes/data/folder.py

@@ -15,10 +15,12 @@
 import pathlib
 import typing as t
 
-from aiida.repository import File
-
 from .data import Data
 
+if t.TYPE_CHECKING:
+    from aiida.repository import File
+
+
 __all__ = ('FolderData',)
 
 FilePath = t.Union[str, pathlib.PurePosixPath]

src/aiida/orm/nodes/node.py

@@ -38,13 +38,13 @@
 from .caching import NodeCaching
 from .comments import NodeComments
 from .links import NodeLinks
-from .repository import NodeRepository
 
 if TYPE_CHECKING:
     from importlib_metadata import EntryPoint
 
     from ..implementation import StorageBackend
     from ..implementation.nodes import BackendNode  # noqa: F401
+    from .repository import NodeRepository
 
 __all__ = ('Node',)
 
@@ -107,6 +107,8 @@ def __init__(self, node: 'Node') -> None:
     @cached_property
     def repository(self) -> 'NodeRepository':
         """Return the repository for this node."""
+        from .repository import NodeRepository
+
         return NodeRepository(self._node)
 
     @cached_property

src/aiida/orm/nodes/repository.py

@@ -12,10 +12,10 @@
 
 from aiida.common import exceptions
 from aiida.manage import get_config_option
-from aiida.repository import File, Repository
-from aiida.repository.backend import SandboxRepositoryBackend
 
 if t.TYPE_CHECKING:
+    from aiida.repository import File, Repository
+
     from .node import Node
 
 __all__ = ('NodeRepository',)
@@ -77,6 +77,9 @@ def _repository(self) -> Repository:
 
         :return: the file repository instance.
         """
+        from aiida.repository import Repository
+        from aiida.repository.backend import SandboxRepositoryBackend
+
         if self._repository_instance is None:
             if self._node.is_stored:
                 backend = self._node.backend.get_repository()
@@ -100,6 +103,9 @@ def _repository(self, repository: Repository) -> None:
 
     def _store(self) -> None:
         """Store the repository in the backend."""
+        from aiida.repository import Repository
+        from aiida.repository.backend import SandboxRepositoryBackend
+
         if isinstance(self._repository.backend, SandboxRepositoryBackend):
             # Only if the backend repository is a sandbox do we have to clone its contents to the permanent repository.
             repository_backend = self._node.backend.get_repository()

src/aiida/storage/psql_dos/backend.py

@@ -74,7 +74,7 @@ class PsqlDosBackend(StorageBackend):
     The `django` backend was removed, to consolidate access to this storage.
     """
 
-    class Model(BaseModel):
+    class Model(BaseModel, defer_build=True):
         """Model describing required information to configure an instance of the storage."""
 
         database_engine: str = Field(

src/aiida/storage/sqlite_dos/backend.py

@@ -100,7 +100,7 @@ class SqliteDosStorage(PsqlDosBackend):
 
     migrator = SqliteDosMigrator
 
-    class Model(BaseModel):
+    class Model(BaseModel, defer_build=True):
         """Model describing required information to configure an instance of the storage."""
 
         filepath: str = Field(

src/aiida/storage/sqlite_temp/backend.py

@@ -43,7 +43,7 @@ class SqliteTempBackend(StorageBackend):
     and destroys it when it is garbage collected.
     """
 
-    class Model(BaseModel):
+    class Model(BaseModel, defer_build=True):
         filepath: str = Field(
             title='Temporary directory',
             description='Temporary directory in which to store data for this backend.',

src/aiida/storage/sqlite_zip/backend.py

@@ -20,7 +20,6 @@
 from typing import BinaryIO, Iterable, Iterator, Optional, Sequence, Tuple, cast
 from zipfile import ZipFile, is_zipfile
 
-from archive_path import ZipPath, extract_file_in_zip
 from pydantic import BaseModel, Field, field_validator
 from sqlalchemy.orm import Session
 
@@ -33,7 +32,6 @@
 from aiida.repository.backend.abstract import AbstractRepositoryBackend
 
 from . import orm
-from .migrator import get_schema_version_head, migrate, validate_storage
 from .utils import (
     DB_FILENAME,
     META_FILENAME,
@@ -68,7 +66,7 @@ class SqliteZipBackend(StorageBackend):
     read_only = True
     """This plugin is read only and data cannot be created or mutated."""
 
-    class Model(BaseModel):
+    class Model(BaseModel, defer_build=True):
         """Model describing required information to configure an instance of the storage."""
 
         filepath: str = Field(title='Filepath of the archive', description='Filepath of the archive.')
@@ -83,6 +81,8 @@ def filepath_exists_and_is_absolute(cls, value: str) -> str:
 
     @classmethod
     def version_head(cls) -> str:
+        from .migrator import get_schema_version_head
+
         return get_schema_version_head()
 
     @staticmethod
@@ -111,9 +111,13 @@ def initialise(cls, profile: 'Profile', reset: bool = False) -> bool:
             tests having run.
         :returns: ``True`` if the storage was initialised by the function call, ``False`` if it was already initialised.
         """
+        from archive_path import ZipPath
+
         filepath_archive = Path(profile.storage_config['filepath'])
 
         if filepath_archive.exists() and not reset:
+            from .migrator import migrate
+
             # The archive exists but ``reset == False``, so we try to migrate to the latest schema version. If the
             # migration works, we replace the original archive with the migrated one.
             with tempfile.TemporaryDirectory() as dirpath:
@@ -162,6 +166,8 @@ def migrate(cls, profile: Profile):
         raise NotImplementedError('use the :func:`aiida.storage.sqlite_zip.migrator.migrate` function directly.')
 
     def __init__(self, profile: Profile):
+        from .migrator import validate_storage
+
         super().__init__(profile)
         self._path = Path(profile.storage_config['filepath'])
         validate_storage(self._path)
@@ -194,6 +200,8 @@ def close(self):
 
     def get_session(self) -> Session:
         """Return an SQLAlchemy session."""
+        from archive_path import extract_file_in_zip
+
         if self._closed:
             raise ClosedStorage(str(self))
         if self._session is None:

src/aiida/storage/sqlite_zip/migrations/legacy_to_main.py

@@ -27,7 +27,6 @@
 from aiida.storage.log import MIGRATE_LOGGER
 
 from ..utils import DB_FILENAME, REPO_FOLDER, create_sqla_engine
-from . import v1_db_schema as v1_schema
 from .utils import update_metadata
 
 _NODE_ENTITY_NAME = 'Node'
@@ -45,15 +44,6 @@
     _COMMENT_ENTITY_NAME: {'dbnode': 'dbnode_id', 'user': 'user_id'},
 }
 
-aiida_orm_to_backend = {
-    _USER_ENTITY_NAME: v1_schema.DbUser,
-    _GROUP_ENTITY_NAME: v1_schema.DbGroup,
-    _NODE_ENTITY_NAME: v1_schema.DbNode,
-    _COMMENT_ENTITY_NAME: v1_schema.DbComment,
-    _COMPUTER_ENTITY_NAME: v1_schema.DbComputer,
-    _LOG_ENTITY_NAME: v1_schema.DbLog,
-}
-
 LEGACY_TO_MAIN_REVISION = 'main_0000'
 
 
@@ -138,6 +128,17 @@ def _json_to_sqlite(
     """Convert a JSON archive format to SQLite."""
     from aiida.tools.archive.common import batch_iter
 
+    from . import v1_db_schema as v1_schema
+
+    aiida_orm_to_backend = {
+        _USER_ENTITY_NAME: v1_schema.DbUser,
+        _GROUP_ENTITY_NAME: v1_schema.DbGroup,
+        _NODE_ENTITY_NAME: v1_schema.DbNode,
+        _COMMENT_ENTITY_NAME: v1_schema.DbComment,
+        _COMPUTER_ENTITY_NAME: v1_schema.DbComputer,
+        _LOG_ENTITY_NAME: v1_schema.DbLog,
+    }
+
     MIGRATE_LOGGER.report('Converting DB to SQLite')
 
     engine = create_sqla_engine(outpath)

src/aiida/storage/sqlite_zip/utils.py

@@ -9,12 +9,10 @@
 """Utilities for this backend."""
 
 import json
-import tarfile
 import zipfile
 from pathlib import Path
 from typing import Any, Dict, Optional, Union
 
-from archive_path import read_file_in_tar, read_file_in_zip
 from sqlalchemy import event
 from sqlalchemy.future.engine import Engine, create_engine
 
@@ -64,6 +62,10 @@ def extract_metadata(path: Union[str, Path], *, search_limit: Optional[int] = 10
 
     :param search_limit: the maximum number of records to search for the metadata file in a zip file.
     """
+    import tarfile
+
+    from archive_path import read_file_in_tar, read_file_in_zip
+
     path = Path(path)
     if not path.exists():
         raise UnreachableStorage(f'path not found: {path}')