1+ import os
2+ import psycopg2
3+ from psycopg2 .extras import RealDictCursor
4+
5+ class UserDB :
6+ def __init__ (self , connection_string = None ):
7+ if connection_string is None :
8+ connection_string = os .getenv ('DATABASE_URL' , 'postgresql://postgres:postgres@localhost:5432/testdb' )
9+ self .connection_string = connection_string
10+
11+ def connect (self ):
12+ return psycopg2 .connect (self .connection_string )
13+
14+ def setup_database (self ):
15+ with self .connect () as conn :
16+ with conn .cursor () as cur :
17+ cur .execute ("""
18+ CREATE TABLE IF NOT EXISTS users (
19+ id SERIAL PRIMARY KEY,
20+ username VARCHAR(100) UNIQUE NOT NULL,
21+ email VARCHAR(255) UNIQUE NOT NULL,
22+ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
23+ )
24+ """ )
25+ conn .commit ()
26+
27+ def add_user (self , username , email ):
28+ with self .connect () as conn :
29+ with conn .cursor (cursor_factory = RealDictCursor ) as cur :
30+ cur .execute ("""
31+ INSERT INTO users (username, email)
32+ VALUES (%s, %s)
33+ RETURNING id, username, email, created_at
34+ """ , (username , email ))
35+ conn .commit ()
36+ return cur .fetchone ()
37+
38+ def get_user (self , username ):
39+ with self .connect () as conn :
40+ with conn .cursor (cursor_factory = RealDictCursor ) as cur :
41+ cur .execute ("SELECT * FROM users WHERE username = %s" , (username ,))
42+ return cur .fetchone ()
43+
44+ def delete_user (self , username ):
45+ with self .connect () as conn :
46+ with conn .cursor () as cur :
47+ cur .execute ("DELETE FROM users WHERE username = %s" , (username ,))
48+ conn .commit ()
49+ return cur .rowcount > 0
0 commit comments