ayaka.depend.db

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