database.py

import asyncio
from typing import Any
from sqlalchemy import JSON
from sqlalchemy import select,update,delete,func,insert
from sqlalchemy.orm import Mapped,DeclarativeBase,mapped_column
from sqlalchemy.ext.asyncio import AsyncAttrs,async_sessionmaker,create_async_engine
from pathlib import Path
import logging

log=logging.getLogger()

engine=create_async_engine(url='sqlite+aiosqlite:///Storage/db.sqlite3', echo=True)
async_session=async_sessionmaker(bind=engine,expire_on_commit=False)

class Base(AsyncAttrs,DeclarativeBase):
    type_annotation_map = {
        dict[str, int]: JSON,
        dict: JSON 
    }

class Client(Base):
    __tablename__='Clients'
    id:Mapped[int] = mapped_column(primary_key=True)
    chat_id:Mapped[int] = mapped_column(primary_key=True)
    info_discount:Mapped[dict[str, int]] = mapped_column()
    is_admin:Mapped[bool]=mapped_column()
    username:Mapped[str]=mapped_column()


async def init_db():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


async def add_user(id:int,chat_id:int,username:str)->bool:
    async with async_session() as s:
        user=await s.scalar(select(Client).where(Client.id==id ,Client.chat_id==chat_id))
        if user:
            log.info(f"User:{id}, is already connected to bot")
            return False
        newUser=Client(id=id,chat_id=chat_id,info_discount={},is_admin=False,username=username)
        s.add(newUser)
        await s.commit()
        await s.refresh(newUser)
        await s.close()
        log.info(f"User:{id}, connected to bot")
        return True
    
async def get_user(id:int,get_all=True,get_admin=False)->dict[str,int]:
    async with async_session() as s:
        if get_all==False:
            if get_admin==False:
                user=await s.scalar(select(Client).where(Client.id==id))
            else:
                user=await s.scalar(select(Client).where(Client.id==id,Client.is_admin==True))
            id_return=user.id
            discount_return=user.info_discount
            return {id_return:discount_return}
        elif get_all==True:
            if get_admin==False:
                user=await s.scalars(select(Client))
            else:
                user=await s.scalars(select(Client).where(Client.is_admin==True))
            if user:
                for row in user:
                    id_return=row.id
                    discount_return=row.info_discount
                    return {id_return:discount_return}

async def get_users_is_admin(is_admin=False)->dict:
    async with async_session() as s:
        if is_admin==False:
            users=await s.scalars(select(Client))
        elif is_admin==True:
            users=await s.scalars(select(Client).where(Client.is_admin==True))
        if users:
            dict_to_return={}
            for row in users:
                dict_to_return[row.id] = row.chat_id
            log.info(f"Got users with is_admin column: {is_admin}")
            return dict_to_return
        
async def get_usernames_list(get_admins=False)->list:
    async with async_session() as s:
        if get_admins==False:
            users=await s.scalars(select(Client).where(Client.is_admin==False))
        elif get_admins==True:
            users=await s.scalars(select(Client).where(Client.is_admin==True))
        print(users)
        list_return=[]
        for row in users:
            list_return.append(row.username)
        log.info(f"Got usernames list with setting of get_admins: {get_admins}")
        return list_return

async def get_user_by_username(username:str)->int:
    async with async_session() as s:
        user=await s.scalar(select(Client).where(Client.username==username))
        log.info(f"Got user id: {user.id} from username: {username}")
        return user.id    

async def change_admin_status(id:int,make_admin:bool):
    async with async_session() as s:
        await s.execute(update(Client).where(Client.id==id).values(is_admin=make_admin))
        await s.commit()
        await s.close()
        log.info(f"Changed admin status to: {id} to status: {make_admin}")

async def discount_get(id:int)->dict:
    async with async_session() as s:
        user=await s.scalar(select(Client).where(Client.id==id))
        log.info(f"Got discount from user with id: {id}")
        return user.info_discount
    
async def discount_use(id:int,use_name:str,cost:float)->float:
    discounts = await discount_get(id)
    cost=cost*discounts[use_name]
    log.info(f"Used discount with name:{use_name} to user with id:{id}")
    return cost

async def add_discount(id:int,discount:dict):
    async with async_session() as s:
        await s.execute(update(Client).where(Client.id==id).values(info_discount=discount.values))
        await s.commit()
        await s.close()
        log.info(f"Added discount to user with id: {id}")

async def remove_discount(id:int,discount_to_remove:str):
    async with async_session() as s:
        user = await s.scalar(select(Client).where(Client.id==id))
        user.info_discount.pop(discount_to_remove)
        await s.execute(update(Client).where(Client.id==id).values(info_discount=user.info_discount))
        await s.commit()
        await s.close()
        log.info(f"Removed discount with name: {discount_to_remove} from user with id: {id}")