updated something :)

deeppavlov · Feb 13, 2025 · af64ae2 · af64ae2
1 parent 0caa820
commit af64ae2
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 10 deletions.
diff --git a/chatsky/context_storages/ydb.py b/chatsky/context_storages/ydb.py
@@ -17,7 +17,6 @@
 
 try:
     from ydb import (
-        SerializableReadWrite,
         SchemeError,
         TableDescription,
         Column,
@@ -146,7 +145,7 @@ async def callee(session: QuerySession) -> Optional[ContextMainInfo]:
                 FROM {self.main_table}
                 WHERE {NameConfig._id_column} = ${NameConfig._id_column};
                 """  # noqa: E501
-            result_sets = await session.transaction(SerializableReadWrite()).execute(
+            result_sets = await session.transaction().execute(
                 await session.prepare(query),
                 {
                     f"${NameConfig._id_column}": ctx_id,
@@ -167,7 +166,7 @@ async def _update_context(
         self, ctx_id: str, ctx_info: Optional[ContextMainInfo], field_info: List[Tuple[str, List[Tuple[int, Optional[bytes]]]]]
     ) -> None:
         async def callee(session: QuerySession) -> None:
-            transaction = await session.transaction(SerializableReadWrite()).begin()
+            transaction = await session.transaction().begin()
             if ctx_info is not None:
                 ctx_info_dump = ctx_info.model_dump(mode="python")
                 query = f"""
@@ -230,7 +229,7 @@ async def callee(session: QuerySession) -> None:
                     DELETE FROM {table_name}
                     WHERE {NameConfig._id_column} = ${NameConfig._id_column};
                     """  # noqa: E501
-                await session.transaction(SerializableReadWrite()).execute(
+                await session.transaction().execute(
                     await session.prepare(query),
                     {
                         f"${NameConfig._id_column}": ctx_id,
@@ -268,7 +267,7 @@ async def callee(session: QuerySession) -> List[Tuple[int, bytes]]:
                 WHERE {NameConfig._id_column} = ${NameConfig._id_column} AND {field_name} IS NOT NULL {key}
                 ORDER BY {NameConfig._key_column} DESC {limit};
                 """  # noqa: E501
-            result_sets = await session.transaction(SerializableReadWrite()).execute(
+            result_sets = await session.transaction().execute(
                 await session.prepare(query),
                 {
                     f"${NameConfig._id_column}": ctx_id,
@@ -293,7 +292,7 @@ async def callee(session: QuerySession) -> List[int]:
                 FROM {self.turns_table}
                 WHERE {NameConfig._id_column} = ${NameConfig._id_column} AND {field_name} IS NOT NULL;
                 """  # noqa: E501
-            result_sets = await session.transaction(SerializableReadWrite()).execute(
+            result_sets = await session.transaction().execute(
                 await session.prepare(query),
                 {
                     f"${NameConfig._id_column}": ctx_id,
@@ -319,7 +318,7 @@ async def callee(session: QuerySession) -> List[Tuple[int, bytes]]:
                 WHERE {NameConfig._id_column} = ${NameConfig._id_column} AND {field_name} IS NOT NULL
                 AND {NameConfig._key_column} IN ({", ".join(prepare.keys())});
                 """  # noqa: E501
-            result_sets = await session.transaction(SerializableReadWrite()).execute(
+            result_sets = await session.transaction().execute(
                 await session.prepare(query),
                 {
                     f"${NameConfig._id_column}": ctx_id,
@@ -342,7 +341,7 @@ async def callee(session: QuerySession) -> None:
                     PRAGMA TablePathPrefix("{self.database}");
                     DELETE FROM {table_name};
                     """  # noqa: E501
-                await session.transaction(SerializableReadWrite()).execute(
+                await session.transaction().execute(
                     await session.prepare(query), dict(), commit_tx=True
                 )
 

diff --git a/chatsky/utils/testing/cleanup_db.py b/chatsky/utils/testing/cleanup_db.py
@@ -82,8 +82,12 @@ async def delete_ydb(storage: YDBContextStorage):
     if not ydb_available:
         raise Exception("Can't delete ydb database - ydb provider unavailable!")
 
-    async def callee(session: Any) -> None:
+    from ydb.aio import SessionPool
+    from ydb.table import Session
+
+    async def callee(session: Session) -> None:
         for table in [storage.main_table, storage.turns_table]:
             await session.drop_table("/".join([storage.database, table]))
 
-    await storage.pool.retry_operation(callee)
+    async with SessionPool(storage._driver, size=10) as pool:
+        await pool.retry_operation(callee)