@gencow/core 0.1.21 → 0.1.23

package/src/rls-db.ts

@@ -1,27 +1,294 @@
+import { AsyncLocalStorage } from "node:async_hooks";
 import { sql } from "drizzle-orm";
 import type { PgDatabase } from "drizzle-orm/pg-core";
 
 /**
- * Wraps Drizzle so `transaction()` runs `set_config('app.current_user_id', ...)` first (RLS policies).
- * Supabase-style session GUC injection.
+ * RLS DB wrapper — execution paths for `withRlsConnection`:
+ * 1. **Reuse outer Drizzle transaction** (`reuseOuterConnection`): same connection, apply GUCs then run `fn`.
+ * 2. **Bun SQL** (`session.client.begin`): short transaction around `fn`.
+ * 3. **PGlite / drivers with `transaction()`**: delegate to driver helper.
+ * 4. **node-pg Pool** (`connect` + `query`): lease a client, BEGIN → set_config → `fn` → COMMIT, or ROLLBACK on error, then `release()`.
+ */
+
+/**
+ * Values pushed into PostgreSQL session GUCs (`set_config(..., true)`) for RLS / tenant checks.
+ * Optional fields become `app.*` settings; add arbitrary allowed keys via `vars`.
+ */
+export type RlsSessionContext = {
+    /** `app.current_user_id` — required by `ownerRls()` policies in this repo. */
+    userId: string;
+    /** When set: `app.current_user_role`. */
+    role?: string;
+    /** When set: `app.tenant_id`. */
+    tenantId?: string;
+    /**
+     * Extra transaction-local settings: keys must be validated `app.*` GUC names
+     * and must not duplicate `app.current_user_id` / `app.current_user_role` / `app.tenant_id`
+     * (use the top-level fields for those).
+     */
+    vars?: Record<string, string>;
+};
+
+const gucNameRe = /^app\.[a-z][a-z0-9_]*(?:\.[a-z][a-z0-9_]*)*$/;
+
+const RESERVED_VARS_KEYS = new Set([
+    "app.current_user_id",
+    "app.current_user_role",
+    "app.tenant_id",
+]);
+
+function assertSafeGucName(key: string): void {
+    if (!gucNameRe.test(key)) {
+        throw new Error(
+            `createRlsDb: GUC name "${key}" is invalid — use lowercase app.* names (e.g. app.org_id)`,
+        );
+    }
+}
+
+/** Ordered `(name, value)` pairs for `set_config(name, value, true)`. */
+function rlsSetConfigPairs(rls: RlsSessionContext): [string, string][] {
+    const pairs: [string, string][] = [["app.current_user_id", rls.userId]];
+    if (rls.role !== undefined) {
+        pairs.push(["app.current_user_role", rls.role]);
+    }
+    if (rls.tenantId !== undefined) {
+        pairs.push(["app.tenant_id", rls.tenantId]);
+    }
+    if (rls.vars) {
+        for (const [key, value] of Object.entries(rls.vars)) {
+            assertSafeGucName(key);
+            if (RESERVED_VARS_KEYS.has(key)) {
+                throw new Error(
+                    `createRlsDb: vars must not set "${key}" — use userId, role, or tenantId on the context object`,
+                );
+            }
+            pairs.push([key, value]);
+        }
+    }
+    return pairs;
+}
+
+async function forEachSetConfig(
+    rls: RlsSessionContext,
+    setOne: (name: string, value: string) => Promise<void>,
+): Promise<void> {
+    for (const [name, value] of rlsSetConfigPairs(rls)) {
+        await setOne(name, value);
+    }
+}
+
+/**
+ * While a prepared query runs, nested queries reuse this client (same outer `begin` / tx).
+ */
+const rlsExecClient = new AsyncLocalStorage<unknown>();
+
+function isDrizzleTransactionDb(db: unknown): boolean {
+    const d = db as { constructor?: { name?: string }; rollback?: unknown };
+    if (typeof d?.rollback === "function") {
+        return true;
+    }
+    const name = d?.constructor?.name;
+    return typeof name === "string" && name.includes("Transaction");
+}
+
+async function execSetConfig(client: any, name: string, value: string): Promise<void> {
+    if (typeof client?.unsafe === "function") {
+        await client.unsafe(`select set_config($1::text, $2::text, true)`, [name, value]);
+        return;
+    }
+    if (typeof client?.query === "function") {
+        await client.query(`select set_config($1::text, $2::text, true)`, [name, value]);
+        return;
+    }
+    throw new Error(
+        "createRlsDb: unsupported SQL driver (expected Bun SQL, node-pg client/pool, or PGlite)",
+    );
+}
+
+async function applyRlsSessionVars(client: any, rls: RlsSessionContext): Promise<void> {
+    await forEachSetConfig(rls, (name, value) => execSetConfig(client, name, value));
+}
+
+async function injectRlsVarsOnTx(tx: any, rls: RlsSessionContext): Promise<void> {
+    await forEachSetConfig(rls, (name, value) =>
+        tx.execute(sql`SELECT set_config(${name}, ${value}, true)`),
+    );
+}
+
+/**
+ * pg `Pool.connect()`-style client: BEGIN → apply RLS GUCs → `fn` → COMMIT on success,
+ * or ROLLBACK if anything fails (including failed COMMIT). Always `release()` in `finally`.
+ *
+ * @internal Exported only for unit tests — not re-exported from `@gencow/core` public API.
+ */
+export async function withRlsLeasedConnection<T>(
+    leased: { query: (sql: string) => Promise<unknown>; release: () => void },
+    rls: RlsSessionContext,
+    fn: (client: { query: (sql: string) => Promise<unknown>; release: () => void }) => Promise<T>,
+): Promise<T> {
+    try {
+        await leased.query("begin");
+        await applyRlsSessionVars(leased, rls);
+        const result = await rlsExecClient.run(leased, () => fn(leased));
+        await leased.query("commit");
+        return result;
+    } catch (e) {
+        try {
+            await leased.query("rollback");
+        } catch {
+            /* ignore */
+        }
+        throw e;
+    } finally {
+        leased.release();
+    }
+}
+
+/**
+ * Runs `fn` with a connection that has transaction-local RLS session variables set.
+ * Pool / autocommit: opens a short transaction (Bun `begin`, PGlite `transaction`).
+ * When `db` is already a Drizzle transaction object, reuses that connection (no nested top-level tx).
+ */
+async function withRlsConnection<T>(
+    session: any,
+    rls: RlsSessionContext,
+    reuseOuterConnection: boolean,
+    fn: (client: any) => Promise<T>,
+): Promise<T> {
+    if (reuseOuterConnection) {
+        const c = session.client;
+        return rlsExecClient.run(c, async () => {
+            await applyRlsSessionVars(c, rls);
+            return fn(c);
+        });
+    }
+    const c = session.client;
+    const runInner = async (client: any) => {
+        await applyRlsSessionVars(client, rls);
+        return rlsExecClient.run(client, () => fn(client));
+    };
+    if (typeof c.begin === "function") {
+        return c.begin(runInner);
+    }
+    if (typeof c.transaction === "function") {
+        return c.transaction(runInner);
+    }
+    if (typeof c.connect === "function" && typeof c.query === "function") {
+        const leased = await c.connect();
+        return withRlsLeasedConnection(leased, rls, fn);
+    }
+    return runInner(c);
+}
+
+function wrapPreparedQuery(
+    pq: any,
+    session: any,
+    rls: RlsSessionContext,
+    reuseOuterConnection: boolean,
+) {
+    const origExecute = pq.execute.bind(pq);
+    const origAll = pq.all.bind(pq);
+
+    pq.execute = async (placeholderValues?: unknown) => {
+        const active = rlsExecClient.getStore();
+        if (active) {
+            const prev = pq.client;
+            pq.client = active;
+            try {
+                return await origExecute(placeholderValues);
+            } finally {
+                pq.client = prev;
+            }
+        }
+        return withRlsConnection(session, rls, reuseOuterConnection, async (client) => {
+            const prev = pq.client;
+            pq.client = client;
+            try {
+                return await origExecute(placeholderValues);
+            } finally {
+                pq.client = prev;
+            }
+        });
+    };
+
+    pq.all = async (placeholderValues?: unknown) => {
+        const active = rlsExecClient.getStore();
+        if (active) {
+            const prev = pq.client;
+            pq.client = active;
+            try {
+                return await origAll(placeholderValues);
+            } finally {
+                pq.client = prev;
+            }
+        }
+        return withRlsConnection(session, rls, reuseOuterConnection, async (client) => {
+            const prev = pq.client;
+            pq.client = client;
+            try {
+                return await origAll(placeholderValues);
+            } finally {
+                pq.client = prev;
+            }
+        });
+    };
+}
+
+function wrapSession(session: any, rls: RlsSessionContext, reuseOuterConnection: boolean) {
+    return new Proxy(session, {
+        get(sTarget, sProp, sRecv) {
+            if (sProp === "prepareQuery") {
+                return (...args: unknown[]) => {
+                    const pq = sTarget.prepareQuery(...args);
+                    wrapPreparedQuery(pq, sTarget, rls, reuseOuterConnection);
+                    return pq;
+                };
+            }
+            const v = Reflect.get(sTarget, sProp, sRecv);
+            return typeof v === "function" ? v.bind(sTarget) : v;
+        },
+    });
+}
+
+/**
+ * RLS-scoped Drizzle handle: every query path (`select`, `insert`, `execute`, …) runs
+ * `set_config` for the given context in the same DB transaction as the query
+ * (PgBouncer–safe transaction-local GUCs).
  *
- * `set_config(..., true)` is transaction-local (safe with PgBouncer transaction pooling).
+ * `db.transaction()` still injects the same variables at the start of the callback transaction.
  */
-export function createRlsDb(db: PgDatabase<any, any, any>, userId: string) {
+export function createRlsDb(
+    db: PgDatabase<any, any, any>,
+    rls: RlsSessionContext,
+): PgDatabase<any, any, any> {
+    const reuseOuterConnection = isDrizzleTransactionDb(db);
+    const baseSession = (db as unknown as { session: any }).session;
+    const wrappedSession = wrapSession(baseSession, rls, reuseOuterConnection);
+
     return new Proxy(db, {
         get(target, prop, receiver) {
+            if (prop === "session") {
+                return wrappedSession;
+            }
+            if (prop === "_") {
+                const inner = target._;
+                return new Proxy(inner, {
+                    get(i, p, r) {
+                        if (p === "session") return wrappedSession;
+                        return Reflect.get(i, p, r);
+                    },
+                });
+            }
             if (prop === "transaction") {
                 return async (callback: any, ...rest: any[]) => {
-                    return await target.transaction(async (tx) => {
-                        await tx.execute(
-                            sql`SELECT set_config('app.current_user_id', ${userId}, true)`
-                        );
+                    return await target.transaction(async (tx: any) => {
+                        await injectRlsVarsOnTx(tx, rls);
                         return await callback(tx);
                     }, ...rest as any);
                 };
             }
             const value = Reflect.get(target, prop, receiver);
-            return typeof value === "function" ? value.bind(target) : value;
-        }
+            return typeof value === "function" ? value.bind(receiver) : value;
+        },
     });
 }

package/src/rls.ts

@@ -43,7 +43,7 @@ export function registerOwnerRls(table: PgTable, meta: OwnerRlsMeta): void {
  *   모든 CRUD 쿼리에 WHERE userId = auth.userId 자동 주입.
  *   PGlite + PostgreSQL 양쪽에서 동작.
  * - Layer 2 (DB 레벨): PostgreSQL RLS 정책으로 심층 방어.
- *   createRlsDb()의 transaction() 내에서 set_config 주입.
+ *   createRlsDb()가 RlsSessionContext로 set_config 주입.
  *
  * @example
  * ```ts

package/src/scheduler.ts

@@ -10,6 +10,30 @@ export interface ScheduleOptions {
     onError?: string;
 }
 
+/** scheduled job의 DB 영속화를 위한 콜백 */
+export interface ScheduledJobRecord {
+    id: string;
+    action: string;
+    args: unknown;
+    runAt: Date;
+    onErrorAction?: string;
+}
+
+/**
+ * createScheduler 옵션.
+ *
+ * persistJob/removeJob이 주어지면 runAfter()가 setTimeout 대신 DB에 영속화.
+ * 외부 폴러(Platform Cron Runner)가 주기적으로 due된 job을 조회하여 실행.
+ *
+ * 콜백 미설정 시 기존 인메모리 setTimeout 동작 유지 (로컬 dev).
+ */
+export interface CreateSchedulerOptions {
+    /** Job을 DB에 영속화 (INSERT) */
+    persistJob?: (job: ScheduledJobRecord) => Promise<void>;
+    /** Job을 DB에서 제거 (DELETE) — cancel/완료 시 */
+    removeJob?: (jobId: string) => Promise<boolean>;
+}
+
 /** 실패한 작업의 기록 */
 export interface FailedJob {
     id: string;
@@ -39,18 +63,26 @@ export interface Scheduler {
 /**
  * Create a scheduler instance — Convex scheduler 패턴 재현
  *
+ * @param options  persistJob/removeJob 콜백을 전달하면 durable mode 활성화.
+ *                 BaaS 모드에서는 서버가 DB 콜백을 전달하여 sleep-safe 보장.
+ *                 로컬 dev에서는 콜백 없이 호출하여 기존 setTimeout 동작.
+ *
  * @example
+ * // 로컬 dev (인메모리)
  * const scheduler = createScheduler();
  *
+ * // BaaS (DB-backed durable)
+ * const scheduler = createScheduler({
+ *     persistJob: async (job) => { await rawSql('INSERT INTO ...', ...) },
+ *     removeJob: async (id) => { await rawSql('DELETE FROM ...', [id]) },
+ * });
+ *
  * // Register actions
  * scheduler.registerAction('emails.send', async (args) => { ... });
  *
- * // Schedule (Convex-style)
+ * // Schedule (Convex-style) — durable mode에서는 DB에 영속화
  * scheduler.runAfter(5 * 60 * 1000, 'emails.send', { to: 'user@test.com' });
  *
- * // Schedule with error callback (dead-letter 패턴)
- * scheduler.runAfter(0, 'pipeline.step2', args, { onError: 'pipeline.onStepError' });
- *
  * // Cron (Convex-style)
  * scheduler.cron('daily-cleanup', '0 2 * * *', async () => { ... });
  */
@@ -84,8 +116,9 @@ export function getSchedulerInfo() {
     };
 }
 
-export function createScheduler(): Scheduler {
+export function createScheduler(options?: CreateSchedulerOptions): Scheduler {
     const timers = new Map<string, NodeJS.Timeout>();
+    const isDurable = !!(options?.persistJob && options?.removeJob);
     const cronJobs = new Map<string, cron.ScheduledTask>();
     const actions = new Map<string, ActionHandler>();
 
@@ -120,52 +153,80 @@ export function createScheduler(): Scheduler {
         console.log(`[scheduler] Action "${action}" completed`);
     }
 
+    /** @internal 인메모리 setTimeout 기반 스케줄링 (로컬 dev + durable fallback) */
+    function scheduleInMemory(id: string, ms: number, action: string, args: any, scheduleOpts: ScheduleOptions | undefined, jobEntry: PendingJobEntry): void {
+        const timer = setTimeout(async () => {
+            jobEntry.status = "running";
+            try {
+                await executeAction(action, args);
+                jobEntry.status = "completed";
+            } catch (error) {
+                const err = error instanceof Error ? error : new Error(String(error));
+                jobEntry.status = "failed";
+                console.error(`[scheduler] Action "${action}" failed (job: ${id}):`, err.message);
+
+                // 실패 기록 보관 (dead-letter)
+                recordFailure(id, action, args, err);
+
+                // onError 콜백 실행 — 다른 action에 에러 정보 전달
+                if (scheduleOpts?.onError) {
+                    try {
+                        await executeAction(scheduleOpts.onError, {
+                            failedAction: action,
+                            failedJobId: id,
+                            error: err.message,
+                            originalArgs: args,
+                        });
+                        console.log(`[scheduler] onError handler "${scheduleOpts.onError}" completed for "${action}"`);
+                    } catch (onErrorErr) {
+                        console.error(`[scheduler] onError handler "${scheduleOpts.onError}" also failed:`, onErrorErr);
+                    }
+                }
+            } finally {
+                timers.delete(id);
+                // completed/failed 기록은 잠시 유지 후 정리 (status 조회 가능하도록)
+                setTimeout(() => {
+                    const idx = _pendingJobs.findIndex((j) => j.id === id);
+                    if (idx >= 0) _pendingJobs.splice(idx, 1);
+                }, 60_000); // 1분 후 정리
+            }
+        }, ms);
+        timers.set(id, timer);
+        console.log(
+            `[scheduler] Scheduled "${action}" to run after ${ms}ms (id: ${id})${scheduleOpts?.onError ? ` [onError: ${scheduleOpts.onError}]` : ""}`
+        );
+    }
+
     return {
-        runAfter(ms: number, action: string, args?: any, options?: ScheduleOptions): string {
+        runAfter(ms: number, action: string, args?: any, scheduleOpts?: ScheduleOptions): string {
             const id = generateId();
             const jobEntry: PendingJobEntry = { id, action, scheduledAt: new Date().toISOString(), status: "pending" };
             _pendingJobs.push(jobEntry);
 
-            const timer = setTimeout(async () => {
-                jobEntry.status = "running";
-                try {
-                    await executeAction(action, args);
-                    jobEntry.status = "completed";
-                } catch (error) {
-                    const err = error instanceof Error ? error : new Error(String(error));
-                    jobEntry.status = "failed";
-                    console.error(`[scheduler] Action "${action}" failed (job: ${id}):`, err.message);
-
-                    // 실패 기록 보관 (dead-letter)
-                    recordFailure(id, action, args, err);
-
-                    // onError 콜백 실행 — 다른 action에 에러 정보 전달
-                    if (options?.onError) {
-                        try {
-                            await executeAction(options.onError, {
-                                failedAction: action,
-                                failedJobId: id,
-                                error: err.message,
-                                originalArgs: args,
-                            });
-                            console.log(`[scheduler] onError handler "${options.onError}" completed for "${action}"`);
-                        } catch (onErrorErr) {
-                            console.error(`[scheduler] onError handler "${options.onError}" also failed:`, onErrorErr);
-                        }
-                    }
-                } finally {
-                    timers.delete(id);
-                    // completed/failed 기록은 잠시 유지 후 정리 (status 조회 가능하도록)
-                    setTimeout(() => {
-                        const idx = _pendingJobs.findIndex((j) => j.id === id);
-                        if (idx >= 0) _pendingJobs.splice(idx, 1);
-                    }, 60_000); // 1분 후 정리
-                }
-            }, ms);
-            timers.set(id, timer);
-            console.log(
-                `[scheduler] Scheduled "${action}" to run after ${ms}ms (id: ${id})${options?.onError ? ` [onError: ${options.onError}]` : ""}`
-            );
+            // ── Durable mode: DB에 영속화, 실행은 외부 폴러에 위임 ──
+            if (isDurable) {
+                const runAt = new Date(Date.now() + ms);
+                options!.persistJob!({
+                    id,
+                    action,
+                    args: args ?? {},
+                    runAt,
+                    onErrorAction: scheduleOpts?.onError,
+                }).then(() => {
+                    console.log(
+                        `[scheduler] Persisted "${action}" to run at ${runAt.toISOString()} (id: ${id}, durable)` +
+                        `${scheduleOpts?.onError ? ` [onError: ${scheduleOpts.onError}]` : ""}`
+                    );
+                }).catch((err) => {
+                    console.error(`[scheduler] Failed to persist job ${id}:`, err instanceof Error ? err.message : err);
+                    // DB persist 실패 → 인메모리 fallback
+                    scheduleInMemory(id, ms, action, args, scheduleOpts, jobEntry);
+                });
+                return id;
+            }
+
+            // ── 인메모리 mode: 기존 setTimeout 동작 ──
+            scheduleInMemory(id, ms, action, args, scheduleOpts, jobEntry);
             return id;
         },
 
@@ -177,6 +238,23 @@ export function createScheduler(): Scheduler {
         },
 
         cancel(jobId: string): boolean {
+            // ── Durable mode: DB에서도 제거 ──
+            if (isDurable) {
+                const idx = _pendingJobs.findIndex((j) => j.id === jobId);
+                if (idx >= 0) {
+                    _pendingJobs.splice(idx, 1);
+                    options!.removeJob!(jobId).catch((err) => {
+                        console.error(`[scheduler] Failed to remove persisted job ${jobId}:`, err instanceof Error ? err.message : err);
+                    });
+                    console.log(`[scheduler] Cancelled job ${jobId} (durable)`);
+                    return true;
+                }
+                // pendingJobs에 없어도 DB에는 있을 수 있으므로 삭제 시도
+                options!.removeJob!(jobId).catch(() => {});
+                return false;
+            }
+
+            // ── 인메모리 mode ──
             const timer = timers.get(jobId);
             if (timer) {
                 clearTimeout(timer);

package/src/server.ts

@@ -5,8 +5,8 @@
  * executing server. Excluded from client-side core (`index.ts`) so they aren't
  * bundled into user functions which run in Firecracker.
  */
-export { createDb } from "./db";
-export { createStorage, storageRoutes } from "./storage";
-export type { StorageImageTierConfig } from "./storage";
-export { createScheduler, getSchedulerInfo } from "./scheduler";
-export { authMiddleware, authRoutes, getUsers } from "./auth";
+export { createDb } from "./db.js";
+export { createStorage, storageRoutes } from "./storage.js";
+export type { StorageImageTierConfig } from "./storage.js";
+export { createScheduler, getSchedulerInfo } from "./scheduler.js";
+export { authMiddleware, authRoutes, getUsers } from "./auth.js";