ORM: Use pydantic to specify a schema for each ORM entity (aiidateam#6255)

This PR addresses [AEP
010](https://github.com/aiidateam/AEP/blob/983a645c9285ba65c7cf07fe6064c23e7e994c06/010_orm_schema/readme.md)
by introducing a schema interface for AiiDA's ORM entities using `pydantic`. The
goal is to make the structure and content of ORM classes introspectable and
serializable, enabling external applications—such as web APIs or CLIs—to
interface with AiiDA in a programmatic, schema-driven way. Each core ORM class
now defines a `Model` subclass that exposes its fields using rich metadata
annotations. These models support conversion to and from ORM instances
(`_from_model`, `_to_model`), as well as a user API yielding JSON-compatible
formats (`from_serialized`, `serialize`). The serialization mechanism is
experimental and is stated to the user as such on use of the serialization
methods.

In the process of implementing the original proposal, the `MetadataField`
utility was extended with additional metadata such as `exclude_to_orm`,
`exclude_from_cli`, `orm_class`, `orm_to_model`, and `model_to_orm`, enabling
precise control over how fields behave across layers. A helper function
`get_metadata` was added to facilitate metadata retrieval from `FieldInfo`
objects.

The PR refactors the `AuthInfo`, `Comment`, and `PortableCode` classes to adopt
this new schema interface. For `AuthInfo` and `Comment`, equality checks are now
defined based on the models to improve testability and object comparison. The
`PortableCode` now also supports `PurePath` for the `filepath_executable` on
initialization, though the type of its model's `filepath_files` is corrected to
`str` only, reflecting how it is stored in the database. The serialization of
repository files is unified for the `filepath_files` field and yaml-export
functionality.

For objects with repository-stored files, on serialization, files are written to
a user-specified directory (or a temporary one if not provided) and the paths
are stored in the serialized model. The `from_serialized` method can then be
used to reconstruct the object, including its files, from the serialized data.

Furthemore, CLI support was also expanded: dynamic option generation now
respects `exclude_from_cli` and allows prioritizing or customizing CLI arguments
using schema metadata. The `create_code` command and dynamic group listing logic
were updated accordingly to use the new models for constructing instances.

Additionally, entry points were registered for core ORM classes to improve
plugin discoverability, and a new `FilePath` typing alias was introduced for
shared use across the codebase.

---------

Co-authored-by: Edan Bainglass <[email protected]>
Co-authored-by: Alexander Goscinski <[email protected]>
Co-authored-by: Alexander Goscinski <[email protected]>

Author: 3 people

docs/source/nitpick-exceptions

@@ -140,7 +140,9 @@ py:meth click.Option.get_default
 py:meth fail
 
 py:class ComputedFieldInfo
+py:class BaseModel
 py:class pydantic.fields.Field
+py:class pydantic.fields.FieldInfo
 py:class pydantic.main.BaseModel
 py:class PluggableSchemaValidator
 
@@ -157,6 +159,7 @@ py:class frozenset
 
 py:class numpy.bool_
 py:class numpy.ndarray
+py:class np.ndarray
 py:class ndarray
 
 py:class paramiko.proxy.ProxyCommand

pyproject.toml

@@ -139,6 +139,17 @@ requires-python = '>=3.9'
 'process.workflow.workchain' = 'aiida.orm.nodes.process.workflow.workchain:WorkChainNode'
 'process.workflow.workfunction' = 'aiida.orm.nodes.process.workflow.workfunction:WorkFunctionNode'
 
+[project.entry-points.'aiida.orm']
+'core.auth_info' = 'aiida.orm.authinfos:AuthInfo'
+'core.comment' = 'aiida.orm.comments:Comment'
+'core.computer' = 'aiida.orm.computers:Computer'
+'core.data' = 'aiida.orm.nodes.data.data:Data'
+'core.entity' = 'aiida.orm.entities:Entity'
+'core.group' = 'aiida.orm.groups:Group'
+'core.log' = 'aiida.orm.logs:Log'
+'core.node' = 'aiida.orm.nodes.node:Node'
+'core.user' = 'aiida.orm.users:User'
+
 [project.entry-points.'aiida.parsers']
 'core.arithmetic.add' = 'aiida.parsers.plugins.arithmetic.add:ArithmeticAddParser'
 'core.templatereplacer' = 'aiida.parsers.plugins.templatereplacer.parser:TemplatereplacerParser'

src/aiida/cmdline/commands/cmd_code.py

@@ -30,10 +30,10 @@ def verdi_code():
     """Setup and manage codes."""
 
 
-def create_code(ctx: click.Context, cls, non_interactive: bool, **kwargs):
+def create_code(ctx: click.Context, cls, **kwargs):
     """Create a new `Code` instance."""
     try:
-        instance = cls(**kwargs)
+        instance = cls._from_model(cls.Model(**kwargs))
     except (TypeError, ValueError) as exception:
         echo.echo_critical(f'Failed to create instance `{cls}`: {exception}')
 
@@ -243,9 +243,7 @@ def show(code):
 @with_dbenv()
 def export(code, output_file, overwrite, sort):
     """Export code to a yaml file. If no output file is given, default name is created based on the code label."""
-
     other_args = {'sort': sort}
-
     fileformat = 'yaml'
 
     if output_file is None:

src/aiida/cmdline/commands/cmd_profile.py

@@ -29,7 +29,6 @@ def verdi_profile():
 def command_create_profile(
     ctx: click.Context,
     storage_cls,
-    non_interactive: bool,
     profile: Profile,
     set_as_default: bool = True,
     email: str | None = None,

src/aiida/cmdline/groups/dynamic.py

@@ -88,29 +88,25 @@ def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None
             command = super().get_command(ctx, cmd_name)
         return command
 
-    def call_command(self, ctx, cls, **kwargs):
+    def call_command(self, ctx, cls, non_interactive, **kwargs):
         """Call the ``command`` after validating the provided inputs."""
         from pydantic import ValidationError
 
         if hasattr(cls, 'Model'):
             # The plugin defines a pydantic model: use it to validate the provided arguments
             try:
-                model = cls.Model(**kwargs)
+                cls.Model(**kwargs)
             except ValidationError as exception:
                 param_hint = [
                     f'--{loc.replace("_", "-")}'  # type: ignore[union-attr]
                     for loc in exception.errors()[0]['loc']
                 ]
-                message = '\n'.join([str(e['ctx']['error']) for e in exception.errors()])
+                message = '\n'.join([str(e['msg']) for e in exception.errors()])
                 raise click.BadParameter(
                     message,
-                    param_hint=param_hint or 'multiple parameters',  # type: ignore[arg-type]
+                    param_hint=param_hint or 'one or more parameters',  # type: ignore[arg-type]
                 ) from exception
 
-            # Update the arguments with the dictionary representation of the model. This will include any type coercions
-            # that may have been applied with validators defined for the model.
-            kwargs.update(**model.model_dump())
-
         return self._command(ctx, cls, **kwargs)
 
     def create_command(self, ctx: click.Context, entry_point: str) -> click.Command:
@@ -154,6 +150,8 @@ def list_options(self, entry_point: str) -> list:
         """
         from pydantic_core import PydanticUndefined
 
+        from aiida.common.pydantic import get_metadata
+
         cls = self.factory(entry_point)
 
         if not hasattr(cls, 'Model'):
@@ -170,6 +168,9 @@ def list_options(self, entry_point: str) -> list:
         options_spec = {}
 
         for key, field_info in cls.Model.model_fields.items():
+            if get_metadata(field_info, 'exclude_from_cli'):
+                continue
+
             default = field_info.default_factory if field_info.default is PydanticUndefined else field_info.default
 
             # If the annotation has the ``__args__`` attribute it is an instance of a type from ``typing`` and the real
@@ -194,7 +195,8 @@ def list_options(self, entry_point: str) -> list:
             }
             for metadata in field_info.metadata:
                 for metadata_key, metadata_value in metadata.items():
-                    options_spec[key][metadata_key] = metadata_value
+                    if metadata_key in ('priority', 'short_name', 'option_cls'):
+                        options_spec[key][metadata_key] = metadata_value
 
         options_ordered = []
 

src/aiida/common/pydantic.py

@@ -3,16 +3,44 @@
 from __future__ import annotations
 
 import typing as t
+from pathlib import Path
 
 from pydantic import Field
+from pydantic_core import PydanticUndefined
+
+if t.TYPE_CHECKING:
+    from pydantic import BaseModel
+
+    from aiida.orm import Entity
+
+
+def get_metadata(field_info, key: str, default: t.Any | None = None):
+    """Return a the metadata of the given field for a particular key.
+
+    :param field_info: The field from which to retrieve the metadata.
+    :param key: The metadata name.
+    :param default: Optional default value to return in case the metadata is not defined on the field.
+    :returns: The metadata if defined, otherwise the default.
+    """
+    for element in field_info.metadata:
+        if key in element:
+            return element[key]
+    return default
 
 
 def MetadataField(  # noqa: N802
-    default: t.Any | None = None,
+    default: t.Any = PydanticUndefined,
     *,
     priority: int = 0,
     short_name: str | None = None,
     option_cls: t.Any | None = None,
+    orm_class: type['Entity'] | str | None = None,
+    orm_to_model: t.Callable[['Entity', Path], t.Any] | None = None,
+    model_to_orm: t.Callable[['BaseModel'], t.Any] | None = None,
+    exclude_to_orm: bool = False,
+    exclude_from_cli: bool = False,
+    is_attribute: bool = True,
+    is_subscriptable: bool = False,
     **kwargs,
 ):
     """Return a :class:`pydantic.fields.Field` instance with additional metadata.
@@ -37,11 +65,36 @@ class Model(BaseModel):
     :param priority: Used to order the list of all fields in the model. Ordering is done from small to large priority.
     :param short_name: Optional short name to use for an option on a command line interface.
     :param option_cls: The :class:`click.Option` class to use to construct the option.
+    :param orm_class: The class, or entry point name thereof, to which the field should be converted. If this field is
+        defined, the value of this field should acccept an integer which will automatically be converted to an instance
+        of said ORM class using ``orm_class.collection.get(id={field_value})``. This is useful, for example, where a
+        field represents an instance of a different entity, such as an instance of ``User``. The serialized data would
+        store the ``pk`` of the user, but the ORM entity instance would receive the actual ``User`` instance with that
+        primary key.
+    :param orm_to_model: Optional callable to convert the value of a field from an ORM instance to a model instance.
+    :param model_to_orm: Optional callable to convert the value of a field from a model instance to an ORM instance.
+    :param exclude_to_orm: When set to ``True``, this field value will not be passed to the ORM entity constructor
+        through ``Entity.from_model``.
+    :param exclude_to_orm: When set to ``True``, this field value will not be exposed on the CLI command that is
+        dynamically generated to create a new instance.
+    :param is_attribute: Whether the field is stored as an attribute.
+    :param is_subscriptable: Whether the field can be indexed like a list or dictionary.
     """
     field_info = Field(default, **kwargs)
 
-    for key, value in (('priority', priority), ('short_name', short_name), ('option_cls', option_cls)):
-        if value is not None and field_info is not None:
+    for key, value in (
+        ('priority', priority),
+        ('short_name', short_name),
+        ('option_cls', option_cls),
+        ('orm_class', orm_class),
+        ('orm_to_model', orm_to_model),
+        ('model_to_orm', model_to_orm),
+        ('exclude_to_orm', exclude_to_orm),
+        ('exclude_from_cli', exclude_from_cli),
+        ('is_attribute', is_attribute),
+        ('is_subscriptable', is_subscriptable),
+    ):
+        if value is not None:
             field_info.metadata.append({key: value})
 
     return field_info

src/aiida/common/typing.py

@@ -0,0 +1,19 @@
+###########################################################################
+# Copyright (c), The AiiDA team. All rights reserved.                     #
+# This file is part of the AiiDA code.                                    #
+#                                                                         #
+# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
+# For further information on the license, see the LICENSE.txt file        #
+# For further information please visit http://www.aiida.net               #
+###########################################################################
+"""Module to define commonly used data structures."""
+
+from __future__ import annotations
+
+import pathlib
+from typing import Union
+
+__all__ = ('FilePath',)
+
+
+FilePath = Union[str, pathlib.PurePath]

src/aiida/orm/authinfos.py

@@ -8,17 +8,20 @@
 ###########################################################################
 """Module for the `AuthInfo` ORM class."""
 
+from __future__ import annotations
+
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 
 from aiida.common import exceptions
+from aiida.common.pydantic import MetadataField
 from aiida.manage import get_manager
 from aiida.plugins import TransportFactory
 
 from . import entities, users
-from .fields import add_field
+from .computers import Computer
+from .users import User
 
 if TYPE_CHECKING:
-    from aiida.orm import Computer, User
     from aiida.orm.implementation import StorageBackend
     from aiida.orm.implementation.authinfos import BackendAuthInfo  # noqa: F401
     from aiida.transports import Transport
@@ -45,51 +48,60 @@ class AuthInfo(entities.Entity['BackendAuthInfo', AuthInfoCollection]):
     """ORM class that models the authorization information that allows a `User` to connect to a `Computer`."""
 
     _CLS_COLLECTION = AuthInfoCollection
+    PROPERTY_WORKDIR = 'workdir'
 
-    __qb_fields__ = [
-        add_field(
-            'enabled',
-            dtype=bool,
+    class Model(entities.Entity.Model):
+        computer: int = MetadataField(
+            description='The PK of the computer',
             is_attribute=False,
-            doc='Whether the instance is enabled',
-        ),
-        add_field(
-            'auth_params',
-            dtype=Dict[str, Any],
+            orm_class=Computer,
+            orm_to_model=lambda auth_info, _: auth_info.computer.pk,  # type: ignore[attr-defined]
+        )
+        user: int = MetadataField(
+            description='The PK of the user',
             is_attribute=False,
-            doc='Dictionary of authentication parameters',
-        ),
-        add_field(
-            'metadata',
-            dtype=Dict[str, Any],
+            orm_class=User,
+            orm_to_model=lambda auth_info, _: auth_info.user.pk,  # type: ignore[attr-defined]
+        )
+        enabled: bool = MetadataField(
+            True,
+            description='Whether the instance is enabled',
             is_attribute=False,
-            doc='Dictionary of metadata',
-        ),
-        add_field(
-            'computer_pk',
-            dtype=int,
+        )
+        auth_params: Dict[str, Any] = MetadataField(
+            default_factory=dict,
+            description='Dictionary of authentication parameters',
             is_attribute=False,
-            doc='The PK of the computer',
-        ),
-        add_field(
-            'user_pk',
-            dtype=int,
+        )
+        metadata: Dict[str, Any] = MetadataField(
+            default_factory=dict,
+            description='Dictionary of metadata',
             is_attribute=False,
-            doc='The PK of the user',
-        ),
-    ]
-
-    PROPERTY_WORKDIR = 'workdir'
-
-    def __init__(self, computer: 'Computer', user: 'User', backend: Optional['StorageBackend'] = None) -> None:
+        )
+
+    def __init__(
+        self,
+        computer: 'Computer',
+        user: 'User',
+        enabled: bool = True,
+        auth_params: Dict[str, Any] | None = None,
+        metadata: Dict[str, Any] | None = None,
+        backend: Optional['StorageBackend'] = None,
+    ) -> None:
         """Create an `AuthInfo` instance for the given computer and user.
 
         :param computer: a `Computer` instance
         :param user: a `User` instance
         :param backend: the backend to use for the instance, or use the default backend if None
         """
         backend = backend or get_manager().get_profile_storage()
-        model = backend.authinfos.create(computer=computer.backend_entity, user=user.backend_entity)
+        model = backend.authinfos.create(
+            computer=computer.backend_entity,
+            user=user.backend_entity,
+            enabled=enabled,
+            auth_params=auth_params or {},
+            metadata=metadata or {},
+        )
         super().__init__(model)
 
     def __str__(self) -> str:
@@ -98,6 +110,18 @@ def __str__(self) -> str:
 
         return f'AuthInfo for {self.user.email} on {self.computer.label} [DISABLED]'
 
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, AuthInfo):
+            return False
+
+        return (
+            self.user.pk == other.user.pk
+            and self.computer.pk == other.computer.pk
+            and self.enabled == other.enabled
+            and self.auth_params == other.auth_params
+            and self.metadata == other.metadata
+        )
+
     @property
     def enabled(self) -> bool:
         """Return whether this instance is enabled.
@@ -126,6 +150,14 @@ def user(self) -> 'User':
         """Return the user associated with this instance."""
         return entities.from_backend_entity(users.User, self._backend_entity.user)
 
+    @property
+    def auth_params(self) -> Dict[str, Any]:
+        return self._backend_entity.get_auth_params()
+
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        return self._backend_entity.get_metadata()
+
     def get_auth_params(self) -> Dict[str, Any]:
         """Return the dictionary of authentication parameters