mfinder/db/broadcast_sql.py

import threading
from sqlalchemy import create_engine, Column, TEXT, BigInteger
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy.pool import StaticPool
from sqlalchemy.exc import IntegrityError
from mfinder import DB_URL

BASE = declarative_base()

class Broadcast(BASE):
    __tablename__ = "broadcast"
    user_id = Column(BigInteger, primary_key=True)
    user_name = Column(TEXT)

    def __init__(self, user_id, user_name):
        self.user_id = user_id
        self.user_name = user_name

def start() -> scoped_session:
    connect_args = {}
    if DB_URL and DB_URL.startswith(("postgres://", "postgresql://")):
        connect_args["sslmode"] = "require"
    engine = create_engine(
        DB_URL,
        connect_args=connect_args,
        client_encoding="utf8",
        poolclass=StaticPool,
        pool_pre_ping=True
    )
    BASE.metadata.bind = engine
    BASE.metadata.create_all(engine)
    return scoped_session(sessionmaker(bind=engine, autoflush=False))

SESSION = start()
INSERTION_LOCK = threading.RLock()

async def add_user(user_id, user_name):
    with INSERTION_LOCK:
        try:
            usr = SESSION.query(Broadcast).filter_by(user_id=user_id).first()
            if usr:
                if usr.user_name != user_name:
                    usr.user_name = user_name
                    SESSION.commit()
                return
            
            usr = Broadcast(user_id=user_id, user_name=user_name)
            SESSION.add(usr)
            SESSION.commit()
        except Exception as e:
            SESSION.rollback()
            if isinstance(e, IntegrityError) or "UniqueViolation" in str(e):
                pass
            else:
                raise e
        finally:
            SESSION.close()

async def is_user(user_id):
    with INSERTION_LOCK:
        try:
            usr = SESSION.query(Broadcast).filter_by(user_id=user_id).first()
            return usr.user_id if usr else False
        except Exception as e:
            SESSION.rollback()
            raise e
        finally:
            SESSION.close()

async def query_msg():
    try:
        query = SESSION.query(Broadcast.user_id).order_by(Broadcast.user_id)
        return query.all()
    except Exception as e:
        SESSION.rollback()
        raise e
    finally:
        SESSION.close()

async def del_user(user_id):
    with INSERTION_LOCK:
        try:
            usr = SESSION.query(Broadcast).filter_by(user_id=user_id).first()
            if usr:
                SESSION.delete(usr)
                SESSION.commit()
        except Exception as e:
            SESSION.rollback()
            raise e
        finally:
            SESSION.close()