route-commerce/db/client.ts

/**
 * Drizzle client + brand-scoped query helper.
 *
 * The app connects to Postgres directly via the `pg` driver. Drizzle sits
 * on top, providing typed queries. The `withBrand` wrapper is the only
 * sanctioned way to run a brand-scoped query — it sets the
 * `app.current_brand_id` GUC transaction-locally, and the database's
 * RLS policies enforce brand isolation even if application code forgets
 * a `WHERE brand_id = $1`.
 *
 * Usage (read):
 *   const products = await withBrand(brandId, (db) =>
 *     db.select().from(productsTable).where(eq(productsTable.active, true)),
 *   );
 *
 * Usage (platform admin — sees all brands):
 *   const allBrands = await withPlatformAdmin((db) =>
 *     db.select().from(brandsTable),
 *   );
 *
 * Usage (no brand — for the rare case the query isn't brand-scoped):
 *   const plans = await withDb((db) => db.select().from(plansTable));
 */

import "server-only";
import { Pool, type PoolClient } from "pg";
import { drizzle, type NodePgDatabase } from "drizzle-orm/node-postgres";
import * as brands from "./schema/brands";
import * as billing from "./schema/billing";
import * as products from "./schema/products";
import * as stops from "./schema/stops";
import * as customers from "./schema/customers";
import * as orders from "./schema/orders";
import * as brand from "./schema/brand";
import * as wholesale from "./schema/wholesale";
import * as waterLog from "./schema/water-log";
import * as communications from "./schema/communications";
import * as marketing from "./schema/marketing";
import * as timeTracking from "./schema/time-tracking";
import * as shipping from "./schema/shipping";
import * as support from "./schema/support";
import * as files from "./schema/files";

const schema = {
  ...brands,
  ...billing,
  ...products,
  ...stops,
  ...customers,
  ...orders,
  ...brand,
  ...wholesale,
  ...waterLog,
  ...communications,
  ...marketing,
  ...timeTracking,
  ...shipping,
  ...support,
  ...files,
};

type Schema = typeof schema;
export type Db = NodePgDatabase<Schema>;

let _pool: Pool | null = null;

function getPool(): Pool {
  if (_pool) return _pool;
  const connectionString = process.env.DATABASE_URL;
  if (!connectionString) {
    throw new Error(
      "DATABASE_URL is not set. Add it to .env.local (see .env.example).",
    );
  }
  _pool = new Pool({
    connectionString,
    max: parseInt(process.env.PG_POOL_MAX ?? "50", 10),
    idleTimeoutMillis: parseInt(process.env.PG_POOL_IDLE_MS ?? "30000", 10),
    connectionTimeoutMillis: parseInt(
      process.env.PG_POOL_CONN_TIMEOUT_MS ?? "5000",
      10,
    ),
    allowExitOnIdle: false,
  });
  _pool.on("error", (err) => {
    console.error("[db] idle client error", err);
  });
  return _pool;
}

/**
 * Run `fn` with a Drizzle client. No brand context is set — the caller
 * is responsible for RLS bypass (e.g. for the `plans` and `add_ons` tables,
 * which are not brand-scoped). For brand-scoped reads, prefer
 * `withBrand` or `withPlatformAdmin`.
 */
export async function withDb<T>(fn: (db: Db) => Promise<T>): Promise<T> {
  const client = await getPool().connect();
  try {
    const db = drizzle(client, { schema });
    return await fn(db);
  } finally {
    client.release();
  }
}

/**
 * Run `fn` inside a transaction with the current brand id set as a
 * transaction-local GUC. RLS policies on brand-scoped tables will allow
 * reads/writes only for rows where `brand_id` matches. Pass `null` to
 * fail open (don't set the GUC) — only useful for the migrations
 * themselves, never for app code.
 */
export async function withBrand<T>(
  brandId: string,
  fn: (db: Db) => Promise<T>,
): Promise<T> {
  return runInTransaction(async (client) => {
    // set_config(setting, value, is_local) — is_local=true makes it
    // transaction-local so it auto-resets at COMMIT/ROLLBACK and never
    // leaks across pooled connections.
    await client.query("SELECT set_config('app.current_brand_id', $1, true)", [
      brandId,
    ]);
    await client.query("SELECT set_config('app.platform_admin', 'false', true)");
    const db = drizzle(client, { schema });
    return fn(db);
  });
}

/**
 * Run `fn` as platform admin. RLS policies permit access to all brands.
 * Use sparingly — typically only in the /admin/platform routes.
 */
export async function withPlatformAdmin<T>(
  fn: (db: Db) => Promise<T>,
): Promise<T> {
  return runInTransaction(async (client) => {
    await client.query("SELECT set_config('app.current_brand_id', '', true)");
    await client.query("SELECT set_config('app.platform_admin', 'true', true)");
    const db = drizzle(client, { schema });
    return fn(db);
  });
}

async function runInTransaction<T>(
  fn: (client: PoolClient) => Promise<T>,
): Promise<T> {
  const client = await getPool().connect();
  try {
    await client.query("BEGIN");
    const result = await fn(client);
    await client.query("COMMIT");
    return result;
  } catch (err) {
    try {
      await client.query("ROLLBACK");
    } catch {
      // ignore secondary rollback failure
    }
    throw err;
  } finally {
    client.release();
  }
}

export { schema };