cms-SiteForge/app/pg_compat.py

"""Conservative PostgreSQL DB-API bridge for the existing SQLite-shaped application.

The row adapter intentionally supports BOTH row['column'] and row[0], because
several existing queue paths use positional aggregate results while API paths
serialize rows as dictionaries. PostgreSQL data migration remains explicit.
"""
from __future__ import annotations
import re

class PGRow(dict):
    """Mapping-compatible row with SQLite-style positional indexing."""
    __slots__ = ('_values',)

    def __init__(self, columns, values):
        values = tuple(values)
        super().__init__(zip(columns, values))
        self._values = values

    def __getitem__(self, key):
        if isinstance(key, int):
            return self._values[key]
        return super().__getitem__(key)


def row_factory(cursor):
    columns = [col.name for col in (cursor.description or ())]
    return lambda values: PGRow(columns, values)


def placeholders(sql: str) -> str:
    """Translate ? placeholders without modifying quoted strings or comments."""
    out=[]; i=0; state='code'
    while i<len(sql):
        ch=sql[i]; nxt=sql[i+1] if i+1<len(sql) else ''
        if state=='code':
            if ch in ("'",'"'):
                state=ch;out.append(ch)
            elif ch=='-' and nxt=='-':
                state='line';out.extend((ch,nxt));i+=1
            elif ch=='/' and nxt=='*':
                state='block';out.extend((ch,nxt));i+=1
            elif ch=='?':out.append('%s')
            else:out.append(ch)
        elif state in ("'",'"'):
            out.append(ch)
            if ch==state:
                if nxt==state:out.append(nxt);i+=1
                else:state='code'
        elif state=='line':
            out.append(ch)
            if ch=='\n':state='code'
        elif state=='block':
            out.append(ch)
            if ch=='*' and nxt=='/':out.append(nxt);i+=1;state='code'
        i+=1
    return ''.join(out)


def statements(script: str):
    """Split migration DDL at top-level semicolons, not inside quoted literals."""
    start=0;i=0;state='code'
    while i<len(script):
        ch=script[i];nxt=script[i+1] if i+1<len(script) else ''
        if state=='code':
            if ch in ("'",'"'):state=ch
            elif ch=='-' and nxt=='-':state='line';i+=1
            elif ch=='/' and nxt=='*':state='block';i+=1
            elif ch==';':
                candidate=script[start:i].strip()
                if candidate:yield candidate
                start=i+1
        elif state in ("'",'"'):
            if ch==state:
                if nxt==state:i+=1
                else:state='code'
        elif state=='line':
            if ch=='\n':state='code'
        elif state=='block':
            if ch=='*' and nxt=='/':i+=1;state='code'
        i+=1
    candidate=script[start:].strip()
    if candidate:yield candidate


def normalize(sql: str) -> str:
    sql=sql.strip()
    if re.match(r'^INSERT\s+OR\s+IGNORE\s+INTO\b',sql,re.I):
        sql=re.sub(r'^INSERT\s+OR\s+IGNORE\s+INTO\b','INSERT INTO',sql,flags=re.I)
        sql=sql.rstrip(';')+' ON CONFLICT DO NOTHING'
    return placeholders(sql)


class PGConnection:
    def __init__(self,dsn):
        import psycopg
        self.connection=psycopg.connect(dsn,row_factory=row_factory,connect_timeout=5)
        self.total_changes=0

    def execute(self,sql,params=()):
        stripped=sql.strip()
        if stripped.upper().startswith('PRAGMA TABLE_INFO('):
            table=stripped.split('(',1)[1].split(')',1)[0].strip().strip('"\'')
            if not re.fullmatch(r'[A-Za-z_][A-Za-z0-9_]*',table):raise ValueError('Invalid table name')
            return self.connection.execute('SELECT column_name AS name FROM information_schema.columns WHERE table_schema=current_schema() AND table_name=%s',(table,))
        if stripped.upper()=='BEGIN IMMEDIATE':
            # Serialize short queue claim/limit transactions, not the worker execution.
            # This prevents per-site races between concurrent enqueuers and claimers.
            return self.connection.execute('SELECT pg_advisory_xact_lock(823947231)')
        if stripped.upper().startswith('PRAGMA '):return _EmptyResult()
        cursor=self.connection.execute(normalize(sql),params)
        if cursor.rowcount and cursor.rowcount>0 and stripped.split(None,1)[0].upper() in ('UPDATE','INSERT','DELETE'):
            self.total_changes+=cursor.rowcount
        return cursor

    def executescript(self,script):
        for statement in statements(script):
            if not statement.lstrip().startswith('--') or '\n' in statement:
                self.execute(statement)

    def __enter__(self):return self
    def __exit__(self,typ,val,tb):
        try:
            if typ:self.connection.rollback()
            else:self.connection.commit()
        finally:self.connection.close()
        return False

class _EmptyResult:
    rowcount=0
    def fetchone(self):return None
    def fetchall(self):return []


def connect(dsn):return PGConnection(dsn)