feat: implement connection pooling and automatic schema migration for performance
This commit is contained in:
Cipher Vance
parent
45a1b2f234
commit
66011bcd0f
1 changed files with 201 additions and 45 deletions
246
database.py
246
database.py
|
|
@ -1,66 +1,222 @@
|
||||||
import os
|
import os
|
||||||
import psycopg2
|
import psycopg2
|
||||||
from psycopg2 import IntegrityError
|
from psycopg2 import pool, IntegrityError
|
||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
|
import logging
|
||||||
|
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
|
|
||||||
|
# Global connection pool
|
||||||
|
_connection_pool = None
|
||||||
|
|
||||||
|
def get_connection_pool():
|
||||||
|
"""Initialize and return the connection pool"""
|
||||||
|
global _connection_pool
|
||||||
|
if _connection_pool is None:
|
||||||
|
try:
|
||||||
|
_connection_pool = psycopg2.pool.ThreadedConnectionPool(
|
||||||
|
minconn=2,
|
||||||
|
maxconn=20,
|
||||||
|
host=os.getenv("PG_HOST"),
|
||||||
|
port=os.getenv("PG_PORT", 5432),
|
||||||
|
dbname=os.getenv("PG_DATABASE"),
|
||||||
|
user=os.getenv("PG_USER"),
|
||||||
|
password=os.getenv("PG_PASSWORD"),
|
||||||
|
connect_timeout=5
|
||||||
|
)
|
||||||
|
logging.info("Database connection pool created successfully")
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(f"Error creating connection pool: {e}")
|
||||||
|
raise
|
||||||
|
return _connection_pool
|
||||||
|
|
||||||
def get_connection():
|
def get_connection():
|
||||||
"""Return a database connection."""
|
"""Get a connection from the pool"""
|
||||||
return psycopg2.connect(
|
try:
|
||||||
host=os.getenv("PG_HOST"),
|
pool = get_connection_pool()
|
||||||
port=os.getenv("PG_PORT"),
|
conn = pool.getconn()
|
||||||
dbname=os.getenv("PG_DATABASE"),
|
if conn.closed:
|
||||||
user=os.getenv("PG_USER"),
|
# Connection is closed, remove it and get a new one
|
||||||
password=os.getenv("PG_PASSWORD"),
|
pool.putconn(conn, close=True)
|
||||||
connect_timeout=10
|
conn = pool.getconn()
|
||||||
)
|
return conn
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(f"Error getting connection from pool: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
def return_connection(conn):
|
||||||
|
"""Return a connection to the pool"""
|
||||||
|
try:
|
||||||
|
pool = get_connection_pool()
|
||||||
|
pool.putconn(conn)
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(f"Error returning connection to pool: {e}")
|
||||||
|
|
||||||
|
def close_all_connections():
|
||||||
|
"""Close all connections in the pool"""
|
||||||
|
global _connection_pool
|
||||||
|
if _connection_pool:
|
||||||
|
_connection_pool.closeall()
|
||||||
|
_connection_pool = None
|
||||||
|
logging.info("All database connections closed")
|
||||||
|
|
||||||
|
def column_exists(cursor, table_name, column_name):
|
||||||
|
"""Check if a column exists in a table"""
|
||||||
|
cursor.execute("""
|
||||||
|
SELECT EXISTS (
|
||||||
|
SELECT 1
|
||||||
|
FROM information_schema.columns
|
||||||
|
WHERE table_name = %s AND column_name = %s
|
||||||
|
)
|
||||||
|
""", (table_name, column_name))
|
||||||
|
return cursor.fetchone()[0]
|
||||||
|
|
||||||
|
def index_exists(cursor, index_name):
|
||||||
|
"""Check if an index exists"""
|
||||||
|
cursor.execute("""
|
||||||
|
SELECT EXISTS (
|
||||||
|
SELECT 1 FROM pg_class c
|
||||||
|
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||||
|
WHERE c.relname = %s AND n.nspname = 'public'
|
||||||
|
)
|
||||||
|
""", (index_name,))
|
||||||
|
return cursor.fetchone()[0]
|
||||||
|
|
||||||
def init_db():
|
def init_db():
|
||||||
conn = get_connection()
|
"""Initialize database tables and indexes"""
|
||||||
cursor = conn.cursor()
|
conn = None
|
||||||
cursor.execute("""
|
try:
|
||||||
CREATE TABLE IF NOT EXISTS subscribers (
|
conn = get_connection()
|
||||||
id SERIAL PRIMARY KEY,
|
cursor = conn.cursor()
|
||||||
email TEXT UNIQUE NOT NULL
|
|
||||||
)
|
|
||||||
""")
|
|
||||||
|
|
||||||
cursor.execute("""
|
|
||||||
CREATE TABLE IF NOT EXISTS newsletters(
|
|
||||||
id SERIAL PRIMARY KEY,
|
|
||||||
subject TEXT NOT NULL,
|
|
||||||
body TEXT NOT NULL,
|
|
||||||
sent_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
||||||
)
|
|
||||||
""")
|
|
||||||
|
|
||||||
conn.commit()
|
# Create subscribers table
|
||||||
cursor.close()
|
cursor.execute("""
|
||||||
conn.close()
|
CREATE TABLE IF NOT EXISTS subscribers (
|
||||||
|
id SERIAL PRIMARY KEY,
|
||||||
|
email TEXT UNIQUE NOT NULL
|
||||||
|
)
|
||||||
|
""")
|
||||||
|
|
||||||
|
# Add created_at column if it doesn't exist
|
||||||
|
if not column_exists(cursor, 'subscribers', 'created_at'):
|
||||||
|
cursor.execute("""
|
||||||
|
ALTER TABLE subscribers
|
||||||
|
ADD COLUMN created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||||
|
""")
|
||||||
|
logging.info("Added created_at column to subscribers table")
|
||||||
|
|
||||||
|
# Create newsletters table
|
||||||
|
cursor.execute("""
|
||||||
|
CREATE TABLE IF NOT EXISTS newsletters(
|
||||||
|
id SERIAL PRIMARY KEY,
|
||||||
|
subject TEXT NOT NULL,
|
||||||
|
body TEXT NOT NULL,
|
||||||
|
sent_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||||
|
)
|
||||||
|
""")
|
||||||
|
|
||||||
|
# Create indexes only if they don't exist
|
||||||
|
if not index_exists(cursor, 'idx_newsletters_sent_at'):
|
||||||
|
cursor.execute("CREATE INDEX idx_newsletters_sent_at ON newsletters(sent_at DESC)")
|
||||||
|
logging.info("Created index idx_newsletters_sent_at")
|
||||||
|
|
||||||
|
if not index_exists(cursor, 'idx_subscribers_email'):
|
||||||
|
cursor.execute("CREATE INDEX idx_subscribers_email ON subscribers(email)")
|
||||||
|
logging.info("Created index idx_subscribers_email")
|
||||||
|
|
||||||
|
if not index_exists(cursor, 'idx_subscribers_created_at'):
|
||||||
|
cursor.execute("CREATE INDEX idx_subscribers_created_at ON subscribers(created_at DESC)")
|
||||||
|
logging.info("Created index idx_subscribers_created_at")
|
||||||
|
|
||||||
|
conn.commit()
|
||||||
|
cursor.close()
|
||||||
|
logging.info("Database tables and indexes initialized successfully")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(f"Error initializing database: {e}")
|
||||||
|
if conn:
|
||||||
|
conn.rollback()
|
||||||
|
raise
|
||||||
|
finally:
|
||||||
|
if conn:
|
||||||
|
return_connection(conn)
|
||||||
|
|
||||||
def add_email(email):
|
def add_email(email):
|
||||||
|
"""Add email to subscribers with connection pooling"""
|
||||||
|
conn = None
|
||||||
try:
|
try:
|
||||||
with get_connection() as conn:
|
conn = get_connection()
|
||||||
with conn.cursor() as 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()
|
||||||
|
logging.info(f"Email added successfully: {email}")
|
||||||
return True
|
return True
|
||||||
|
|
||||||
except IntegrityError:
|
except IntegrityError:
|
||||||
|
# Email already exists
|
||||||
|
if conn:
|
||||||
|
conn.rollback()
|
||||||
|
logging.info(f"Email already exists: {email}")
|
||||||
return False
|
return False
|
||||||
except psycopg2.OperationalError as e:
|
|
||||||
print(f"Error: {e}")
|
except Exception as e:
|
||||||
|
if conn:
|
||||||
|
conn.rollback()
|
||||||
|
logging.error(f"Error adding email {email}: {e}")
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
finally:
|
||||||
|
if conn:
|
||||||
|
return_connection(conn)
|
||||||
|
|
||||||
def remove_email(email):
|
def remove_email(email):
|
||||||
|
"""Remove email from subscribers with connection pooling"""
|
||||||
|
conn = None
|
||||||
try:
|
try:
|
||||||
with get_connection() as conn:
|
conn = get_connection()
|
||||||
with conn.cursor() as cursor:
|
cursor = conn.cursor()
|
||||||
cursor.execute("DELETE FROM subscribers WHERE email = %s", (email,))
|
cursor.execute("DELETE FROM subscribers WHERE email = %s", (email,))
|
||||||
conn.commit()
|
conn.commit()
|
||||||
if cursor.rowcount > 0:
|
rows_affected = cursor.rowcount
|
||||||
return True
|
cursor.close()
|
||||||
return False
|
|
||||||
|
if rows_affected > 0:
|
||||||
|
logging.info(f"Email removed successfully: {email}")
|
||||||
|
return True
|
||||||
|
else:
|
||||||
|
logging.info(f"Email not found for removal: {email}")
|
||||||
|
return False
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"Error removing email: {e}")
|
if conn:
|
||||||
return False
|
conn.rollback()
|
||||||
|
logging.error(f"Error removing email {email}: {e}")
|
||||||
|
return False
|
||||||
|
|
||||||
|
finally:
|
||||||
|
if conn:
|
||||||
|
return_connection(conn)
|
||||||
|
|
||||||
|
def get_subscriber_count():
|
||||||
|
"""Get total number of subscribers"""
|
||||||
|
conn = None
|
||||||
|
try:
|
||||||
|
conn = get_connection()
|
||||||
|
cursor = conn.cursor()
|
||||||
|
cursor.execute("SELECT COUNT(*) FROM subscribers")
|
||||||
|
count = cursor.fetchone()[0]
|
||||||
|
cursor.close()
|
||||||
|
return count
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(f"Error getting subscriber count: {e}")
|
||||||
|
return 0
|
||||||
|
|
||||||
|
finally:
|
||||||
|
if conn:
|
||||||
|
return_connection(conn)
|
||||||
|
|
||||||
|
# Cleanup function for graceful shutdown
|
||||||
|
import atexit
|
||||||
|
atexit.register(close_all_connections)
|
||||||
Loading…
Add table
Add a link
Reference in a new issue