Files
nonebot-plugin-tetris-stats/nonebot_plugin_tetris_stats/utils/database.py
2022-10-09 13:37:16 +08:00

89 lines
2.6 KiB
Python

import os
from asyncio import gather
from sqlite3 import Connection, connect
from nonebot import get_driver
from nonebot.log import logger
from .config import Config
driver = get_driver()
config = Config.parse_obj(get_driver().config)
@driver.on_startup
async def _():
await DataBase.init_db()
@driver.on_shutdown
async def _():
await DataBase.close_db()
class DataBase():
'''数据库交互类'''
_db: Connection | None = None
@classmethod
async def init_db(cls) -> Connection:
'''初始化数据库'''
if not os.path.exists(os.path.dirname(config.db_path)):
os.makedirs(os.path.dirname(config.db_path))
cls._db = connect(config.db_path)
cursor = cls._db.cursor()
cursor.execute('''CREATE TABLE IF NOT EXISTS IOBIND
(QQ INTEGER NOT NULL,
USER TEXT NOT NULL)''')
cursor.execute('''CREATE TABLE IF NOT EXISTS TOPBIND
(QQ INTEGER NOT NULL,
USER TEXT NOT NULL)''')
cls._db.commit()
logger.info('数据库初始化完成')
return cls._db
@classmethod
async def _get_db(cls) -> Connection:
'''获取数据库对象'''
return cls._db or await cls.init_db()
@classmethod
async def query_bind_info(cls, qq_number: str | int, game_type: str) -> str | None:
'''查询绑定信息'''
db = await cls._get_db()
cursor = db.cursor()
cursor.execute(
f'SELECT USER FROM {game_type}BIND WHERE QQ = {qq_number}')
user = cursor.fetchone()
if user is None:
return None
return user[0]
@classmethod
async def write_bind_info(cls, qq_number: str | int, user: str, game_type: str) -> str:
'''写入绑定信息'''
bind_info, db = await gather(
cls.query_bind_info(qq_number=qq_number, game_type=game_type),
cls._get_db()
)
cursor = db.cursor()
if bind_info is not None:
cursor.execute(
f'UPDATE {game_type}BIND SET USER = ? WHERE QQ = ?', (user, qq_number))
message = '更新成功'
elif bind_info is None:
cursor.execute(
f'INSERT INTO {game_type}BIND (QQ, USER) VALUES (?, ?)', (qq_number, user))
message = '绑定成功'
else:
raise ValueError('预期外行为, 请上报GitHub')
db.commit()
return message
@classmethod
async def close_db(cls) -> None:
'''关闭数据库对象'''
if isinstance(cls._db, Connection):
cls._db.close()