diff --git a/database.py b/database.py index b7f1494..0ee6567 100644 --- a/database.py +++ b/database.py @@ -1,45 +1,56 @@ -import sqlite3 import os +import psycopg2 +from psycopg2 import IntegrityError from dotenv import load_dotenv load_dotenv() -DATABASE_URL = os.getenv("DATABASE_FILE") +def get_connection(): + """Return a database connection.""" + return psycopg2.connect( + host=os.getenv("PG_HOST"), + port=os.getenv("PG_PORT"), + dbname=os.getenv("PG_DATABASE"), + user=os.getenv("PG_USER"), + password=os.getenv("PG_PASSWORD"), + connect_timeout=10 + ) def init_db(): - conn = sqlite3.connect(DATABASE_URL) + conn = get_connection() cursor = conn.cursor() cursor.execute(""" CREATE TABLE IF NOT EXISTS subscribers ( - id INTEGER PRIMARY KEY AUTOINCREMENT, + id SERIAL PRIMARY KEY, email TEXT UNIQUE NOT NULL ) """) conn.commit() + cursor.close() conn.close() def add_email(email): try: - with sqlite3.connect(DATABASE_URL, timeout=10) as conn: - cursor = conn.cursor() - cursor.execute("INSERT INTO subscribers (email) VALUES (?)", (email,)) - conn.commit() + with get_connection() as conn: + with conn.cursor() as cursor: + cursor.execute("INSERT INTO subscribers (email) VALUES (%s)", (email,)) + conn.commit() return True - except sqlite3.IntegrityError: + except IntegrityError: return False - except sqlite3.OperationalError as e: - print(f"Operational Error: {e}") + except psycopg2.OperationalError as e: + print(f"Error: {e}") return False def remove_email(email): try: - with sqlite3.connect(DATABASE_URL, timeout=10) as conn: - cursor = conn.cursor() - cursor.execute("DELETE FROM subscribers WHERE email = ?", (email,)) - conn.commit() - if cursor.rowcount > 0: - return True - return False + with get_connection() as conn: + with conn.cursor() as cursor: + cursor.execute("DELETE FROM subscribers WHERE email = %s", (email,)) + conn.commit() + if cursor.rowcount > 0: + return True + return False except Exception as e: print(f"Error removing email: {e}") return False \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 5a4a2b9..e4b25c7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ flask -python-dotenv \ No newline at end of file +python-dotenv +psycopg2-binary \ No newline at end of file