feat: Add pg-cluster package (#240)

Author: avaly

README.md

@@ -45,6 +45,7 @@ Package Name | Version | Docs
 @databases/mysql-typed | [![NPM version](https://img.shields.io/npm/v/@databases/mysql-typed?style=for-the-badge)](https://www.npmjs.com/package/@databases/mysql-typed) | [https://www.atdatabases.org/docs/mysql-typed](https://www.atdatabases.org/docs/mysql-typed)
 @databases/pg | [![NPM version](https://img.shields.io/npm/v/@databases/pg?style=for-the-badge)](https://www.npmjs.com/package/@databases/pg) | [https://www.atdatabases.org/docs/pg](https://www.atdatabases.org/docs/pg)
 @databases/pg-bulk | [![NPM version](https://img.shields.io/npm/v/@databases/pg-bulk?style=for-the-badge)](https://www.npmjs.com/package/@databases/pg-bulk) | [https://www.atdatabases.org/docs/pg-bulk](https://www.atdatabases.org/docs/pg-bulk)
+@databases/pg-cluster | [![NPM version](https://img.shields.io/npm/v/@databases/pg-cluster?style=for-the-badge)](https://www.npmjs.com/package/@databases/pg-cluster) | [https://www.atdatabases.org/docs/pg-cluster](https://www.atdatabases.org/docs/pg-cluster)
 @databases/pg-migrations | [![NPM version](https://img.shields.io/npm/v/@databases/pg-migrations?style=for-the-badge)](https://www.npmjs.com/package/@databases/pg-migrations) | [https://www.atdatabases.org/docs/pg-migrations](https://www.atdatabases.org/docs/pg-migrations)
 @databases/pg-test | [![NPM version](https://img.shields.io/npm/v/@databases/pg-test?style=for-the-badge)](https://www.npmjs.com/package/@databases/pg-test) | [https://www.atdatabases.org/docs/pg-test](https://www.atdatabases.org/docs/pg-test)
 @databases/pg-typed | [![NPM version](https://img.shields.io/npm/v/@databases/pg-typed?style=for-the-badge)](https://www.npmjs.com/package/@databases/pg-typed) | [https://www.atdatabases.org/docs/pg-typed](https://www.atdatabases.org/docs/pg-typed)

docs/pg-cluster.md

@@ -0,0 +1,102 @@
+---
+id: pg-cluster
+title: Postgres Node.js Cluster Connection
+sidebar_label: Cluster
+---
+
+The `Cluster` object represents group of physical connections or connection pools to the underlying database in a primary & replicas setup.
+
+### `Cluster.query(SQLQuery): Promise<any[]>`
+
+Run an SQL Query and get a promise for an array of results. If your query contains multiple statements, only the results of the final statement are returned.
+
+Write queries are executed in the primary connection, and read-only queries are executed in the replica connections.
+
+```ts
+// query executed on a secondary connection
+const result = await cluster.query(sql`SELECT 1 + 1 AS a`);
+result[0].a;
+// => 2
+
+// query executed on the primary connection
+await cluster.query(sql`UPDATE users SET active = true`);
+```
+
+### `Cluster.query(SQLQuery[]): Promise<any[]>`
+
+If you pass an array of SQLQueries, you will get an array in response where each element of the array is the results of one of the queries.
+
+If there is at least one write query in the argument list, then the queries are executed in the primary connection, otherwise in the replica connections.
+
+```ts
+// query executed on a secondary connection
+const [resultA, resultB] = await cluster.query([
+  sql`SELECT 1 + 1 AS a`,
+  sql`SELECT 1 + 1 AS b`,
+]);
+resultA[0].a + resultB[0].b;
+// => 4
+```
+
+### `Cluster.queryStream(SQLQuery): AsyncIterable<any>`
+
+Run an SQL Query and get an async iterable of the results. e.g.
+
+```js
+for await (const record of cluster.queryStream(
+  sql`SELECT * FROM massive_table`,
+)) {
+  console.log(result);
+}
+```
+
+Write queries are executed in the primary connection, and read-only queries are executed in the replica connections.
+
+### `Cluster.queryNodeStream(SQLQuery): ReadableStream`
+
+Run an SQL Query and get a node.js readable stream of the results. e.g.
+
+```js
+const Stringifier = require('newline-json').Stringifier;
+
+cluster
+  .queryNodeStream(sql`SELECT * FROM massive_table`)
+  .pipe(new Stringifier())
+  .pipe(process.stdout);
+```
+
+Write queries are executed in the primary connection, and read-only queries are executed in the replica connections.
+
+### `Cluster.tx<T>(fn: (tx: Transaction) => Promise<T>, options?): Promise<T>`
+
+Executes the callback `fn` within a transaction on the primary connection or replica connections (depending on `options.readOnly`).
+
+A transaction wraps a regular task with 3 additional queries:
+
+1. it executes `BEGIN` just before invoking the callback function
+2. it executes `ROLLBACK`, if the callback throws an error or returns a rejected promise
+3. it executes `COMMIT`, if the callback does throw any error and does not return a rejected promise
+
+```ts
+// transaction executed on a replica connection
+const result = await cluster.tx(
+  async (tx) => {
+    const resultA = await tx.query(sql`SELECT 1 + 1 AS a`);
+    const resultB = await tx.query(sql`SELECT 1 + 1 AS b`);
+    return resultA[0].a + resultB[0].b;
+  },
+  {readOnly: true},
+);
+
+// transaction executed on the primary connection
+const resultPrimary = await cluster.tx(async (tx) => {
+  const resultA = await tx.query(sql`SELECT 1 + 1 AS a`);
+  const resultB = await tx.query(sql`SELECT 1 + 1 AS b`);
+  return resultA[0].a + resultB[0].b;
+});
+// => 4
+```
+
+### `Cluster.task(fn): Promise<T>`
+
+This method exists to mimic the API in `ConnectionPool.task`. It executes the `task` method on the primary connection.

docs/pg-connection-pool.md

@@ -10,7 +10,7 @@ A ConnectionPool represents a set of automatically managed physical connections
 
 Acquires a connection from the pool. If the pool is 'full' and all connections are currently checked out, this will wait in a FIFO queue until a connection becomes available by it being released back to the pool.
 
-Once a connetion has been acquired, `fn` is called with that connection.
+Once a connection has been acquired, `fn` is called with that connection.
 
 When `fn` returns, the connection is returned to the pool.
 

docs/pg-guide-connections.md

@@ -81,6 +81,25 @@ process.once('SIGTERM', () => {
 });
 ```
 
+## Cluster connections
+
+In a highly available Postgres cluster with a primary node & multiple replicas nodes (e.g. AWS RDS Aurora), write queries should be sent to the primary node, while read queries that do not require strong consistency can be distributed equally to the read replicas nodes.
+
+To connect to such a cluster you can use the `@databases/pg-cluster` package like this:
+
+```typescript
+// database.ts
+
+import createConnectionPool from '@databases/pg';
+import createCluster from '@databases/pg-cluster';
+
+const primary = createConnectionPool(process.env.MY_CUSTOM_PRIMARY_ENV_VAR);
+const replicas = [createConnectionPool(process.env.MY_CUSTOM_REPLICA_ENV_VAR)];
+
+const db = createCluster(primary, replicas);
+export default db;
+```
+
 ## Connecting to/from cloud providers
 
 You can normally follow the instructions from the cloud providers, but we have prepared the following guides to make things easier for these common platforms:

docs/pg-queryable.md

@@ -4,13 +4,14 @@ title: Postgres Queryable
 sidebar_label: Queryable
 ---
 
-There are three types of `Queryable` in postgres:
+There are four types of `Queryable` in postgres:
 
+- `Cluster` - represents a group of primary and replica connections to a database cluster
 - `ConnectionPool` - represents a set of automatically managed connections to the database
 - `Connection` - represents a single physical connection to the database
 - `Transaction` - represents a transaction (or nested transaction) on a single physical connection to the database
 
-All three share a common API, allowing you to write methods that can be used both inside and outside a transaction. e.g.
+All four share a common API, allowing you to write methods that can be used both inside and outside a transaction. e.g.
 
 ```ts
 import {Queryable} from '@databases/pg';

docs/pg-typed.md

@@ -44,6 +44,34 @@ module.exports = {users, posts};
 
 ## Table
 
+### initializer
+
+The objects returned by the `tables` function are initialized with an optional argument represeting the connection(s) for the queries.
+
+The initializer argument is a single `Queryable` (i.e. `ConnectionPool`, `Connection`, `Transaction` or `Cluster`).
+
+```typescript
+import db, {users} from './database';
+
+export async function initialize() {
+  // These 2 queries are not run in a transaction
+  await users(db).find().all();
+  await users(db).insert({
+    email: `[email protected]`,
+    favorite_color: `blue`,
+  });
+
+  await db.tx(async db => {
+    // These 2 queries are run in the same transaction
+    await users(db).find().all();
+    await users(db).insert({
+      email: `[email protected]`,
+      favorite_color: `blue`,
+    });
+  });
+}
+```
+
 ### insert(...records)
 
 Inserts records into the database table. If you pass multiple records to `insert`, they will all be added "atomically", i.e. either all of the records will be added, or none of them will be added.

docs/sidebars.json

@@ -27,6 +27,7 @@
         "ids": [
           "pg-options",
           "pg-queryable",
+          "pg-cluster",
           "pg-connection-pool",
           "pg-connection",
           "pg-transaction",

packages/pg-cluster/README.md

@@ -0,0 +1,3 @@
+# @databases/pg-cluster
+
+For documentation, see https://www.atdatabases.org/docs/pg-cluster

packages/pg-cluster/package.json

@@ -0,0 +1,25 @@
+{
+  "name": "@databases/pg-cluster",
+  "version": "0.0.0",
+  "description": "",
+  "main": "./lib/index.js",
+  "types": "./lib/index.d.ts",
+  "dependencies": {},
+  "devDependencies": {
+    "@databases/pg": "^0.0.0"
+  },
+  "peerDependencies": {
+    "@databases/pg": "*"
+  },
+  "scripts": {},
+  "repository": "https://github.com/ForbesLindesay/atdatabases/tree/master/packages/pg-cluster",
+  "bugs": "https://github.com/ForbesLindesay/atdatabases/issues",
+  "license": "MIT",
+  "publishConfig": {
+    "access": "public"
+  },
+  "files": [
+    "lib/"
+  ],
+  "homepage": "https://www.atdatabases.org/docs/pg-cluster"
+}

packages/pg-cluster/src/Cluster.ts

@@ -0,0 +1,110 @@
+import {Readable} from 'stream';
+import {
+  Connection,
+  pgFormat,
+  Queryable,
+  QueryableType,
+  Transaction,
+  sql,
+  SQLQuery,
+} from '@databases/pg';
+import AbortSignal from '@databases/pg/lib/types/AbortSignal';
+import TransactionOptions from '@databases/pg/lib/types/TransactionOptions';
+
+export default class Cluster implements Queryable {
+  public readonly type: QueryableType = QueryableType.Cluster;
+  public readonly sql = sql;
+
+  private readonly _primary: Queryable;
+  private readonly _replicas: Queryable[];
+
+  constructor(primary: Queryable, replicas: Queryable[]) {
+    if (!replicas.length) {
+      throw new Error(
+        'You must provide at least one replica when using pg-cluster',
+      );
+    }
+    this._primary = primary;
+    this._replicas = replicas;
+  }
+
+  private _getReplica(): Queryable {
+    return this._replicas[Math.floor(Math.random() * this._replicas.length)];
+  }
+
+  async query(query: SQLQuery): Promise<any[]>;
+  async query(query: SQLQuery[]): Promise<any[][]>;
+  async query(query: SQLQuery | SQLQuery[]): Promise<any[]> {
+    if (Array.isArray(query)) {
+      if (query.length === 0) {
+        return [];
+      }
+      const hasWriteableQueries = query.some(isQueryWriteable);
+      if (hasWriteableQueries) {
+        return this._primary.query(query);
+      } else {
+        return this._getReplica().query(query);
+      }
+    } else {
+      if (isQueryWriteable(query)) {
+        return this._primary.query(query);
+      } else {
+        return this._getReplica().query(query);
+      }
+    }
+  }
+
+  queryStream(
+    query: SQLQuery,
+    options: {batchSize?: number | undefined; signal?: AbortSignal | undefined},
+  ): AsyncIterable<any> {
+    if (isQueryWriteable(query)) {
+      return this._primary.queryStream(query, options);
+    } else {
+      return this._getReplica().queryStream(query, options);
+    }
+  }
+
+  queryNodeStream(
+    query: SQLQuery,
+    options?: {
+      highWaterMark?: number | undefined;
+      batchSize?: number | undefined;
+    },
+  ): Readable {
+    if (isQueryWriteable(query)) {
+      return this._primary.queryNodeStream(query, options);
+    } else {
+      return this._getReplica().queryNodeStream(query, options);
+    }
+  }
+
+  async task<T>(
+    fn: (connection: Connection | Transaction) => Promise<T>,
+  ): Promise<T> {
+    return this._primary.task(fn);
+  }
+
+  async tx<T>(
+    fn: (connection: Transaction) => Promise<T>,
+    options?: TransactionOptions,
+  ): Promise<T> {
+    if (options?.readOnly) {
+      return this._getReplica().tx(fn, options);
+    } else {
+      return this._primary.tx(fn, options);
+    }
+  }
+
+  async addPostCommitStep(fn: () => Promise<void>): Promise<void> {
+    await fn();
+  }
+}
+
+const WRITEABLE_REGEX =
+  /\b(alter|create|delete|drop|insert|truncate|update|vacuum)\b/i;
+
+function isQueryWriteable(query: SQLQuery): boolean {
+  const formatted = query.format(pgFormat);
+  return WRITEABLE_REGEX.test(formatted.text);
+}