bsition/backend/postgres.py

from os import getenv

import psycopg2
from dotenv import load_dotenv
from psycopg2 import sql


def get_connection():
    return psycopg2.connect(
        host=getenv("POSTGRES_HOST"),
        database=getenv("POSTGRES_DB"),
        user=getenv("POSTGRES_USER"),
        password=getenv("POSTGRES_PASSWORD"),
    )


def create_table(name):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        sql.SQL("CREATE TABLE {table} (row_number SERIAL PRIMARY KEY)").format(
            table=sql.Identifier(name)
        )
    )
    conn.commit()


def add_column(name, column, type):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        sql.SQL("ALTER TABLE {table} ADD {column}" + type).format(
            table=sql.Identifier(name), column=sql.Identifier(column)
        )
    )
    conn.commit()


def insert_columns(name, data):
    conn = get_connection()
    cur = conn.cursor()
    str = "(" + "DEFAULT, %s," * (len(data) - 1) + "%s" + ")"  # TODO: change.
    cur.execute(
        sql.SQL("INSERT INTO {table} VALUES" + str).format(table=sql.Identifier(name)),
        data,
    )
    conn.commit()


def edit_columns(name, columns, data, id):
    conn = get_connection()
    cur = conn.cursor()
    i = 0
    for column in columns:
        cur.execute(
            sql.SQL("UPDATE {table} SET {col} = %s WHERE row_number = " + id).format(
                table=sql.Identifier(name), col=sql.Identifier(column)
            ),
            [data[i]],
        )
        i += 1
    conn.commit()


def remove_column(name, column):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        sql.SQL("ALTER TABLE {table} DROP COLUMN {column}").format(
            table=sql.Identifier(name), column=sql.Identifier(column)
        )
    )
    conn.commit()


def create_sort(name):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        sql.SQL(
            "CREATE TABLE {table} (property TEXT, _order CHAR(3), priority int)"
        ).format(table=sql.Identifier(name + "_sort"))
    )
    conn.commit()


def add_sort(name, property, order, priority):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        sql.SQL("INSERT INTO {table} VALUES (%s, %s, %s)").format(
            table=sql.Identifier(name + "_sort")
        ),
        (property, order, priority),
    )
    conn.commit()


def sort(name):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        sql.SQL("SELECT * FROM {table} ORDER BY priority").format(
            table=sql.Identifier(name + "_sort")
        ),
    )
    order_clause = "ORDER BY "
    i = 0
    for sort in cur:
        if i > 0:
            order_clause += ", "
        order_clause += sort[0] + " " + sort[1]
        i += 1
    cur.execute(
        sql.SQL("SELECT * FROM {table} " + order_clause).format(
            table=sql.Identifier(name)
        ),
    )
    return cur


def add_function():
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        """
            CREATE OR REPLACE FUNCTION trigger_function()
                RETURNS TRIGGER
                LANGUAGE PLPGSQL
            AS $$
            DECLARE
                name text := TG_ARGV[0]::text;
            BEGIN
                IF NEW.property NOT IN (
                    SELECT column_name
                    FROM INFORMATION_SCHEMA.COLUMNS
                    WHERE TABLE_NAME = name)
                THEN
                    RAISE EXCEPTION 'ERROR %', NEW.property;

                END IF;

                RETURN NEW;
            END;
            $$;
        """
    )
    conn.commit()


def add_filter_trigger(name):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        sql.SQL(
            """
            CREATE TRIGGER {filter}
                BEFORE INSERT OR UPDATE
                ON {filter}
                FOR EACH ROW
                    EXECUTE PROCEDURE trigger_function({table});
            """
        ).format(table=sql.Identifier(name), filter=sql.Identifier(name + "_filter"))
    )
    conn.commit()


def create_filter(name):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        sql.SQL(
            """
        CREATE TABLE {table} (
            property TEXT,
            value TEXT,
            function TEXT CHECK (function IN ('c', 'e', 'n'))
        )
        """
        ).format(table=sql.Identifier(name + "_filter"))
    )
    conn.commit()


def add_filter(name, property, value, function):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        sql.SQL("INSERT INTO {table} VALUES (%s, %s, %s)").format(
            table=sql.Identifier(name + "_filter")
        ),
        (property, value, function),
    )
    conn.commit()


def filter(name):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        sql.SQL("SELECT * FROM {table}").format(table=sql.Identifier(name + "_filter")),
    )
    filter_clause = "WHERE "
    i = 0
    for sort in cur:
        if i > 0:
            filter_clause += " AND "
        filter_clause += sort[0]
        match sort[2]:
            case "e":
                filter_clause += " = '" + sort[1] + "'"
            case "ne":
                filter_clause += " <> '" + sort[1] + "'"
            case "le":
                filter_clause += " <= " + sort[1]
            case "ge":
                filter_clause += " >= " + sort[1]
            case "l":
                filter_clause += " < " + sort[1]
            case "g":
                filter_clause += " > " + sort[1]
            case "c":
                filter_clause += " ILIKE '%" + sort[1] + "'"
            case "_":
                raise "Invalid filter function"
        i += 1
    cur.execute(
        sql.SQL("SELECT * FROM {table} " + filter_clause).format(
            table=sql.Identifier(name)
        ),
    )
    return cur