drinks-manager/lib/upgrade-db.py

#!/usr/bin/env python3

import os, sys

from pathlib import Path

from psycopg2 import connect
from psycopg2._psycopg import cursor as _cursor
from psycopg2._psycopg import connection as _connection
from psycopg2 import Error
from psycopg2 import IntegrityError
from psycopg2 import errorcodes


# setup or upgrade the database

def log(s, error=False):
    if error:
        print(f"{s}", file=sys.stderr)
    else:
        print(f"{s}", file=sys.stdout)


def execute_sql_statement(cursor:_cursor, connection:_connection, sql_statement):
    try:
        cursor.execute(sql_statement)
        connection.commit()
    except IntegrityError as ie:
        if ie.pgcode == errorcodes.UNIQUE_VIOLATION:
            log("Skipping one row that already exists.")
            connection.rollback()
        else:
            log(f"An integrity error occured:\n{ie}\nRolling back...", error=True)
            connection.rollback()
    except Error as e:
        log(f"An SQL statement failed while upgrading the database at {os.environ['PGDB_HOST']}:\n{e}", error=True)
        connection.rollback()


if __name__ == "__main__":
    exit_code = 1
    conn = connect(
        user = os.environ["PGDB_USER"],
        password = os.environ["PGDB_PASSWORD"],
        host = os.environ["PGDB_HOST"],
        port = os.environ["PGDB_PORT"],
        database = os.environ["PGDB_DB"]
    )
    cur = conn.cursor()
    try:
        log("\nSetting up/upgrading database...")
        # # # # #
        log("Not deleting register_balance. You can delete it via the Admin Panel (Globals -> register_balance), as it is no more used.")
        execute_sql_statement(cur, conn, """
            insert into app_global
            values ('global_message', 'Here you can set a global message that will be shown to every user.', 0.0, '');
        """)
        execute_sql_statement(cur, conn, """
            insert into app_global
            values ('admin_info', 'Here you can set am infotext that will be displayed on the admin panel.', 0.0, '');
        """)
        execute_sql_statement(cur, conn, """
            create or replace view app_userdeposits_view as
                select * from app_registertransaction
                where is_user_deposit = true;
        """)
        # # # # #
        # set app_version in file and database
        # database
        try:
            cur.execute("""
                select value from application_info
                where key = 'app_version';
            """)
            result = cur.fetchone()
            if result == None:
                cur.execute(f"""
                    insert into application_info values ('app_version', '{os.environ['APP_VERSION']}');
                """)
                conn.commit()
            else:
                cur.execute(f"""
                    update application_info set value = '{os.environ['APP_VERSION']}' where key = 'app_version';
                """)
                conn.commit()
        except Error as err:
            if err.pgcode == errorcodes.UNDEFINED_TABLE:
                try:
                    conn.rollback()
                    cur.execute("""
                        create table application_info (
                            key varchar(32) primary key,
                            value text
                        );
                    """)
                    cur.execute(f"""
                        insert into application_info values ('app_version', '{os.environ['APP_VERSION']}');
                    """)
                    conn.commit()
                except Error as err2:
                    log(f"An error occurred while setting app_version in table application_info: {err}", error=True)
                    exit_code = 1
            else:
                log(f"An error occurred while setting app_version in table application_info: {err}", error=True)
                exit_code = 1
        # file
        Path("./config/db_app_version.txt").write_text(os.environ["APP_VERSION"])
        # done
        exit_code = 0
        log("done with db setup/upgrade.")
    except (Error, Exception) as err:
        log(f"An error occured while upgrading the database at {os.environ['PGDB_HOST']}:\n{err}", error=True)
        exit_code = 1
    finally:
        cur.close()
        conn.close()
        exit(exit_code)