get Pooler, HTTP, and Serverless connectors for Neon

Author: atilafassina

src/connectors/neon.ts renamed to src/connectors/neon/http.ts

@@ -5,7 +5,7 @@ import {
 } from "@neondatabase/serverless";
 import type { Connector, Primitive } from "db0";
 
-import { BoundableStatement } from "./_internal/statement.ts";
+import { BoundableStatement } from "../_internal/statement.ts";
 
 export type ConnectorOptions = {
   connectionString: string;

src/connectors/neon/index.ts

@@ -0,0 +1,82 @@
+import {
+  neon,
+  type FullQueryResults,
+  type NeonQueryFunction,
+} from "@neondatabase/serverless";
+import type { Connector, Primitive, Statement } from "db0";
+
+import { BoundableStatement } from "../_internal/statement.ts";
+
+export type ConnectorOptions = {
+  connectionString: string;
+};
+
+type InternalQuery = (
+  sql: string,
+  params?: Primitive[],
+) => Promise<FullQueryResults<false>>;
+
+export default function neonServerlessConnector(
+  opts: ConnectorOptions,
+): Connector<NeonQueryFunction<false, true>> {
+  let _client: undefined | NeonQueryFunction<false, true>;
+
+  function getClient() {
+    if (_client) {
+      return _client;
+    }
+
+    _client = neon(opts.connectionString, { fullResults: true });
+
+    return _client;
+  }
+
+  const query: InternalQuery = async (sql, params) => {
+    const client = getClient();
+
+    return client.query(normalizeParams(sql), params);
+  };
+
+  return {
+    name: "neon",
+    dialect: "postgresql",
+    getInstance: (): NeonQueryFunction<false, true> => getClient(),
+    exec: (sql: string) => query(sql),
+    prepare: (sql: string): Statement => new StatementWrapper(sql, query),
+  };
+}
+
+// // 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;
+
+  constructor(sql: string, query: InternalQuery) {
+    super();
+    this.#sql = 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/neon/ws.ts

@@ -0,0 +1,100 @@
+import {
+  neonConfig,
+  Pool as NeonPool,
+  type QueryResult,
+  type WebSocketConstructor,
+} from "@neondatabase/serverless";
+import type { Connector, Primitive } from "db0";
+
+import { BoundableStatement } from "../_internal/statement.ts";
+
+export type ConnectorOptions = {
+  connectionString: string;
+  webSocketConstructor?: WebSocketConstructor;
+};
+
+type InternalQuery = (
+  sql: string,
+  params?: Primitive[],
+) => Promise<QueryResult>;
+
+/**
+ * @description Creates a new Neon pool connector.
+ * @param opts
+ * @returns
+ */
+export default function neonPoolConnector(
+  opts: ConnectorOptions,
+): Connector<NeonPool> {
+  let _client: undefined | NeonPool | Promise<NeonPool>;
+  function getClient() {
+    if (_client) {
+      return _client;
+    }
+
+    const client = new NeonPool({ connectionString: opts.connectionString });
+    _client = client.connect().then(() => {
+      /**
+       * @description Allow to override the WebSocket constructor or provide one when platform does not support it.
+       * @see https://github.com/neondatabase/serverless?tab=readme-ov-file#pool-and-client
+       */
+      if (opts.webSocketConstructor) {
+        neonConfig.webSocketConstructor = opts.webSocketConstructor;
+      }
+
+      _client = client;
+      return _client;
+    });
+
+    return _client;
+  }
+
+  const query: InternalQuery = async (sql, params) => {
+    const client = getClient();
+
+    return (await client).query(normalizeParams(sql), params);
+  };
+
+  return {
+    name: "neon",
+    dialect: "postgresql",
+    getInstance: () => getClient(),
+    exec: (sql) => query(sql),
+    prepare: (sql) => new StatementWrapper(sql, query),
+  };
+}
+
+// // 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;
+
+  constructor(sql: string, query: InternalQuery) {
+    super();
+    this.#sql = 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];
+  }
+}

test/connectors/neon.test.ts

@@ -0,0 +1,12 @@
+import { describe } from "vitest";
+import { testConnector } from "../_tests";
+import neonClientConnector from "../../../src/connectors/neon/http";
+
+describe.runIf(process.env.NEON_URL_HTTP)("connectors: neon.test", () => {
+  testConnector({
+    dialect: "postgresql",
+    connector: neonClientConnector({
+      connectionString: process.env.NEON_URL_HTTP!,
+    }),
+  });
+});

test/connectors/neon/http.test.ts

@@ -0,0 +1,15 @@
+import { describe } from "vitest";
+import connector from "../../../src/connectors/neon";
+import { testConnector } from "../_tests";
+
+describe.runIf(process.env.NEON_URL_SERVERLESS)(
+  "connectors: neon serverless (index)",
+  () => {
+    testConnector({
+      dialect: "postgresql",
+      connector: connector({
+        connectionString: process.env.NEON_URL_SERVERLESS!,
+      }),
+    });
+  },
+);

test/connectors/neon/index.test.ts

@@ -0,0 +1,12 @@
+import { describe } from "vitest";
+import connector from "../../../src/connectors/neon/ws";
+import { testConnector } from "../_tests";
+
+describe.runIf(process.env.NEON_URL_WS)("connectors: neon pool (ws)", () => {
+  testConnector({
+    dialect: "postgresql",
+    connector: connector({
+      connectionString: process.env.NEON_URL_WS!,
+    }),
+  });
+});