(refactor): Enhance database connection and initialization with error
handling and logging Implemented robust logging using the logging module Improved database connection with try...except...finally blocks for better error handling and resource management Enhanced database initialization with proper error handling and logging
This commit is contained in:
138
database.py
138
database.py
@@ -1,44 +1,90 @@
|
|||||||
import os
|
import os
|
||||||
|
import logging
|
||||||
import psycopg2
|
import psycopg2
|
||||||
from psycopg2 import IntegrityError
|
from psycopg2 import IntegrityError
|
||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
from werkzeug.security import generate_password_hash
|
from werkzeug.security import generate_password_hash
|
||||||
|
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
|
|
||||||
|
# Logging setup
|
||||||
|
logging.basicConfig(
|
||||||
|
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
|
||||||
|
)
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
def get_connection():
|
def get_connection():
|
||||||
"""Return a new connection to the PostgreSQL database."""
|
"""Return a new connection to the PostgreSQL database."""
|
||||||
return psycopg2.connect(
|
try:
|
||||||
|
conn = psycopg2.connect(
|
||||||
host=os.getenv("PG_HOST"),
|
host=os.getenv("PG_HOST"),
|
||||||
port=os.getenv("PG_PORT"),
|
port=os.getenv("PG_PORT"),
|
||||||
dbname=os.getenv("PG_DATABASE"),
|
dbname=os.getenv("PG_DATABASE"),
|
||||||
user=os.getenv("PG_USER"),
|
user=os.getenv("PG_USER"),
|
||||||
password=os.getenv("PG_PASSWORD"),
|
password=os.getenv("PG_PASSWORD"),
|
||||||
connect_timeout=10
|
connect_timeout=10,
|
||||||
)
|
)
|
||||||
|
return conn
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Database connection error: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
def init_db():
|
def init_db():
|
||||||
"""Initialize the database tables."""
|
"""Initialize the database tables."""
|
||||||
|
conn = None
|
||||||
|
try:
|
||||||
conn = get_connection()
|
conn = get_connection()
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
|
|
||||||
# Create subscribers table (if not exists)
|
# Create subscribers table (if not exists)
|
||||||
cursor.execute("""
|
cursor.execute(
|
||||||
|
"""
|
||||||
CREATE TABLE IF NOT EXISTS subscribers (
|
CREATE TABLE IF NOT EXISTS subscribers (
|
||||||
id SERIAL PRIMARY KEY,
|
id SERIAL PRIMARY KEY,
|
||||||
email TEXT UNIQUE NOT NULL
|
email TEXT UNIQUE NOT NULL
|
||||||
)
|
)
|
||||||
""")
|
"""
|
||||||
|
)
|
||||||
|
|
||||||
# Create admin_users table (if not exists)
|
# Create admin_users table (if not exists)
|
||||||
cursor.execute("""
|
cursor.execute(
|
||||||
|
"""
|
||||||
CREATE TABLE IF NOT EXISTS admin_users (
|
CREATE TABLE IF NOT EXISTS admin_users (
|
||||||
id SERIAL PRIMARY KEY,
|
id SERIAL PRIMARY KEY,
|
||||||
username TEXT UNIQUE NOT NULL,
|
username TEXT UNIQUE NOT NULL,
|
||||||
password TEXT NOT NULL
|
password TEXT NOT NULL
|
||||||
)
|
)
|
||||||
""")
|
"""
|
||||||
|
)
|
||||||
|
|
||||||
|
# Newsletter storage
|
||||||
|
cursor.execute(
|
||||||
|
"""
|
||||||
|
CREATE TABLE IF NOT EXISTS newsletters (
|
||||||
|
id SERIAL PRIMARY KEY,
|
||||||
|
subject TEXT NOT NULL,
|
||||||
|
body TEXT NOT NULL,
|
||||||
|
sent_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
|
||||||
conn.commit()
|
conn.commit()
|
||||||
|
logger.info("Database initialized successfully.")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Database initialization error: {e}")
|
||||||
|
if conn:
|
||||||
|
conn.rollback() # Rollback if there's an error
|
||||||
|
|
||||||
|
raise
|
||||||
|
finally:
|
||||||
|
if conn:
|
||||||
cursor.close()
|
cursor.close()
|
||||||
conn.close()
|
conn.close()
|
||||||
|
|
||||||
|
|
||||||
def get_all_emails():
|
def get_all_emails():
|
||||||
"""Return a list of all subscriber emails."""
|
"""Return a list of all subscriber emails."""
|
||||||
try:
|
try:
|
||||||
@@ -46,78 +92,112 @@ def get_all_emails():
|
|||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
cursor.execute("SELECT email FROM subscribers")
|
cursor.execute("SELECT email FROM subscribers")
|
||||||
results = cursor.fetchall()
|
results = cursor.fetchall()
|
||||||
|
emails = [row[0] for row in results]
|
||||||
|
logger.debug(f"Retrieved emails: {emails}")
|
||||||
|
return emails
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error retrieving emails: {e}")
|
||||||
|
return []
|
||||||
|
finally:
|
||||||
|
if conn:
|
||||||
cursor.close()
|
cursor.close()
|
||||||
conn.close()
|
conn.close()
|
||||||
return [row[0] for row in results]
|
|
||||||
except Exception as e:
|
|
||||||
print(f"Error retrieving emails: {e}")
|
|
||||||
return []
|
|
||||||
|
|
||||||
def add_email(email):
|
def add_email(email):
|
||||||
"""Insert an email into the subscribers table."""
|
"""Insert an email into the subscribers table."""
|
||||||
|
conn = None
|
||||||
try:
|
try:
|
||||||
conn = get_connection()
|
conn = get_connection()
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
cursor.execute("INSERT INTO subscribers (email) VALUES (%s)", (email,))
|
cursor.execute("INSERT INTO subscribers (email) VALUES (%s)", (email,))
|
||||||
conn.commit()
|
conn.commit()
|
||||||
cursor.close()
|
logger.info(f"Email {email} added successfully.")
|
||||||
conn.close()
|
|
||||||
return True
|
return True
|
||||||
except IntegrityError:
|
except IntegrityError:
|
||||||
|
logger.warning(f"Attempted to add duplicate email: {email}")
|
||||||
return False
|
return False
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"Error adding email: {e}")
|
logger.error(f"Error adding email {email}: {e}")
|
||||||
return False
|
return False
|
||||||
|
finally:
|
||||||
|
if conn:
|
||||||
|
cursor.close()
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
|
||||||
def remove_email(email):
|
def remove_email(email):
|
||||||
"""Remove an email from the subscribers table."""
|
"""Remove an email from the subscribers table."""
|
||||||
|
conn = None
|
||||||
try:
|
try:
|
||||||
conn = get_connection()
|
conn = get_connection()
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
cursor.execute("DELETE FROM subscribers WHERE email = %s", (email,))
|
cursor.execute("DELETE FROM subscribers WHERE email = %s", (email,))
|
||||||
conn.commit()
|
|
||||||
rowcount = cursor.rowcount
|
rowcount = cursor.rowcount
|
||||||
cursor.close()
|
conn.commit()
|
||||||
conn.close()
|
logger.info(f"Email {email} removed successfully.")
|
||||||
return rowcount > 0
|
return rowcount > 0
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"Error removing email: {e}")
|
logger.error(f"Error removing email {email}: {e}")
|
||||||
return False
|
return False
|
||||||
|
finally:
|
||||||
|
if conn:
|
||||||
|
cursor.close()
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
|
||||||
def get_admin(username):
|
def get_admin(username):
|
||||||
"""Retrieve admin credentials for a given username.
|
"""Retrieve admin credentials for a given username.
|
||||||
Returns a tuple (username, password_hash) if found, otherwise None.
|
Returns a tuple (username, password_hash) if found, otherwise None.
|
||||||
"""
|
"""
|
||||||
|
conn = None
|
||||||
try:
|
try:
|
||||||
conn = get_connection()
|
conn = get_connection()
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
cursor.execute("SELECT username, password FROM admin_users WHERE username = %s", (username,))
|
cursor.execute(
|
||||||
|
"SELECT username, password FROM admin_users WHERE username = %s",
|
||||||
|
(username,),
|
||||||
|
)
|
||||||
result = cursor.fetchone()
|
result = cursor.fetchone()
|
||||||
cursor.close()
|
|
||||||
conn.close()
|
|
||||||
return result # (username, password_hash)
|
return result # (username, password_hash)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"Error retrieving admin: {e}")
|
logger.error(f"Error retrieving admin: {e}")
|
||||||
return None
|
return None
|
||||||
|
finally:
|
||||||
|
if conn:
|
||||||
|
cursor.close()
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
|
||||||
def create_default_admin():
|
def create_default_admin():
|
||||||
"""Create a default admin user if one doesn't already exist."""
|
"""Create a default admin user if one doesn't already exist."""
|
||||||
default_username = os.getenv("ADMIN_USERNAME", "admin")
|
default_username = os.getenv("ADMIN_USERNAME", "admin")
|
||||||
default_password = os.getenv("ADMIN_PASSWORD", "changeme")
|
default_password = os.getenv("ADMIN_PASSWORD", "changeme")
|
||||||
hashed = generate_password_hash(default_password, method="pbkdf2:sha256")
|
hashed_password = generate_password_hash(default_password, method="pbkdf2:sha256")
|
||||||
|
conn = None
|
||||||
try:
|
try:
|
||||||
conn = get_connection()
|
conn = get_connection()
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
|
|
||||||
# Check if the admin already exists
|
# Check if the admin already exists
|
||||||
cursor.execute("SELECT id FROM admin_users WHERE username = %s", (default_username,))
|
cursor.execute(
|
||||||
|
"SELECT id FROM admin_users WHERE username = %s", (default_username,)
|
||||||
|
)
|
||||||
if cursor.fetchone() is None:
|
if cursor.fetchone() is None:
|
||||||
cursor.execute("INSERT INTO admin_users (username, password) VALUES (%s, %s)",
|
cursor.execute(
|
||||||
(default_username, hashed))
|
"INSERT INTO admin_users (username, password) VALUES (%s, %s)",
|
||||||
|
(default_username, hashed_password),
|
||||||
|
)
|
||||||
conn.commit()
|
conn.commit()
|
||||||
print("Default admin created successfully")
|
logger.info("Default admin created successfully")
|
||||||
else:
|
else:
|
||||||
print("Default admin already exists")
|
logger.info("Default admin already exists")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error creating default admin: {e}")
|
||||||
|
if conn:
|
||||||
|
conn.rollback()
|
||||||
|
finally:
|
||||||
|
if conn:
|
||||||
cursor.close()
|
cursor.close()
|
||||||
conn.close()
|
conn.close()
|
||||||
except Exception as e:
|
|
||||||
print(f"Error creating default admin: {e}")
|
|
||||||
|
|||||||
Reference in New Issue
Block a user