@prairielearn/postgres 1.6.0 → 1.7.0

package/src/pool.test.ts

@@ -1,8 +1,17 @@
 import chai from 'chai';
 import chaiAsPromised from 'chai-as-promised';
+import { Writable } from 'node:stream';
+import { pipeline } from 'node:stream/promises';
 import { z, ZodError } from 'zod';
 
-import { queryAsync, queryCursor, queryValidatedCursor } from './default-pool';
+import {
+  queryAsync,
+  queryRows,
+  queryRow,
+  queryOptionalRow,
+  queryCursor,
+  queryValidatedCursor,
+} from './default-pool';
 import { makePostgresTestUtils } from './test-utils';
 
 chai.use(chaiAsPromised);
@@ -12,24 +21,118 @@ const postgresTestUtils = makePostgresTestUtils({
   database: 'prairielearn_postgres',
 });
 
+const WorkspaceSchema = z.object({
+  id: z.string(),
+  created_at: z.date(),
+});
+
 describe('@prairielearn/postgres', function () {
   before(async () => {
     await postgresTestUtils.createDatabase();
-    await queryAsync('CREATE TABLE workspaces (id BIGSERIAL PRIMARY KEY, state TEXT);', {});
-    await queryAsync("INSERT INTO workspaces (id, state) VALUES (1,'uninitialized');", {});
-    await queryAsync("INSERT INTO workspaces (id, state) VALUES (2, 'stopped');", {});
-    await queryAsync("INSERT INTO workspaces (id, state) VALUES (3, 'launching');", {});
-    await queryAsync("INSERT INTO workspaces (id, state) VALUES (4, 'running');", {});
+    await queryAsync(
+      'CREATE TABLE workspaces (id BIGSERIAL PRIMARY KEY, created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP);',
+      {}
+    );
+    await queryAsync('INSERT INTO workspaces (id) SELECT s FROM generate_series(1, 100) AS s', {});
   });
 
   after(async () => {
     await postgresTestUtils.dropDatabase();
   });
 
+  describe('queryRows', () => {
+    it('handles single column', async () => {
+      const rows = await queryRows('SELECT id FROM workspaces WHERE id <= 10;', z.string());
+      assert.lengthOf(rows, 10);
+      assert.equal(rows[0], '1');
+    });
+
+    it('handles multiple columns', async () => {
+      const rows = await queryRows('SELECT * FROM workspaces WHERE id <= 10;', WorkspaceSchema);
+      assert.lengthOf(rows, 10);
+      assert.equal(rows[0].id, '1');
+      assert.isNotNull(rows[0].created_at);
+    });
+
+    it('handles parameters', async () => {
+      const rows = await queryRows(
+        'SELECT * FROM workspaces WHERE id <= $1;',
+        [10],
+        WorkspaceSchema
+      );
+      assert.lengthOf(rows, 10);
+    });
+  });
+
+  describe('queryRow', () => {
+    it('handles single column', async () => {
+      const row = await queryRow('SELECT id FROM workspaces WHERE id = 1;', z.string());
+      assert.equal(row, '1');
+    });
+
+    it('handles multiple columns', async () => {
+      const row = await queryRow('SELECT * FROM workspaces WHERE id = 1;', WorkspaceSchema);
+      assert.equal(row.id, '1');
+      assert.isNotNull(row.created_at);
+    });
+
+    it('handles parameters', async () => {
+      const row = await queryRow('SELECT * FROM workspaces WHERE id = $1;', [1], WorkspaceSchema);
+      assert.equal(row.id, '1');
+    });
+
+    it('rejects results with zero rows', async () => {
+      const rows = queryRow('SELECT * FROM workspaces WHERE id = -1;', WorkspaceSchema);
+      await assert.isRejected(rows, 'Incorrect rowCount: 0');
+    });
+
+    it('rejects results with multiple rows', async () => {
+      const rows = queryRow('SELECT * FROM workspaces', WorkspaceSchema);
+      await assert.isRejected(rows, 'Incorrect rowCount: 100');
+    });
+  });
+
+  describe('queryOptionalRow', () => {
+    it('handles single column', async () => {
+      const row = await queryRow('SELECT id FROM workspaces WHERE id = 1;', z.string());
+      assert.equal(row, '1');
+    });
+
+    it('handles multiple columns', async () => {
+      const row = await queryOptionalRow('SELECT * FROM workspaces WHERE id = 1;', WorkspaceSchema);
+      assert.isNotNull(row);
+      assert.equal(row?.id, '1');
+      assert.isNotNull(row?.created_at);
+    });
+
+    it('handles parameters', async () => {
+      const row = await queryOptionalRow(
+        'SELECT * FROM workspaces WHERE id = $1;',
+        [1],
+        WorkspaceSchema
+      );
+      assert.isNotNull(row);
+      assert.equal(row?.id, '1');
+    });
+
+    it('handles missing result', async () => {
+      const row = await queryOptionalRow(
+        'SELECT * FROM workspaces WHERE id = -1;',
+        WorkspaceSchema
+      );
+      assert.isNull(row);
+    });
+
+    it('rejects with multiple rows', async () => {
+      const rows = queryOptionalRow('SELECT * FROM workspaces', WorkspaceSchema);
+      await assert.isRejected(rows, 'Incorrect rowCount: 100');
+    });
+  });
+
   describe('queryCursor', () => {
     it('returns zero rows', async () => {
-      const cursor = await queryCursor('SELECT * FROM workspaces WHERE id = 5;', {});
-      let rowBatches = [];
+      const cursor = await queryCursor('SELECT * FROM workspaces WHERE id = 10000;', {});
+      const rowBatches = [];
       for await (const rows of cursor.iterate(10)) {
         rowBatches.push(rows);
       }
@@ -48,13 +151,13 @@ describe('@prairielearn/postgres', function () {
     });
 
     it('returns all rows at once', async () => {
-      const cursor = queryCursor('SELECT * FROM workspaces;', {});
+      const cursor = queryCursor('SELECT * FROM workspaces WHERE id <= 10;', {});
       const rowBatches = [];
       for await (const rows of (await cursor).iterate(10)) {
         rowBatches.push(rows);
       }
       assert.lengthOf(rowBatches, 1);
-      assert.lengthOf(rowBatches[0], 4);
+      assert.lengthOf(rowBatches[0], 10);
     });
 
     it('handles errors', async () => {
@@ -80,88 +183,111 @@ describe('@prairielearn/postgres', function () {
   });
 
   describe('queryValidatedCursor', () => {
-    it('validates with provided schema', async () => {
-      const WorkspaceSchema = z.object({
-        id: z.string(),
-      });
-      const cursor = await queryValidatedCursor(
-        'SELECT * FROM workspaces ORDER BY id ASC;',
-        {},
-        WorkspaceSchema
-      );
-      const allRows = [];
-      for await (const rows of cursor.iterate(10)) {
-        allRows.push(...rows);
-      }
-      assert.lengthOf(allRows, 4);
-      const workspace = allRows[0] as any;
-      assert.equal(workspace.id, '1');
-      assert.isUndefined(workspace.state);
+    const WorkspaceSchema = z.object({
+      id: z.string(),
     });
 
-    it('throws error when validation fails', async () => {
-      const BadWorkspaceSchema = z.object({
-        badProperty: z.string(),
-      });
-      const cursor = await queryValidatedCursor(
-        'SELECT * FROM workspaces ORDER BY id ASC;',
-        {},
-        BadWorkspaceSchema
-      );
+    const BadWorkspaceSchema = z.object({
+      badProperty: z.string(),
+    });
 
-      async function readAllRows() {
+    describe('iterator', () => {
+      it('validates with provided schema', async () => {
+        const cursor = await queryValidatedCursor(
+          'SELECT * FROM workspaces WHERE id <= 10 ORDER BY id ASC;',
+          {},
+          WorkspaceSchema
+        );
         const allRows = [];
         for await (const rows of cursor.iterate(10)) {
           allRows.push(...rows);
         }
-        return allRows;
-      }
-
-      const maybeError = await readAllRows().catch((err) => err);
-      assert.instanceOf(maybeError, ZodError);
-      assert.lengthOf(maybeError.errors, 4);
-    });
-
-    it('returns a stream', async () => {
-      const WorkspaceSchema = z.object({
-        id: z.string(),
+        assert.lengthOf(allRows, 10);
+        const workspace = allRows[0] as any;
+        assert.equal(workspace.id, '1');
+        assert.isUndefined(workspace.state);
       });
-      const cursor = await queryValidatedCursor(
-        'SELECT * FROM workspaces ORDER BY id ASC;',
-        {},
-        WorkspaceSchema
-      );
-      const stream = cursor.stream(1);
-      const allRows = [];
-      for await (const row of stream) {
-        allRows.push(row);
-      }
 
-      assert.lengthOf(allRows, 4);
-    });
+      it('throws error when validation fails', async () => {
+        const cursor = await queryValidatedCursor(
+          'SELECT * FROM workspaces WHERE id <= 10 ORDER BY id ASC;',
+          {},
+          BadWorkspaceSchema
+        );
 
-    it('emits an error when validation fails', async () => {
-      const BadWorkspaceSchema = z.object({
-        badProperty: z.string(),
+        async function readAllRows() {
+          const allRows = [];
+          for await (const rows of cursor.iterate(10)) {
+            allRows.push(...rows);
+          }
+          return allRows;
+        }
+
+        const maybeError = await readAllRows().catch((err) => err);
+        assert.instanceOf(maybeError, ZodError);
+        assert.lengthOf(maybeError.errors, 10);
       });
-      const cursor = await queryValidatedCursor(
-        'SELECT * FROM workspaces ORDER BY id ASC;',
-        {},
-        BadWorkspaceSchema
-      );
-      const stream = cursor.stream(1);
+    });
 
-      async function readAllRows() {
+    describe('stream', () => {
+      it('validates with provided schema', async () => {
+        const cursor = await queryValidatedCursor(
+          'SELECT * FROM workspaces WHERE id <= 10 ORDER BY id ASC;',
+          {},
+          WorkspaceSchema
+        );
+        const stream = cursor.stream(1);
         const allRows = [];
         for await (const row of stream) {
           allRows.push(row);
         }
-        return allRows;
-      }
 
-      const maybeError = await readAllRows().catch((err) => err);
-      assert.instanceOf(maybeError, ZodError);
-      assert.lengthOf(maybeError.errors, 1);
+        assert.lengthOf(allRows, 10);
+      });
+
+      it('emits an error when validation fails', async () => {
+        const cursor = await queryValidatedCursor(
+          'SELECT * FROM workspaces ORDER BY id ASC;',
+          {},
+          BadWorkspaceSchema
+        );
+        const stream = cursor.stream(1);
+
+        async function readAllRows() {
+          const allRows = [];
+          for await (const row of stream) {
+            allRows.push(row);
+          }
+          return allRows;
+        }
+
+        const maybeError = await readAllRows().catch((err) => err);
+        assert.instanceOf(maybeError, ZodError);
+        assert.lengthOf(maybeError.errors, 1);
+      });
+
+      it('closes the cursor when the stream is closed', async () => {
+        const cursor = await queryValidatedCursor('SELECT * FROM workspaces;', {}, WorkspaceSchema);
+        const stream = cursor.stream(1);
+
+        const rows: any[] = [];
+        const ac = new AbortController();
+        const writable = new Writable({
+          objectMode: true,
+          write: function (chunk, _encoding, callback) {
+            rows.push(chunk);
+
+            // After receiving the first row, abort the stream. This lets us test
+            // that the underlying cursor is closed. If it is *not* closed, this
+            // `after` hook will fail with a timeout.
+            ac.abort();
+            callback();
+          },
+        });
+
+        await assert.isRejected(pipeline(stream, writable, { signal: ac.signal }));
+        assert.lengthOf(rows, 1);
+      });
     });
   });
 });

package/src/pool.ts

@@ -1,7 +1,6 @@
 import _ from 'lodash';
 import pg, { QueryResult } from 'pg';
 import Cursor from 'pg-cursor';
-import path from 'node:path';
 import debugFactory from 'debug';
 import { callbackify } from 'node:util';
 import { AsyncLocalStorage } from 'node:async_hooks';
@@ -16,7 +15,7 @@ export interface CursorIterator<T> {
   stream: (batchSize: number) => NodeJS.ReadWriteStream;
 }
 
-const debug = debugFactory('prairielib:' + path.basename(__filename, '.js'));
+const debug = debugFactory('@prairielearn/postgres');
 const lastQueryMap: WeakMap<pg.PoolClient, string> = new WeakMap();
 const searchSchemaMap: WeakMap<pg.PoolClient, string> = new WeakMap();
 
@@ -239,7 +238,7 @@ export class PostgresPool {
 
     // If we're inside a transaction, we'll reuse the same client to avoid a
     // potential deadlock.
-    let client = this.alsClient.getStore() ?? (await this.pool.connect());
+    const client = this.alsClient.getStore() ?? (await this.pool.connect());
 
     // If we're configured to use a particular schema, we'll store whether or
     // not the search path has already been configured for this particular
@@ -475,7 +474,7 @@ export class PostgresPool {
   async runInTransactionAsync<T>(fn: (client: pg.PoolClient) => Promise<T>): Promise<T> {
     // Check if we're already inside a transaction. If so, we won't start another one,
     // as Postgres doesn't support nested transactions.
-    let client = this.alsClient.getStore();
+    const client = this.alsClient.getStore();
     const isNestedTransaction = client !== undefined;
     const transactionClient = client ?? (await this.beginTransactionAsync());
 
@@ -754,7 +753,7 @@ export class PostgresPool {
     model: Model
   ): Promise<z.infer<Model> | null> {
     const results = await this.queryZeroOrOneRowAsync(query, params);
-    if (results.rows.length == 0) {
+    if (results.rows.length === 0) {
       return null;
     } else {
       return model.parse(results.rows[0]);
@@ -772,7 +771,7 @@ export class PostgresPool {
     model: Model
   ): Promise<z.infer<Model>[]> {
     const results = await this.queryAsync(query, params);
-    if (results.fields.length != 1) {
+    if (results.fields.length !== 1) {
       throw new Error(`Expected one column, got ${results.fields.length}`);
     }
     const columnName = results.fields[0].name;
@@ -791,7 +790,7 @@ export class PostgresPool {
     model: Model
   ): Promise<z.infer<Model>> {
     const results = await this.queryOneRowAsync(query, params);
-    if (results.fields.length != 1) {
+    if (results.fields.length !== 1) {
       throw new Error(`Expected one column, got ${results.fields.length}`);
     }
     const columnName = results.fields[0].name;
@@ -809,10 +808,10 @@ export class PostgresPool {
     model: Model
   ): Promise<z.infer<Model> | null> {
     const results = await this.queryZeroOrOneRowAsync(query, params);
-    if (results.fields.length != 1) {
+    if (results.fields.length !== 1) {
       throw new Error(`Expected one column, got ${results.fields.length}`);
     }
-    if (results.rows.length == 0) {
+    if (results.rows.length === 0) {
       return null;
     } else {
       const columnName = results.fields[0].name;
@@ -856,13 +855,85 @@ export class PostgresPool {
     model: Model
   ): Promise<z.infer<Model> | null> {
     const results = await this.callZeroOrOneRowAsync(sprocName, params);
-    if (results.rows.length == 0) {
+    if (results.rows.length === 0) {
       return null;
     } else {
       return model.parse(results.rows[0]);
     }
   }
 
+  async queryRows<Model extends z.ZodTypeAny>(sql: string, model: Model): Promise<z.infer<Model>[]>;
+  async queryRows<Model extends z.ZodTypeAny>(
+    sql: string,
+    params: QueryParams,
+    model: Model
+  ): Promise<z.infer<Model>[]>;
+  async queryRows<Model extends z.ZodTypeAny>(
+    sql: string,
+    paramsOrSchema: QueryParams | Model,
+    maybeModel?: Model
+  ) {
+    const params = maybeModel === undefined ? {} : (paramsOrSchema as QueryParams);
+    const model = maybeModel === undefined ? (paramsOrSchema as Model) : maybeModel;
+    const results = await this.queryAsync(sql, params);
+    if (results.fields.length === 1) {
+      const columnName = results.fields[0].name;
+      const rawData = results.rows.map((row) => row[columnName]);
+      return z.array(model).parse(rawData);
+    } else {
+      return z.array(model).parse(results.rows);
+    }
+  }
+
+  async queryRow<Model extends z.ZodTypeAny>(sql: string, model: Model): Promise<z.infer<Model>>;
+  async queryRow<Model extends z.ZodTypeAny>(
+    sql: string,
+    params: QueryParams,
+    model: Model
+  ): Promise<z.infer<Model>>;
+  async queryRow<Model extends z.ZodTypeAny>(
+    sql: string,
+    paramsOrSchema: QueryParams | Model,
+    maybeModel?: Model
+  ) {
+    const params = maybeModel === undefined ? {} : (paramsOrSchema as QueryParams);
+    const model = maybeModel === undefined ? (paramsOrSchema as Model) : maybeModel;
+    const results = await this.queryOneRowAsync(sql, params);
+    if (results.fields.length === 1) {
+      const columnName = results.fields[0].name;
+      return model.parse(results.rows[0][columnName]);
+    } else {
+      return model.parse(results.rows[0]);
+    }
+  }
+
+  async queryOptionalRow<Model extends z.ZodTypeAny>(
+    sql: string,
+    model: Model
+  ): Promise<z.infer<Model> | null>;
+  async queryOptionalRow<Model extends z.ZodTypeAny>(
+    sql: string,
+    params: QueryParams,
+    model: Model
+  ): Promise<z.infer<Model> | null>;
+  async queryOptionalRow<Model extends z.ZodTypeAny>(
+    sql: string,
+    paramsOrSchema: QueryParams | Model,
+    maybeModel?: Model
+  ) {
+    const params = maybeModel === undefined ? {} : (paramsOrSchema as QueryParams);
+    const model = maybeModel === undefined ? (paramsOrSchema as Model) : maybeModel;
+    const results = await this.queryZeroOrOneRowAsync(sql, params);
+    if (results.rows.length === 0) {
+      return null;
+    } else if (results.fields.length === 1) {
+      const columnName = results.fields[0].name;
+      return model.parse(results.rows[0][columnName]);
+    } else {
+      return model.parse(results.rows[0]);
+    }
+  }
+
   /**
    * Returns a {@link Cursor} for the given query. The cursor can be used to
    * read results in batches, which is useful for large result sets.
@@ -938,7 +1009,11 @@ export class PostgresPool {
         } catch (err: any) {
           throw enhanceError(err, sql, params);
         } finally {
-          client.release();
+          try {
+            await cursor.close();
+          } finally {
+            client.release();
+          }
         }
       },
       stream: function (batchSize: number) {
@@ -954,7 +1029,20 @@ export class PostgresPool {
         });
 
         // TODO: use native `node:stream#compose` once it's stable.
-        return multipipe(Readable.from(iterator.iterate(batchSize)), transform);
+        const generator = iterator.iterate(batchSize);
+        const pipe = multipipe(Readable.from(generator), transform);
+
+        // When the underlying stream is closed, we need to make sure that the
+        // cursor is also closed. We do this by calling `return()` on the generator,
+        // which will trigger its `finally` block, which will in turn release
+        // the client and close the cursor. The fact that the stream is already
+        // closed by this point means that someone reading from the stream will
+        // never actually see the `null` value that's returned.
+        pipe.once('close', () => {
+          generator.return(null);
+        });
+
+        return pipe;
       },
     };
     return iterator;

package/tsconfig.json

@@ -1,5 +1,5 @@
 {
-  "extends": "@prairielearn/tsconfig",
+  "extends": "@prairielearn/tsconfig/tsconfig.package.json",
   "compilerOptions": {
     "outDir": "./dist",
     "rootDir": "./src",