upd project structure

src/qbot/fsm/db_storage.py

@@ -0,0 +1,84 @@
+from aiogram.fsm.state import State
+from aiogram.fsm.storage.base import (
+    BaseStorage,
+    StorageKey,
+    StateType,
+    DefaultKeyBuilder,
+    KeyBuilder,
+)
+from sqlmodel import select
+from typing import Any, Dict
+import ujson as json
+
+from ..db import async_session
+from ..model.fsm_storage import FSMStorage
+
+
+class DbStorage(BaseStorage):
+    def __init__(self, key_builder: KeyBuilder | None = None) -> None:
+        if key_builder is None:
+            key_builder = DefaultKeyBuilder()
+        self.key_builder = key_builder
+
+    async def set_state(self, key: StorageKey, state: StateType = None) -> None:
+        db_key = self.key_builder.build(key, "state")
+        async with async_session() as session:
+            db_state = (
+                await session.exec(select(FSMStorage).where(FSMStorage.key == db_key))
+            ).first()
+
+            if db_state:
+                if state is None:
+                    await session.delete(db_state)
+                else:
+                    db_state.value = state.state if isinstance(state, State) else state
+            elif state is not None:
+                db_state = FSMStorage(
+                    key=db_key, value=state.state if isinstance(state, State) else state
+                )
+                session.add(db_state)
+            else:
+                return
+
+            await session.commit()
+
+    async def get_state(self, key: StorageKey) -> str | None:
+        db_key = self.key_builder.build(key, "state")
+        async with async_session() as session:
+            db_state = (
+                await session.exec(select(FSMStorage).where(FSMStorage.key == db_key))
+            ).first()
+            return db_state.value if db_state else None
+
+    async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None:
+        db_key = self.key_builder.build(key, "data")
+        async with async_session() as session:
+            db_data = (
+                await session.exec(select(FSMStorage).where(FSMStorage.key == db_key))
+            ).first()
+
+            if db_data:
+                if not data:
+                    await session.delete(db_data)
+                else:
+                    db_data.value = json.dumps(data, ensure_ascii=False)
+            elif data:
+                db_data = FSMStorage(
+                    key=db_key, value=json.dumps(data, ensure_ascii=False)
+                )
+                session.add(db_data)
+            else:
+                return
+
+            await session.commit()
+
+    async def get_data(self, key: StorageKey) -> Dict[str, Any]:
+        db_key = self.key_builder.build(key, "data")
+        async with async_session() as session:
+            db_data = (
+                await session.exec(select(FSMStorage).where(FSMStorage.key == db_key))
+            ).first()
+            return json.loads(db_data.value) if db_data else {}
+
+    async def close(self):
+        return await super().close()