import os import pickle import sqlite3 from types import NoneType from typing import Any, Callable from packaging.version import parse import inspect import nonebot import pydantic from pydantic import BaseModel class LiteModel(BaseModel): TABLE_NAME: str = None id: int = None def dump(self, *args, **kwargs): if parse(pydantic.__version__) < parse("2.0.0"): return self.dict(*args, **kwargs) else: return self.model_dump(*args, **kwargs) class Database: def __init__(self, db_name: str): if os.path.dirname(db_name) != "" and not os.path.exists(os.path.dirname(db_name)): os.makedirs(os.path.dirname(db_name)) self.db_name = db_name self.conn = sqlite3.connect(db_name) self.cursor = self.conn.cursor() self._on_save_callbacks = [] def where_one(self, model: LiteModel, condition: str = "", *args: Any, default: Any = None) -> LiteModel | Any | None: """查询第一个 Args: model: 数据模型实例 condition: 查询条件，不给定则查询所有 *args: 参数化查询参数 default: 默认值 Returns: """ all_results = self.where_all(model, condition, *args) return all_results[0] if all_results else default def where_all(self, model: LiteModel, condition: str = "", *args: Any, default: Any = None) -> list[LiteModel | Any] | None: """查询所有 Args: model: 数据模型实例 condition: 查询条件，不给定则查询所有 *args: 参数化查询参数 default: 默认值 Returns: """ table_name = model.TABLE_NAME model_type = type(model) nonebot.logger.debug(f"Selecting {model.TABLE_NAME} WHERE {condition.replace('?', '%s') % args}") if not table_name: raise ValueError(f"数据模型{model_type.__name__}未提供表名") # condition = f"WHERE {condition}" # print(f"SELECT * FROM {table_name} {condition}", args) # if len(args) == 0: # results = self.cursor.execute(f"SELECT * FROM {table_name} {condition}").fetchall() # else: # results = self.cursor.execute(f"SELECT * FROM {table_name} {condition}", args).fetchall() if condition: results = self.cursor.execute(f"SELECT * FROM {table_name} WHERE {condition}", args).fetchall() else: results = self.cursor.execute(f"SELECT * FROM {table_name}").fetchall() fields = [description[0] for description in self.cursor.description] if not results: return default else: return [model_type(**self._load(dict(zip(fields, result)))) for result in results] def save(self, *args: LiteModel): """增/改操作 Args: *args: Returns: """ table_list = [item[0] for item in self.cursor.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()] for model in args: nonebot.logger.debug(f"Upserting {model}") if not model.TABLE_NAME: raise ValueError(f"数据模型 {model.__class__.__name__} 未提供表名") elif model.TABLE_NAME not in table_list: raise ValueError(f"数据模型 {model.__class__.__name__} 表 {model.TABLE_NAME} 不存在，请先迁移") else: self._save(model.dump(by_alias=True)) for callback in self._on_save_callbacks: callback(model) def _save(self, obj: Any) -> Any: # obj = copy.deepcopy(obj) if isinstance(obj, dict): table_name = obj.get("TABLE_NAME") row_id = obj.get("id") new_obj = {} for field, value in obj.items(): if isinstance(value, self.ITERABLE_TYPE): new_obj[self._get_stored_field_prefix(value) + field] = self._save(value) # self._save(value) # -> bytes elif isinstance(value, self.BASIC_TYPE): new_obj[field] = value else: raise ValueError(f"数据模型{table_name}包含不支持的数据类型，字段：{field} 值：{value} 值类型：{type(value)}") if table_name: fields, values = [], [] for n_field, n_value in new_obj.items(): if n_field not in ["TABLE_NAME", "id"]: fields.append(n_field) values.append(n_value) # 移除TABLE_NAME和id fields = list(fields) values = list(values) if row_id is not None: # 如果 _id 不为空，将 'id' 插入到字段列表的开始 fields.insert(0, 'id') # 将 _id 插入到值列表的开始 values.insert(0, row_id) fields = ', '.join([f'"{field}"' for field in fields]) placeholders = ', '.join('?' for _ in values) self.cursor.execute(f"INSERT OR REPLACE INTO {table_name}({fields}) VALUES ({placeholders})", tuple(values)) self.conn.commit() foreign_id = self.cursor.execute("SELECT last_insert_rowid()").fetchone()[0] return f"{self.FOREIGN_KEY_PREFIX}{foreign_id}@{table_name}" # -> FOREIGN_KEY_123456@{table_name} id@{table_name} else: return pickle.dumps(new_obj) # -> bytes elif isinstance(obj, (list, set, tuple)): obj_type = type(obj) # 到时候转回去 new_obj = [] for item in obj: if isinstance(item, self.ITERABLE_TYPE): new_obj.append(self._save(item)) elif isinstance(item, self.BASIC_TYPE): new_obj.append(item) else: raise ValueError(f"数据模型包含不支持的数据类型，值：{item} 值类型：{type(item)}") return pickle.dumps(obj_type(new_obj)) # -> bytes else: raise ValueError(f"数据模型包含不支持的数据类型，值：{obj} 值类型：{type(obj)}") def _load(self, obj: Any) -> Any: if isinstance(obj, dict): new_obj = {} for field, value in obj.items(): field: str if field.startswith(self.BYTES_PREFIX): if isinstance(value, bytes): new_obj[field.replace(self.BYTES_PREFIX, "")] = self._load(pickle.loads(value)) else: # 从value字段可能为None，fix at 2024/6/13 pass # 暂时不作处理，后面再修 elif field.startswith(self.FOREIGN_KEY_PREFIX): new_obj[field.replace(self.FOREIGN_KEY_PREFIX, "")] = self._load(self._get_foreign_data(value)) else: new_obj[field] = value return new_obj elif isinstance(obj, (list, set, tuple)): new_obj = [] for item in obj: if isinstance(item, bytes): # 对bytes进行尝试解析，解析失败则返回原始bytes try: new_obj.append(self._load(pickle.loads(item))) except Exception as e: new_obj.append(self._load(item)) elif isinstance(item, str) and item.startswith(self.FOREIGN_KEY_PREFIX): new_obj.append(self._load(self._get_foreign_data(item))) else: new_obj.append(self._load(item)) return new_obj else: return obj def delete(self, model: LiteModel, condition: str, *args: Any, allow_empty: bool = False): """ 删除满足条件的数据 Args: allow_empty: 允许空条件删除整个表 model: condition: *args: Returns: """ table_name = model.TABLE_NAME nonebot.logger.debug(f"Deleting {model} WHERE {condition} {args}") if not table_name: raise ValueError(f"数据模型{model.__class__.__name__}未提供表名") if model.id is not None: condition = f"id = {model.id}" if not condition and not allow_empty: raise ValueError("删除操作必须提供条件") self.cursor.execute(f"DELETE FROM {table_name} WHERE {condition}", args) self.conn.commit() def auto_migrate(self, *args: LiteModel): """ 自动迁移模型 Args: *args: 模型类实例化对象，支持空默认值，不支持嵌套迁移 Returns: """ for model in args: if not model.TABLE_NAME: raise ValueError(f"数据模型{type(model).__name__}未提供表名") # 若无则创建表 self.cursor.execute( f'CREATE TABLE IF NOT EXISTS "{model.TABLE_NAME}" (id INTEGER PRIMARY KEY AUTOINCREMENT)' ) # 获取表结构,field -> SqliteType new_structure = {} for n_field, n_value in model.dump(by_alias=True).items(): if n_field not in ["TABLE_NAME", "id"]: new_structure[self._get_stored_field_prefix(n_value) + n_field] = self._get_stored_type(n_value) # 原有的字段列表 existing_structure = dict([(column[1], column[2]) for column in self.cursor.execute(f'PRAGMA table_info({model.TABLE_NAME})').fetchall()]) # 检测缺失字段，由于SQLite是动态类型，所以不需要检测类型 for n_field, n_type in new_structure.items(): if n_field not in existing_structure.keys() and n_field.lower() not in ["id", "table_name"]: default_value = self.DEFAULT_MAPPING.get(n_type, 'NULL') self.cursor.execute( f"ALTER TABLE '{model.TABLE_NAME}' ADD COLUMN {n_field} {n_type} DEFAULT {self.DEFAULT_MAPPING.get(n_type, default_value)}" ) # 检测多余字段进行删除 for e_field in existing_structure.keys(): if e_field not in new_structure.keys() and e_field.lower() not in ['id']: self.cursor.execute( f'ALTER TABLE "{model.TABLE_NAME}" DROP COLUMN "{e_field}"' ) self.conn.commit() # 已完成 def _get_stored_field_prefix(self, value) -> str: """根据类型获取存储字段前缀，一定在后加上字段名 * -> "" Args: value: 储存的值 Returns: Sqlite3存储字段 """ if isinstance(value, LiteModel) or isinstance(value, dict) and "TABLE_NAME" in value: return self.FOREIGN_KEY_PREFIX elif type(value) in self.ITERABLE_TYPE: return self.BYTES_PREFIX return "" def _get_stored_type(self, value) -> str: """获取存储类型 Args: value: 储存的值 Returns: Sqlite3存储类型 """ if isinstance(value, dict) and "TABLE_NAME" in value: # 是一个模型字典，储存外键 return "INTEGER" return self.TYPE_MAPPING.get(type(value), "TEXT") def _get_foreign_data(self, foreign_value: str) -> dict: """ 获取外键数据 Args: foreign_value: Returns: """ foreign_value = foreign_value.replace(self.FOREIGN_KEY_PREFIX, "") table_name = foreign_value.split("@")[-1] foreign_id = foreign_value.split("@")[0] fields = [description[1] for description in self.cursor.execute(f"PRAGMA table_info({table_name})").fetchall()] result = self.cursor.execute(f"SELECT * FROM {table_name} WHERE id = ?", (foreign_id,)).fetchone() return dict(zip(fields, result)) def on_save(self, func: Callable[[LiteModel | Any], None]): """ 装饰一个可调用对象使其在储存数据模型时被调用 Args: func: Returns: """ def wrapper(model): # 检查被装饰函数声明的model类型和传入的model类型是否一致 sign = inspect.signature(func) if param := sign.parameters.get("model"): if isinstance(model, param.annotation): pass else: return else: return result = func(model) for callback in self._on_save_callbacks: callback(result) return result self._on_save_callbacks.append(wrapper) return wrapper TYPE_MAPPING = { int : "INTEGER", float : "REAL", str : "TEXT", bool : "INTEGER", bytes : "BLOB", NoneType : "NULL", # dict : "TEXT", # list : "TEXT", # tuple : "TEXT", # set : "TEXT", dict : "BLOB", # LITEYUKIDICT{key_name} list : "BLOB", # LITEYUKILIST{key_name} tuple : "BLOB", # LITEYUKITUPLE{key_name} set : "BLOB", # LITEYUKISET{key_name} LiteModel: "TEXT" # FOREIGN_KEY_{table_name} } DEFAULT_MAPPING = { "TEXT" : "''", "INTEGER": 0, "REAL" : 0.0, "BLOB" : None, "NULL" : None } # 基础类型 BASIC_TYPE = (int, float, str, bool, bytes, NoneType) # 可序列化类型 ITERABLE_TYPE = (dict, list, tuple, set, LiteModel) # 外键前缀 FOREIGN_KEY_PREFIX = "FOREIGN_KEY_" # 转换为的字节前缀 BYTES_PREFIX = "PICKLE_BYTES_" # transaction tx 事务操作 def first(self, model: LiteModel) -> "Database": pass def where(self, condition: str, *args) -> "Database": pass def limit(self, limit: int) -> "Database": pass def order(self, order: str) -> "Database": pass def check_sqlite_keyword(name): sqlite_keywords = [ "ABORT", "ACTION", "ADD", "AFTER", "ALL", "ALTER", "ANALYZE", "AND", "AS", "ASC", "ATTACH", "AUTOINCREMENT", "BEFORE", "BEGIN", "BETWEEN", "BY", "CASCADE", "CASE", "CAST", "CHECK", "COLLATE", "COLUMN", "COMMIT", "CONFLICT", "CONSTRAINT", "CREATE", "CROSS", "CURRENT_DATE", "CURRENT_TIME", "CURRENT_TIMESTAMP", "DATABASE", "DEFAULT", "DEFERRABLE", "DEFERRED", "DELETE", "DESC", "DETACH", "DISTINCT", "DROP", "EACH", "ELSE", "END", "ESCAPE", "EXCEPT", "EXCLUSIVE", "EXISTS", "EXPLAIN", "FAIL", "FOR", "FOREIGN", "FROM", "FULL", "GLOB", "GROUP", "HAVING", "IF", "IGNORE", "IMMEDIATE", "IN", "INDEX", "INDEXED", "INITIALLY", "INNER", "INSERT", "INSTEAD", "INTERSECT", "INTO", "IS", "ISNULL", "JOIN", "KEY", "LEFT", "LIKE", "LIMIT", "MATCH", "NATURAL", "NO", "NOT", "NOTNULL", "NULL", "OF", "OFFSET", "ON", "OR", "ORDER", "OUTER", "PLAN", "PRAGMA", "PRIMARY", "QUERY", "RAISE", "RECURSIVE", "REFERENCES", "REGEXP", "REINDEX", "RELEASE", "RENAME", "REPLACE", "RESTRICT", "RIGHT", "ROLLBACK", "ROW", "SAVEPOINT", "SELECT", "SET", "TABLE", "TEMP", "TEMPORARY", "THEN", "TO", "TRANSACTION", "TRIGGER", "UNION", "UNIQUE", "UPDATE", "USING", "VACUUM", "VALUES", "VIEW", "VIRTUAL", "WHEN", "WHERE", "WITH", "WITHOUT" ] return True # if name.upper() in sqlite_keywords: # raise ValueError(f"'{name}' 是SQLite保留字，不建议使用，请更换名称")