diff --git a/CHANGES/1718.feature.rst b/CHANGES/1718.feature.rst new file mode 100644 index 00000000..648f832a --- /dev/null +++ b/CHANGES/1718.feature.rst @@ -0,0 +1 @@ +Added SqliteStorage which makes local development and small deployments a bit easier. \ No newline at end of file diff --git a/aiogram/fsm/storage/sqlite.py b/aiogram/fsm/storage/sqlite.py new file mode 100644 index 00000000..d65c4990 --- /dev/null +++ b/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 diff --git a/docs/dispatcher/finite_state_machine/storages.rst b/docs/dispatcher/finite_state_machine/storages.rst index ef39c4c3..66c259ab 100644 --- a/docs/dispatcher/finite_state_machine/storages.rst +++ b/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 ------------ diff --git a/pyproject.toml b/pyproject.toml index f4bfe873..b8039c76 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/conftest.py b/tests/conftest.py index 6a0c37f4..36942126 100644 --- a/tests/conftest.py +++ b/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() diff --git a/tests/test_fsm/storage/test_storages.py b/tests/test_fsm/storage/test_storages.py index 884f6874..5e6f8e68 100644 --- a/tests/test_fsm/storage/test_storages.py +++ b/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: