fleetmesh-node 0.2.0__py3-none-any.whl

node/storage/postgres.py

@@ -0,0 +1,294 @@
+"""PostgreSQL, as an option beside the SQLite default.
+
+Optional in the convention this repo already uses for wasmtime: the driver is
+imported lazily, a node without it still starts and serves, and a
+`postgres://` URL is refused by name rather than failing at import. CI sets
+FLEETMESH_REQUIRE_POSTGRES=1 so a missing driver cannot turn a suite green by
+skipping.
+
+The thirteen stores write SQLite's dialect, which is the house dialect, and
+`sql()` here renders it into PostgreSQL's. The census that says this is a
+short list rather than a rewrite is docs/storage-backends.md section 4; what
+it actually comes to:
+
+  placeholders   `?` -> `%s`, outside string literals
+  upserts        `INSERT OR REPLACE` -> `ON CONFLICT (pk) DO UPDATE SET ...`
+                 `INSERT OR IGNORE`  -> `ON CONFLICT DO NOTHING`
+  types          nothing to do. The node's DDL uses TEXT and INTEGER and
+                 nothing else across all 44 tables, and both are PostgreSQL
+                 types.
+  WITHOUT ROWID  dropped; it is a SQLite storage hint.
+  PRAGMA         never reaches here; `tune` sets what a server needs instead.
+
+The conflict target comes from the table's own DDL. Every `CREATE TABLE` in
+the tree passes through `sql()` during `_ensure_schema`, before any INSERT can,
+so the dialect learns each table's columns and primary key on the way past and
+has them when an upsert arrives. All 44 tables declare a primary key, so
+`DO UPDATE` always has a target - checked, not assumed.
+
+One connection behind one lock, the shape BaseStore already had. That is what
+makes this backend need no retry policy: with a single writer there is nothing
+to serialize against, so no transaction is ever told to retry.
+"""
+from __future__ import annotations
+
+import re
+import threading
+from typing import Any, Optional
+
+from fleetmesh_core.storage.store_postgres import ERROR_HINT, available, driver
+
+__all__ = ['ERROR_HINT', 'PostgresDialect', 'available', 'driver', 'placeholders']
+
+
+_CREATE = re.compile(r'CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?([A-Za-z_]\w*)\s*\((.*)\)\s*'
+                     r'(?:WITHOUT\s+ROWID)?\s*;?\s*$', re.I | re.S)
+_INLINE_PK = re.compile(r'^\s*([A-Za-z_]\w*)\s+\w+[^,]*\bPRIMARY\s+KEY\b', re.I)
+_TABLE_PK = re.compile(r'\bPRIMARY\s+KEY\s*\(([^)]*)\)', re.I)
+_UPSERT = re.compile(r'^(\s*)INSERT\s+OR\s+(REPLACE|IGNORE)\s+INTO\s+([A-Za-z_]\w*)\s*'
+                     r'(\(([^)]*)\))?\s*(VALUES\b.*)$', re.I | re.S)
+
+
+def _split_columns(body: str) -> list[str]:
+    """Top-level comma split of a CREATE TABLE body, ignoring nested parens."""
+    out, depth, cur = [], 0, ''
+    for ch in body:
+        if ch == '(':
+            depth += 1
+        elif ch == ')':
+            depth -= 1
+        if ch == ',' and depth == 0:
+            out.append(cur)
+            cur = ''
+        else:
+            cur += ch
+    if cur.strip():
+        out.append(cur)
+    return out
+
+
+def placeholders(statement: str) -> str:
+    """`?` -> `%s`, leaving anything inside a single-quoted literal alone.
+
+    The node's SQL carries literals like `status='verified'`; a blind replace
+    would corrupt one that happened to contain a question mark.
+    """
+    out, in_lit, i = [], False, 0
+    while i < len(statement):
+        ch = statement[i]
+        if ch == "'":
+            # '' inside a literal is an escaped quote, not the end of it
+            if in_lit and i + 1 < len(statement) and statement[i + 1] == "'":
+                out.append("''")
+                i += 2
+                continue
+            in_lit = not in_lit
+            out.append(ch)
+        elif ch == '?' and not in_lit:
+            out.append('%s')
+        else:
+            out.append(ch)
+        i += 1
+    return ''.join(out)
+
+
+class PostgresDialect:
+    """PostgreSQL. Learns each table's shape from the DDL it renders."""
+
+    name: str = 'postgres'
+    rowid_expr: Optional[str] = None   # no insertion-order expression; see Dialect
+
+    def __init__(self) -> None:
+        self._cols: dict[str, list[str]] = {}
+        self._pk: dict[str, list[str]] = {}
+
+    # ---- connection -------------------------------------------------------
+
+    def connect(self, target: str, timeout: float, readonly: bool = False) -> Any:
+        pg = driver()
+        conn = pg.connect(target, connect_timeout=int(max(1, timeout)), autocommit=False)
+        if readonly:
+            conn.read_only = True
+        return conn
+
+    def tune(self, db: Any, target: str, lock: threading.RLock) -> None:
+        """What a server needs in place of SQLite's PRAGMAs: a statement
+        timeout so a wedged query cannot hold the one connection for ever."""
+        with lock:
+            cur = db.cursor()
+            try:
+                cur.execute("SET statement_timeout = '30s'")
+                cur.execute("SET idle_in_transaction_session_timeout = '60s'")
+            finally:
+                cur.close()
+            db.commit()
+
+    def learn_table(self, table: Any) -> None:
+        name = getattr(table, 'name', '')
+        if not name:
+            return
+        cols = [c.name for c in getattr(table, 'columns', ()) if hasattr(c, 'name')]
+        pk = list(getattr(table, 'primary_key', ()))
+        if not pk:
+            pk = [c.name for c in getattr(table, 'columns', ()) if getattr(c, 'primary_key', False)]
+        self._cols[name.lower()] = cols
+        if pk:
+            self._pk[name.lower()] = pk
+
+    def _learn(self, statement: str) -> None:
+        m = _CREATE.match(statement.strip())
+        if not m:
+            return
+        table, body = m.group(1), m.group(2)
+        cols, pk = [], []
+        for part in _split_columns(body):
+            tpk = _TABLE_PK.search(part)
+            if tpk and not re.match(r'^\s*[A-Za-z_]\w*\s+\w', part):
+                pk = [c.strip() for c in tpk.group(1).split(',') if c.strip()]
+                continue
+            name = re.match(r'^\s*([A-Za-z_]\w*)', part)
+            if not name:
+                continue
+            if name.group(1).upper() in ('PRIMARY', 'UNIQUE', 'FOREIGN', 'CHECK', 'CONSTRAINT'):
+                tpk = _TABLE_PK.search(part)
+                if tpk:
+                    pk = [c.strip() for c in tpk.group(1).split(',') if c.strip()]
+                continue
+            cols.append(name.group(1))
+            if _INLINE_PK.match(part):
+                pk = [name.group(1)]
+        self._cols[table.lower()] = cols
+        if pk:
+            self._pk[table.lower()] = pk
+
+    def sql(self, statement: Any) -> str:
+        if hasattr(statement, 'to_sql'):
+            self.learn_table(statement)
+            statement = statement.to_sql('postgres')
+        else:
+            self._learn(statement)
+        out = re.sub(r'\s+WITHOUT\s+ROWID\b', '', statement, flags=re.I)
+        m = _UPSERT.match(out)
+        if m:
+            indent, mode, table, _, named, tail = m.groups()
+            key = table.lower()
+            if mode.upper() == 'IGNORE':
+                out = '%sINSERT INTO %s %s %s ON CONFLICT DO NOTHING' % (
+                    indent, table, '(%s)' % named if named else '', tail)
+            else:
+                pk = self._pk.get(key, [])
+                cols = [c.strip() for c in named.split(',')] if named else self._cols.get(key, [])
+                sets = [c for c in cols if c.lower() not in {p.lower() for p in pk}]
+                if not pk:
+                    raise RuntimeError(
+                        'INSERT OR REPLACE INTO %s needs a primary key to become an upsert, '
+                        'and none was declared in its CREATE TABLE' % table)
+                assign = ', '.join('%s = EXCLUDED.%s' % (c, c) for c in sets) or None
+                action = ('DO UPDATE SET %s' % assign) if assign else 'DO NOTHING'
+                out = '%sINSERT INTO %s %s %s ON CONFLICT (%s) %s' % (
+                    indent, table, '(%s)' % named if named else '', tail,
+                    ', '.join(pk), action)
+        return placeholders(out)
+
+    # ---- transaction control ----------------------------------------------
+    #
+    # The driver owns the transaction. psycopg opens one on the first
+    # statement and tracks it; issuing raw BEGIN/COMMIT/ROLLBACK alongside
+    # that leaves the two disagreeing about whether a transaction is open,
+    # and work committed by one path gets rolled back by the other. That is
+    # how the first Postgres run recorded schema version 2 while the column
+    # version 1 adds was missing: the ALTER applied and was then undone,
+    # while the version write survived on the other side of a raw COMMIT.
+
+    def begin(self, db: Any) -> None:
+        return None              # psycopg begins on the first statement
+
+    def commit(self, db: Any) -> None:
+        db.commit()
+
+    def rollback(self, db: Any) -> None:
+        db.rollback()
+
+    # ---- schema version ---------------------------------------------------
+
+    def get_version(self, db: Any, store: str) -> int:
+        """The schema version of ONE store, keyed by its name.
+
+        SQLite gets this for free: each store is its own file, so
+        `PRAGMA user_version` is per store. Here all fourteen share one
+        database, and a single version row made them collide - AccountStore
+        wrote 2, GrantStore read 2, concluded its own migrations were applied
+        and skipped the ALTER that adds `grants.env_created_at`. The table was
+        then one column short of what its INSERT supplies, and the first
+        write failed with "INSERT has more expressions than target columns".
+        Only a multi-store test on a server could find that; one store alone
+        passes.
+        """
+        cur = db.cursor()
+        try:
+            cur.execute('CREATE TABLE IF NOT EXISTS fleetmesh_schema_version ('
+                        ' store TEXT PRIMARY KEY, version INTEGER NOT NULL)')
+            cur.execute('SELECT version FROM fleetmesh_schema_version WHERE store = %s', (store,))
+            row = cur.fetchone()
+        finally:
+            cur.close()
+        return int(row[0]) if row else 0
+
+    def set_version(self, db: Any, store: str, version: int) -> None:
+        cur = db.cursor()
+        try:
+            cur.execute('INSERT INTO fleetmesh_schema_version VALUES (%s, %s) '
+                        'ON CONFLICT (store) DO UPDATE SET version = EXCLUDED.version',
+                        (store, int(version)))
+        finally:
+            cur.close()
+
+    # ---- reflection and maintenance ---------------------------------------
+
+    def has_table(self, db: Any, table: str) -> bool:
+        cur = db.cursor()
+        try:
+            cur.execute('SELECT 1 FROM information_schema.tables '
+                        'WHERE table_schema = current_schema() AND table_name = %s', (table,))
+            return cur.fetchone() is not None
+        finally:
+            cur.close()
+
+    def has_column(self, db: Any, table: str, column: str) -> bool:
+        cur = db.cursor()
+        try:
+            cur.execute('SELECT 1 FROM information_schema.columns '
+                        'WHERE table_schema = current_schema() AND table_name = %s '
+                        'AND column_name = %s', (table, column))
+            return cur.fetchone() is not None
+        finally:
+            cur.close()
+
+    def is_duplicate_column(self, exc: BaseException) -> bool:
+        return getattr(getattr(exc, 'sqlstate', None), 'strip', lambda: '')() == '42701' \
+            or 'already exists' in str(exc).lower()
+
+    def vacuum(self, db: Any, pages: int) -> int:
+        """ANALYZE, and nothing more. Postgres reclaims by autovacuum; a manual
+        VACUUM here would block the one connection the node has, and
+        `pages` has no meaning on a server. Returns 0 - nothing was reclaimed
+        by this call, and saying otherwise would be a fake count."""
+        try:
+            db.commit()
+            with db.cursor() as cur:
+                cur.execute('ANALYZE')
+            db.commit()
+        except Exception:
+            return 0
+        return 0
+
+    def integrity_check(self, db: Any) -> Optional[str]:
+        """None. File integrity is a file-backed idea; a server answers for its
+        own soundness and `PRAGMA quick_check` has no counterpart here."""
+        return None
+
+    def errors(self) -> tuple[type[BaseException], ...]:
+        try:
+            return (driver().Error,)
+        except RuntimeError:
+            return ()