mirror of
https://github.com/Mayuri-Chan/pyrofork.git
synced 2025-12-29 12:04:51 +00:00
Move to aiosqlite and enable WAL
Signed-off-by: wulan17 <wulan17@nusantararom.org>
This commit is contained in:
Hitalo
wulan17
parent
c034a2da27
commit
afd2e3f0bc
4 changed files with 73 additions and 75 deletions
|
|
@ -18,7 +18,7 @@
|
|||
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
import aiosqlite
|
||||
from pathlib import Path
|
||||
|
||||
from .sqlite_storage import SQLiteStorage
|
||||
|
|
@ -34,39 +34,38 @@ class FileStorage(SQLiteStorage):
|
|||
|
||||
self.database = workdir / (self.name + self.FILE_EXTENSION)
|
||||
|
||||
def update(self):
|
||||
version = self.version()
|
||||
async def update(self):
|
||||
version = await self.version()
|
||||
|
||||
if version == 1:
|
||||
with self.lock, self.conn:
|
||||
self.conn.execute("DELETE FROM peers")
|
||||
await self.conn.execute("DELETE FROM peers")
|
||||
|
||||
version += 1
|
||||
|
||||
if version == 2:
|
||||
with self.lock, self.conn:
|
||||
self.conn.execute("ALTER TABLE sessions ADD api_id INTEGER")
|
||||
await self.conn.execute("ALTER TABLE sessions ADD api_id INTEGER")
|
||||
|
||||
version += 1
|
||||
|
||||
self.version(version)
|
||||
await self.version(version)
|
||||
|
||||
async def open(self):
|
||||
path = self.database
|
||||
file_exists = path.is_file()
|
||||
|
||||
self.conn = sqlite3.connect(str(path), timeout=1, check_same_thread=False)
|
||||
self.conn = await aiosqlite.connect(str(path), timeout=1)
|
||||
|
||||
await self.conn.execute("PRAGMA journal_mode=WAL")
|
||||
|
||||
if not file_exists:
|
||||
self.create()
|
||||
await self.create()
|
||||
else:
|
||||
self.update()
|
||||
await self.update()
|
||||
|
||||
with self.conn:
|
||||
try: # Python 3.6.0 (exactly this version) is bugged and won't successfully execute the vacuum
|
||||
self.conn.execute("VACUUM")
|
||||
except sqlite3.OperationalError:
|
||||
pass
|
||||
try: # Python 3.6.0 (exactly this version) is bugged and won't successfully execute the vacuum
|
||||
await self.conn.execute("VACUUM")
|
||||
except aiosqlite.OperationalError:
|
||||
pass
|
||||
|
||||
async def delete(self):
|
||||
os.remove(self.database)
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@
|
|||
|
||||
import base64
|
||||
import logging
|
||||
import sqlite3
|
||||
import aiosqlite
|
||||
import struct
|
||||
|
||||
from .sqlite_storage import SQLiteStorage
|
||||
|
|
@ -33,8 +33,8 @@ class MemoryStorage(SQLiteStorage):
|
|||
self.session_string = session_string
|
||||
|
||||
async def open(self):
|
||||
self.conn = sqlite3.connect(":memory:", check_same_thread=False)
|
||||
self.create()
|
||||
self.conn = await aiosqlite.connect(":memory:")
|
||||
await self.create()
|
||||
|
||||
if self.session_string:
|
||||
# Old format
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@
|
|||
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
import inspect
|
||||
import sqlite3
|
||||
import aiosqlite
|
||||
import time
|
||||
from threading import Lock
|
||||
from typing import List, Tuple, Any
|
||||
|
|
@ -97,22 +97,20 @@ class SQLiteStorage(Storage):
|
|||
def __init__(self, name: str):
|
||||
super().__init__(name)
|
||||
|
||||
self.conn = None # type: sqlite3.Connection
|
||||
self.lock = Lock()
|
||||
self.conn = None # type: aiosqlite.Connection
|
||||
|
||||
def create(self):
|
||||
with self.lock, self.conn:
|
||||
self.conn.executescript(SCHEMA)
|
||||
async def create(self):
|
||||
await self.conn.executescript(SCHEMA)
|
||||
|
||||
self.conn.execute(
|
||||
"INSERT INTO version VALUES (?)",
|
||||
(self.VERSION,)
|
||||
)
|
||||
await self.conn.execute(
|
||||
"INSERT INTO version VALUES (?)",
|
||||
(self.VERSION,)
|
||||
)
|
||||
|
||||
self.conn.execute(
|
||||
"INSERT INTO sessions VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
(2, None, None, None, 0, None, None)
|
||||
)
|
||||
await self.conn.execute(
|
||||
"INSERT INTO sessions VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
(2, None, None, None, 0, None, None)
|
||||
)
|
||||
|
||||
async def open(self):
|
||||
raise NotImplementedError
|
||||
|
|
@ -120,29 +118,27 @@ class SQLiteStorage(Storage):
|
|||
async def save(self):
|
||||
await self.date(int(time.time()))
|
||||
|
||||
with self.lock:
|
||||
self.conn.commit()
|
||||
await self.conn.commit()
|
||||
|
||||
async def close(self):
|
||||
with self.lock:
|
||||
self.conn.close()
|
||||
await self.conn.close()
|
||||
|
||||
async def delete(self):
|
||||
raise NotImplementedError
|
||||
|
||||
async def update_peers(self, peers: List[Tuple[int, int, str, str, str]]):
|
||||
with self.lock:
|
||||
self.conn.executemany(
|
||||
"REPLACE INTO peers (id, access_hash, type, username, phone_number)"
|
||||
"VALUES (?, ?, ?, ?, ?)",
|
||||
peers
|
||||
)
|
||||
await self.conn.executemany(
|
||||
"REPLACE INTO peers (id, access_hash, type, username, phone_number)"
|
||||
"VALUES (?, ?, ?, ?, ?)",
|
||||
peers
|
||||
)
|
||||
|
||||
async def get_peer_by_id(self, peer_id: int):
|
||||
r = self.conn.execute(
|
||||
q = await self.conn.execute(
|
||||
"SELECT id, access_hash, type FROM peers WHERE id = ?",
|
||||
(peer_id,)
|
||||
).fetchone()
|
||||
)
|
||||
r = await q.fetchone()
|
||||
|
||||
if r is None:
|
||||
raise KeyError(f"ID not found: {peer_id}")
|
||||
|
|
@ -150,11 +146,12 @@ class SQLiteStorage(Storage):
|
|||
return get_input_peer(*r)
|
||||
|
||||
async def get_peer_by_username(self, username: str):
|
||||
r = self.conn.execute(
|
||||
q = await self.conn.execute(
|
||||
"SELECT id, access_hash, type, last_update_on FROM peers WHERE username = ?"
|
||||
"ORDER BY last_update_on DESC",
|
||||
(username,)
|
||||
).fetchone()
|
||||
)
|
||||
r = await q.fetchone()
|
||||
|
||||
if r is None:
|
||||
raise KeyError(f"Username not found: {username}")
|
||||
|
|
@ -165,64 +162,65 @@ class SQLiteStorage(Storage):
|
|||
return get_input_peer(*r[:3])
|
||||
|
||||
async def get_peer_by_phone_number(self, phone_number: str):
|
||||
r = self.conn.execute(
|
||||
q = await self.conn.execute(
|
||||
"SELECT id, access_hash, type FROM peers WHERE phone_number = ?",
|
||||
(phone_number,)
|
||||
).fetchone()
|
||||
)
|
||||
r = await q.fetchone()
|
||||
|
||||
if r is None:
|
||||
raise KeyError(f"Phone number not found: {phone_number}")
|
||||
|
||||
return get_input_peer(*r)
|
||||
|
||||
def _get(self):
|
||||
async def _get(self):
|
||||
attr = inspect.stack()[2].function
|
||||
|
||||
return self.conn.execute(
|
||||
q = await self.conn.execute(
|
||||
f"SELECT {attr} FROM sessions"
|
||||
).fetchone()[0]
|
||||
)
|
||||
return (await q.fetchone())[0]
|
||||
|
||||
def _set(self, value: Any):
|
||||
async def _set(self, value: Any):
|
||||
attr = inspect.stack()[2].function
|
||||
|
||||
with self.lock, self.conn:
|
||||
self.conn.execute(
|
||||
f"UPDATE sessions SET {attr} = ?",
|
||||
(value,)
|
||||
)
|
||||
await self.conn.execute(
|
||||
f"UPDATE sessions SET {attr} = ?",
|
||||
(value,)
|
||||
)
|
||||
|
||||
def _accessor(self, value: Any = object):
|
||||
return self._get() if value == object else self._set(value)
|
||||
async def _accessor(self, value: Any = object):
|
||||
return await self._get() if value == object else await self._set(value)
|
||||
|
||||
async def dc_id(self, value: int = object):
|
||||
return self._accessor(value)
|
||||
return await self._accessor(value)
|
||||
|
||||
async def api_id(self, value: int = object):
|
||||
return self._accessor(value)
|
||||
return await self._accessor(value)
|
||||
|
||||
async def test_mode(self, value: bool = object):
|
||||
return self._accessor(value)
|
||||
return await self._accessor(value)
|
||||
|
||||
async def auth_key(self, value: bytes = object):
|
||||
return self._accessor(value)
|
||||
return await self._accessor(value)
|
||||
|
||||
async def date(self, value: int = object):
|
||||
return self._accessor(value)
|
||||
return await self._accessor(value)
|
||||
|
||||
async def user_id(self, value: int = object):
|
||||
return self._accessor(value)
|
||||
return await self._accessor(value)
|
||||
|
||||
async def is_bot(self, value: bool = object):
|
||||
return self._accessor(value)
|
||||
return await self._accessor(value)
|
||||
|
||||
def version(self, value: int = object):
|
||||
async def version(self, value: int = object):
|
||||
if value == object:
|
||||
return self.conn.execute(
|
||||
q = await self.conn.execute(
|
||||
"SELECT number FROM version"
|
||||
).fetchone()[0]
|
||||
)
|
||||
return (await q.fetchone())[0]
|
||||
else:
|
||||
with self.lock, self.conn:
|
||||
self.conn.execute(
|
||||
"UPDATE version SET number = ?",
|
||||
(value,)
|
||||
)
|
||||
await self.conn.execute(
|
||||
"UPDATE version SET number = ?",
|
||||
(value,)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
pyaes==1.6.1
|
||||
pymediainfo==6.0.1
|
||||
pysocks==1.7.1
|
||||
aiosqlite>=0.16.0,<0.18.0
|
||||
|
|
|
|||
Loading…
Reference in a new issue