finnandClaude Opus 5.5
4a159006aa
app.logger.info("message", key=value) appends the keyword arguments as
JSON encoded fields, so callers no longer format the values themselves.
All log calls use fixed messages and put variable values into fields.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
109 lines
2.6 KiB
Python
109 lines
2.6 KiB
Python
import os
|
|
from functools import wraps
|
|
|
|
import psycopg2
|
|
|
|
|
|
class DatabaseError(Exception):
|
|
pass
|
|
|
|
|
|
class ProxyCursor(object):
|
|
def __init__(self, cur):
|
|
self.cur = cur
|
|
|
|
original_execute = object.__getattribute__(self, "cur").execute
|
|
|
|
@wraps(original_execute)
|
|
def execute(query, args=None):
|
|
query = query.strip()
|
|
database.add_query(query, args)
|
|
return original_execute(query, args)
|
|
|
|
self.execute = execute
|
|
|
|
def __getattribute__(self, name):
|
|
__dict__ = object.__getattribute__(self, "__dict__")
|
|
if name in __dict__ and name != "cur":
|
|
return object.__getattribute__(self, name)
|
|
else:
|
|
return getattr(object.__getattribute__(self, "cur"), name)
|
|
|
|
|
|
def get_env(name, default=None):
|
|
return os.environ.get(name, default)
|
|
|
|
|
|
def ensure_connection(fn):
|
|
@wraps(fn)
|
|
def wrapper(*args, **kwargs):
|
|
while True:
|
|
try:
|
|
result = fn(*args, **kwargs)
|
|
except psycopg2.InterfaceError:
|
|
continue
|
|
return result
|
|
|
|
return wrapper
|
|
|
|
|
|
def with_cursor(fn):
|
|
@wraps(fn)
|
|
@ensure_connection
|
|
def wrapper(*args, **kwargs):
|
|
conn = database.get_connection()
|
|
with conn:
|
|
with conn.cursor() as cur:
|
|
cur = ProxyCursor(cur)
|
|
return fn(cur, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
class Database:
|
|
app = None
|
|
|
|
def __init__(self):
|
|
self.connection = None
|
|
self.query_count = 0
|
|
|
|
def init(self, app):
|
|
self.app = app
|
|
|
|
def get_connection(self):
|
|
if not self.app:
|
|
raise RuntimeError()
|
|
if not self.connection:
|
|
self.connection = self.create_connection()
|
|
return self.connection
|
|
|
|
def create_connection(self):
|
|
try:
|
|
return psycopg2.connect(
|
|
host=get_env("PGSQL_HOST"),
|
|
port=int(get_env("PGSQL_PORT", 5432)),
|
|
database=get_env("PGSQL_NAME"),
|
|
user=get_env("PGSQL_USER"),
|
|
password=get_env("PGSQL_PASSWORD"),
|
|
)
|
|
except psycopg2.Error as e:
|
|
raise DatabaseError(f"Database connect error {e.pgcode}: {e.pgerror}")
|
|
|
|
def add_query(self, query, args=None):
|
|
if args is None:
|
|
args = []
|
|
self.app.logger.debug("query", query=query, args=args)
|
|
self.query_count += 1
|
|
|
|
def collect_query_count(self):
|
|
val = self.query_count
|
|
self.query_count = 0
|
|
return val
|
|
|
|
def shutdown(self):
|
|
if self.connection:
|
|
self.connection.close()
|
|
|
|
|
|
database = Database()
|