uwu-me/uwume/lib/databaseMethods.py

from sqlalchemy import create_engine, insert, update, select, MetaData, Table, Column, Integer, String, Boolean
from os import environ

pg_pass_file = open(environ['POSTGRES_PASSWORD_FILE'])
pg_pass = pg_pass_file.readline().strip('\n')
engine = create_engine(
    f"postgresql+psycopg2://{environ['POSTGRES_USER']}:{pg_pass}@{environ['POSTGRES_HOST']}:{environ['POSTGRES_PORT']}/{environ['POSTGRES_DB']}")
conn = engine.connect()
md = MetaData()


def initialize_database():
    users = Table('users', md,
                  Column('id', Integer()),
                  Column('name', String(255)),
                  Column('pass', String(255)),
                  Column('count', Integer()))
    try:
        md.create_all(engine)
    except Exception as e:
        print(f'{e}')


def create_user(idn, name, password, count):
    users = Table('users', md, autoload=True, autoload_with=engine)
    query = insert(users)
    values = [{
        'id': idn,
        'name': name,
        'password': password,
        'count': count
    }]
    output = conn.execute(query, values)


def get_table():
    users = Table('users', md, autoload=True, autoload_with=engine)
    return users.columns.keys()


def get_count(user):
    users = Table('users', md, autoload=True, autoload_with=engine)
    query = select([users.columns.count]).where(users.columns.name == user)
    output = conn.execute(query)
    return output.fetchone()[0]


def update_count(user, count):
    users = Table('users', md, autoload=True, autoload_with=engine)
    query = update(users).values(count=count).where(users.columns.name == user)
    return conn.execute(query)