b3-prosaic-harness-postgres 0.2.1__py3-none-any.whl

prosaic_harness_postgres/schema.py

@@ -0,0 +1,135 @@
+"""Explicit, transactional version-1 setup and read-only readiness checks."""
+from prosaic_harness import StoreUninitialized, StoreIncompatible
+from prosaic_harness.run_store import STORAGE_CONTRACT_VERSION
+
+SCHEMA_VERSION = 1
+MIGRATION_LOCK = 0x50524f53414943
+
+DDL = """
+CREATE SCHEMA IF NOT EXISTS prosaic_harness;
+CREATE TABLE prosaic_harness.schema_version (
+ singleton BOOLEAN PRIMARY KEY CHECK (singleton), schema_version INTEGER NOT NULL,
+ storage_contract_version INTEGER NOT NULL);
+CREATE TABLE prosaic_harness.leases (
+ namespace TEXT NOT NULL CHECK (namespace ~ '^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$'),
+ run_id TEXT NOT NULL CHECK (run_id ~ '^[0-9a-f]{32}$'),
+ generation BIGINT NOT NULL DEFAULT 0 CHECK (generation >= 0), owner TEXT,
+ expires_at TIMESTAMPTZ, PRIMARY KEY(namespace, run_id),
+ CHECK ((owner IS NULL) = (expires_at IS NULL)));
+CREATE TABLE prosaic_harness.runs (
+ namespace TEXT NOT NULL, run_id TEXT NOT NULL,
+ revision BIGINT NOT NULL CHECK(revision > 0),
+ document TEXT NOT NULL CHECK(octet_length(document) <= 8388608),
+ PRIMARY KEY(namespace, run_id),
+ FOREIGN KEY(namespace, run_id) REFERENCES prosaic_harness.leases(namespace, run_id));
+CREATE TABLE prosaic_harness.receipts (
+ namespace TEXT NOT NULL, run_id TEXT NOT NULL,
+ invocation_id TEXT NOT NULL CHECK(invocation_id ~ '^[0-9a-f]{32}$'),
+ document TEXT NOT NULL CHECK(octet_length(document) <= 8388608),
+ sha256 TEXT NOT NULL CHECK(sha256 ~ '^[0-9a-f]{64}$'),
+ PRIMARY KEY(namespace, run_id, invocation_id),
+ FOREIGN KEY(namespace, run_id) REFERENCES prosaic_harness.runs(namespace, run_id));
+INSERT INTO prosaic_harness.schema_version VALUES(TRUE, 1, 1);
+"""
+
+COLUMNS = {
+    'schema_version': {'singleton': 'boolean', 'schema_version': 'integer', 'storage_contract_version': 'integer'},
+    'leases': {'namespace': 'text', 'run_id': 'text', 'generation': 'bigint', 'owner': 'text', 'expires_at': 'timestamp with time zone'},
+    'runs': {'namespace': 'text', 'run_id': 'text', 'revision': 'bigint', 'document': 'text'},
+    'receipts': {'namespace': 'text', 'run_id': 'text', 'invocation_id': 'text', 'document': 'text', 'sha256': 'text'},
+}
+
+CONSTRAINTS = {
+    'schema_version': {'CHECK (singleton)', 'PRIMARY KEY (singleton)'},
+    'leases': {
+        "CHECK ((namespace ~ '^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$'::text))",
+        "CHECK ((run_id ~ '^[0-9a-f]{32}$'::text))",
+        'CHECK ((generation >= 0))', 'CHECK (((owner IS NULL) = (expires_at IS NULL)))',
+        'PRIMARY KEY (namespace, run_id)'},
+    'runs': {'CHECK ((revision > 0))', 'CHECK ((octet_length(document) <= 8388608))',
+        'PRIMARY KEY (namespace, run_id)',
+        'FOREIGN KEY (namespace, run_id) REFERENCES prosaic_harness.leases(namespace, run_id)'},
+    'receipts': {"CHECK ((invocation_id ~ '^[0-9a-f]{32}$'::text))",
+        "CHECK ((sha256 ~ '^[0-9a-f]{64}$'::text))",
+        'CHECK ((octet_length(document) <= 8388608))',
+        'PRIMARY KEY (namespace, run_id, invocation_id)',
+        'FOREIGN KEY (namespace, run_id) REFERENCES prosaic_harness.runs(namespace, run_id)'},
+}
+
+
+async def ready_on_connection(conn):
+    exists = await (await conn.execute("SELECT pg_catalog.to_regclass('prosaic_harness.schema_version')")).fetchone()
+    if not exists or exists[0] is None:
+        raise StoreUninitialized('run init with migration credentials before accepting work')
+    tables = await (await conn.execute("""SELECT c.relname, c.relpersistence, c.relkind
+        FROM pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid=c.relnamespace
+        WHERE n.nspname='prosaic_harness' AND c.relname = ANY(%s)""", (list(COLUMNS),))).fetchall()
+    if {r[0] for r in tables} != set(COLUMNS) or any(r[1:] != ('p', 'r') for r in tables):
+        raise StoreIncompatible('owned schema must contain permanent logged adapter tables')
+    columns = await (await conn.execute("""SELECT c.relname, a.attname, pg_catalog.format_type(a.atttypid,a.atttypmod), a.attnotnull
+        FROM pg_catalog.pg_attribute a JOIN pg_catalog.pg_class c ON c.oid=a.attrelid
+        JOIN pg_catalog.pg_namespace n ON n.oid=c.relnamespace
+        WHERE n.nspname='prosaic_harness' AND c.relname = ANY(%s)
+        AND a.attnum>0 AND NOT a.attisdropped""", (list(COLUMNS),))).fetchall()
+    actual = {name: {} for name in COLUMNS}
+    for table, column, kind, notnull in columns:
+        actual[table][column] = kind
+        if notnull != ((table, column) not in {('leases', 'owner'), ('leases', 'expires_at')}):
+            raise StoreIncompatible('adapter schema nullability is incompatible')
+    if actual != COLUMNS:
+        raise StoreIncompatible('adapter schema columns are incompatible')
+    constraints = await (await conn.execute("""SELECT c.relname,
+        pg_catalog.pg_get_constraintdef(k.oid), k.convalidated, k.contype
+        FROM pg_catalog.pg_constraint k JOIN pg_catalog.pg_class c ON c.oid=k.conrelid
+        JOIN pg_catalog.pg_namespace n ON n.oid=c.relnamespace
+        WHERE n.nspname='prosaic_harness' AND c.relname = ANY(%s)""", (list(COLUMNS),))).fetchall()
+    actual_constraints = {name: set() for name in COLUMNS}
+    for table, definition, validated, kind in constraints:
+        if not validated:
+            raise StoreIncompatible('adapter schema constraints must be validated')
+        # PostgreSQL 18 catalogs NOT NULL separately; attnotnull above checks
+        # the identical contract on both supported major versions.
+        if kind == 'n':
+            continue
+        actual_constraints[table].add(definition)
+    if actual_constraints != CONSTRAINTS:
+        raise StoreIncompatible('adapter schema constraints are incompatible')
+    version = await (await conn.execute('SELECT schema_version, storage_contract_version FROM prosaic_harness.schema_version')).fetchall()
+    if version != [(SCHEMA_VERSION, STORAGE_CONTRACT_VERSION)]:
+        raise StoreIncompatible('deploy a compatible adapter/schema; no automatic downgrade')
+    for table in COLUMNS:
+        privileges = ['SELECT']
+        if table in {'leases', 'runs', 'receipts'}:
+            privileges.append('INSERT')
+        if table in {'leases', 'runs'}:
+            privileges.append('UPDATE')
+        for privilege in privileges:
+            allowed = await (await conn.execute('SELECT pg_catalog.has_table_privilege(%s,%s)',
+                                                ('prosaic_harness.' + table, privilege))).fetchone()
+            if not allowed or not allowed[0]:
+                raise StoreIncompatible('grant required runtime table privileges before accepting work')
+
+
+def check_ready(transport):
+    transport.call(ready_on_connection)
+
+
+def initialize(transport):
+    async def migrate(conn):
+        await conn.execute("SELECT pg_catalog.set_config('lock_timeout','25000ms',true), "
+                           "pg_catalog.set_config('statement_timeout','28000ms',true)")
+        await conn.execute('SELECT pg_catalog.pg_advisory_xact_lock(%s)', (MIGRATION_LOCK,))
+        namespace = await (await conn.execute("""SELECT n.oid,
+            pg_catalog.pg_has_role(n.nspowner,'USAGE') FROM pg_catalog.pg_namespace n
+            WHERE n.nspname='prosaic_harness'""")).fetchone()
+        if namespace:
+            if not namespace[1]:
+                raise StoreIncompatible('migration role must own the adapter schema')
+            count = await (await conn.execute('SELECT count(*) FROM pg_catalog.pg_class WHERE relnamespace=%s',
+                                              (namespace[0],))).fetchone()
+            if count[0]:
+                await ready_on_connection(conn)
+                return
+        await conn.execute(DDL)
+        await ready_on_connection(conn)
+    transport.call(migrate, timeout_s=30)

prosaic_harness_postgres/store.py

@@ -0,0 +1,277 @@
+"""Namespace-bound RunStore with short transactions and no inference retries."""
+from contextlib import contextmanager
+from datetime import timedelta
+import hashlib
+import os
+import re
+import threading
+import time
+
+from prosaic_harness import (RunSnapshot, StoreBusy, StoreCorrupt, StoreUnavailable,
+                            LeaseLost, RunBusy, RunNotFound, RunAlreadyExists,
+                            RevisionConflict, ReceiptConflict)
+from prosaic_harness.run_store import (validate_namespace, validate_id, validate_revision,
+                                      encode_document, decode_document, revision_token)
+from .transport import PostgresTransport, positive
+from . import schema
+from .leases import PostgresLease
+
+
+class PostgresRunStore:
+    @classmethod
+    def from_env(cls, name, *, namespace, **kwargs):
+        if not isinstance(name, str) or not re.fullmatch(r'[A-Za-z_][A-Za-z0-9_]{0,127}', name):
+            raise ValueError('dsn environment name must be a bounded identifier')
+        dsn = os.environ.get(name)
+        if not dsn:
+            raise StoreUnavailable('required database environment variable is missing')
+        return cls(dsn, namespace=namespace, **kwargs)
+
+    def __init__(self, dsn, *, namespace, lease_s=60, renew_s=15,
+                 operation_timeout_s=10, pool_timeout_s=5, connect_timeout_s=5,
+                 lock_timeout_s=1, statement_timeout_s=5, pool_size=8, max_active_leases=4):
+        self.namespace = validate_namespace(namespace)
+        positive(lease_s, 'lease_s')
+        positive(renew_s, 'renew_s')
+        positive(operation_timeout_s, 'operation_timeout_s')
+        if (lease_s < 15 or renew_s < 1 or renew_s > lease_s / 3
+                or renew_s + operation_timeout_s + 1 >= lease_s):
+            raise ValueError('lease timing requires duration >=15, renewal >=1 and adequate deadline margin')
+        self.lease_s, self.renew_s = lease_s, renew_s
+        self._pid = os.getpid()
+        self._mutex = threading.RLock()
+        self._sessions = set()
+        self._reservations = set()
+        self._max_active_leases = max_active_leases
+        self._transport = PostgresTransport(dsn, operation_timeout_s=operation_timeout_s,
+            pool_timeout_s=pool_timeout_s, connect_timeout_s=connect_timeout_s,
+            lock_timeout_s=lock_timeout_s, statement_timeout_s=statement_timeout_s,
+            pool_size=pool_size, max_active_leases=max_active_leases)
+
+    def initialize(self):
+        schema.initialize(self._transport)
+
+    def check_ready(self):
+        schema.check_ready(self._transport)
+
+    def _local(self, run_id, lease):
+        self._transport._check_process()
+        validate_id(run_id)
+        if (not isinstance(lease, PostgresLease) or lease.store is not self
+                or lease.run_id != run_id or lease not in self._sessions):
+            raise LeaseLost('lease does not own this store/run')
+        lease.check_valid()
+
+    async def _lease_row(self, conn, run_id):
+        return await (await conn.execute('SELECT owner,generation,expires_at FROM prosaic_harness.leases '
+            'WHERE namespace=%s AND run_id=%s FOR UPDATE', (self.namespace, run_id))).fetchone()
+
+    async def _guard_lease(self, conn, lease):
+        self._local(lease.run_id, lease)
+        row = await self._lease_row(conn, lease.run_id)
+        now = (await (await conn.execute('SELECT pg_catalog.clock_timestamp()')).fetchone())[0]
+        if (not row or row[0] != lease.owner or row[1] != lease.generation
+                or row[2] is None or row[2] <= now):
+            lease._lost = True
+            raise LeaseLost('database execution ownership expired or changed')
+        lease.check_valid()
+        return now
+
+    async def _acquire(self, conn, lease):
+        await conn.execute('INSERT INTO prosaic_harness.leases(namespace,run_id) VALUES(%s,%s) '
+                           'ON CONFLICT DO NOTHING', (self.namespace, lease.run_id))
+        row = await self._lease_row(conn, lease.run_id)
+        now = (await (await conn.execute('SELECT pg_catalog.clock_timestamp()')).fetchone())[0]
+        if row[0] is not None and row[2] > now:
+            raise RunBusy('run is owned by another execution')
+        generation = row[1] + 1
+        await conn.execute('UPDATE prosaic_harness.leases SET owner=%s,generation=%s,expires_at=%s '
+                           'WHERE namespace=%s AND run_id=%s',
+                           (lease.owner, generation, now + timedelta(seconds=self.lease_s),
+                            self.namespace, lease.run_id))
+        return generation
+
+    async def _renew(self, conn, lease):
+        now = await self._guard_lease(conn, lease)
+        cursor = await conn.execute('UPDATE prosaic_harness.leases SET expires_at=%s '
+            'WHERE namespace=%s AND run_id=%s AND owner=%s AND generation=%s',
+            (now + timedelta(seconds=self.lease_s), self.namespace, lease.run_id, lease.owner, lease.generation))
+        if cursor.rowcount != 1:
+            lease._lost = True
+            raise LeaseLost('database renewal ownership changed')
+
+    async def _release(self, conn, lease):
+        await conn.execute('UPDATE prosaic_harness.leases SET owner=NULL,expires_at=NULL '
+            'WHERE namespace=%s AND run_id=%s AND owner=%s AND generation=%s',
+            (self.namespace, lease.run_id, lease.owner, lease.generation))
+
+    @contextmanager
+    def lease(self, run_id):
+        self._transport._check_process()
+        validate_id(run_id)
+        reservation = object()
+        session = None
+        primary = None
+        try:
+            with self._mutex:
+                if len(self._reservations) >= self._max_active_leases:
+                    raise StoreBusy('active execution capacity exhausted')
+                # One token owns capacity throughout setup, execution and
+                # teardown. Registration never transfers or decrements it.
+                self._reservations.add(reservation)
+            session = PostgresLease(self, run_id, time.monotonic())
+            session.generation = self._transport.call(lambda conn: self._acquire(conn, session))
+            session.check_valid()
+            with self._mutex:
+                self._sessions.add(session)
+            # Retain the coroutine even if submission is interrupted before it
+            # returns. stop() closes an unstarted runner on its owning loop or
+            # drains a started renewal before ownership can be released.
+            session._renewal = session.maintain()
+            self._transport.submit(session._renewal)
+            yield session
+        except BaseException as exc:
+            primary = exc
+            raise
+        finally:
+            try:
+                if session is not None and session.generation is not None:
+                    if session._renewal is None:
+                        session._stopped.set()
+                    else:
+                        session.stop()
+                    self._transport.call(lambda conn: self._release(conn, session), renewal=True)
+            except BaseException:
+                session._lost = True
+                if primary is None:
+                    raise
+                primary.add_note('database lease cleanup failed; reconcile ownership before retrying')
+            finally:
+                if session is not None:
+                    session._active = False
+                with self._mutex:
+                    if (session is None or session.generation is None
+                            or session._stopped.is_set()):
+                        self._reservations.discard(reservation)
+                    if session is not None and session._stopped.is_set():
+                        self._sessions.discard(session)
+
+    def _token(self, run_id, revision):
+        return revision_token('postgres', self.namespace, run_id, revision)
+
+    async def _read_run(self, conn, run_id, *, lock=False):
+        row = await (await conn.execute('SELECT revision,CASE WHEN octet_length(document)<=8388608 '
+            'THEN document ELSE NULL END FROM prosaic_harness.runs WHERE namespace=%s AND run_id=%s'
+            + (' FOR UPDATE' if lock else ''), (self.namespace, run_id))).fetchone()
+        if not row:
+            raise RunNotFound('run document not found')
+        state = decode_document(row[1])
+        if state.get('run_id') != run_id:
+            raise StoreCorrupt('checkpoint identity does not match its address')
+        return RunSnapshot(state, self._token(run_id, row[0])), row[0]
+
+    def load_run(self, run_id):
+        validate_id(run_id)
+        async def read(conn):
+            return (await self._read_run(conn, run_id))[0]
+        return self._transport.call(read)
+
+    def _document(self, run_id, state):
+        content = encode_document(state)
+        if state.get('run_id') != run_id:
+            raise StoreCorrupt('checkpoint identity does not match its address')
+        return content.decode('utf-8')
+
+    def _mutate(self, run_id, lease, callback):
+        self._local(run_id, lease)
+        async def guarded(conn):
+            await self._guard_lease(conn, lease)
+            return await callback(conn)
+        try:
+            result = self._transport.call(guarded)
+            lease.check_valid()
+            return result
+        except StoreUnavailable:
+            lease._lost = True
+            raise
+
+    def create_run(self, run_id, state, lease):
+        self._local(run_id, lease)
+        document = self._document(run_id, state)
+        async def create(conn):
+            row = await (await conn.execute('INSERT INTO prosaic_harness.runs(namespace,run_id,revision,document) '
+                'VALUES(%s,%s,1,%s) ON CONFLICT DO NOTHING RETURNING revision',
+                (self.namespace, run_id, document))).fetchone()
+            if not row:
+                raise RunAlreadyExists('run already exists; use resume or a new identity')
+            return self._token(run_id, row[0])
+        return self._mutate(run_id, lease, create)
+
+    def save_run(self, run_id, state, expected_revision, lease):
+        self._local(run_id, lease)
+        validate_revision(expected_revision)
+        document = self._document(run_id, state)
+        async def save(conn):
+            snapshot, revision = await self._read_run(conn, run_id, lock=True)
+            if snapshot.revision != expected_revision:
+                raise RevisionConflict('checkpoint changed; refresh status')
+            row = await (await conn.execute('UPDATE prosaic_harness.runs SET document=%s,revision=revision+1 '
+                'WHERE namespace=%s AND run_id=%s AND revision=%s RETURNING revision',
+                (document, self.namespace, run_id, revision))).fetchone()
+            if not row:
+                raise RevisionConflict('checkpoint changed; refresh status')
+            return self._token(run_id, row[0])
+        return self._mutate(run_id, lease, save)
+
+    async def _read_receipt(self, conn, run_id, invocation_id):
+        row = await (await conn.execute('SELECT CASE WHEN octet_length(document)<=8388608 THEN document '
+            'ELSE NULL END,sha256 FROM prosaic_harness.receipts WHERE namespace=%s AND run_id=%s AND invocation_id=%s',
+            (self.namespace, run_id, invocation_id))).fetchone()
+        if row is None:
+            return None
+        receipt = decode_document(row[0])
+        if (receipt.get('id') != invocation_id or receipt.get('run_id') != run_id
+                or hashlib.sha256(encode_document(receipt)).hexdigest() != row[1]):
+            raise StoreCorrupt('receipt integrity or identity does not match its address')
+        return receipt
+
+    def load_receipt(self, run_id, invocation_id):
+        validate_id(run_id)
+        validate_id(invocation_id)
+        return self._transport.call(lambda conn: self._read_receipt(conn, run_id, invocation_id))
+
+    def save_receipt(self, run_id, invocation_id, receipt, lease):
+        self._local(run_id, lease)
+        validate_id(invocation_id)
+        content = encode_document(receipt)
+        if receipt.get('id') != invocation_id or receipt.get('run_id') != run_id:
+            raise StoreCorrupt('receipt identity does not match its address')
+        async def save(conn):
+            await self._read_run(conn, run_id, lock=True)
+            await conn.execute('INSERT INTO prosaic_harness.receipts(namespace,run_id,invocation_id,document,sha256) '
+                'VALUES(%s,%s,%s,%s,%s) ON CONFLICT DO NOTHING',
+                (self.namespace, run_id, invocation_id, content.decode('utf-8'), hashlib.sha256(content).hexdigest()))
+            existing = await self._read_receipt(conn, run_id, invocation_id)
+            if encode_document(existing) != content:
+                raise ReceiptConflict('immutable receipt differs from committed evidence')
+        self._mutate(run_id, lease, save)
+
+    def close(self):
+        if self._pid != os.getpid():
+            raise StoreBusy('construct database stores after worker fork')
+        with self._mutex:
+            if self._sessions or self._reservations:
+                raise StoreBusy('drain active executions before closing the store')
+            self._transport.close()
+
+    def __enter__(self):
+        self._transport._check_process()
+        return self
+
+    def __exit__(self, exc_type, exc, traceback):
+        try:
+            self.close()
+        except Exception:
+            if exc is None:
+                raise
+            exc.add_note('database store cleanup failed; drain pending operations')