Added serializer, removed bug with duplicate key

Signed-off-by: chandr-andr (Kiselev Aleksandr) <[email protected]>

Author: chandr-andr

taskiq_psqlpy/queries.py

@@ -11,6 +11,9 @@
 
 INSERT_RESULT_QUERY = """
 INSERT INTO {} VALUES ($1, $2)
+ON CONFLICT (task_id)
+DO UPDATE
+SET result = $2
 """
 
 IS_RESULT_EXISTS_QUERY = """

taskiq_psqlpy/result_backend.py

@@ -1,4 +1,3 @@
-import pickle
 from typing import (
     Any,
     Final,
@@ -11,6 +10,9 @@
 from psqlpy import ConnectionPool
 from psqlpy.exceptions import RustPSQLDriverPyBaseError
 from taskiq import AsyncResultBackend, TaskiqResult
+from taskiq.abc.serializer import TaskiqSerializer
+from taskiq.compat import model_dump, model_validate
+from taskiq.serializers import PickleSerializer
 
 from taskiq_psqlpy.exceptions import ResultIsMissingError
 from taskiq_psqlpy.queries import (
@@ -34,19 +36,24 @@ def __init__(
         keep_results: bool = True,
         table_name: str = "taskiq_results",
         field_for_task_id: Literal["VarChar", "Text"] = "VarChar",
+        serializer: Optional[TaskiqSerializer] = None,
         **connect_kwargs: Any,
     ) -> None:
         """Construct new result backend.
 
         :param dsn: connection string to PostgreSQL.
         :param keep_results: flag to not remove results from Redis after reading.
+        :param table_name: name of the table to store results.
+        :param field_for_task_id: type of the field to store task_id.
+        :param serializer: serializer class to serialize/deserialize result from task.
         :param connect_kwargs: additional arguments for nats `ConnectionPool` class.
         """
         self.dsn: Final = dsn
         self.keep_results: Final = keep_results
         self.table_name: Final = table_name
         self.field_for_task_id: Final = field_for_task_id
         self.connect_kwargs: Final = connect_kwargs
+        self.serializer = serializer or PickleSerializer()
 
         self._database_pool: ConnectionPool
 
@@ -93,7 +100,7 @@ async def set_result(
             ),
             parameters=[
                 task_id,
-                pickle.dumps(result),
+                self.serializer.dumpb(model_dump(result)),
             ],
         )
 
@@ -149,8 +156,9 @@ async def get_result(
                 parameters=[task_id],
             )
 
-        taskiq_result: Final = pickle.loads(  # noqa: S301
-            result_in_bytes,
+        taskiq_result: Final = model_validate(
+            TaskiqResult[_ReturnType],
+            self.serializer.loadb(result_in_bytes),
         )
 
         if not with_logs:

tests/conftest.py

@@ -45,7 +45,7 @@ def postgresql_dsn() -> str:
     """
     return (
         os.environ.get("POSTGRESQL_URL")
-        or "postgresql://postgres:postgres@localhost:5432/taskiqpsqlpy"
+        or "postgresql://akiselev:12345@localhost:5432/taskiqpsqlpy"
     )
 
 

tests/test_result_backend.py

@@ -2,6 +2,7 @@
 from typing import Any, TypeVar
 
 import pytest
+from pydantic import BaseModel
 from taskiq import TaskiqResult
 
 from taskiq_psqlpy.exceptions import ResultIsMissingError
@@ -11,12 +12,10 @@
 pytestmark = pytest.mark.anyio
 
 
-class ResultForTest:
+class ResultForTest(BaseModel):
     """Just test class for testing."""
 
-    def __init__(self) -> None:
-        """Generates test class for result testing."""
-        self.test_arg = uuid.uuid4()
+    test_arg: str = uuid.uuid4().hex
 
 
 @pytest.fixture
@@ -137,7 +136,7 @@ async def test_success_backend_custom_result(
     result = await psqlpy_result_backend.get_result(task_id=task_id)
 
     assert (
-        result.return_value.test_arg  # type: ignore
+        result.return_value["test_arg"]  # type: ignore
         == custom_taskiq_result.return_value.test_arg  # type: ignore
     )
     assert result.is_err == custom_taskiq_result.is_err
@@ -158,3 +157,41 @@ async def test_success_backend_is_result_ready(
     )
 
     assert await psqlpy_result_backend.is_result_ready(task_id=task_id)
+
+
+async def test_test_success_backend_custom_result_set_same_task_id(
+    psqlpy_result_backend: PSQLPyResultBackend[_ReturnType],
+    custom_taskiq_result: TaskiqResult[_ReturnType],
+    task_id: str,
+) -> None:
+    await psqlpy_result_backend.set_result(
+        task_id=task_id,
+        result=custom_taskiq_result,
+    )
+    result = await psqlpy_result_backend.get_result(task_id=task_id)
+
+    assert (
+        result.return_value["test_arg"]  # type: ignore
+        == custom_taskiq_result.return_value.test_arg  # type: ignore
+    )
+
+    await psqlpy_result_backend.set_result(
+        task_id=task_id,
+        result=custom_taskiq_result,
+    )
+
+    another_taskiq_res_uuid = uuid.uuid4().hex
+    another_taskiq_res = TaskiqResult(
+        is_err=False,
+        log=None,
+        return_value=ResultForTest(test_arg=another_taskiq_res_uuid),
+        execution_time=0.1,
+    )
+
+    await psqlpy_result_backend.set_result(
+        task_id=task_id,
+        result=another_taskiq_res,  # type: ignore[arg-type]
+    )
+    result = await psqlpy_result_backend.get_result(task_id=task_id)
+
+    assert result.return_value["test_arg"] == another_taskiq_res_uuid  # type: ignore