MikrotikManager/backend/src/services/partition-write-transaction.ts

import { drizzle, type NodePgDatabase } from "drizzle-orm/node-postgres"
import type { PoolClient } from "pg"
import * as schema from "../db/schema.js"

// Independent of the optional trigger proposal/migration and maintenance
// runtime: collectors can acquire this key before their first relation access.
export const PARTITION_WRITER_ADVISORY_KEY = 8723110
const admittedClients = new WeakMap<object, () => boolean>()
/** Identity proof for the guarded client owned by this module. A caller's
 * query property or serialized marker can never manufacture this capability. */
export function isPartitionWriteClient(value: unknown): value is PoolClient {
  return !!value && typeof value === "object" && admittedClients.get(value)?.() === true
}
export interface PartitionWriteOptions {
  authorize?: () => Promise<void>
  signal?: AbortSignal
  deadlineAt?: number
  connect?: () => Promise<PoolClient>
}
export interface PartitionWriteOwner { client: PoolClient; db: NodePgDatabase<typeof schema> }
export class PartitionWriteError extends Error {
  readonly automaticReplay = false
  readonly code: "PARTITION_WRITE_UNAVAILABLE" | "PARTITION_WRITE_UNKNOWN"
  constructor(readonly outcomeUnknown = false) {
    super(outcomeUnknown ? "Partition write outcome unknown; automatic replay forbidden" : "Partition write transaction unavailable or closed")
    this.code = outcomeUnknown ? "PARTITION_WRITE_UNKNOWN" : "PARTITION_WRITE_UNAVAILABLE"
  }
}
export const isPartitionWriteOutcomeUnknown = (error: unknown): error is PartitionWriteError => error instanceof PartitionWriteError && error.outcomeUnknown
const boundary = /^\s*(?:BEGIN|COMMIT|END|ROLLBACK|SAVEPOINT|RELEASE|START\s+TRANSACTION|PREPARE\s+TRANSACTION)\b/i
function settings(options: PartitionWriteOptions) {
  const now = Date.now(), deadlineAt = options.deadlineAt ?? now + 30000
  if (!Number.isSafeInteger(deadlineAt) || deadlineAt <= now || deadlineAt - now > 120000 || options.signal?.aborted) throw new PartitionWriteError()
  return { ...options, deadlineAt }
}
function remaining(deadlineAt: number) {
  const value = deadlineAt - Date.now()
  if (value <= 0) throw new PartitionWriteError()
  return Math.max(1, Math.min(5000, value))
}

/** Caller MUST already own an open SQL transaction on this exact backend and
 * MUST call before any partition/table query or secondary writer lock. This
 * helper never starts/commits a caller's transaction and never retries SQL.
 * Its bounded server timeout also applies when a surrounding Drizzle tx owns
 * the lifecycle. The surrounding owner must roll back/destroy on rejection.
 */
export async function admitPartitionWriter(client: { query(text: string, values?: [number]): Promise<unknown> }, options: Omit<PartitionWriteOptions, "connect"> = {}): Promise<void> {
  const current = settings(options)
  const check = async () => {
    if (current.signal?.aborted) throw new PartitionWriteError()
    remaining(current.deadlineAt); await current.authorize?.()
    if (current.signal?.aborted) throw new PartitionWriteError()
    remaining(current.deadlineAt)
  }
  await check()
  await client.query(`SET LOCAL lock_timeout='${remaining(current.deadlineAt)}ms'`)
  await check()
  await client.query(`SET LOCAL statement_timeout='${remaining(current.deadlineAt)}ms'`)
  await check()
  await client.query("SELECT pg_advisory_xact_lock_shared($1::bigint)", [PARTITION_WRITER_ADVISORY_KEY])
  await check()
}

/** Owns exactly one checked-out backend for raw SQL AND Drizzle. Admission is
 * BEGIN → bounded LOCAL timeouts → shared writer key → target relation SQL.
 * Query handles expire before COMMIT; detached work cannot use a pooled client
 * after release. Any failed/uncertain boundary destroys the owner. No retries
 * and no ROLLBACK is sent after an uncertain BEGIN/COMMIT acknowledgement.
 */
export async function withPartitionWrite<T>(work: (owner: PartitionWriteOwner) => Promise<T>, options: PartitionWriteOptions = {}): Promise<T> {
  const current = settings(options)
  let client: PoolClient | undefined, released = false, active = false, accepting = false
  let beginSent = false, beginAcknowledged = false, commitSent = false, committed = false, ended = false
  let queryFailed = false
  const inflight = new Set<Promise<unknown>>()
  const release = (destroy: boolean) => { active = false; accepting = false; if (client && !released) { released = true; client.release(destroy) } }
  const check = async () => {
    if (ended || current.signal?.aborted) throw new PartitionWriteError(beginSent && !beginAcknowledged || commitSent)
    remaining(current.deadlineAt); await Promise.race([current.authorize?.() ?? Promise.resolve(), stopped])
    if (ended || current.signal?.aborted) throw new PartitionWriteError(beginSent && !beginAcknowledged || commitSent)
    remaining(current.deadlineAt)
  }
  let stop!: (error: PartitionWriteError) => void
  const stopped = new Promise<never>((_resolve, reject) => { stop = reject })
  void stopped.catch(() => {})
  const destroy = () => { ended = true; release(true); stop(new PartitionWriteError(beginSent && !beginAcknowledged || commitSent)) }
  const timer = setTimeout(destroy, Math.max(1, current.deadlineAt - Date.now()))
  current.signal?.addEventListener("abort", destroy, { once: true })
  try {
    const connect = current.connect ?? (async () => (await import("../db/index.js")).pool.connect())
    const pending = connect()
    client = await Promise.race([pending, stopped, new Promise<never>((_resolve, reject) => {
      const admissionTimer = setTimeout(() => reject(new PartitionWriteError()), Math.min(5000, remaining(current.deadlineAt)))
      pending.finally(() => clearTimeout(admissionTimer)).catch(() => {})
    })]).catch(error => { ended = true; void pending.then(late => late.release(true), () => {}); throw error })
    if (ended) { release(true); throw new PartitionWriteError() }
    await check()
    beginSent = true
    await Promise.race([client.query("BEGIN"), stopped])
    beginAcknowledged = true
    await check()
    await Promise.race([admitPartitionWriter(client, current), stopped])
    active = true; accepting = true
    const rawClient = client
    const guardedQuery = (...args: unknown[]) => {
      const callback = typeof args.at(-1) === "function" ? args.pop() as (error: Error | null, value?: unknown) => void : undefined
      const run = async () => {
        if (!active || !accepting || released) throw new PartitionWriteError()
        const text = typeof args[0] === "string" ? args[0] : args[0] && typeof args[0] === "object" && "text" in args[0] ? String(args[0].text) : ""
        // Scope clients only run application statements. The owner keeps all
        // transaction boundaries and cannot be escaped through Drizzle SQL.
        const statement = text.replace(/\/\*[\s\S]*?\*\/|--[^\r\n]*/g, " ").trim()
        if (!statement || boundary.test(statement) || /;\s*\S/.test(statement)) throw new PartitionWriteError()
        await check()
        if (!active || !accepting || released) throw new PartitionWriteError()
        const result = await Reflect.apply(rawClient.query, rawClient, args)
        await check()
        return result
      }
      const pending = run(); inflight.add(pending)
      void pending.then(() => inflight.delete(pending), () => { queryFailed = true; inflight.delete(pending) })
      if (callback) { void pending.then(value => callback(null, value), error => callback(error)); return }
      return pending
    }
    const guarded = new Proxy(rawClient, { get(target, property) {
      if (property === "query") return guardedQuery
      if (property === "release" || property === "end" || property === "connect" || property === "connection") throw new PartitionWriteError()
      const value = Reflect.get(target, property, target)
      return typeof value === "function" ? value.bind(target) : value
    } })
    const owner: PartitionWriteOwner = { client: guarded, db: drizzle(guarded, { schema }) }
    admittedClients.set(guarded, () => active && accepting && !released && !ended && current.deadlineAt > Date.now() && !current.signal?.aborted)
    const result = await Promise.race([work(owner), stopped])
    accepting = false
    // Await already-started work, then close both SQL handles BEFORE COMMIT.
    await Promise.race([Promise.all([...inflight]), stopped]); active = false
    if (queryFailed) throw new PartitionWriteError()
    await check(); commitSent = true
    const commitResult = await Promise.race([client.query("COMMIT"), stopped])
    if (commitResult.command !== "COMMIT") throw new PartitionWriteError(true)
    committed = true
    await check(); release(false)
    return result
  } catch (error) {
    const unknown = beginSent && !beginAcknowledged || commitSent || committed
    destroy()
    if (error instanceof PartitionWriteError && error.outcomeUnknown === unknown) throw error
    throw new PartitionWriteError(unknown)
  } finally {
    clearTimeout(timer); current.signal?.removeEventListener("abort", destroy)
    if (!released) release(true)
  }
}