├── AiogramStorages ├── __init__.py └── storages.py ├── banner.png ├── banner.psd ├── LICENSE └── README.md /AiogramStorages/__init__.py: -------------------------------------------------------------------------------- 1 | from . import storages -------------------------------------------------------------------------------- /banner.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Aiogram-Templates/AiogramStorages/HEAD/banner.png -------------------------------------------------------------------------------- /banner.psd: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Aiogram-Templates/AiogramStorages/HEAD/banner.psd -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 DIMFLIX 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | ![enter image description here](https://github.com/DIMFLIX-OFFICIAL/aiogram-storages/blob/main/banner.png?raw=true) 3 | # Save your data! 4 | 5 | **Aiogram-storages** was created to extend the standard fsm_storage options in **aiogram**. 6 | Our library supports such databases as: 7 | 8 | # Installation 9 | 10 | pip install AiogramStorages==1.0.0 11 | 12 | # PostgreSQL 13 | 14 | Support for storage with the **PostgreSQL** database is due to the **asyncpg** asynchronous library, which gives a **huge data processing speed**, and, accordingly, the bot itself. 15 | 16 | In order to use it, you need to create an instance of the **PGStorage** class, to which you need to pass the **required parameters (user, password, db_name).** You can also specify additional parameters (host, port). 17 | 18 | Next, this instance must be passed to the **Dispatcher.** 19 | ## Example 20 | 21 | from AiogramStorages.storages import PGStorage 22 | 23 | storage = PGStorage(username='YourUser', password='YourPassword', db_name='YourDbName') 24 | dp = Dispatcher(bot, storage=storage) 25 | 26 | ## Warning 27 | 28 | By default, **PGStorage** creates three tables in your database named: **aiogram-states**, **aiogram-data**, **aiogram-buckets**. 29 | 30 | We **strongly recommend** that you do **not use these names as the name of the table**, otherwise there may be disagreements. 31 | 32 | 33 | 34 | # SQLiteStorage 35 | 36 | 37 | Support for storage with the **SQLite** database is due to the **aiosqlite** asynchronous library, which gives a **huge data processing speed**, and, accordingly, the bot itself. 38 | 39 | In order to use it, you need to create an instance of the **SQLiteStorage** class, to which you need to pass the **required parameters (db_path).** 40 | 41 | Next, this instance must be passed to the **Dispatcher.** 42 | ## Example 43 | 44 | from AiogramStorages.storages import SQLiteStorage 45 | 46 | storage = SQLiteStorage(db_path='your_path') 47 | dp = Dispatcher(bot, storage=storage) 48 | 49 | ## Warning 50 | 51 | By default, **SQLiteStorage** creates three tables in your database named: **aiogram-states**, **aiogram-data**, **aiogram-buckets**. 52 | 53 | We **strongly recommend** that you do **not use these names as the name of the table**, otherwise there may be disagreements. 54 | 55 | -------------------------------------------------------------------------------- /AiogramStorages/storages.py: -------------------------------------------------------------------------------- 1 | from typing import Union, Dict, Optional, List, Tuple, AnyStr 2 | 3 | import aiosqlite 4 | import asyncpg 5 | import jsonpickle 6 | from aiogram.dispatcher.storage import BaseStorage 7 | 8 | 9 | class SQLiteStorage(BaseStorage): 10 | """ 11 | AioSqlite-based storage for FSM. 12 | 13 | Usage: 14 | 15 | storage = SQLiteStorage(db_path='data/database.db') 16 | dp = Dispatcher(bot, storage=storage) 17 | """ 18 | 19 | def __init__(self, db_path): 20 | self._path_db = db_path 21 | self._db = None 22 | 23 | async def close(self): 24 | if isinstance(self._db, aiosqlite.Connection): 25 | await self._db.close() 26 | 27 | async def get_db(self) -> aiosqlite.Connection: 28 | if isinstance(self._db, aiosqlite.Connection): 29 | return self._db 30 | 31 | self._db = await aiosqlite.connect(database=self._path_db) 32 | await self._db.execute("""CREATE TABLE IF NOT EXISTS "aiogram_state"( 33 | "user" BIGINT NOT NULL PRIMARY KEY, 34 | "chat" BIGINT NOT NULL, 35 | "state" TEXT NOT NULL)""") 36 | await self._db.execute("""CREATE TABLE IF NOT EXISTS "aiogram_data"( 37 | "user" BIGINT NOT NULL PRIMARY KEY, 38 | "chat" BIGINT NOT NULL, 39 | "data" TEXT)""") 40 | await self._db.execute("""CREATE TABLE IF NOT EXISTS "aiogram_bucket"( 41 | "user" BIGINT NOT NULL PRIMARY KEY, 42 | "chat" BIGINT NOT NULL, 43 | "bucket" TEXT NOT NULL)""") 44 | 45 | return self._db 46 | 47 | async def wait_closed(self): 48 | return True 49 | 50 | async def set_state(self, *, chat: Union[str, int, None] = None, 51 | user: Union[str, int, None] = None, 52 | state: Optional[AnyStr] = None): 53 | chat, user = map(int, self.check_address(chat=chat, user=user)) 54 | db = await self.get_db() 55 | 56 | if state is not None: 57 | await db.execute("""INSERT INTO "aiogram_state" VALUES(?, ?, ?)""" 58 | """ON CONFLICT ("user") DO UPDATE SET "state" = ?""", 59 | (user, chat, self.resolve_state(state), self.resolve_state(state))) 60 | await db.commit() 61 | else: 62 | await db.execute("""DELETE FROM "aiogram_state" WHERE chat=? AND "user"=?""", (chat, user)) 63 | await db.commit() 64 | 65 | async def get_state(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 66 | default: Optional[str] = None) -> Optional[str]: 67 | chat, user = map(int, self.check_address(chat=chat, user=user)) 68 | db = await self.get_db() 69 | async with db.execute("""SELECT "state" FROM "aiogram_state" WHERE "chat"=? AND "user"=?""", (chat, user)) as cursor: 70 | result = await cursor.fetchone() 71 | return result[0] if result else self.resolve_state(default) 72 | 73 | async def set_data(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 74 | data: Dict = None): 75 | chat, user = map(int, self.check_address(chat=chat, user=user)) 76 | db = await self.get_db() 77 | 78 | if data: 79 | await db.execute("""INSERT INTO "aiogram_data" VALUES(?, ?, ?)""" 80 | """ON CONFLICT ("user") DO UPDATE SET "data" = ?""", 81 | (user, chat, jsonpickle.encode(data), jsonpickle.encode(data))) 82 | await db.commit() 83 | else: 84 | await db.execute("""DELETE FROM "aiogram_data" WHERE "chat"=? AND "user"=?""", (chat, user)) 85 | await db.commit() 86 | 87 | async def get_data(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 88 | default: Optional[dict] = None) -> Dict: 89 | chat, user = map(int, self.check_address(chat=chat, user=user)) 90 | db = await self.get_db() 91 | async with db.execute("""SELECT "data" FROM "aiogram_data" WHERE "chat"=? AND "user"=?""", (chat, user)) as cursor: 92 | result = await cursor.fetchone() 93 | 94 | return jsonpickle.decode(result[0]) if result else default or {} 95 | 96 | async def update_data(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 97 | data: Dict = None, **kwargs): 98 | if data is None: 99 | data = {} 100 | temp_data = await self.get_data(chat=chat, user=user, default={}) 101 | temp_data.update(data, **kwargs) 102 | await self.set_data(chat=chat, user=user, data=temp_data) 103 | 104 | def has_bucket(self): 105 | return True 106 | 107 | async def get_bucket(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 108 | default: Optional[dict] = None) -> Dict: 109 | chat, user = map(int, self.check_address(chat=chat, user=user)) 110 | db = await self.get_db() 111 | async with db.execute("""SELECT "bucket" FROM "aiogram_bucket" WHERE "chat"=? AND "user"=?""", 112 | (chat, user)) as cursor: 113 | result = await cursor.fetchone() 114 | return jsonpickle.decode(result[0]) if result else default or {} 115 | 116 | async def set_bucket(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 117 | bucket: Dict = None): 118 | chat, user = map(int, self.check_address(chat=chat, user=user)) 119 | db = await self.get_db() 120 | await db.execute("""INSERT INTO "aiogram_bucket" VALUES(?, ?, ?)""" 121 | """ON CONFLICT ("user") DO UPDATE SET "bucket" = ?""", 122 | (user, chat, jsonpickle.encode(bucket), jsonpickle.encode(bucket))) 123 | await db.commit() 124 | 125 | async def update_bucket(self, *, chat: Union[str, int, None] = None, 126 | user: Union[str, int, None] = None, 127 | bucket: Dict = None, **kwargs): 128 | if bucket is None: 129 | bucket = {} 130 | temp_bucket = await self.get_bucket(chat=chat, user=user) 131 | temp_bucket.update(bucket, **kwargs) 132 | await self.set_bucket(chat=chat, user=user, bucket=temp_bucket) 133 | 134 | async def reset_all(self, full=True): 135 | db = await self.get_db() 136 | await db.execute("DROP TABLE aiogram_state") 137 | if full: 138 | await db.execute("DROP TABLE aiogram_data") 139 | await db.execute("DROP TABLE aiogram_bucket") 140 | 141 | await db.commit() 142 | 143 | async def get_states_list(self) -> List[Tuple[int, int]]: 144 | db = await self.get_db() 145 | async with db.execute("SELECT * FROM aiogram_state") as cursor: 146 | items = await cursor.fetchall() 147 | return [(int(item[1]), int(item[0])) for item in items] 148 | 149 | 150 | class PGStorage(BaseStorage): 151 | """ 152 | AsyncPG-based storage for FSM. 153 | 154 | Usage: 155 | 156 | storage = PGStorage(host='localhost', port=5432, db_name='aiogram_fsm') 157 | dp = Dispatcher(bot, storage=storage) 158 | """ 159 | 160 | def __init__(self, username: str, password: str, host='localhost', port=5432, db_name='aiogram_fsm'): 161 | self._host = host 162 | self._port = port 163 | self._db_name: str = db_name 164 | self._username = username 165 | self._password = password 166 | self._db = None 167 | 168 | async def close(self): 169 | if isinstance(self._db, asyncpg.Connection): 170 | await self._db.close() 171 | 172 | async def get_db(self) -> asyncpg.Connection: 173 | if isinstance(self._db, asyncpg.Connection): 174 | return self._db 175 | 176 | self._db = await asyncpg.connect( 177 | user=self._username, 178 | password=self._password, 179 | host=self._host, 180 | port=self._port, 181 | database=self._db_name 182 | ) 183 | await self._db.execute("""CREATE TABLE IF NOT EXISTS "aiogram_state"( 184 | "user" BIGINT NOT NULL PRIMARY KEY, 185 | "chat" BIGINT NOT NULL, 186 | "state" TEXT NOT NULL)""") 187 | await self._db.execute("""CREATE TABLE IF NOT EXISTS "aiogram_data"( 188 | "user" BIGINT NOT NULL PRIMARY KEY, 189 | "chat" BIGINT NOT NULL, 190 | "data" JSON)""") 191 | await self._db.execute("""CREATE TABLE IF NOT EXISTS "aiogram_bucket"( 192 | "user" BIGINT NOT NULL PRIMARY KEY, 193 | "chat" BIGINT NOT NULL, 194 | "bucket" JSON NOT NULL)""") 195 | 196 | return self._db 197 | 198 | async def wait_closed(self): 199 | return True 200 | 201 | async def set_state(self, *, chat: Union[str, int, None] = None, 202 | user: Union[str, int, None] = None, 203 | state: Optional[AnyStr] = None): 204 | chat, user = map(int, self.check_address(chat=chat, user=user)) 205 | db = await self.get_db() 206 | if state is not None: 207 | await db.execute("""INSERT INTO "aiogram_state" VALUES($1, $2, $3)""" 208 | """ON CONFLICT ("user") DO UPDATE SET "state" = $3""", 209 | user, chat, self.resolve_state(state)) 210 | else: 211 | await db.execute("""DELETE FROM "aiogram_state" WHERE chat=$1 AND "user"=$2""", chat, user) 212 | 213 | async def get_state(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 214 | default: Optional[str] = None) -> Optional[str]: 215 | chat, user = map(int, self.check_address(chat=chat, user=user)) 216 | db = await self.get_db() 217 | result = await db.fetchval("""SELECT "state" FROM "aiogram_state" WHERE "chat"=$1 AND "user"=$2""", chat, user) 218 | return result if result else self.resolve_state(default) 219 | 220 | async def set_data(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 221 | data: Dict = None): 222 | chat, user = map(int, self.check_address(chat=chat, user=user)) 223 | db = await self.get_db() 224 | if data: 225 | await db.execute("""INSERT INTO "aiogram_data" VALUES($1, $2, $3)""" 226 | """ON CONFLICT ("user") DO UPDATE SET "data" = $3""", 227 | user, chat, jsonpickle.encode(data)) 228 | else: 229 | await db.execute("""DELETE FROM "aiogram_data" WHERE "chat"=$1 AND "user"=$2""", chat, user) 230 | 231 | async def get_data(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 232 | default: Optional[dict] = None) -> Dict: 233 | chat, user = map(int, self.check_address(chat=chat, user=user)) 234 | db = await self.get_db() 235 | result = await db.fetchval("""SELECT "data" FROM "aiogram_data" WHERE "chat"=$1 AND "user"=$2""", chat, user) 236 | return jsonpickle.decode(result) if result else default or {} 237 | 238 | async def update_data(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 239 | data: Dict = None, **kwargs): 240 | if data is None: 241 | data = {} 242 | temp_data = await self.get_data(chat=chat, user=user, default={}) 243 | temp_data.update(data, **kwargs) 244 | await self.set_data(chat=chat, user=user, data=temp_data) 245 | 246 | def has_bucket(self): 247 | return True 248 | 249 | async def get_bucket(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 250 | default: Optional[dict] = None) -> Dict: 251 | chat, user = map(int, self.check_address(chat=chat, user=user)) 252 | db = await self.get_db() 253 | result = await db.fetchval("""SELECT "bucket" FROM "aiogram_bucket" WHERE "chat"=$1 AND "user"=$2""", chat, user) 254 | return jsonpickle.decode(result) if result else default or {} 255 | 256 | async def set_bucket(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, 257 | bucket: Dict = None): 258 | chat, user = map(int, self.check_address(chat=chat, user=user)) 259 | db = await self.get_db() 260 | await db.execute("""INSERT INTO "aiogram_bucket" VALUES($1, $2, $3)""" 261 | """ON CONFLICT ("user") DO UPDATE SET "bucket" = $3""", 262 | user, chat, jsonpickle.encode(bucket)) 263 | 264 | async def update_bucket(self, *, chat: Union[str, int, None] = None, 265 | user: Union[str, int, None] = None, 266 | bucket: Dict = None, **kwargs): 267 | if bucket is None: 268 | bucket = {} 269 | temp_bucket = await self.get_bucket(chat=chat, user=user) 270 | temp_bucket.update(bucket, **kwargs) 271 | await self.set_bucket(chat=chat, user=user, bucket=temp_bucket) 272 | 273 | async def reset_all(self, full=True): 274 | db = await self.get_db() 275 | await db.execute("DROP TABLE aiogram_state") 276 | if full: 277 | await db.execute("DROP TABLE aiogram_data") 278 | await db.execute("DROP TABLE aiogram_bucket") 279 | 280 | async def get_states_list(self) -> List[Tuple[int, int]]: 281 | db = await self.get_db() 282 | items = await db.fetch("SELECT * FROM aiogram_state") 283 | return [(int(item['chat']), int(item['user'])) for item in map(dict, items)] 284 | --------------------------------------------------------------------------------