database.py

from datetime import datetime
import os
from flask import g
import psycopg2

DB_NAME = os.getenv("DB_NAME")
DB_USER = os.getenv("DB_USER")
DB_PASSWORD = os.getenv("DB_PASSWORD")
DB_HOST = os.getenv("DB_HOST")
DB_PORT = os.getenv("DB_PORT")

# test data variables
CUR_MONTH = datetime.now().month



# init_db run through a cli command to initialize the database tables
def init_db():
    conn = db_connection()
    cur = conn.cursor()
    
    # pgcrypto for gen_random_uuid()
    cur.execute("CREATE EXTENSION IF NOT EXISTS pgcrypto;")
    
    # users table
    cur.execute("""
    CREATE TABLE IF NOT EXISTS users (
        id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
        email text UNIQUE NOT NULL,
        password_hash TEXT NOT NULL,
        created_at timestamptz DEFAULT now()
    );
    """)

    # categories table
    cur.execute("""
    CREATE TABLE IF NOT EXISTS categories (
        id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
        user_id uuid NOT NULL,
        name text NOT NULL,
        color varchar(7) NOT NULL,
        created_at timestamptz DEFAULT now(),
        updated_at timestamptz DEFAULT now(),
        CONSTRAINT fk_user FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
    );
    """)
    
    # events table
    cur.execute("""
    CREATE TABLE IF NOT EXISTS events (
        id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
        user_id uuid NOT NULL,
        title text NOT NULL,
        description text,
        all_day boolean NOT NULL,
        start_time timestamptz,
        end_time timestamptz,
        created_at timestamptz DEFAULT now(),
        updated_at timestamptz DEFAULT now(),
        category_id uuid,
        color_override varchar(7),
        location text,
        CONSTRAINT fk_user FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE,
        CONSTRAINT fk_category FOREIGN KEY (category_id) REFERENCES categories(id) ON DELETE SET NULL
    );
    """)
    
    conn.commit()
    cur.close()
    conn.close()
    
# TODO: create function to create various test events for the current month
#def create_test_data():
    

# TODO: consider using connection pool in the future
def db_connection():
    return psycopg2.connect(
        dbname = DB_NAME,
        user = DB_USER,
        password = DB_PASSWORD,
        host = DB_HOST,
        port = DB_PORT
)   

def get_db():
    if 'db' not in g:
        g.db = db_connection()
    return g.db


def close_db(exception):
    db = g.pop('db', None)

    if db is not None:
        db.close()