Merge 66402e9cc0 into 0c6a705310

CHANGES/1718.feature.rst

@@ -0,0 +1 @@
+Added SqliteStorage which makes local development and small deployments a bit easier.

aiogram/fsm/storage/sqlite.py

@@ -0,0 +1,158 @@
+import json
+from typing import Any, Dict, Mapping, Optional, cast
+
+from aiosqlite import Connection, connect
+
+from aiogram.exceptions import DataNotDictLikeError
+from aiogram.fsm.state import State
+from aiogram.fsm.storage.base import (
+    BaseEventIsolation,
+    BaseStorage,
+    DefaultKeyBuilder,
+    KeyBuilder,
+    StateType,
+    StorageKey,
+)
+
+SQLITE_FILENAME = "aiogram_fsm.sqlite"
+
+
+class SqliteStorage(BaseStorage):
+    """
+    SQLite storage required :code:`aiosqlite` package installed (:code:`pip install aiosqlite`)
+    """
+
+    def __init__(
+        self,
+        connection: Connection,
+        key_builder: Optional[KeyBuilder] = None,
+    ) -> None:
+        """
+        Create an instance of :class:`SqliteStorage` with provided SQLite connection
+
+        :param key_builder: builder that helps to convert contextual key to string
+        :param db_filename: name of the SQLite database file for FSM
+        """
+        if key_builder is None:
+            key_builder = DefaultKeyBuilder()
+        self._key_builder = key_builder
+        self._connection = connection
+
+    @classmethod
+    async def connect(cls, db_filename: str = SQLITE_FILENAME) -> "SqliteStorage":
+        """
+        Create an instance of :class:`SqliteStorage` with specifying the DB filename
+
+        :param db_filename: for example :code:`aiogram_fsm.sqlite`
+        :param connection_kwargs: see :code:`motor` docs
+        :param kwargs: arguments to be passed to :class:`MongoStorage`
+        :return: an instance of :class:`MongoStorage`
+        """
+        connection = await connect(db_filename)
+        await connection.execute(
+            f"""CREATE TABLE IF NOT EXISTS aiogram_fsm (
+                    id TEXT PRIMARY KEY,
+                    state TEXT,
+                    data TEXT)"""
+        )
+        # db optimization on start
+        await connection.execute(f"""VACUUM""")
+        await connection.commit()
+        return cls(connection=connection)
+
+    async def close(self) -> None:
+        await self._connection.close()
+
+    async def set_state(self, key: StorageKey, state: StateType = None) -> None:
+        id = self._key_builder.build(key)
+        state = cast(str, state.state if isinstance(state, State) else state)
+
+        cursor = await self._connection.execute(
+            f"""SELECT data
+                FROM aiogram_fsm
+                WHERE id = ?""",
+            (id,),
+        )
+        row = await cursor.fetchone()
+
+        if not state and (not row or not row[0]):
+            # db clean up on the go
+            await self._connection.execute(
+                f"""DELETE FROM aiogram_fsm
+                    where id = ?""",
+                (id,),
+            )
+        else:
+            await self._connection.execute(
+                f"""INSERT INTO aiogram_fsm (id, state)
+                    VALUES (?, ?)
+                    ON CONFLICT (id)
+                    DO UPDATE SET state = ?""",
+                (id, state, state),
+            )
+
+        await self._connection.commit()
+
+    async def get_state(self, key: StorageKey) -> Optional[str]:
+        id = self._key_builder.build(key)
+
+        cursor = await self._connection.execute(
+            f"""SELECT state
+                FROM aiogram_fsm
+                WHERE id = ?""",
+            (id,),
+        )
+
+        row = await cursor.fetchone()
+        return row[0] if row else None
+
+    async def set_data(self, key: StorageKey, data: Mapping[str, Any]) -> None:
+        if not isinstance(data, dict):
+            raise DataNotDictLikeError(
+                f"Data must be a dict or dict-like object, got {type(data).__name__}"
+            )
+
+        id = self._key_builder.build(key)
+        data_cell = json.dumps(data) if data else None
+
+        cursor = await self._connection.execute(
+            f"""SELECT state
+                FROM aiogram_fsm
+                WHERE id = ?""",
+            (id,),
+        )
+        row = await cursor.fetchone()
+
+        if not data and not row:
+            # db clean up on the go
+            await self._connection.execute(
+                f"""DELETE FROM aiogram_fsm
+                    where id = ?""",
+                (id,),
+            )
+        else:
+            await self._connection.execute(
+                f"""INSERT INTO aiogram_fsm (id, data)
+                    VALUES (?, ?)
+                    ON CONFLICT (id)
+                    DO UPDATE SET data = ?""",
+                (id, data_cell, data_cell),
+            )
+
+        await self._connection.commit()
+
+    async def get_data(self, key: StorageKey) -> Dict[str, Any]:
+        id = self._key_builder.build(key)
+        data_cell = {}
+
+        cursor = await self._connection.execute(
+            f"""SELECT data
+                FROM aiogram_fsm
+                WHERE id = ?""",
+            (id,),
+        )
+        row = await cursor.fetchone()
+
+        if row and row[0]:
+            data_cell = cast(Dict[str, Any], json.loads(row[0]))
+        return data_cell

docs/dispatcher/finite_state_machine/storages.rst

@@ -30,6 +30,13 @@ MongoStorage
     :members: __init__, from_url
     :member-order: bysource
 
+SqliteStorage
+------------
+
+.. autoclass:: aiogram.fsm.storage.sqlite.SqliteStorage
+    :members: __init__, connect
+    :member-order: bysource
+
 KeyBuilder
 ------------
 

pyproject.toml

@@ -65,6 +65,9 @@ mongo = [
     "motor>=3.3.2,<3.7.0",
     "pymongo>4.5,<4.11",
 ]
+sqlite = [
+    "aiosqlite>=0.21.0",
+]
 proxy = [
     "aiohttp-socks~=0.8.3",
 ]
@@ -125,6 +128,7 @@ features = [
     "fast",
     "redis",
     "mongo",
+    "sqlite",
     "proxy",
     "i18n",
     "cli",
@@ -145,6 +149,7 @@ features = [
     "fast",
     "redis",
     "mongo",
+    "sqlite",
     "proxy",
     "i18n",
     "docs",
@@ -160,6 +165,7 @@ features = [
     "fast",
     "redis",
     "mongo",
+    "sqlite",
     "proxy",
     "i18n",
     "test",
@@ -178,6 +184,7 @@ features = [
     "fast",
     "redis",
     "mongo",
+    "sqlite",
     "proxy",
     "i18n",
     "test",

tests/conftest.py

@@ -19,6 +19,7 @@ from aiogram.fsm.storage.memory import (
 from aiogram.fsm.storage.mongo import MongoStorage
 from aiogram.fsm.storage.pymongo import PyMongoStorage
 from aiogram.fsm.storage.redis import RedisStorage
+from aiogram.fsm.storage.sqlite import SqliteStorage
 from tests.mocked_bot import MockedBot
 
 DATA_DIR = Path(__file__).parent / "data"
@@ -142,6 +143,16 @@ async def memory_storage():
         await storage.close()
 
 
+@pytest.fixture()
+async def sqlite_storage():
+    storage = await SqliteStorage.connect("aiogram_fsm_test.sqlite")
+    try:
+        yield storage
+    finally:
+        await storage.close()
+        Path("aiogram_fsm_test.sqlite").unlink()
+
+
 @pytest.fixture()
 async def redis_isolation(redis_storage):
     isolation = redis_storage.create_isolation()

tests/test_fsm/storage/test_storages.py

@@ -13,6 +13,7 @@ from aiogram.fsm.storage.base import BaseStorage, StorageKey
         pytest.lazy_fixture("mongo_storage"),
         pytest.lazy_fixture("pymongo_storage"),
         pytest.lazy_fixture("memory_storage"),
+        pytest.lazy_fixture("sqlite_storage"),
     ],
 )
 class TestStorages: