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_by_app方法
 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对象
 23```
 24'''
 25    __table_name__ = ""
 26    __primary_key__ = PrimaryKey
 27    __json_key__ = JsonKey
 28    __created__ = False
 29
 30    def __init__(self, **data) -> None:
 31        self.create_table()
 32        super().__init__(**data)
 33
 34    @classmethod
 35    def _create_by_db_data(cls, data: dict):
 36        props = cls.props()
 37
 38        # 特殊处理json
 39        for k, v in props.items():
 40            extra: dict = v.get("extra", {})
 41            if extra.get("json"):
 42                if k in data:
 43                    data[k] = json.loads(data[k])
 44
 45        return cls(**data)
 46
 47    def dict(self, **params):
 48        data = super().dict(**params)
 49        props = self.props()
 50
 51        # 特殊处理json
 52        for k, v in props.items():
 53            extra: dict = v.get("extra", {})
 54            if extra.get("json"):
 55                if k in data:
 56                    data[k] = json.dumps(data[k], ensure_ascii=0)
 57        return data
 58
 59    @classmethod
 60    def drop_table(cls):
 61        drop_table(cls.__table_name__)
 62
 63    @classmethod
 64    def create_table(cls):
 65        '''根据数据类型自动创建表'''
 66        if not cls.__created__:
 67            if not cls.__table_name__:
 68                raise Exception("__table_name__不可为空")
 69            cls.__created__ = True
 70            create_table(cls.__table_name__, cls)
 71
 72    @classmethod
 73    def replace(cls, data: Self):
 74        cls.create_table()
 75        insert_or_replace(cls.__table_name__, data, "replace")
 76
 77    def save(self):
 78        '''写入数据库'''
 79        self.replace(self)
 80
 81    @classmethod
 82    def replace_many(cls, datas: List[Self]):
 83        cls.create_table()
 84        insert_or_replace_many(cls.__table_name__, datas, "replace")
 85
 86    @classmethod
 87    def insert(cls, data: Self):
 88        cls.create_table()
 89        insert_or_replace(cls.__table_name__, data, "insert")
 90
 91    @classmethod
 92    def insert_many(cls, datas: List[Self]):
 93        cls.create_table()
 94        insert_or_replace_many(cls.__table_name__, datas, "insert")
 95
 96    @classmethod
 97    def select_many(cls, **params) -> List[Self]:
 98        '''按照params的值搜索数据，返回数据列表，若没有符合的数据则返回空列表'''
 99        where = "1"
100        if params:
101            where = " and ".join(f"{k}={wrap(v)}" for k, v in params.items())
102        cls.create_table()
103        return select_many(cls.__table_name__, cls, where)
104
105    @classmethod
106    def select_one(cls, **params):
107        '''按照params的值搜索数据，返回一项数据，若不存在，则自动根据params创建，创建后自动写入数据库'''
108        datas = cls.select_many(**params)
109        if datas:
110            return datas[0]
111        data = cls(**params)
112        data.save()
113        return data
114
115    @classmethod
116    def get_db(cls):
117        return db
118
119
120class AyakaGroupDB(AyakaDB):
121    '''继承时要书写`__table_name__`
122
123    主键有且仅有 group_id'''
124    group_id: int = Field(extra=AyakaDB.__primary_key__)
125
126    @classmethod
127    async def _create_by_app(cls, app: "AyakaApp"):
128        return cls.select_one(group_id=app.group_id)
129
130
131class AyakaUserDB(AyakaDB):
132    '''继承时要书写`__table_name__`
133
134    主键有且仅有 group_id, user_id'''
135    group_id: int = Field(extra=AyakaDB.__primary_key__)
136    user_id: int = Field(extra=AyakaDB.__primary_key__)
137
138    @classmethod
139    async def _create_by_app(cls, app: "AyakaApp"):
140        return cls.select_one(
141            group_id=app.group_id,
142            user_id=app.user_id
143        )