Migrate code to use Pydantic V2

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
   "aiohttp-jinja2==1.6",
   "aioipfs~=0.7.1",
   "alembic==1.15.1",
-  "aleph-message==0.6.1",
+  "aleph-message @ git+https://github.com/aleph-im/aleph-message@108-upgrade-pydantic-version#egg=aleph-message",
   "aleph-nuls2==0.1",
   "aleph-p2p-client @ git+https://github.com/aleph-im/p2p-service-client-python@cbfebb871db94b2ca580e66104a67cd730c5020c",
   "asyncpg==0.30",
@@ -158,9 +158,8 @@ dependencies = [
   "isort==5.13.2",
   "check-sdist==0.1.3",
   "sqlalchemy[mypy]==1.4.41",
-  "yamlfix==1.16.1",
-  # because of aleph messages otherwise yamlfix install a too new version
-  "pydantic>=1.10.5,<2.0.0",
+  "yamlfix>=1.17",
+  "pydantic>=2,<3.0.0",
   "pyproject-fmt==2.2.1",
   "types-aiofiles",
   "types-protobuf",

src/aleph/chains/chain_data_service.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 from typing import Any, Dict, List, Mapping, Optional, Self, Set, Type, Union, cast
 
 import aio_pika.abc
@@ -66,10 +67,12 @@ async def prepare_sync_event_payload(
             protocol=ChainSyncProtocol.ON_CHAIN_SYNC,
             version=1,
             content=OnChainContent(
-                messages=[OnChainMessage.from_orm(message) for message in messages]
+                messages=[
+                    OnChainMessage.model_validate(message) for message in messages
+                ]
             ),
         )
-        archive_content: bytes = archive.json().encode("utf-8")
+        archive_content: bytes = archive.model_dump_json().encode("utf-8")
 
         ipfs_cid = await self.storage_service.add_file(
             session=session, file_content=archive_content, engine=ItemType.ipfs
@@ -166,7 +169,9 @@ def _get_tx_messages_smart_contract_protocol(tx: ChainTxDb) -> List[Dict[str, An
         )
 
         try:
-            payload = cast(GenericMessageEvent, payload_model.parse_obj(tx.content))
+            payload = cast(
+                GenericMessageEvent, payload_model.model_validate(tx.content)
+            )
         except ValidationError:
             raise InvalidContent(f"Incompatible tx content for {tx.chain}/{tx.hash}")
 
@@ -189,7 +194,7 @@ def _get_tx_messages_smart_contract_protocol(tx: ChainTxDb) -> List[Dict[str, An
                 item_hash=ItemHash(payload.content),
                 metadata=None,
             )
-            item_content = content.json(exclude_none=True)
+            item_content = json.dumps(content.model_dump(exclude_none=True))
         else:
             item_content = payload.content
 

src/aleph/chains/ethereum.py

@@ -342,7 +342,7 @@ async def packer(self, config: Config):
                     account,
                     int(gas_price * 1.1),
                     nonce,
-                    sync_event_payload.json(),
+                    sync_event_payload.model_dump_json(),
                 )
                 LOGGER.info("Broadcast %r on %s" % (response, CHAIN_NAME))
 

src/aleph/chains/indexer_reader.py

@@ -93,7 +93,7 @@ def make_events_query(
         model = SyncEvent
         event_type_str = "syncEvents"
 
-    fields = "\n".join(model.__fields__.keys())
+    fields = "\n".join(model.__annotations__.keys())
     params: Dict[str, Any] = {
         "blockchain": f'"{blockchain.value}"',
         "limit": limit,
@@ -147,7 +147,7 @@ async def _query(self, query: str, model: Type[T]) -> T:
         response = await self.http_session.post("/", json={"query": query})
         response.raise_for_status()
         response_json = await response.json()
-        return model.parse_obj(response_json)
+        return model.model_validate(response_json)
 
     async def fetch_account_state(
         self,
@@ -196,7 +196,7 @@ def indexer_event_to_chain_tx(
     if isinstance(indexer_event, MessageEvent):
         protocol = ChainSyncProtocol.SMART_CONTRACT
         protocol_version = 1
-        content = indexer_event.dict()
+        content = indexer_event.model_dump()
     else:
         sync_message = aleph_json.loads(indexer_event.message)
 

src/aleph/chains/nuls2.py

@@ -210,7 +210,7 @@ async def packer(self, config: Config):
                 # Required to apply update to the files table in get_chaindata
                 session.commit()
 
-                content = sync_event_payload.json()
+                content = sync_event_payload.model_dump_json()
                 tx = await prepare_transfer_tx(
                     address,
                     [(target_addr, CHEAP_UNIT_FEE)],

src/aleph/chains/tezos.py

@@ -162,7 +162,7 @@ async def fetch_messages(
     response.raise_for_status()
     response_json = await response.json()
 
-    return IndexerResponse[IndexerMessageEvent].parse_obj(response_json)
+    return IndexerResponse[IndexerMessageEvent].model_validate(response_json)
 
 
 def indexer_event_to_chain_tx(
@@ -176,7 +176,7 @@ def indexer_event_to_chain_tx(
         publisher=indexer_event.source,
         protocol=ChainSyncProtocol.SMART_CONTRACT,
         protocol_version=1,
-        content=indexer_event.payload.dict(),
+        content=indexer_event.payload.model_dump(),
     )
 
     return chain_tx

src/aleph/db/models/messages.py

@@ -14,7 +14,6 @@
     StoreContent,
 )
 from pydantic import ValidationError
-from pydantic.error_wrappers import ErrorWrapper
 from sqlalchemy import (
     ARRAY,
     TIMESTAMP,
@@ -62,14 +61,14 @@ def validate_message_content(
     content_dict: Dict[str, Any],
 ) -> BaseContent:
     content_type = CONTENT_TYPE_MAP[message_type]
-    content = content_type.parse_obj(content_dict)
+    content = content_type.model_validate(content_dict)
     # Validate that the content time can be converted to datetime. This will
     # raise a ValueError and be caught
     # TODO: move this validation in aleph-message
     try:
         _ = dt.datetime.fromtimestamp(content_dict["time"])
     except ValueError as e:
-        raise ValidationError([ErrorWrapper(e, loc="time")], model=content_type) from e
+        raise ValidationError(str(e)) from e
 
     return content
 

src/aleph/handlers/content/vm.py

@@ -195,7 +195,7 @@ def vm_message_to_db(message: MessageDb) -> VmBaseDb:
 
         if content.on.message:
             vm.message_triggers = [
-                subscription.dict() for subscription in content.on.message
+                subscription.model_dump() for subscription in content.on.message
             ]
 
         vm.code_volume = CodeVolumeDb(

src/aleph/schemas/api/accounts.py

@@ -3,7 +3,7 @@
 from typing import Dict, List, Optional
 
 from aleph_message.models import Chain
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from aleph.types.files import FileType
 from aleph.types.sort_order import SortOrder
@@ -53,16 +53,15 @@ class GetBalancesChainsQueryParams(BaseModel):
     )
     min_balance: int = Field(default=0, ge=1, description="Minimum Balance needed")
 
-    @validator("chains", pre=True)
+    @field_validator("chains", mode="before")
     def split_str(cls, v):
         if isinstance(v, str):
             return v.split(LIST_FIELD_SEPARATOR)
         return v
 
 
 class AddressBalanceResponse(BaseModel):
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
     address: str
     balance: str
@@ -78,8 +77,7 @@ class GetAccountFilesResponseItem(BaseModel):
 
 
 class GetAccountFilesResponse(BaseModel):
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
     address: str
     total_size: int

src/aleph/schemas/api/costs.py

@@ -1,27 +1,25 @@
 from typing import List
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, ConfigDict, field_validator
 
 from aleph.toolkit.costs import format_cost_str
 
 
 class EstimatedCostDetailResponse(BaseModel):
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
     type: str
     name: str
     cost_hold: str
     cost_stream: str
 
-    @validator("cost_hold", "cost_stream")
+    @field_validator("cost_hold", "cost_stream", mode="after")
     def check_format_price(cls, v):
         return format_cost_str(v)
 
 
 class EstimatedCostsResponse(BaseModel):
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
     required_tokens: float
     payment_type: str

src/aleph/schemas/api/messages.py

@@ -25,10 +25,8 @@
     ProgramContent,
     StoreContent,
 )
-from pydantic import BaseModel, Field
-from pydantic.generics import GenericModel
+from pydantic import BaseModel, ConfigDict, Field, field_serializer
 
-import aleph.toolkit.json as aleph_json
 from aleph.db.models import MessageDb
 from aleph.types.message_status import ErrorCode, MessageStatus
 
@@ -39,26 +37,26 @@
 class MessageConfirmation(BaseModel):
     """Format of the result when a message has been confirmed on a blockchain"""
 
-    class Config:
-        orm_mode = True
-        json_encoders = {dt.datetime: lambda d: d.timestamp()}
+    model_config = ConfigDict(from_attributes=True)
 
     chain: Chain
     height: int
     hash: str
+    time: dt.datetime
+
+    @field_serializer("time")
+    def serialize_time(self, dt: dt.datetime, _info) -> float:
+        return dt.timestamp()
 
 
-class BaseMessage(GenericModel, Generic[MType, ContentType]):
-    class Config:
-        orm_mode = True
-        json_loads = aleph_json.loads
-        json_encoders = {dt.datetime: lambda d: d.timestamp()}
+class BaseMessage(BaseModel, Generic[MType, ContentType]):
+    model_config = ConfigDict(from_attributes=True)
 
     sender: str
     chain: Chain
-    signature: Optional[str]
+    signature: Optional[str] = None
     type: MType
-    item_content: Optional[str]
+    item_content: Optional[str] = None
     item_type: ItemType
     item_hash: str
     time: dt.datetime
@@ -67,6 +65,10 @@ class Config:
     confirmed: bool
     confirmations: List[MessageConfirmation]
 
+    @field_serializer("time")
+    def serialize_time(self, dt: dt.datetime, _info) -> float:
+        return dt.timestamp()
+
 
 class AggregateMessage(
     BaseMessage[Literal[MessageType.aggregate], AggregateContent]  # type: ignore
@@ -127,13 +129,13 @@ def format_message(message: MessageDb) -> AlephMessage:
     message_type = message.type
 
     message_cls = MESSAGE_CLS_DICT[message_type]
-    return message_cls.from_orm(message)
+    return message_cls.model_validate(message)
 
 
 def format_message_dict(message: Dict[str, Any]) -> AlephMessage:
     message_type = message.get("type")
     message_cls = MESSAGE_CLS_DICT[message_type]
-    return message_cls.parse_obj(message)
+    return message_cls.model_validate(message)
 
 
 class BaseMessageStatus(BaseModel):
@@ -145,45 +147,41 @@ class BaseMessageStatus(BaseModel):
 # We already have a model for the validation of pending messages, but this one
 # is only used for formatting and does not try to be smart.
 class PendingMessage(BaseModel):
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
     sender: str
     chain: Chain
-    signature: Optional[str]
+    signature: Optional[str] = None
     type: MessageType
-    item_content: Optional[str]
+    item_content: Optional[str] = None
     item_type: ItemType
     item_hash: str
     time: dt.datetime
     channel: Optional[str] = None
-    content: Optional[Dict[str, Any]]
+    content: Optional[Dict[str, Any]] = None
     reception_time: dt.datetime
 
 
 class PendingMessageStatus(BaseMessageStatus):
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
     status: MessageStatus = MessageStatus.PENDING
     messages: List[PendingMessage]
 
 
 class ProcessedMessageStatus(BaseMessageStatus):
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
     status: MessageStatus = MessageStatus.PROCESSED
     message: AlephMessage
 
 
 class ForgottenMessage(BaseModel):
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
     sender: str
     chain: Chain
-    signature: Optional[str]
+    signature: Optional[str] = None
     type: MessageType
     item_type: ItemType
     item_hash: str
@@ -201,18 +199,15 @@ class RejectedMessageStatus(BaseMessageStatus):
     status: MessageStatus = MessageStatus.REJECTED
     message: Mapping[str, Any]
     error_code: ErrorCode
-    details: Any
+    details: Any = None
 
 
 class MessageStatusInfo(BaseMessageStatus):
-    class Config:
-        orm_mode = True
-        fields = {"item_hash": {"exclude": True}}
+    model_config = ConfigDict(from_attributes=True)
 
 
 class MessageHashes(BaseMessageStatus):
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
 
 MessageWithStatus = Union[
@@ -224,12 +219,13 @@ class Config:
 
 
 class MessageListResponse(BaseModel):
-    class Config:
-        json_encoders = {dt.datetime: lambda d: d.timestamp()}
-        json_loads = aleph_json.loads
-
     messages: List[AlephMessage]
     pagination_page: int
     pagination_total: int
     pagination_per_page: int
     pagination_item: Literal["messages"] = "messages"
+    time: dt.datetime
+
+    @field_serializer("time")
+    def serialize_time(self, dt: dt.datetime, _info) -> float:
+        return dt.timestamp()