116 lines
3.3 KiB
TypeScript
116 lines
3.3 KiB
TypeScript
import dotenv from "dotenv";
|
|
|
|
dotenv.config();
|
|
|
|
import { schemas as schema } from "@autumn/shared";
|
|
import { instrumentDrizzleClient } from "@kubiks/otel-drizzle";
|
|
|
|
import type { SQLWrapper } from "drizzle-orm";
|
|
import { drizzle } from "drizzle-orm/node-postgres";
|
|
import pg, { type PoolConfig } from "pg";
|
|
import { otelConfig } from "../utils/otel/otelConfig.js";
|
|
|
|
type AutumnDb = Omit<ReturnType<typeof drizzle<typeof schema>>, "execute"> & {
|
|
execute: <TRow = Record<string, unknown>>(
|
|
query: string | SQLWrapper,
|
|
) => Promise<TRow[]>;
|
|
};
|
|
|
|
const normalizeExecuteRows = <TRow>(result: unknown): TRow[] => {
|
|
if (result && typeof result === "object" && "rows" in result) {
|
|
return (result as { rows: TRow[] }).rows;
|
|
}
|
|
|
|
return result as TRow[];
|
|
};
|
|
|
|
const normalizeDbExecute = <
|
|
TDb extends { execute: (query: string | SQLWrapper) => Promise<unknown> },
|
|
>(
|
|
db: TDb,
|
|
) => {
|
|
const execute = db.execute.bind(db);
|
|
return Object.assign(db, {
|
|
execute: async <TRow = Record<string, unknown>>(
|
|
query: string | SQLWrapper,
|
|
) => normalizeExecuteRows<TRow>(await execute(query)),
|
|
});
|
|
};
|
|
|
|
/** Creates a Drizzle pool with the given configuration. */
|
|
export const initDrizzle = ({
|
|
maxConnections = 10,
|
|
replica = false,
|
|
connectTimeout = 5,
|
|
databaseUrl,
|
|
poolConfig,
|
|
}: {
|
|
maxConnections?: number;
|
|
replica?: boolean;
|
|
/** Connect timeout in seconds */
|
|
connectTimeout?: number | null;
|
|
databaseUrl?: string;
|
|
poolConfig?: PoolConfig;
|
|
} = {}) => {
|
|
const envDbUrl = replica
|
|
? process.env.DATABASE_REPLICA_URL
|
|
: process.env.DATABASE_URL;
|
|
|
|
const dbUrl = databaseUrl || envDbUrl || "";
|
|
|
|
const client = new pg.Pool({
|
|
connectionString: dbUrl,
|
|
...poolConfig,
|
|
max: maxConnections,
|
|
connectionTimeoutMillis:
|
|
connectTimeout === null ? undefined : connectTimeout * 1000,
|
|
});
|
|
|
|
const drizzleDb = drizzle(client, { schema });
|
|
const transaction = drizzleDb.transaction.bind(drizzleDb);
|
|
const db = normalizeDbExecute(drizzleDb) as unknown as AutumnDb;
|
|
const normalizedTransaction: typeof drizzleDb.transaction = ((fn, config) =>
|
|
transaction(
|
|
(tx) => fn(normalizeDbExecute(tx) as typeof tx),
|
|
config,
|
|
)) as typeof drizzleDb.transaction;
|
|
db.transaction = normalizedTransaction as typeof db.transaction;
|
|
|
|
if (otelConfig.drizzle) {
|
|
instrumentDrizzleClient(db);
|
|
}
|
|
|
|
return { db, client };
|
|
};
|
|
|
|
export const { db: dbCritical, client: clientCritical } = initDrizzle({
|
|
maxConnections: 5,
|
|
connectTimeout: 2,
|
|
databaseUrl: process.env.DATABASE_CRITICAL_URL,
|
|
poolConfig: {
|
|
application_name: "autumn-critical",
|
|
// Server-side timeout is configured on the DATABASE_CRITICAL_URL role.
|
|
query_timeout: 2_000,
|
|
},
|
|
});
|
|
|
|
// -- General pool: used by all other endpoints --
|
|
export const { db: dbGeneral, client: clientGeneral } = initDrizzle({
|
|
// connectTimeout: 5,
|
|
});
|
|
|
|
// -- Replica pool: used as fallback when primary is degraded --
|
|
// Only created if DATABASE_REPLICA_URL is configured.
|
|
const replicaResult = process.env.DATABASE_REPLICA_URL
|
|
? initDrizzle({ replica: true, maxConnections: 5, connectTimeout: null })
|
|
: null;
|
|
export const dbReplica = replicaResult?.db ?? null;
|
|
export const clientReplica = replicaResult?.client ?? null;
|
|
|
|
// Backward-compatible exports — existing code that imports `db` or `client`
|
|
// gets the general pool automatically.
|
|
export const client = clientGeneral;
|
|
export const db = dbGeneral;
|
|
|
|
export type DrizzleCli = ReturnType<typeof initDrizzle>["db"];
|