ayaka.depend.db

  1import json
  2from pydantic import Field
  3from typing import List, TYPE_CHECKING
  4from typing_extensions import Self
  5from .depend import AyakaDepend
  6from .sql import PrimaryKey, JsonKey, insert_or_replace, create_table, drop_table, insert_or_replace_many, select_many, wrap, db
  7
  8if TYPE_CHECKING:
  9    from .. import AyakaApp
 10
 11
 12class AyakaDB(AyakaDepend):
 13    '''
 14```
 151. 继承时要书写 __table_name__
 162. 如果要把该类放入回调函数的参数表中，则还要编写classmethod async def create方法
 173. 设置主键需要使用
 18    <name>:<type> = Field(extra=AyakaDB.__primary_key__)
 194. 一些特殊类型的数据请设置其为json形式存取 
 20    <name>:<type> = Field(extra=AyakaDB.__json_key__)
 21    AyakaDB在写入时会自动序列化该数据为字符串，写入数据库，读取时则相反
 225. 若需要编写自定义读写数据方法，可以使用AyakaDB.get_db()方法获取sqlite3.Connection对象
 236. 使用前先调用classmethod def create_table
 24```
 25'''
 26    __table_name__ = ""
 27    __primary_key__ = PrimaryKey
 28    __json_key__ = JsonKey
 29
 30    def __init__(self, **data) -> None:
 31        if not self.__table_name__:
 32            raise Exception("__table_name__不可为空")
 33        super().__init__(**data)
 34
 35    @classmethod
 36    def _create_by_db_data(cls, data: dict):
 37        props = cls.props()
 38
 39        # 特殊处理json
 40        for k, v in props.items():
 41            extra: dict = v.get("extra", {})
 42            if extra.get("json"):
 43                if k in data:
 44                    data[k] = json.loads(data[k])
 45
 46        return cls(**data)
 47
 48    def dict(self, **params):
 49        data = super().dict(**params)
 50        props = self.props()
 51
 52        # 特殊处理json
 53        for k, v in props.items():
 54            extra: dict = v.get("extra", {})
 55            if extra.get("json"):
 56                if k in data:
 57                    data[k] = json.dumps(data[k], ensure_ascii=0)
 58        return data
 59
 60    @classmethod
 61    def drop_table(cls):
 62        drop_table(cls.__table_name__)
 63
 64    @classmethod
 65    def create_table(cls):
 66        '''根据数据类型自动创建表'''
 67        create_table(cls.__table_name__, cls)
 68
 69    @classmethod
 70    def replace(cls, data: Self):
 71        insert_or_replace(cls.__table_name__, data, "replace")
 72
 73    def save(self):
 74        '''写入数据库'''
 75        self.replace(self)
 76
 77    @classmethod
 78    def replace_many(cls, datas: List[Self]):
 79        insert_or_replace_many(cls.__table_name__, datas, "replace")
 80
 81    @classmethod
 82    def insert(cls, data: Self):
 83        insert_or_replace(cls.__table_name__, data, "insert")
 84
 85    @classmethod
 86    def insert_many(cls, datas: List[Self]):
 87        insert_or_replace_many(cls.__table_name__, datas, "insert")
 88
 89    @classmethod
 90    def select_many(cls, **params) -> List[Self]:
 91        '''按照params的值搜索数据，返回数据列表'''
 92        where = "1"
 93        if params:
 94            where = " and ".join(f"{k}={wrap(v)}" for k, v in params.items())
 95        return select_many(cls.__table_name__, cls, where)
 96
 97    @classmethod
 98    def select_one(cls, **params):
 99        '''按照params的值搜索数据，返回一项数据，若不存在，则自动根据params创建'''
100        datas = cls.select_many(**params)
101        if datas:
102            return datas[0]
103        return cls(**params)
104
105    @classmethod
106    def get_db(cls):
107        return db
108
109
110class AyakaGroupDB(AyakaDB):
111    '''继承时要书写`__table_name__`
112
113    主键有且仅有 group_id
114
115    使用前先调用classmethod def create_table'''
116    group_id: int = Field(extra=AyakaDB.__primary_key__)
117
118    @classmethod
119    async def _create_by_app(cls, app: "AyakaApp"):
120        return cls.select_one(group_id=app.group_id)
121
122
123class AyakaUserDB(AyakaDB):
124    '''继承时要书写`__table_name__`
125
126    主键有且仅有 group_id, user_id
127
128    使用前先调用classmethod def create_table'''
129    group_id: int = Field(extra=AyakaDB.__primary_key__)
130    user_id: int = Field(extra=AyakaDB.__primary_key__)
131
132    @classmethod
133    async def _create_by_app(cls, app: "AyakaApp"):
134        return cls.select_one(
135            group_id=app.group_id,
136            user_id=app.user_id
137        )