253 lines
8.1 KiB
Python
253 lines
8.1 KiB
Python
"""PostgreSQL compatibility layer for the legacy sqlite3-style data access.
|
|
|
|
The rest of the codebase talks to the database through hand-written SQL that
|
|
was originally written for sqlite3 (``?`` placeholders, ``datetime('now')``,
|
|
``INSERT OR IGNORE``, ``cursor.lastrowid``, ``conn.total_changes`` ...).
|
|
This module exposes a sqlite3-like connection/cursor interface on top of
|
|
psycopg2 and transparently translates the dialect differences, so the
|
|
business-layer SQL keeps working with minimal changes.
|
|
|
|
Translations applied inside :meth:`PgCursor.execute`:
|
|
- ``?`` placeholders -> psycopg2 ``%s`` placeholders
|
|
- ``datetime('now')`` -> ``to_char(now(), 'YYYY-MM-DD HH24:MI:SS')``
|
|
(UTC, byte-identical format to before)
|
|
- ``INSERT OR IGNORE INTO`` -> ``INSERT INTO ... ON CONFLICT DO NOTHING``
|
|
- ``BEGIN IMMEDIATE`` -> ``BEGIN``
|
|
- ``INSERT`` (without RETURNING) -> appends ``RETURNING <pk>`` so that
|
|
``cursor.lastrowid`` keeps working; the
|
|
primary key is looked up once from
|
|
information_schema and cached.
|
|
|
|
Rows are returned as a dict subclass that supports both ``row["col"]`` and
|
|
``row[0]`` (positional) access to match the old ``sqlite3.Row`` behaviour.
|
|
|
|
The connection timezone is pinned to UTC so ``now()`` output matches the
|
|
previous ``datetime('now')`` (UTC) semantics.
|
|
"""
|
|
import re
|
|
import threading
|
|
|
|
import psycopg2
|
|
import psycopg2.extras
|
|
|
|
_DATETIME_NOW_RE = re.compile(r"datetime\('now'\)", re.IGNORECASE)
|
|
_INSERT_OR_IGNORE_RE = re.compile(r"INSERT\s+OR\s+IGNORE\s+INTO", re.IGNORECASE)
|
|
_INSERT_TABLE_RE = re.compile(
|
|
r"^\s*INSERT\s+(?:OR\s+IGNORE\s+)?INTO\s+([\"'`\w]+)", re.IGNORECASE
|
|
)
|
|
_BEGIN_IMMEDIATE_RE = re.compile(r"^\s*BEGIN\s+IMMEDIATE", re.IGNORECASE)
|
|
|
|
_PK_CACHE: dict[str, str | None] = {}
|
|
_PK_LOCK = threading.Lock()
|
|
|
|
|
|
class _Row(dict):
|
|
"""dict-like row supporting both column-name and positional access."""
|
|
|
|
__slots__ = ("_cols",)
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._cols = None
|
|
|
|
def __getitem__(self, key):
|
|
if isinstance(key, int):
|
|
if self._cols is None or key < 0 or key >= len(self._cols):
|
|
raise IndexError(key)
|
|
return dict.__getitem__(self, self._cols[key])
|
|
return dict.__getitem__(self, key)
|
|
|
|
def get_by_index(self, i):
|
|
return self[i]
|
|
|
|
|
|
def _wrap_row(row):
|
|
"""Convert a psycopg2 result row into a `_Row` (positional access enabled)."""
|
|
if row is None or isinstance(row, _Row):
|
|
return row
|
|
if isinstance(row, dict):
|
|
r = _Row(row)
|
|
r._cols = list(row.keys())
|
|
return r
|
|
return row
|
|
|
|
|
|
def _primary_key(conn, table: str) -> str | None:
|
|
"""Return the primary-key column of `table` (cached per process)."""
|
|
with _PK_LOCK:
|
|
if table in _PK_CACHE:
|
|
return _PK_CACHE[table]
|
|
pk: str | None = None
|
|
try:
|
|
cur = conn.cursor()
|
|
cur.execute(
|
|
"""
|
|
SELECT kcu.column_name
|
|
FROM information_schema.table_constraints tc
|
|
JOIN information_schema.key_column_usage kcu
|
|
ON tc.constraint_name = kcu.constraint_name
|
|
AND tc.table_schema = kcu.table_schema
|
|
WHERE tc.constraint_type = 'PRIMARY KEY'
|
|
AND tc.table_name = %s
|
|
AND tc.table_schema = current_schema()
|
|
ORDER BY kcu.ordinal_position
|
|
LIMIT 1
|
|
""",
|
|
(table,),
|
|
)
|
|
row = cur.fetchone()
|
|
pk = row["column_name"] if row else None
|
|
cur.close()
|
|
except Exception:
|
|
pk = None
|
|
with _PK_LOCK:
|
|
_PK_CACHE[table] = pk
|
|
return pk
|
|
|
|
|
|
def _translate_sql(sql: str) -> tuple[str, bool]:
|
|
"""Return (translated_sql, has_insert_or_ignore)."""
|
|
sql = _DATETIME_NOW_RE.sub("to_char(now(), 'YYYY-MM-DD HH24:MI:SS')", sql)
|
|
sql = _BEGIN_IMMEDIATE_RE.sub("BEGIN", sql)
|
|
has_ignore = bool(_INSERT_OR_IGNORE_RE.search(sql))
|
|
if has_ignore:
|
|
sql = _INSERT_OR_IGNORE_RE.sub("INSERT INTO", sql)
|
|
# `?` -> `%s`. The codebase SQL contains no literal `?` inside strings.
|
|
sql = sql.replace("?", "%s")
|
|
return sql, has_ignore
|
|
|
|
|
|
class PgCursor:
|
|
"""A psycopg2 cursor wrapper exposing the sqlite3 cursor surface."""
|
|
|
|
def __init__(self, proxy: "PgConnection", cursor):
|
|
self._proxy = proxy
|
|
self._cursor = cursor
|
|
self._lastrowid = None
|
|
|
|
def execute(self, sql: str, params=None) -> "PgCursor":
|
|
sql, has_ignore = _translate_sql(sql)
|
|
conn = self._proxy._conn
|
|
|
|
m = _INSERT_TABLE_RE.match(sql)
|
|
table = m.group(1).strip('"`') if m else None
|
|
returning_pk = None
|
|
if table is not None and " RETURNING " not in sql.upper():
|
|
returning_pk = _primary_key(conn, table)
|
|
|
|
sql = sql.rstrip().rstrip(";")
|
|
if has_ignore:
|
|
sql += " ON CONFLICT DO NOTHING"
|
|
if returning_pk:
|
|
sql += f" RETURNING {returning_pk}"
|
|
|
|
try:
|
|
self._cursor.execute(sql, params if params is not None else ())
|
|
if returning_pk:
|
|
row = _wrap_row(self._cursor.fetchone())
|
|
self._lastrowid = row[0] if row is not None else None
|
|
rc = self._cursor.rowcount
|
|
if rc and rc > 0:
|
|
self._proxy._total_changes += rc
|
|
except Exception:
|
|
self._proxy._safe_rollback()
|
|
raise
|
|
return self
|
|
|
|
# ── sqlite3 cursor surface ───────────────────────────────
|
|
@property
|
|
def lastrowid(self):
|
|
return self._lastrowid
|
|
|
|
@property
|
|
def rowcount(self):
|
|
return self._cursor.rowcount
|
|
|
|
def fetchone(self):
|
|
return _wrap_row(self._cursor.fetchone())
|
|
|
|
def fetchall(self):
|
|
return [_wrap_row(r) for r in self._cursor.fetchall()]
|
|
|
|
def fetchmany(self, size=None):
|
|
return [_wrap_row(r) for r in self._cursor.fetchmany(size)]
|
|
|
|
def close(self):
|
|
try:
|
|
self._cursor.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
class PgConnection:
|
|
"""A sqlite3-like connection backed by psycopg2.
|
|
|
|
Note: unlike a raw psycopg2 connection, a failed statement automatically
|
|
rolls the connection back (sqlite allows continuing after an error; this
|
|
keeps the connection usable under the same call pattern).
|
|
"""
|
|
|
|
def __init__(self, dsn: str):
|
|
self._conn = psycopg2.connect(dsn, cursor_factory=psycopg2.extras.RealDictCursor)
|
|
self._conn.autocommit = False
|
|
self._total_changes = 0
|
|
self._closed = False
|
|
cur = self._conn.cursor()
|
|
cur.execute("SET TIME ZONE 'UTC'")
|
|
cur.close()
|
|
|
|
def execute(self, sql: str, params=None) -> PgCursor:
|
|
return self.cursor().execute(sql, params)
|
|
|
|
def cursor(self) -> PgCursor:
|
|
return PgCursor(self, self._conn.cursor())
|
|
|
|
def executescript(self, sql: str):
|
|
"""Run a multi-statement DDL script (SCHEMA). No translation."""
|
|
cur = self._conn.cursor()
|
|
try:
|
|
cur.execute(sql)
|
|
except Exception:
|
|
self._safe_rollback()
|
|
raise
|
|
finally:
|
|
try:
|
|
cur.close()
|
|
except Exception:
|
|
pass
|
|
|
|
def commit(self):
|
|
self._conn.commit()
|
|
|
|
def rollback(self):
|
|
self._conn.rollback()
|
|
|
|
def close(self):
|
|
if not self._closed:
|
|
try:
|
|
self._conn.rollback()
|
|
self._conn.close()
|
|
except Exception:
|
|
pass
|
|
self._closed = True
|
|
|
|
def _safe_rollback(self):
|
|
try:
|
|
self._conn.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
# ── sqlite3 connection surface ───────────────────────────
|
|
@property
|
|
def total_changes(self):
|
|
return self._total_changes
|
|
|
|
@property
|
|
def closed(self):
|
|
return self._closed
|
|
|
|
|
|
def connect(dsn: str) -> PgConnection:
|
|
"""Open a PostgreSQL connection with the sqlite3-compatible wrapper."""
|
|
return PgConnection(dsn)
|