87 lines
2.5 KiB
Python
87 lines
2.5 KiB
Python
import logging
|
|
import psycopg2
|
|
import psycopg2.extras
|
|
import sqlite3
|
|
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class PostgresCursorWrapper:
|
|
def __init__(self, cursor):
|
|
self.cursor = cursor
|
|
|
|
def execute(self, query, args=None):
|
|
# Convert SQLite ? placeholders to Postgres %s
|
|
if query:
|
|
# Replace ? with %s
|
|
pg_query = query.replace('?', '%s')
|
|
|
|
# Handle BLOB to BYTEA differences if necessary
|
|
# (psycopg2 usually handles bytes objects correctly as BYTEA)
|
|
else:
|
|
pg_query = query
|
|
|
|
try:
|
|
return self.cursor.execute(pg_query, args)
|
|
except Exception as e:
|
|
logger.error(f"SQL Error: {e} | Query: {pg_query} | Args: {args}")
|
|
raise e
|
|
|
|
def fetchall(self):
|
|
return self.cursor.fetchall()
|
|
|
|
def fetchone(self):
|
|
return self.cursor.fetchone()
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self.cursor, name)
|
|
|
|
class PostgresConnectionWrapper:
|
|
def __init__(self, conn):
|
|
self.conn = conn
|
|
# Enable autocommit or manage it manually?
|
|
# sqlite3 connection in app.py context manager does not seem to set autocommit,
|
|
# but the context manager does commit() or rollback().
|
|
|
|
def cursor(self):
|
|
# Use RealDictCursor to mimic sqlite3.Row (dictionary access)
|
|
return PostgresCursorWrapper(self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor))
|
|
|
|
def commit(self):
|
|
return self.conn.commit()
|
|
|
|
def rollback(self):
|
|
return self.conn.rollback()
|
|
|
|
def close(self):
|
|
return self.conn.close()
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self.conn, name)
|
|
|
|
def get_db_connection(database_url):
|
|
"""
|
|
Factory function to return a connection object.
|
|
Supports 'sqlite://' (or just file path) and 'postgresql://'.
|
|
"""
|
|
if database_url.startswith('postgres://') or database_url.startswith('postgresql://'):
|
|
try:
|
|
conn = psycopg2.connect(database_url)
|
|
return PostgresConnectionWrapper(conn)
|
|
except Exception as e:
|
|
logger.error(f"Failed to connect to Postgres: {e}")
|
|
raise e
|
|
else:
|
|
# Fallback to SQLite
|
|
# If it starts with sqlite:///, strip it
|
|
if database_url.startswith('sqlite:///'):
|
|
db_path = database_url.replace('sqlite:///', '')
|
|
else:
|
|
db_path = database_url
|
|
|
|
conn = sqlite3.connect(db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|