feat: connection pooling (wip)

examples/drizzle/index.ts

@@ -1,7 +1,7 @@
 import { sqliteTable, text, numeric } from "drizzle-orm/sqlite-core";
 
 import { createDatabase } from "../../src";
-import { drizzle } from "../../src/integrations/drizzle"
+import { drizzle } from "../../src/integrations/drizzle";
 
 import sqlite from "../../src/connectors/better-sqlite3";
 

package.json

@@ -53,6 +53,7 @@
     "automd": "^0.4.2",
     "better-sqlite3": "^12.4.1",
     "changelogen": "^0.6.2",
+    "ciorent": "^1.0.10",
     "db0": "link:.",
     "dotenv": "^17.2.3",
     "drizzle-orm": "^0.44.7",
@@ -62,6 +63,7 @@
     "mlly": "^1.8.0",
     "mysql2": "^3.15.3",
     "obuild": "^0.4.2",
+    "pathe": "^2.0.3",
     "pg": "^8.16.3",
     "prettier": "^3.6.2",
     "scule": "^1.3.0",

scripts/gen-connectors.ts

@@ -1,15 +1,13 @@
 import { readFile, readdir, writeFile } from "node:fs/promises";
-import { join } from "node:path";
-import { fileURLToPath } from "node:url";
+import { join, resolve } from "pathe";
 import { findTypeExports } from "mlly";
 import { camelCase, upperFirst } from "scule";
 
-const connectorsDir = fileURLToPath(
-  new URL("../src/connectors", import.meta.url),
-);
+const connectorsDir = resolve(import.meta.dirname, "../src/connectors");
 
-const connectorsMetaFile = fileURLToPath(
-  new URL("../src/_connectors.ts", import.meta.url),
+const connectorsMetaFile = resolve(
+  import.meta.dirname,
+  "../src/_connectors.ts",
 );
 
 const aliases = {
@@ -38,7 +36,7 @@ async function getConnectorFiles(dir: string): Promise<string[]> {
 
 const connectorFiles = await getConnectorFiles(connectorsDir);
 const connectorEntries = connectorFiles.map((file) =>
-  file.replace(connectorsDir + "/", ""),
+  file.slice(connectorsDir.length + 1),
 );
 
 const connectors: {
@@ -52,7 +50,7 @@ const connectors: {
 
 for (const entry of connectorEntries) {
   const pathName = entry.replace(/\.ts$/, "");
-  const name = pathName.replace(/\/|\\/g, "-");
+  const name = pathName.replace(/[/\\]/g, "-");
   const subpath = `db0/connectors/${pathName}`;
   const fullPath = join(connectorsDir, `${pathName}.ts`);
 

src/_connectors.ts

@@ -11,12 +11,13 @@ import type { ConnectorOptions as LibSQLNodeOptions } from "db0/connectors/libsq
 import type { ConnectorOptions as LibSQLWebOptions } from "db0/connectors/libsql/web";
 import type { ConnectorOptions as MySQL2Options } from "db0/connectors/mysql2";
 import type { ConnectorOptions as NodeSQLiteOptions } from "db0/connectors/node-sqlite";
+import type { ConnectorOptions as PgPoolOptions } from "db0/connectors/pg-pool";
 import type { ConnectorOptions as PgliteOptions } from "db0/connectors/pglite";
 import type { ConnectorOptions as PlanetscaleOptions } from "db0/connectors/planetscale";
 import type { ConnectorOptions as PostgreSQLOptions } from "db0/connectors/postgresql";
 import type { ConnectorOptions as SQLite3Options } from "db0/connectors/sqlite3";
 
-export type ConnectorName = "better-sqlite3" | "bun-sqlite" | "bun" | "cloudflare-d1" | "cloudflare-hyperdrive-mysql" | "cloudflare-hyperdrive-postgresql" | "libsql-core" | "libsql-http" | "libsql-node" | "libsql" | "libsql-web" | "mysql2" | "node-sqlite" | "sqlite" | "pglite" | "planetscale" | "postgresql" | "sqlite3";
+export type ConnectorName = "better-sqlite3" | "bun-sqlite" | "bun" | "cloudflare-d1" | "cloudflare-hyperdrive-mysql" | "cloudflare-hyperdrive-postgresql" | "libsql-core" | "libsql-http" | "libsql-node" | "libsql" | "libsql-web" | "mysql2" | "node-sqlite" | "sqlite" | "pg-pool" | "pglite" | "planetscale" | "postgresql" | "sqlite3";
 
 export type ConnectorOptions = {
   "better-sqlite3": BetterSQLite3Options;
@@ -36,6 +37,7 @@ export type ConnectorOptions = {
   "node-sqlite": NodeSQLiteOptions;
   /** alias of node-sqlite */
   "sqlite": NodeSQLiteOptions;
+  "pg-pool": PgPoolOptions;
   "pglite": PgliteOptions;
   "planetscale": PlanetscaleOptions;
   "postgresql": PostgreSQLOptions;
@@ -60,6 +62,7 @@ export const connectors: Record<ConnectorName, string> = Object.freeze({
   "node-sqlite": "db0/connectors/node-sqlite",
   /** alias of node-sqlite */
   "sqlite": "db0/connectors/node-sqlite",
+  "pg-pool": "db0/connectors/pg-pool",
   "pglite": "db0/connectors/pglite",
   "planetscale": "db0/connectors/planetscale",
   "postgresql": "db0/connectors/postgresql",

src/connectors/_internal/postgresql.ts

@@ -0,0 +1,5 @@
+// https://www.postgresql.org/docs/9.3/sql-prepare.html
+export function normalizeParams(sql: string): string {
+  let i = 0;
+  return sql.replace(/\?/g, () => `$${++i}`);
+}

src/connectors/better-sqlite3.ts

@@ -35,6 +35,7 @@ export default function sqliteConnector(
   return {
     name: "sqlite",
     dialect: "sqlite",
+    supportsPooling: false,
     getInstance: () => getDB(),
     exec: (sql) => getDB().exec(sql),
     prepare: (sql) => new StatementWrapper(() => getDB().prepare(sql)),

src/connectors/bun-sqlite.ts

@@ -34,6 +34,7 @@ export default function bunSqliteConnector(
   return {
     name: "sqlite",
     dialect: "sqlite",
+    supportsPooling: false,
     getInstance: () => getDB(),
     exec: (sql) => getDB().exec(sql),
     prepare: (sql) => new StatementWrapper(getDB().prepare(sql)),

src/connectors/cloudflare-hyperdrive-postgresql.ts

@@ -2,6 +2,7 @@ import pg from "pg";
 
 import type { Connector, Primitive } from "db0";
 
+import { normalizeParams } from "./_internal/postgresql.ts";
 import { BoundableStatement } from "./_internal/statement.ts";
 import { getHyperdrive } from "./_internal/cloudflare.ts";
 
@@ -56,12 +57,6 @@ export default function cloudflareHyperdrivePostgresqlConnector(
   };
 }
 
-// https://www.postgresql.org/docs/9.3/sql-prepare.html
-function normalizeParams(sql: string) {
-  let i = 0;
-  return sql.replace(/\?/g, () => `$${++i}`);
-}
-
 class StatementWrapper extends BoundableStatement<void> {
   #query: InternalQuery;
   #sql: string;

src/connectors/node-sqlite.ts

@@ -41,6 +41,7 @@ export default function nodeSqlite3Connector(
   return {
     name: "node-sqlite",
     dialect: "sqlite",
+    supportsPooling: false,
     getInstance: () => getDB(),
     exec(sql: string) {
       getDB().exec(sql);

src/connectors/pg-pool.ts

@@ -0,0 +1,85 @@
+import { Pool, type PoolConfig, type QueryResult } from "pg";
+import type { Connector, Primitive } from "db0";
+import { normalizeParams } from "./_internal/postgresql.ts";
+import { BoundableStatement } from "./_internal/statement.ts";
+import type { ConnectorConnection } from "../types.ts";
+
+export type ConnectorOptions = { url: string } | PoolConfig;
+
+type InternalQuery = (
+  sql: string,
+  params?: Primitive[],
+) => Promise<QueryResult>;
+
+export default function postgresqlPoolConnector(
+  opts: ConnectorOptions,
+): Connector<Pool> {
+  let _pool: undefined | Pool;
+  /*
+    TODO: decide what behavior to use here, the non-pooling connector connects before
+          proceeding with calling functions, while `pg` wants you to run queries
+          and let the library handle creating connections via the pool when it needs to.
+          should the non-pooling connector do the same? should it just be removed altogether?
+   */
+  const getPool = () => {
+    _pool ??= new Pool("url" in opts ? { connectionString: opts.url } : opts);
+    return _pool;
+  };
+
+  const acquireConnection = async (): Promise<ConnectorConnection> => {
+    const client = await getPool().connect();
+    return {
+      dialect: "postgresql",
+      exec: (sql) => client.query(normalizeParams(sql)),
+      prepare: (sql) => new StatementWrapper(sql, client.query.bind(client)),
+      dispose: async () => client.release(),
+      [Symbol.asyncDispose]: async () => client.release(),
+    };
+  };
+
+  const dispose = async () => {
+    await _pool?.end?.();
+    _pool = undefined;
+  };
+
+  return {
+    name: "pg-pool",
+    dialect: "postgresql",
+    supportsPooling: true,
+    getInstance: () => getPool(),
+    exec: (sql) => getPool().query(normalizeParams(sql)),
+    prepare: (sql) => new StatementWrapper(sql, getPool().query.bind(_pool)),
+    acquireConnection,
+    dispose,
+    [Symbol.asyncDispose]: dispose,
+  };
+}
+
+class StatementWrapper extends BoundableStatement<void> {
+  readonly #query: InternalQuery;
+  readonly #sql: string;
+
+  constructor(sql: string, query: InternalQuery) {
+    super();
+    this.#sql = normalizeParams(sql);
+    this.#query = query;
+  }
+
+  async all(...params: Primitive[]) {
+    const res = await this.#query(this.#sql, params);
+    return res.rows;
+  }
+
+  async run(...params: Primitive[]) {
+    const res = await this.#query(this.#sql, params);
+    return {
+      success: true,
+      ...res,
+    };
+  }
+
+  async get(...params: Primitive[]) {
+    const res = await this.#query(this.#sql, params);
+    return res.rows[0];
+  }
+}

src/connectors/pglite.ts

@@ -5,6 +5,7 @@ import type {
 } from "@electric-sql/pglite";
 import { PGlite } from "@electric-sql/pglite";
 import type { Connector, Primitive } from "db0";
+import { normalizeParams } from "./_internal/postgresql.ts";
 import { BoundableStatement } from "./_internal/statement.ts";
 
 export type ConnectorOptions = PGliteOptions;
@@ -46,12 +47,6 @@ export default function pgliteConnector<TOptions extends ConnectorOptions>(
   };
 }
 
-// https://www.postgresql.org/docs/9.3/sql-prepare.html
-function normalizeParams(sql: string) {
-  let i = 0;
-  return sql.replace(/\?/g, () => `$${++i}`);
-}
-
 class StatementWrapper extends BoundableStatement<void> {
   #query: InternalQuery;
   #sql: string;

src/connectors/postgresql.ts

@@ -1,63 +1,79 @@
-import pg from "pg";
+import { type ClientConfig, Client, type QueryResult } from "pg";
+import * as mutex from "ciorent/mutex";
 
 import type { Connector, Primitive } from "db0";
 
+import { normalizeParams } from "./_internal/postgresql.ts";
 import { BoundableStatement } from "./_internal/statement.ts";
+import type { ConnectorConnection } from "../types.ts";
 
-export type ConnectorOptions = { url: string } | pg.ClientConfig;
+export type ConnectorOptions = { url: string } | ClientConfig;
 
 type InternalQuery = (
   sql: string,
   params?: Primitive[],
-) => Promise<pg.QueryResult>;
+) => Promise<QueryResult>;
 
 export default function postgresqlConnector(
   opts: ConnectorOptions,
-): Connector<pg.Client> {
-  let _client: undefined | pg.Client | Promise<pg.Client>;
+): Connector<Client> {
+  let _client: undefined | Client | Promise<Client>;
   function getClient() {
     if (_client) {
       return _client;
     }
-    const client = new pg.Client("url" in opts ? opts.url : opts);
+    const client = new Client("url" in opts ? opts.url : opts);
     _client = client.connect().then(() => {
       _client = client;
       return _client;
     });
     return _client;
   }
 
-  const query: InternalQuery = async (sql, params) => {
+  const connectionMutex = mutex.init();
+  const acquireConnection = async (): Promise<ConnectorConnection> => {
+    const releaseMutex = await mutex.acquire(connectionMutex);
+
     const client = await getClient();
-    return client.query(normalizeParams(sql), params);
+    return {
+      dialect: "postgresql",
+      exec: (sql) => client.query(normalizeParams(sql)),
+      prepare: (sql) => new StatementWrapper(sql, client.query.bind(client)),
+      dispose: async () => releaseMutex(),
+      [Symbol.asyncDispose]: async () => releaseMutex(),
+    };
+  };
+
+  const buildQueryFn =
+    (getClient: () => Client | Promise<Client>): InternalQuery =>
+    async (...params) =>
+      (await getClient()).query(...params);
+
+  const dispose = async () => {
+    (await _client)?.end?.();
+    _client = undefined;
   };
 
   return {
     name: "postgresql",
     dialect: "postgresql",
+    supportsPooling: false,
     getInstance: () => getClient(),
-    exec: (sql) => query(sql),
-    prepare: (sql) => new StatementWrapper(sql, query),
-    dispose: async () => {
-      await (await _client)?.end?.();
-      _client = undefined;
-    },
+    exec: buildQueryFn(getClient),
+    prepare: (sql) => new StatementWrapper(sql, buildQueryFn(getClient)),
+    acquireConnection,
+    dispose,
+    [Symbol.asyncDispose]: dispose,
   };
 }
 
-// https://www.postgresql.org/docs/9.3/sql-prepare.html
-function normalizeParams(sql: string) {
-  let i = 0;
-  return sql.replace(/\?/g, () => `$${++i}`);
-}
-
 class StatementWrapper extends BoundableStatement<void> {
-  #query: InternalQuery;
-  #sql: string;
+  readonly #query: InternalQuery;
+  readonly #sql: string;
 
   constructor(sql: string, query: InternalQuery) {
     super();
-    this.#sql = sql;
+    this.#sql = normalizeParams(sql);
     this.#query = query;
   }