douyin/db/user_oauth.py

import datetime
from typing import Optional
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from db.engine import engine

from sqlmodel import Field, SQLModel,create_engine,Session,select,func,Column,Relationship
from config import DB_URL,logger
# from db.common import engine
from sqlalchemy import UniqueConstraint, Index, asc  
from sqlalchemy.dialects.postgresql import insert
from db.base import DouyinBaseRepository
from db.model import *


# Database manager class
class UserOAuthRepository(DouyinBaseRepository):  
    def __init__(self, engine=engine):  
        super().__init__(UserOAuthToken, engine)  
        self.model:UserOAuthToken

    def dict_to_model(self, data: dict):  
        access_token_expires_in = datetime.datetime.now() + datetime.timedelta(seconds=data.get("expires_in"))  
        refresh_token_expires_in = datetime.datetime.now() + datetime.timedelta(seconds=data.get("refresh_expires_in"))  

        model:UserOAuthToken = super().dict_to_model(data,self.model)
        model.expires_at = access_token_expires_in
        model.refresh_expires_at = refresh_token_expires_in
        return model
    
    # field = UserOAuthToken.expires_at | UserOAuthToken.refresh_expires_at 
    def select_nearest_expire(self, field: Column) -> UserOAuthToken:  
        with Session(bind=self.engine) as session:
            statement = select(self.model).where(field > datetime.datetime.now()).order_by(asc(field)).limit(1)
            res = session.exec(statement)
            return res.first()
        
    def update_refresh_token(self, open_id:str, refresh_token:str, refresh_expires_in:int):
        refresh_token_expires_at = datetime.datetime.now() + datetime.timedelta(seconds=refresh_expires_in)
        model = UserOAuthToken(open_id=open_id, refresh_expires_in=refresh_expires_in, refresh_token=refresh_token,refresh_expires_at=refresh_token_expires_at)
        return self.add_or_update(model)
    
    def update_access_token(self, open_id:str, access_token:str, expires_in:int):
        access_token_expires_at = datetime.datetime.now() + datetime.timedelta(seconds=expires_in)
        model = UserOAuthToken(open_id=open_id, access_token=access_token, expires_in=expires_in,expires_at=access_token_expires_at)
        return self.add_or_update(model)


    async def delete_token(self, token_id: int):  
        async with self.session_factory() as session:  
            statement = select(UserOAuthToken).where(UserOAuthToken.id == token_id)  
            token = await session.execute(statement).scalars().first()  
            if token:  
                await session.delete(token)  
                await session.commit()  
                print(f"Record deleted: ID - {token_id}")  
            else:  
                print(f"Record with ID {token_id} not found") 

    

def test_add(open_id=None):  
    SQLModel.metadata.create_all(engine)
    
    user_oauth = {'access_token': 'act.3.m3kiZmfxxIH95i1bHZ7Bq3Wkv_Xm5TtD3kpLGjtCr3G96WIINBKEvzlsaObrGcH4GaxTQeLZA13jkzoZhpAwPwMRqFxlVuIcxpge_-BpdFib1xHqkcFa4B-LX4zpd2YK3kDFTFfMcJXN_fZ2eByg6oqqa1OieUWcvlaVgw==', 'captcha': '', 'desc_url': '', 'description': '', 'error_code': 0, 'expires_in': 1296000, 'log_id': '202402171353461C33F969CEFB511B216F', 'open_id': '_000LiV_o0FGKMwaZgqGMCvcVSf-UAnNU_kk', 'refresh_expires_in': 2592000, 'refresh_token': 'rft.e4b3da8bd3ef01d880d827b11e826391OEGHiRrTLcp5zsGYP1dh6F9Bo7fg', 'scope': 'user_info,trial.whitelist'}
    user_oauth2 = {'access_token': 'act', 'expires_in': 19290, 'open_id': '55test2', 'refresh_expires_in': 1950, 'refresh_token': 'rft', 'scope': 'user_info,trial.whitelist'}
    user_info = {
    "avatar": "https://p26.douyinpic.com/aweme/100x100/aweme-avatar/tos-cn-i-0813_66c4e34ae8834399bbf967c3d3c919db.jpeg?from=4010531038",
    "avatar_larger": "https://p3.douyinpic.com/aweme/1080x1080/aweme-avatar/tos-cn-i-0813_66c4e34ae8834399bbf967c3d3c919db.jpeg?from=4010531038",
    "captcha": "",
    "city": "",
    "client_key": "55",
    "country": "",
    "desc_url": "",
    "description": "",
    "district": "",
    "e_account_role": "",
    "error_code": 0,
    "gender": 0,
    "log_id": "202401261424326FE877A6CAB03910C553",
    "nickname": "程序员马工",
    "open_id": "_000QadFMhmU1jNCI3JdPnyVDL6XavC70dFy",
    "province": "",
    "union_id": "123-01ae-59bd-978a-1de8566186a8"
  }
    if open_id:
        user_oauth["open_id"] = open_id
        user_info["open_id"] = open_id
        user_info["nickname"] = "user" + open_id[:5]
    else:
        open_id = user_oauth["open_id"]
    db_manager = UserOAuthRepository()
    res = db_manager.save_login_data(user_oauth2)
    res = db_manager.select_nearest_expire(UserOAuthToken.expires_at)
    
    logger.debug(res)
    # db_user_info = UserInfoRepository()
    # res = db_user_info.add_or_update(user_info)
    # logger.debug(db_manager.get_by_open_id(open_id))
    return user_oauth["open_id"]

if __name__ == "__main__":
    test_add()