chore(infrastructure): replace pickle to pydantic

Author: emptybutton

pyproject.toml

@@ -82,7 +82,6 @@ lines-after-imports = 2
 "src/tgdb/entities/*" = ["PLR2004"]
 "src/tgdb/application/*" = ["PLR0917"]
 "src/tgdb/infrastructure/adapters/*" = ["RUF029"]
-"src/tgdb/infrastructure/*" = ["S403", "S301"]
 "tests/*" = ["PLR0124", "PLR0917", "S106", "C901", "PLR2004"]
 "__init__.py" = ["PLC0414"]
 

src/tgdb/infrastructure/adapters/buffer.py

@@ -1,13 +1,18 @@
-import pickle
 from asyncio import Event, wait_for
 from collections import deque
 from collections.abc import AsyncIterator, Sequence
 from dataclasses import dataclass, field
 from types import TracebackType
-from typing import Self
+from typing import ClassVar, Self
+
+from pydantic import TypeAdapter
 
 from tgdb.application.common.ports.buffer import Buffer
 from tgdb.entities.horizon.transaction import Commit, PreparedCommit
+from tgdb.infrastructure.pydantic.horizon.commit import (
+    EncodableCommit,
+    EncodablePreparedCommit,
+)
 from tgdb.infrastructure.telethon.in_telegram_bytes import InTelegramBytes
 
 
@@ -51,22 +56,20 @@ class InTelegramReplicablePreparedCommitBuffer(Buffer[Commit | PreparedCommit]):
     _buffer: Buffer[Commit | PreparedCommit]
     _in_tg_encoded_commits: InTelegramBytes
 
+    _adapter: ClassVar = TypeAdapter(
+        tuple[EncodableCommit | EncodablePreparedCommit, ...]
+    )
+
     async def __aenter__(self) -> Self:
         encoded_commits = await self._in_tg_encoded_commits
 
         if encoded_commits is None:
             return self
 
-        commits = pickle.loads(encoded_commits)
-
-        if not isinstance(commits, list):
-            raise TypeError(str(commits))
+        encodable_commits = self._adapter.validate_json(encoded_commits)
 
-        for commit in commits:
-            if not isinstance(commit, PreparedCommit):
-                raise TypeError(str(commits))
-
-            await self._buffer.add(commit)
+        for commit in encodable_commits:
+            await self._buffer.add(commit.entity())
 
         return self
 
@@ -84,7 +87,18 @@ async def __aiter__(
         self,
     ) -> AsyncIterator[Sequence[Commit | PreparedCommit]]:
         async for commits in self._buffer:
-            encoded_commits = pickle.dumps(list(commits))
+            encodable_commits = tuple(map(self._encodable_commit, commits))
+            encoded_commits = self._adapter.dump_json(encodable_commits)
             await self._in_tg_encoded_commits.set(encoded_commits)
 
             yield commits
+
+    def _encodable_commit(
+        self, commit: Commit | PreparedCommit
+    ) -> EncodableCommit | EncodablePreparedCommit:
+        match commit:
+            case Commit():
+                return EncodableCommit.of(commit)
+
+            case PreparedCommit():
+                return EncodablePreparedCommit.of(commit)

src/tgdb/infrastructure/adapters/relations.py

@@ -1,9 +1,9 @@
-import pickle
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Self
+from typing import ClassVar, Self
 
 from in_memory_db import InMemoryDb
+from pydantic import TypeAdapter
 
 from tgdb.application.relation.ports.relations import (
     NoRelationError,
@@ -12,6 +12,7 @@
 )
 from tgdb.entities.numeration.number import Number
 from tgdb.entities.relation.relation import Relation
+from tgdb.infrastructure.pydantic.relation.relation import EncodableRelation
 from tgdb.infrastructure.telethon.in_telegram_bytes import InTelegramBytes
 
 
@@ -44,6 +45,8 @@ class InTelegramReplicableRelations(Relations):
     _in_tg_encoded_relations: InTelegramBytes
     _cached_relations: InMemoryDb[Relation]
 
+    _adapter: ClassVar = TypeAdapter(tuple[EncodableRelation, ...])
+
     async def __aenter__(self) -> Self:
         for loaded_relation in await self._loaded_relations():
             self._cached_relations.insert(loaded_relation)
@@ -88,22 +91,17 @@ async def add(self, relation: Relation) -> None:
 
         self._cached_relations.insert(relation)
 
-        encoded_cached_relations = pickle.dumps(tuple(self._cached_relations))
-        await self._in_tg_encoded_relations.set(encoded_cached_relations)
+        encodable_relations = (
+            tuple(map(EncodableRelation.of, self._cached_relations))
+        )
+        encoded_relations = self._adapter.dump_json(encodable_relations, by_alias=True)
+        await self._in_tg_encoded_relations.set(encoded_relations)
 
     async def _loaded_relations(self) -> tuple[Relation, ...]:
         encoded_relations = await self._in_tg_encoded_relations
 
         if encoded_relations is None:
             return tuple()
 
-        relations = pickle.loads(encoded_relations)
-
-        if not isinstance(relations, tuple):
-            raise TypeError(relations)
-
-        for relation in relations:
-            if not isinstance(relation, Relation):
-                raise TypeError(relation)
-
-        return relations
+        encodable_relations = self._adapter.validate_json(encoded_relations)
+        return tuple(it.entity() for it in encodable_relations)

src/tgdb/infrastructure/pydantic/__init__.py

@@ -0,0 +1,51 @@
+from typing import Literal
+
+from pydantic import BaseModel
+
+from tgdb.entities.horizon.transaction import XID, Commit, PreparedCommit
+from tgdb.infrastructure.pydantic.horizon.transaction_effect import (
+    EncodableTransactionScalarEffect,
+    encodable_transaction_scalar_effect,
+)
+
+
+class EncodableCommit(BaseModel):
+    type: Literal["commit"] = "commit"
+    xid: XID
+    effect: tuple[EncodableTransactionScalarEffect, ...]
+
+    def entity(self) -> Commit:
+        return Commit(
+            self.xid,
+            frozenset(it.entity() for it in self.effect)
+        )
+
+    @classmethod
+    def of(cls, entity: Commit) -> "EncodableCommit":
+        return EncodableCommit(
+            xid=entity.xid,
+            effect=tuple(
+                map(encodable_transaction_scalar_effect, entity.effect)
+            ),
+        )
+
+
+class EncodablePreparedCommit(BaseModel):
+    type: Literal["prepared_commit"] = "prepared_commit"
+    xid: XID
+    effect: tuple[EncodableTransactionScalarEffect, ...]
+
+    def entity(self) -> Commit:
+        return Commit(
+            self.xid,
+            frozenset(it.entity() for it in self.effect)
+        )
+
+    @classmethod
+    def of(cls, entity: PreparedCommit) -> "EncodablePreparedCommit":
+        return EncodablePreparedCommit(
+            xid=entity.xid,
+            effect=tuple(
+                map(encodable_transaction_scalar_effect, entity.effect)
+            ),
+        )

src/tgdb/infrastructure/pydantic/horizon/__init__.py

@@ -0,0 +1,38 @@
+from tgdb.entities.horizon.transaction import TransactionScalarEffect
+from tgdb.entities.relation.tuple_effect import (
+    DeletedTuple,
+    MigratedTuple,
+    MutatedTuple,
+    NewTuple,
+)
+from tgdb.infrastructure.pydantic.relation.tuple_effect import (
+    EncodableDeletedTuple,
+    EncodableMigratedTuple,
+    EncodableMutatedTuple,
+    EncodableNewTuple,
+)
+
+
+type EncodableTransactionScalarEffect = (
+    EncodableNewTuple
+    | EncodableMutatedTuple
+    | EncodableMigratedTuple
+    | EncodableDeletedTuple
+)
+
+
+def encodable_transaction_scalar_effect(
+    effect: TransactionScalarEffect
+) -> EncodableTransactionScalarEffect:
+    match effect:
+        case NewTuple():
+            return EncodableNewTuple.of(effect)
+
+        case MutatedTuple():
+            return EncodableMutatedTuple.of(effect)
+
+        case MigratedTuple():
+            return EncodableMigratedTuple.of(effect)
+
+        case DeletedTuple():
+            return EncodableDeletedTuple.of(effect)

src/tgdb/infrastructure/pydantic/horizon/commit.py

@@ -0,0 +1,123 @@
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+from tgdb.entities.relation.domain import (
+    BoolDomain,
+    DatetimeDomain,
+    Domain,
+    IntDomain,
+    SetDomain,
+    StrDomain,
+    UuidDomain,
+)
+from tgdb.entities.relation.scalar import Scalar
+
+
+class EncodableIntDomain(BaseModel):
+    type: Literal["int"] = "int"
+    min: int
+    max: int
+    is_nonable: bool
+
+    @classmethod
+    def of(cls, domain: "IntDomain") -> "EncodableIntDomain":
+        return cls(min=domain.min, max=domain.max, is_nonable=domain.is_nonable)
+
+    def entity(self) -> "IntDomain":
+        return IntDomain(self.min, self.max, self.is_nonable)
+
+
+class EncodableStrDomain(BaseModel):
+    type: Literal["str"] = "str"
+    max_len: int
+    is_nonable: bool
+
+    @classmethod
+    def of(cls, domain: "StrDomain") -> "EncodableStrDomain":
+        return cls(
+            max_len=domain.max_len,
+            is_nonable=domain.is_nonable,
+        )
+
+    def entity(self) -> "StrDomain":
+        return StrDomain(self.max_len, self.is_nonable)
+
+
+class EncodableBoolDomain(BaseModel):
+    type: Literal["bool"] = "bool"
+    is_nonable: bool
+
+    @classmethod
+    def of(cls, domain: BoolDomain) -> "EncodableBoolDomain":
+        return cls(is_nonable=domain.is_nonable)
+
+    def entity(self) -> BoolDomain:
+        return BoolDomain(self.is_nonable)
+
+
+class EncodableTimeDomain(BaseModel):
+    type: Literal["time"] = "time"
+    is_nonable: bool
+
+    @classmethod
+    def of(cls, domain: DatetimeDomain) -> "EncodableTimeDomain":
+        return cls(is_nonable=domain.is_nonable)
+
+    def entity(self) -> DatetimeDomain:
+        return DatetimeDomain(self.is_nonable)
+
+
+class EncodableUuidDomain(BaseModel):
+    type: Literal["uuid"] = "uuid"
+    is_nonable: bool
+
+    @classmethod
+    def of(cls, domain: UuidDomain) -> "EncodableUuidDomain":
+        return cls(is_nonable=domain.is_nonable)
+
+    def entity(self) -> UuidDomain:
+        return UuidDomain(self.is_nonable)
+
+
+class EncodableSetDomain(BaseModel):
+    type: Literal["set"] = "set"
+    scalars: tuple[Scalar, ...]
+
+    @classmethod
+    def of(cls, domain: SetDomain) -> "EncodableSetDomain":
+        return cls(scalars=domain)
+
+    def entity(self) -> SetDomain:
+        return self.scalars
+
+
+type EncodableDomain = (
+    EncodableBoolDomain
+    | EncodableIntDomain
+    | EncodableStrDomain
+    | EncodableTimeDomain
+    | EncodableUuidDomain
+    | EncodableSetDomain
+)
+
+
+def encodable_domain(domain: Domain) -> EncodableDomain:
+    match domain:
+        case IntDomain():
+            return EncodableIntDomain.of(domain)
+
+        case StrDomain():
+            return EncodableStrDomain.of(domain)
+
+        case BoolDomain():
+            return EncodableBoolDomain.of(domain)
+
+        case UuidDomain():
+            return EncodableUuidDomain.of(domain)
+
+        case DatetimeDomain():
+            return EncodableTimeDomain.of(domain)
+
+        case tuple():
+            return EncodableSetDomain.of(domain)