init

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)
+            elif data:
+                db_data = FSMStorage(key = db_key, value = json.dumps(data,))
+                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()
+