119 lines
4.6 KiB
Python
119 lines
4.6 KiB
Python
import psycopg2
|
|
from psycopg2 import pool
|
|
from contextlib import contextmanager
|
|
from typing import Generator
|
|
from app.config import settings
|
|
import json
|
|
|
|
# Connection pool for better performance
|
|
connection_pool = pool.SimpleConnectionPool(1, 20, settings.database_url)
|
|
|
|
@contextmanager
|
|
def get_db_connection():
|
|
"""Get a database connection from the pool"""
|
|
conn = connection_pool.getconn()
|
|
try:
|
|
yield conn
|
|
conn.commit()
|
|
except Exception as e:
|
|
conn.rollback()
|
|
raise e
|
|
finally:
|
|
connection_pool.putconn(conn)
|
|
|
|
def init_db():
|
|
"""Initialize database tables"""
|
|
with get_db_connection() as conn:
|
|
cur = conn.cursor()
|
|
|
|
# Create tables
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id SERIAL PRIMARY KEY,
|
|
email VARCHAR(255) UNIQUE NOT NULL,
|
|
hashed_password VARCHAR(255) NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
""")
|
|
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS profiles (
|
|
id SERIAL PRIMARY KEY,
|
|
user_id INTEGER NOT NULL UNIQUE,
|
|
display_name VARCHAR(255) NOT NULL,
|
|
age INTEGER NOT NULL,
|
|
gender VARCHAR(50) NOT NULL,
|
|
location VARCHAR(255) NOT NULL,
|
|
bio TEXT,
|
|
interests JSONB DEFAULT '[]',
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
|
|
);
|
|
""")
|
|
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS photos (
|
|
id SERIAL PRIMARY KEY,
|
|
profile_id INTEGER NOT NULL,
|
|
file_path VARCHAR(255) NOT NULL,
|
|
display_order INTEGER NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (profile_id) REFERENCES profiles(id) ON DELETE CASCADE
|
|
);
|
|
""")
|
|
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS likes (
|
|
id SERIAL PRIMARY KEY,
|
|
liker_id INTEGER NOT NULL,
|
|
liked_id INTEGER NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE(liker_id, liked_id),
|
|
FOREIGN KEY (liker_id) REFERENCES users(id) ON DELETE CASCADE,
|
|
FOREIGN KEY (liked_id) REFERENCES users(id) ON DELETE CASCADE
|
|
);
|
|
""")
|
|
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS conversations (
|
|
id SERIAL PRIMARY KEY,
|
|
user_id_1 INTEGER NOT NULL,
|
|
user_id_2 INTEGER NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE(user_id_1, user_id_2),
|
|
FOREIGN KEY (user_id_1) REFERENCES users(id) ON DELETE CASCADE,
|
|
FOREIGN KEY (user_id_2) REFERENCES users(id) ON DELETE CASCADE
|
|
);
|
|
""")
|
|
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
id SERIAL PRIMARY KEY,
|
|
conversation_id INTEGER NOT NULL,
|
|
sender_id INTEGER NOT NULL,
|
|
content TEXT NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE,
|
|
FOREIGN KEY (sender_id) REFERENCES users(id) ON DELETE CASCADE
|
|
);
|
|
""")
|
|
|
|
# Create indexes for common queries
|
|
cur.execute("CREATE INDEX IF NOT EXISTS idx_profiles_user_id ON profiles(user_id);")
|
|
cur.execute("CREATE INDEX IF NOT EXISTS idx_photos_profile_id ON photos(profile_id);")
|
|
cur.execute("CREATE INDEX IF NOT EXISTS idx_likes_liker_id ON likes(liker_id);")
|
|
cur.execute("CREATE INDEX IF NOT EXISTS idx_likes_liked_id ON likes(liked_id);")
|
|
cur.execute("CREATE INDEX IF NOT EXISTS idx_conversations_users ON conversations(user_id_1, user_id_2);")
|
|
cur.execute("CREATE INDEX IF NOT EXISTS idx_messages_conversation_id ON messages(conversation_id);")
|
|
cur.execute("CREATE INDEX IF NOT EXISTS idx_messages_created_at ON messages(created_at);")
|
|
|
|
conn.commit()
|
|
|
|
def close_db():
|
|
"""Close all database connections"""
|
|
if connection_pool:
|
|
connection_pool.closeall()
|