akari-bot/modules/wiki/dbutils.py

import json

from core.elements import MessageSession
from .orm import session, WikiTargetSetInfo, WikiInfo


class WikiTargetInfo:
    def __init__(self, msg: [MessageSession, str]):
        if isinstance(msg, MessageSession):
            targetId = msg.target.targetId
        else:
            targetId = msg
        self.query = session.query(WikiTargetSetInfo).filter_by(targetId=targetId).first()
        if self.query is None:
            try:
                session.add_all([WikiTargetSetInfo(targetId=targetId, iws='{}', headers='{}')])
                session.commit()
            except Exception:
                session.rollback()
                raise
            self.query = session.query(WikiTargetSetInfo).filter_by(targetId=targetId).first()

    def add_start_wiki(self, url):
        try:
            self.query.link = url
            session.commit()
            session.expire_all()
            return True
        except Exception:
            session.rollback()
            raise

    def get_start_wiki(self):
        if self.query is not None:
            return self.query.link if self.query.link is not None else False
        return False

    def config_interwikis(self, iw: str, iwlink: str = None, let_it=True):
        try:
            interwikis = json.loads(self.query.iws)
            if let_it:
                interwikis[iw] = iwlink
            else:
                if iw in interwikis:
                    del interwikis[iw]
            self.query.iws = json.dumps(interwikis)
            session.commit()
            session.expire_all()
            return True
        except Exception:
            session.rollback()
            raise

    def get_interwikis(self):
        q = self.query.iws
        if q is not None:
            iws = json.loads(q)
            return iws
        else:
            return False

    def config_headers(self, headers, let_it: [bool, None] = True):
        try:
            headers_ = json.loads(self.query.headers)
            if let_it:
                for x in headers:
                    headers_[x] = headers[x]
            elif let_it is None:
                headers_ = {}
            else:
                for x in headers:
                    if x in headers_:
                        del headers_[x]
            self.query.headers = json.dumps(headers_)
            session.commit()
            return True
        except Exception:
            session.rollback()
            raise

    def get_headers(self):
        if self.query is not None:
            q = self.query.headers
            headers = json.loads(q)
        else:
            headers = {'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6'}
        return headers


class WikiSiteInfo:
    def __init__(self, api_link):
        self.api_link = api_link
        self.query = session.query(WikiInfo).filter_by(apiLink=api_link).first()

    def get(self):
        if self.query is not None:
            return self.query.siteInfo, self.query.timestamp
        return False

    def update(self, info: dict):
        try:
            if self.query is None:
                session.add_all([WikiInfo(apiLink=self.api_link, siteInfo=json.dumps(info))])
            else:
                self.query.siteInfo = json.dumps(info)
            session.commit()
            return True
        except Exception:
            session.rollback()
            raise