Spaces:
Paused
Paused
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 |