db/db_utils.py

import os
import psycopg2
from dotenv import load_dotenv

load_dotenv()


def get_connection():
    DATABASE_URL = os.environ['DATABASE_URL']
    conn = psycopg2.connect(DATABASE_URL, sslmode='require')
    # conn = psycopg2.connect(
    #     dbname="brightly_ai",
    #     user="bw",
    #     password="",
    #     host="localhost",
    #     port="5432"
    # )
    initialize_db(conn)
    return conn

def initialize_db(conn):
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS mappings (
            input_word TEXT PRIMARY KEY,
            cleaned_word TEXT,
            matching_word TEXT,
            dictionary_word TEXT,
            similarity_score REAL,
            confidence_score REAL,
            similar_words TEXT,
            is_food BOOLEAN,
            food_nonfood_score REAL,
            reviewed BOOLEAN DEFAULT FALSE,
            flagged BOOLEAN DEFAULT FALSE,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS dictionary (
            fdc_id INTEGER PRIMARY KEY,
            description TEXT,
            food_category TEXT,
            wweia_category TEXT,
            water_content REAL,
            dry_matter_content REAL,
            leakage REAL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()

def get_mapping_from_db(cursor, cleaned_word):
    cursor.execute('SELECT * FROM mappings WHERE cleaned_word = %s', (cleaned_word,))
    row = cursor.fetchone()
    if row:
        columns = [col[0] for col in cursor.description]
        return dict(zip(columns, row))
    return None

def get_dictionary_data_from_db(cursor, dictionary_word):
    cursor.execute('SELECT * FROM dictionary WHERE description = %s', (dictionary_word,))
    row = cursor.fetchone()
    if row:
        columns = [col[0] for col in cursor.description]
        return dict(zip(columns, row))
    return None

def store_mapping_to_db(cursor, conn, mapping):
    try:
        cursor.execute('''
            INSERT INTO mappings (input_word, cleaned_word, matching_word, dictionary_word, similarity_score, confidence_score, similar_words, is_food, food_nonfood_score)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
        ''', (
            mapping['input_word'], 
            mapping['cleaned_word'], 
            mapping['matching_word'], 
            mapping['dictionary_word'],
            mapping['similarity_score'], 
            mapping['confidence_score'], 
            mapping['similar_words'], 
            mapping['is_food'], 
            mapping['food_nonfood_score']
        ))
        conn.commit()
    except Exception as e:
        print(f" - Error storing mapping to db: {e}")
        conn.rollback()
        return False