From 18e1515bba231cb8c05377dbb0ac028337c5faec Mon Sep 17 00:00:00 2001 From: Chris Sewell Date: Fri, 18 Feb 2022 12:46:20 +0100 Subject: [PATCH] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20REFACTOR:=20`backends.Back?= =?UTF-8?q?end`=20->=20`storage=5Fbackend.StorageBackend`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .pre-commit-config.yaml | 2 +- aiida/manage/configuration/profile.py | 4 +- aiida/manage/manager.py | 8 ++-- aiida/orm/authinfos.py | 4 +- aiida/orm/comments.py | 6 ++- aiida/orm/computers.py | 4 +- aiida/orm/entities.py | 20 ++++----- aiida/orm/groups.py | 4 +- aiida/orm/implementation/__init__.py | 4 +- aiida/orm/implementation/entities.py | 10 ++--- aiida/orm/implementation/querybuilder.py | 8 ++-- .../{backends.py => storage_backend.py} | 4 +- aiida/orm/logs.py | 4 +- aiida/orm/nodes/node.py | 17 ++++---- aiida/orm/querybuilder.py | 6 +-- aiida/orm/users.py | 6 +-- aiida/storage/control.py | 9 ++-- aiida/storage/psql_dos/backend.py | 5 +-- aiida/storage/testbase.py | 4 +- aiida/tools/archive/abstract.py | 4 +- aiida/tools/archive/create.py | 10 ++--- .../archive/implementations/sqlite/backend.py | 4 +- aiida/tools/archive/imports.py | 41 ++++++++++--------- aiida/tools/graph/graph_traversers.py | 8 ++-- aiida/tools/visualization/graph.py | 6 +-- tests/tools/archive/test_backend.py | 4 +- 26 files changed, 105 insertions(+), 101 deletions(-) rename aiida/orm/implementation/{backends.py => storage_backend.py} (98%) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3cad7ac827..aa07afaf39 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -76,7 +76,7 @@ repos: aiida/manage/database/delete/nodes.py| aiida/orm/querybuilder.py| aiida/orm/implementation/entities.py| - aiida/orm/implementation/backends.py| + aiida/orm/implementation/storage_backend.py| aiida/orm/implementation/authinfos.py| aiida/orm/implementation/comments.py| aiida/orm/implementation/computers.py| diff --git a/aiida/manage/configuration/profile.py b/aiida/manage/configuration/profile.py index d4698b7d18..19b9aa9270 100644 --- a/aiida/manage/configuration/profile.py +++ b/aiida/manage/configuration/profile.py @@ -20,7 +20,7 @@ from .settings import DAEMON_DIR, DAEMON_LOG_DIR if TYPE_CHECKING: - from aiida.orm.implementation.backends import Backend + from aiida.orm.implementation import StorageBackend __all__ = ('Profile',) @@ -121,7 +121,7 @@ def set_storage(self, name: str, config: Dict[str, Any]) -> None: self._attributes[self.KEY_STORAGE][self.KEY_STORAGE_CONFIG] = config @property - def storage_cls(self) -> Type['Backend']: + def storage_cls(self) -> Type['StorageBackend']: """Return the storage backend class for this profile.""" if self.storage_backend == 'psql_dos': from aiida.storage.psql_dos.backend import PsqlDosBackend diff --git a/aiida/manage/manager.py b/aiida/manage/manager.py index 5d3963f537..4bd6d97631 100644 --- a/aiida/manage/manager.py +++ b/aiida/manage/manager.py @@ -23,7 +23,7 @@ from aiida.engine.runners import Runner from aiida.manage.configuration.config import Config from aiida.manage.configuration.profile import Profile - from aiida.orm.implementation import Backend + from aiida.orm.implementation import StorageBackend __all__ = ('get_manager',) @@ -64,7 +64,7 @@ class Manager: def __init__(self) -> None: # note: the config currently references the global variables self._profile: Optional['Profile'] = None - self._profile_storage: Optional['Backend'] = None + self._profile_storage: Optional['StorageBackend'] = None self._daemon_client: Optional['DaemonClient'] = None self._communicator: Optional['RmqThreadCommunicator'] = None self._process_controller: Optional['RemoteProcessThreadController'] = None @@ -193,7 +193,7 @@ def get_option(self, option_name: str) -> Any: option = get_option(option_name) return option.default - def get_backend(self) -> 'Backend': + def get_backend(self) -> 'StorageBackend': """Return the current profile's storage backend, loading it if necessary. Deprecated: use `get_profile_storage` instead. @@ -202,7 +202,7 @@ def get_backend(self) -> 'Backend': warn('get_backend() is deprecated, use get_profile_storage() instead', AiidaDeprecationWarning) return self.get_profile_storage() - def get_profile_storage(self) -> 'Backend': + def get_profile_storage(self) -> 'StorageBackend': """Return the current profile's storage backend, loading it if necessary.""" from aiida.common import ConfigurationError from aiida.common.log import configure_logging diff --git a/aiida/orm/authinfos.py b/aiida/orm/authinfos.py index d83718ba21..553655928e 100644 --- a/aiida/orm/authinfos.py +++ b/aiida/orm/authinfos.py @@ -19,7 +19,7 @@ if TYPE_CHECKING: from aiida.orm import Computer, User - from aiida.orm.implementation import Backend, BackendAuthInfo + from aiida.orm.implementation import BackendAuthInfo, StorageBackend from aiida.transports import Transport __all__ = ('AuthInfo',) @@ -51,7 +51,7 @@ def objects(cls: Type['AuthInfo']) -> AuthInfoCollection: # type: ignore[misc] PROPERTY_WORKDIR = 'workdir' - def __init__(self, computer: 'Computer', user: 'User', backend: Optional['Backend'] = None) -> None: + def __init__(self, computer: 'Computer', user: 'User', backend: Optional['StorageBackend'] = None) -> None: """Create an `AuthInfo` instance for the given computer and user. :param computer: a `Computer` instance diff --git a/aiida/orm/comments.py b/aiida/orm/comments.py index a8e3a99ea2..0475c4098d 100644 --- a/aiida/orm/comments.py +++ b/aiida/orm/comments.py @@ -18,7 +18,7 @@ if TYPE_CHECKING: from aiida.orm import Node, User - from aiida.orm.implementation import Backend, BackendComment + from aiida.orm.implementation import BackendComment, StorageBackend __all__ = ('Comment',) @@ -72,7 +72,9 @@ class Comment(entities.Entity['BackendComment']): def objects(cls: Type['Comment']) -> CommentCollection: # type: ignore[misc] # pylint: disable=no-self-argument return CommentCollection.get_cached(cls, get_manager().get_profile_storage()) - def __init__(self, node: 'Node', user: 'User', content: Optional[str] = None, backend: Optional['Backend'] = None): + def __init__( + self, node: 'Node', user: 'User', content: Optional[str] = None, backend: Optional['StorageBackend'] = None + ): """Create a Comment for a given node and user :param node: a Node instance diff --git a/aiida/orm/computers.py b/aiida/orm/computers.py index 2a2361aaa8..88200cb3a8 100644 --- a/aiida/orm/computers.py +++ b/aiida/orm/computers.py @@ -21,7 +21,7 @@ if TYPE_CHECKING: from aiida.orm import AuthInfo, User - from aiida.orm.implementation import Backend, BackendComputer + from aiida.orm.implementation import BackendComputer, StorageBackend from aiida.schedulers import Scheduler from aiida.transports import Transport @@ -89,7 +89,7 @@ def __init__( # pylint: disable=too-many-arguments transport_type: str = '', scheduler_type: str = '', workdir: str = None, - backend: Optional['Backend'] = None, + backend: Optional['StorageBackend'] = None, ) -> None: """Construct a new computer.""" backend = backend or get_manager().get_profile_storage() diff --git a/aiida/orm/entities.py b/aiida/orm/entities.py index 62a1a35896..65efd1f3e9 100644 --- a/aiida/orm/entities.py +++ b/aiida/orm/entities.py @@ -21,7 +21,7 @@ from aiida.manage import get_manager if TYPE_CHECKING: - from aiida.orm.implementation import Backend, BackendEntity + from aiida.orm.implementation import BackendEntity, StorageBackend from aiida.orm.querybuilder import FilterType, OrderByType, QueryBuilder __all__ = ('Entity', 'Collection', 'EntityAttributesMixin', 'EntityExtrasMixin', 'EntityTypes') @@ -56,28 +56,28 @@ def _entity_base_cls() -> Type[EntityType]: @classmethod @lru_cache(maxsize=100) - def get_cached(cls, entity_class: Type[EntityType], backend: 'Backend'): + def get_cached(cls, entity_class: Type[EntityType], backend: 'StorageBackend'): """Get the cached collection instance for the given entity class and backend. :param backend: the backend instance to get the collection for """ - from aiida.orm.implementation import Backend - type_check(backend, Backend) + from aiida.orm.implementation import StorageBackend + type_check(backend, StorageBackend) return cls(entity_class, backend=backend) - def __init__(self, entity_class: Type[EntityType], backend: Optional['Backend'] = None) -> None: + def __init__(self, entity_class: Type[EntityType], backend: Optional['StorageBackend'] = None) -> None: """ Construct a new entity collection. :param entity_class: the entity type e.g. User, Computer, etc :param backend: the backend instance to get the collection for, or use the default """ - from aiida.orm.implementation import Backend - type_check(backend, Backend, allow_none=True) + from aiida.orm.implementation import StorageBackend + type_check(backend, StorageBackend, allow_none=True) assert issubclass(entity_class, self._entity_base_cls()) self._backend = backend or get_manager().get_profile_storage() self._entity_type = entity_class - def __call__(self: CollectionType, backend: 'Backend') -> CollectionType: + def __call__(self: CollectionType, backend: 'StorageBackend') -> CollectionType: """Get or create a cached collection using a new backend.""" if backend is self._backend: return self @@ -89,7 +89,7 @@ def entity_type(self) -> Type[EntityType]: return self._entity_type @property - def backend(self) -> 'Backend': + def backend(self) -> 'StorageBackend': """Return the backend.""" return self._backend @@ -244,7 +244,7 @@ def is_stored(self) -> bool: return self._backend_entity.is_stored @property - def backend(self) -> 'Backend': + def backend(self) -> 'StorageBackend': """Get the backend for this entity""" return self._backend_entity.backend diff --git a/aiida/orm/groups.py b/aiida/orm/groups.py index 33d52dfb87..976178acba 100644 --- a/aiida/orm/groups.py +++ b/aiida/orm/groups.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: from aiida.orm import Node, User - from aiida.orm.implementation import Backend, BackendGroup + from aiida.orm.implementation import BackendGroup, StorageBackend __all__ = ('Group', 'AutoGroup', 'ImportGroup', 'UpfFamily') @@ -127,7 +127,7 @@ def __init__( user: Optional['User'] = None, description: str = '', type_string: Optional[str] = None, - backend: Optional['Backend'] = None + backend: Optional['StorageBackend'] = None ): """ Create a new group. Either pass a dbgroup parameter, to reload diff --git a/aiida/orm/implementation/__init__.py b/aiida/orm/implementation/__init__.py index 588a8a7dae..0f02fcbf65 100644 --- a/aiida/orm/implementation/__init__.py +++ b/aiida/orm/implementation/__init__.py @@ -15,7 +15,6 @@ # pylint: disable=wildcard-import from .authinfos import * -from .backends import * from .comments import * from .computers import * from .entities import * @@ -23,11 +22,11 @@ from .logs import * from .nodes import * from .querybuilder import * +from .storage_backend import * from .users import * from .utils import * __all__ = ( - 'Backend', 'BackendAuthInfo', 'BackendAuthInfoCollection', 'BackendCollection', @@ -47,6 +46,7 @@ 'BackendUser', 'BackendUserCollection', 'EntityType', + 'StorageBackend', 'clean_value', 'validate_attribute_extra_key', ) diff --git a/aiida/orm/implementation/entities.py b/aiida/orm/implementation/entities.py index 9e0e4df17e..41f8e8b988 100644 --- a/aiida/orm/implementation/entities.py +++ b/aiida/orm/implementation/entities.py @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Dict, Generic, Iterable, List, Tuple, Type, TypeVar if TYPE_CHECKING: - from aiida.orm.implementation import Backend + from aiida.orm.implementation import StorageBackend __all__ = ('BackendEntity', 'BackendCollection', 'EntityType', 'BackendEntityExtrasMixin') @@ -22,11 +22,11 @@ class BackendEntity(abc.ABC): """An first-class entity in the backend""" - def __init__(self, backend: 'Backend', **kwargs: Any): # pylint: disable=unused-argument + def __init__(self, backend: 'StorageBackend', **kwargs: Any): # pylint: disable=unused-argument self._backend = backend @property - def backend(self) -> 'Backend': + def backend(self) -> 'StorageBackend': """Return the backend this entity belongs to :return: the backend instance @@ -74,7 +74,7 @@ class BackendCollection(Generic[EntityType]): ENTITY_CLASS: ClassVar[Type[EntityType]] # type: ignore[misc] - def __init__(self, backend: 'Backend'): + def __init__(self, backend: 'StorageBackend'): """ :param backend: the backend this collection belongs to """ @@ -82,7 +82,7 @@ def __init__(self, backend: 'Backend'): self._backend = backend @property - def backend(self) -> 'Backend': + def backend(self) -> 'StorageBackend': """Return the backend.""" return self._backend diff --git a/aiida/orm/implementation/querybuilder.py b/aiida/orm/implementation/querybuilder.py index 567b0cfdde..55e649aac3 100644 --- a/aiida/orm/implementation/querybuilder.py +++ b/aiida/orm/implementation/querybuilder.py @@ -16,7 +16,7 @@ from aiida.orm.entities import EntityTypes if TYPE_CHECKING: - from aiida.orm.implementation import Backend + from aiida.orm.implementation import StorageBackend __all__ = ('BackendQueryBuilder',) @@ -80,12 +80,12 @@ class QueryDictType(TypedDict): class BackendQueryBuilder(abc.ABC): """Backend query builder interface""" - def __init__(self, backend: 'Backend'): + def __init__(self, backend: 'StorageBackend'): """ :param backend: the backend """ - from . import backends - type_check(backend, backends.Backend) + from .storage_backend import StorageBackend + type_check(backend, StorageBackend) self._backend = backend @abc.abstractmethod diff --git a/aiida/orm/implementation/backends.py b/aiida/orm/implementation/storage_backend.py similarity index 98% rename from aiida/orm/implementation/backends.py rename to aiida/orm/implementation/storage_backend.py index cf99ea416f..2541925d51 100644 --- a/aiida/orm/implementation/backends.py +++ b/aiida/orm/implementation/storage_backend.py @@ -27,12 +27,12 @@ ) from aiida.repository.backend.abstract import AbstractRepositoryBackend -__all__ = ('Backend',) +__all__ = ('StorageBackend',) TransactionType = TypeVar('TransactionType') -class Backend(abc.ABC): # pylint: disable=too-many-public-methods +class StorageBackend(abc.ABC): # pylint: disable=too-many-public-methods """Abstraction for a backend to read/write persistent data for a profile's provenance graph. AiiDA splits data storage into two sources: diff --git a/aiida/orm/logs.py b/aiida/orm/logs.py index 885ffb7bb1..d1aebb9e5d 100644 --- a/aiida/orm/logs.py +++ b/aiida/orm/logs.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: from aiida.orm import Node - from aiida.orm.implementation import Backend, BackendLog + from aiida.orm.implementation import BackendLog, StorageBackend from aiida.orm.querybuilder import FilterType, OrderByType __all__ = ('Log', 'OrderSpecifier', 'ASCENDING', 'DESCENDING') @@ -142,7 +142,7 @@ def __init__( dbnode_id: int, message: str = '', metadata: Optional[Dict[str, Any]] = None, - backend: Optional['Backend'] = None + backend: Optional['StorageBackend'] = None ): # pylint: disable=too-many-arguments """Construct a new log diff --git a/aiida/orm/nodes/node.py b/aiida/orm/nodes/node.py index 4b416e1b05..0b53719316 100644 --- a/aiida/orm/nodes/node.py +++ b/aiida/orm/nodes/node.py @@ -48,8 +48,7 @@ from .repository import NodeRepositoryMixin if TYPE_CHECKING: - from ..implementation import Backend - from ..implementation.nodes import BackendNode + from ..implementation import BackendNode, StorageBackend __all__ = ('Node',) @@ -157,7 +156,7 @@ def objects(cls: Type[NodeType]) -> NodeCollection[NodeType]: # pylint: disable def __init__( self, - backend: Optional['Backend'] = None, + backend: Optional['StorageBackend'] = None, user: Optional[User] = None, computer: Optional[Computer] = None, **kwargs: Any @@ -292,7 +291,7 @@ def process_type(self, value: str) -> None: :param value: the new value to set """ - self.backend_entity.process_type = value # type: ignore[misc] + self.backend_entity.process_type = value @property def label(self) -> str: @@ -308,7 +307,7 @@ def label(self, value: str) -> None: :param value: the new value to set """ - self.backend_entity.label = value # type: ignore[misc] + self.backend_entity.label = value @property def description(self) -> str: @@ -324,7 +323,7 @@ def description(self, value: str) -> None: :param value: the new value to set """ - self.backend_entity.description = value # type: ignore[misc] + self.backend_entity.description = value @property def repository_metadata(self) -> Dict[str, Any]: @@ -340,7 +339,7 @@ def repository_metadata(self, value: Dict[str, Any]) -> None: :param value: the new value to set """ - self.backend_entity.repository_metadata = value # type: ignore[misc] + self.backend_entity.repository_metadata = value @property def computer(self) -> Optional[Computer]: @@ -361,7 +360,7 @@ def computer(self, computer: Optional[Computer]) -> None: type_check(computer, Computer, allow_none=True) - self.backend_entity.computer = None if computer is None else computer.backend_entity # type: ignore[misc] + self.backend_entity.computer = None if computer is None else computer.backend_entity @property def user(self) -> User: @@ -378,7 +377,7 @@ def user(self, user: User) -> None: raise exceptions.ModificationNotAllowed('cannot set the user on a stored node') type_check(user, User) - self.backend_entity.user = user.backend_entity # type: ignore[misc] + self.backend_entity.user = user.backend_entity @property def ctime(self) -> datetime.datetime: diff --git a/aiida/orm/querybuilder.py b/aiida/orm/querybuilder.py index fcdf6a5fbd..483ca86784 100644 --- a/aiida/orm/querybuilder.py +++ b/aiida/orm/querybuilder.py @@ -53,7 +53,7 @@ if TYPE_CHECKING: # pylint: disable=ungrouped-imports from aiida.engine import Process - from aiida.orm.implementation import Backend + from aiida.orm.implementation import StorageBackend __all__ = ('QueryBuilder',) @@ -94,7 +94,7 @@ class QueryBuilder: def __init__( self, - backend: Optional['Backend'] = None, + backend: Optional['StorageBackend'] = None, *, debug: bool = False, path: Optional[Sequence[Union[str, Dict[str, Any], EntityClsType]]] = (), @@ -190,7 +190,7 @@ def __init__( self.order_by(order_by) @property - def backend(self) -> 'Backend': + def backend(self) -> 'StorageBackend': """Return the backend used by the QueryBuilder.""" return self._backend diff --git a/aiida/orm/users.py b/aiida/orm/users.py index abc56e2e19..309b4e1c7d 100644 --- a/aiida/orm/users.py +++ b/aiida/orm/users.py @@ -17,7 +17,7 @@ from . import entities if TYPE_CHECKING: - from aiida.orm.implementation import Backend, BackendUser + from aiida.orm.implementation import BackendUser, StorageBackend __all__ = ('User',) @@ -29,7 +29,7 @@ class UserCollection(entities.Collection['User']): def _entity_base_cls() -> Type['User']: return User - def __init__(self, entity_class: Type['User'], backend: Optional['Backend'] = None) -> None: + def __init__(self, entity_class: Type['User'], backend: Optional['StorageBackend'] = None) -> None: super().__init__(entity_class=entity_class, backend=backend) self._default_user: Optional[User] = None @@ -82,7 +82,7 @@ def __init__( first_name: str = '', last_name: str = '', institution: str = '', - backend: Optional['Backend'] = None + backend: Optional['StorageBackend'] = None ): """Create a new `User`.""" # pylint: disable=too-many-arguments diff --git a/aiida/storage/control.py b/aiida/storage/control.py index 044424e5f6..96d439f571 100644 --- a/aiida/storage/control.py +++ b/aiida/storage/control.py @@ -18,7 +18,7 @@ from aiida.manage import get_manager if TYPE_CHECKING: - from aiida.orm.implementation import Backend + from aiida.orm.implementation import StorageBackend __all__ = ('MAINTAIN_LOGGER',) @@ -28,7 +28,7 @@ def repository_maintain( full: bool = False, dry_run: bool = False, - backend: Optional['Backend'] = None, + backend: Optional['StorageBackend'] = None, **kwargs, ) -> None: """Performs maintenance tasks on the repository. @@ -56,7 +56,8 @@ def repository_maintain( repository.maintain(live=not full, dry_run=dry_run, **kwargs) -def get_unreferenced_keyset(check_consistency: bool = True, aiida_backend: Optional['Backend'] = None) -> Set[str]: +def get_unreferenced_keyset(check_consistency: bool = True, + aiida_backend: Optional['StorageBackend'] = None) -> Set[str]: """Returns the keyset of objects that exist in the repository but are not tracked by AiiDA. This should be all the soft-deleted files. @@ -92,7 +93,7 @@ def get_unreferenced_keyset(check_consistency: bool = True, aiida_backend: Optio return keyset_repository - keyset_database -def get_repository_info(statistics: bool = False, backend: Optional['Backend'] = None) -> dict: +def get_repository_info(statistics: bool = False, backend: Optional['StorageBackend'] = None) -> dict: """Returns general information on the repository.""" if backend is None: backend = get_manager().get_profile_storage() diff --git a/aiida/storage/psql_dos/backend.py b/aiida/storage/psql_dos/backend.py index 862bd6103b..0ffa7a3035 100644 --- a/aiida/storage/psql_dos/backend.py +++ b/aiida/storage/psql_dos/backend.py @@ -21,8 +21,7 @@ from aiida.manage.configuration.profile import Profile from aiida.orm import User from aiida.orm.entities import EntityTypes -from aiida.orm.implementation.backends import Backend -from aiida.orm.implementation.entities import BackendEntity +from aiida.orm.implementation import BackendEntity, StorageBackend from aiida.storage.psql_dos.migrator import REPOSITORY_UUID_KEY, PsqlDostoreMigrator from aiida.storage.psql_dos.models import base @@ -41,7 +40,7 @@ } -class PsqlDosBackend(Backend): # pylint: disable=too-many-public-methods +class PsqlDosBackend(StorageBackend): # pylint: disable=too-many-public-methods """An AiiDA storage backend that stores data in a PostgreSQL database and disk-objectstore repository. Note, there were originally two such backends, `sqlalchemy` and `django`. diff --git a/aiida/storage/testbase.py b/aiida/storage/testbase.py index 05aa9b7d8e..30a093f658 100644 --- a/aiida/storage/testbase.py +++ b/aiida/storage/testbase.py @@ -16,7 +16,7 @@ from aiida.common.exceptions import TestsNotAllowedError from aiida.common.lang import classproperty from aiida.manage import configuration, get_manager -from aiida.orm.implementation import Backend +from aiida.orm.implementation import StorageBackend TEST_KEYWORD = 'test_' @@ -33,7 +33,7 @@ def check_if_tests_can_run(): class AiidaTestCase(unittest.TestCase): """This is the base class for AiiDA tests, independent of the backend.""" _class_was_setup = False - backend: Optional[Backend] = None + backend: Optional[StorageBackend] = None @classmethod def setUpClass(cls): diff --git a/aiida/tools/archive/abstract.py b/aiida/tools/archive/abstract.py index 4689ecd582..b45eded9a6 100644 --- a/aiida/tools/archive/abstract.py +++ b/aiida/tools/archive/abstract.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: from aiida.orm import QueryBuilder from aiida.orm.entities import Entity, EntityTypes - from aiida.orm.implementation import Backend + from aiida.orm.implementation import StorageBackend from aiida.tools.visualization.graph import Graph SelfType = TypeVar('SelfType') @@ -145,7 +145,7 @@ def get_metadata(self) -> Dict[str, Any]: """ @abstractmethod - def get_backend(self) -> 'Backend': + def get_backend(self) -> 'StorageBackend': """Return a 'read-only' backend for the archive.""" # below are convenience methods for some common use cases diff --git a/aiida/tools/archive/create.py b/aiida/tools/archive/create.py index 34166badb8..edd60d5132 100644 --- a/aiida/tools/archive/create.py +++ b/aiida/tools/archive/create.py @@ -29,7 +29,7 @@ from aiida.common.progress_reporter import get_progress_reporter from aiida.manage import get_manager from aiida.orm.entities import EntityTypes -from aiida.orm.implementation import Backend +from aiida.orm.implementation import StorageBackend from aiida.orm.utils.links import LinkQuadruple from aiida.tools.graph.graph_traversers import get_nodes_export, validate_traversal_rules @@ -59,7 +59,7 @@ def create_archive( batch_size: int = 1000, compression: int = 6, test_run: bool = False, - backend: Optional[Backend] = None, + backend: Optional[StorageBackend] = None, **traversal_rules: bool ) -> Path: """Export AiiDA data to an archive file. @@ -151,7 +151,7 @@ def create_archive( """ # check the backend backend = backend or get_manager().get_profile_storage() - type_check(backend, Backend) + type_check(backend, StorageBackend) # create a function to get a query builder instance for the backend querybuilder = lambda: orm.QueryBuilder(backend=backend) @@ -425,7 +425,7 @@ def _collect_all_entities( def _collect_required_entities( querybuilder: QbType, entity_ids: Dict[EntityTypes, Set[int]], traversal_rules: Dict[str, bool], - include_authinfos: bool, include_comments: bool, include_logs: bool, backend: Backend, batch_size: int + include_authinfos: bool, include_comments: bool, include_logs: bool, backend: StorageBackend, batch_size: int ) -> Tuple[List[list], Set[LinkQuadruple]]: """Collect required entities, given a set of starting entities and provenance graph traversal rules. @@ -565,7 +565,7 @@ def _collect_required_entities( def _stream_repo_files( - key_format: str, writer: ArchiveWriterAbstract, node_ids: Set[int], backend: Backend, batch_size: int + key_format: str, writer: ArchiveWriterAbstract, node_ids: Set[int], backend: StorageBackend, batch_size: int ) -> None: """Collect all repository object keys from the nodes, then stream the files to the archive.""" keys = set(orm.Node.objects(backend).iter_repo_keys(filters={'id': {'in': list(node_ids)}}, batch_size=batch_size)) diff --git a/aiida/tools/archive/implementations/sqlite/backend.py b/aiida/tools/archive/implementations/sqlite/backend.py index 9d93ce7b1d..06afc51907 100644 --- a/aiida/tools/archive/implementations/sqlite/backend.py +++ b/aiida/tools/archive/implementations/sqlite/backend.py @@ -27,7 +27,7 @@ from aiida.common.exceptions import UnreachableStorage from aiida.manage import Profile from aiida.orm.entities import EntityTypes -from aiida.orm.implementation.backends import Backend as BackendAbstract +from aiida.orm.implementation import StorageBackend from aiida.repository.backend.abstract import AbstractRepositoryBackend # we need to import all models, to ensure they are loaded on the SQLA Metadata from aiida.storage.psql_dos.models import authinfo, base, comment, computer, group, log, node, user @@ -245,7 +245,7 @@ def table_groups_nodes(self): return DbGroupNodes.__table__ # type: ignore[attr-defined] # pylint: disable=no-member -class ArchiveReadOnlyBackend(BackendAbstract): # pylint: disable=too-many-public-methods +class ArchiveReadOnlyBackend(StorageBackend): # pylint: disable=too-many-public-methods """A read-only backend for the archive.""" @classmethod diff --git a/aiida/tools/archive/imports.py b/aiida/tools/archive/imports.py index bca6cfaea5..7c0b002bf9 100644 --- a/aiida/tools/archive/imports.py +++ b/aiida/tools/archive/imports.py @@ -22,7 +22,7 @@ from aiida.common.progress_reporter import get_progress_reporter from aiida.manage import get_manager from aiida.orm.entities import EntityTypes -from aiida.orm.implementation import Backend +from aiida.orm.implementation import StorageBackend from aiida.orm.querybuilder import QueryBuilder from aiida.repository import Repository @@ -65,7 +65,7 @@ def import_archive( create_group: bool = True, group: Optional[orm.Group] = None, test_run: bool = False, - backend: Optional[Backend] = None, + backend: Optional[StorageBackend] = None, ) -> Optional[int]: """Import an archive into the AiiDA backend. @@ -116,7 +116,7 @@ def import_archive( type_check(group, orm.Group, allow_none=True) type_check(test_run, bool) backend = backend or get_manager().get_profile_storage() - type_check(backend, Backend) + type_check(backend, StorageBackend) if group and not group.is_stored: group.store() @@ -194,8 +194,8 @@ def import_archive( def _add_new_entities( - etype: EntityTypes, total: int, unique_field: str, backend_unique_id: dict, backend_from: Backend, - backend_to: Backend, batch_size: int, transform: Callable[[dict], dict] + etype: EntityTypes, total: int, unique_field: str, backend_unique_id: dict, backend_from: StorageBackend, + backend_to: StorageBackend, batch_size: int, transform: Callable[[dict], dict] ) -> None: """Add new entities to the output backend and update the mapping of unique field -> id.""" IMPORT_LOGGER.report(f'Adding {total} new {etype.value}(s)') @@ -216,7 +216,7 @@ def _add_new_entities( progress.update(nrows) -def _import_users(backend_from: Backend, backend_to: Backend, batch_size: int) -> Dict[int, int]: +def _import_users(backend_from: StorageBackend, backend_to: StorageBackend, batch_size: int) -> Dict[int, int]: """Import users from one backend to another. :returns: mapping of input backend id to output backend id @@ -254,7 +254,7 @@ def _import_users(backend_from: Backend, backend_to: Backend, batch_size: int) - return {int(i): output_email_id[email] for i, email in input_id_email.items()} -def _import_computers(backend_from: Backend, backend_to: Backend, batch_size: int) -> Dict[int, int]: +def _import_computers(backend_from: StorageBackend, backend_to: StorageBackend, batch_size: int) -> Dict[int, int]: """Import computers from one backend to another. :returns: mapping of input backend id to output backend id @@ -323,7 +323,7 @@ def transform(row: dict) -> dict: def _import_authinfos( - backend_from: Backend, backend_to: Backend, batch_size: int, user_ids_archive_backend: Dict[int, int], + backend_from: StorageBackend, backend_to: StorageBackend, batch_size: int, user_ids_archive_backend: Dict[int, int], computer_ids_archive_backend: Dict[int, int] ) -> None: """Import logs from one backend to another. @@ -402,7 +402,7 @@ def transform(row: dict) -> dict: def _import_nodes( - backend_from: Backend, backend_to: Backend, batch_size: int, user_ids_archive_backend: Dict[int, int], + backend_from: StorageBackend, backend_to: StorageBackend, batch_size: int, user_ids_archive_backend: Dict[int, int], computer_ids_archive_backend: Dict[int, int], import_new_extras: bool, merge_extras: MergeExtrasType ) -> Dict[int, int]: """Import users from one backend to another. @@ -480,8 +480,9 @@ def __call__(self, row: dict) -> dict: return data -def _import_logs(backend_from: Backend, backend_to: Backend, batch_size: int, - node_ids_archive_backend: Dict[int, int]) -> Dict[int, int]: +def _import_logs( + backend_from: StorageBackend, backend_to: StorageBackend, batch_size: int, node_ids_archive_backend: Dict[int, int] +) -> Dict[int, int]: """Import logs from one backend to another. :returns: mapping of input backend id to output backend id @@ -528,7 +529,8 @@ def transform(row: dict) -> dict: def _merge_node_extras( - backend_from: Backend, backend_to: Backend, batch_size: int, backend_uuid_id: Dict[str, int], mode: MergeExtrasType + backend_from: StorageBackend, backend_to: StorageBackend, batch_size: int, backend_uuid_id: Dict[str, int], + mode: MergeExtrasType ) -> None: """Merge extras from the input backend with the ones in the output backend. @@ -672,8 +674,8 @@ def __call__(self, row: dict) -> dict: def _import_comments( - backend_from: Backend, - backend: Backend, + backend_from: StorageBackend, + backend: StorageBackend, batch_size: int, user_ids_archive_backend: Dict[int, int], node_ids_archive_backend: Dict[int, int], @@ -753,7 +755,7 @@ def _transform(row): def _import_links( - backend_from: Backend, backend_to: Backend, batch_size: int, node_ids_archive_backend: Dict[int, int] + backend_from: StorageBackend, backend_to: StorageBackend, batch_size: int, node_ids_archive_backend: Dict[int, int] ) -> None: """Import links from one backend to another.""" @@ -922,7 +924,7 @@ def __call__(self, row: dict) -> dict: def _import_groups( - backend_from: Backend, backend_to: Backend, batch_size: int, user_ids_archive_backend: Dict[int, int], + backend_from: StorageBackend, backend_to: StorageBackend, batch_size: int, user_ids_archive_backend: Dict[int, int], node_ids_archive_backend: Dict[int, int] ) -> Set[str]: """Import groups from the input backend, and add group -> node records. @@ -1002,7 +1004,7 @@ def group_node_transform(row): def _make_import_group( - group: Optional[orm.Group], labels: Set[str], node_ids_archive_backend: Dict[int, int], backend_to: Backend, + group: Optional[orm.Group], labels: Set[str], node_ids_archive_backend: Dict[int, int], backend_to: StorageBackend, batch_size: int ) -> Optional[int]: """Make an import group containing all imported nodes. @@ -1065,7 +1067,8 @@ def _make_import_group( return group_id -def _get_new_object_keys(key_format: str, backend_from: Backend, backend_to: Backend, batch_size: int) -> Set[str]: +def _get_new_object_keys(key_format: str, backend_from: StorageBackend, backend_to: StorageBackend, + batch_size: int) -> Set[str]: """Return the object keys that need to be added to the backend.""" archive_hashkeys: Set[str] = set() query = QueryBuilder(backend=backend_from).append(orm.Node, project='repository_metadata') @@ -1092,7 +1095,7 @@ def _get_new_object_keys(key_format: str, backend_from: Backend, backend_to: Bac return new_hashkeys -def _add_files_to_repo(backend_from: Backend, backend_to: Backend, new_keys: Set[str]) -> None: +def _add_files_to_repo(backend_from: StorageBackend, backend_to: StorageBackend, new_keys: Set[str]) -> None: """Add the new files to the repository.""" if not new_keys: return None diff --git a/aiida/tools/graph/graph_traversers.py b/aiida/tools/graph/graph_traversers.py index 991332a817..22e9ef4c93 100644 --- a/aiida/tools/graph/graph_traversers.py +++ b/aiida/tools/graph/graph_traversers.py @@ -21,7 +21,7 @@ from aiida.tools.graph.age_rules import RuleSaveWalkers, RuleSequence, RuleSetWalkers, UpdateRule if TYPE_CHECKING: - from aiida.orm.implementation import Backend + from aiida.orm.implementation import StorageBackend if sys.version_info >= (3, 8): from typing import TypedDict @@ -38,7 +38,7 @@ def get_nodes_delete( starting_pks: Iterable[int], get_links: bool = False, missing_callback: Optional[Callable[[Iterable[int]], None]] = None, - backend: Optional['Backend'] = None, + backend: Optional['StorageBackend'] = None, **traversal_rules: bool ) -> TraverseGraphOutput: """ @@ -81,7 +81,7 @@ def get_nodes_delete( def get_nodes_export( starting_pks: Iterable[int], get_links: bool = False, - backend: Optional['Backend'] = None, + backend: Optional['StorageBackend'] = None, **traversal_rules: bool ) -> TraverseGraphOutput: """ @@ -196,7 +196,7 @@ def traverse_graph( links_forward: Iterable[LinkType] = (), links_backward: Iterable[LinkType] = (), missing_callback: Optional[Callable[[Iterable[int]], None]] = None, - backend: Optional['Backend'] = None + backend: Optional['StorageBackend'] = None ) -> TraverseGraphOutput: """ This function will return the set of all nodes that can be connected diff --git a/aiida/tools/visualization/graph.py b/aiida/tools/visualization/graph.py index 5803f53f69..37b4fd0b04 100644 --- a/aiida/tools/visualization/graph.py +++ b/aiida/tools/visualization/graph.py @@ -23,7 +23,7 @@ from aiida.tools.graph.graph_traversers import traverse_graph if TYPE_CHECKING: - from aiida.orm.implementation import Backend + from aiida.orm.implementation import StorageBackend __all__ = ('Graph', 'default_link_styles', 'default_node_styles', 'pstate_node_styles', 'default_node_sublabels') @@ -364,7 +364,7 @@ def __init__( node_style_fn=None, node_sublabel_fn=None, node_id_type='pk', - backend: Optional['Backend'] = None + backend: Optional['StorageBackend'] = None ): """a class to create graphviz graphs of the AiiDA node provenance @@ -409,7 +409,7 @@ def __init__( self._origin_node_style = _OVERRIDE_STYLES_DICT['origin_node'] @property - def backend(self) -> 'Backend': + def backend(self) -> 'StorageBackend': """The backend used to create the graph""" return self._backend diff --git a/tests/tools/archive/test_backend.py b/tests/tools/archive/test_backend.py index 7b5d0c8d50..d537c32d09 100644 --- a/tests/tools/archive/test_backend.py +++ b/tests/tools/archive/test_backend.py @@ -13,7 +13,7 @@ from aiida import orm from aiida.common.exceptions import NotExistent -from aiida.orm.implementation.backends import Backend +from aiida.orm.implementation import StorageBackend from aiida.tools.archive import ArchiveFormatSqlZip, ArchiveReaderAbstract from tests.utils.archives import get_archive_file @@ -32,7 +32,7 @@ def archive(tmp_path): def test_get_backend(archive: ArchiveReaderAbstract): """Test retrieving the backend""" backend = archive.get_backend() - assert isinstance(backend, Backend) + assert isinstance(backend, StorageBackend) def test_get(archive: ArchiveReaderAbstract):