refactor(ext-driver-pg): release connection when stream destroyed

Author: oscar60310

packages/extension-driver-pg/src/lib/pgDataSource.ts

@@ -92,34 +92,35 @@ export class PGDataSource extends DataSource<any, PGOptions> {
     const { chunkSize = 100 } = options;
     const cursorRead = this.cursorRead.bind(this);
     const firstChunk = await cursorRead(cursor, chunkSize);
-
+    // save first chunk in buffer for incoming requests
+    let bufferedRows = [...firstChunk.rows];
+    let bufferReadIndex = 0;
+    const fetchNext = async () => {
+      if (bufferReadIndex >= bufferedRows.length) {
+        bufferedRows = (await cursorRead(cursor, chunkSize)).rows;
+        bufferReadIndex = 0;
+      }
+      return bufferedRows[bufferReadIndex++] || null;
+    };
     const stream = new Readable({
       objectMode: true,
       read() {
-        cursorRead(cursor, chunkSize)
-          .then(({ rows }) => {
-            if (rows.length === 0) {
-              this.push(null);
-              // Send done event to notify upstream to release the connection.
-              cursor.emit('done');
-            }
-            for (const row of rows) {
-              this.push(row);
-            }
+        fetchNext()
+          .then((row) => {
+            this.push(row);
           })
           .catch((error) => {
-            this.emit('error', error);
-            // Send done event to notify upstream to release the connection.
-            cursor.emit('done');
+            this.destroy(error);
           });
       },
+      destroy() {
+        // Send done event to notify upstream to release the connection.
+        cursor.emit('done');
+      },
+      // automatically destroy() the stream when it emits 'finish' or errors. Node > 10.16
+      autoDestroy: true,
     });
 
-    // Push the first chunk data.
-    for (const row of firstChunk.rows) {
-      stream.push(row);
-    }
-
     return {
       getColumns: () =>
         firstChunk.result.fields.map((field) => ({

packages/extension-driver-pg/test/pgDataSource.spec.ts

@@ -1,6 +1,7 @@
 import { PGServer } from './pgServer';
 import { PGDataSource, PGOptions } from '../src';
 import { streamToArray } from '@vulcan-sql/core';
+import { Writable } from 'stream';
 
 const pg = new PGServer();
 let dataSource: PGDataSource;
@@ -295,10 +296,58 @@ it('Data source should return correct column types', async () => {
     operations: {} as any,
   });
   const column = getColumns();
-  // We need to consume the data stream or the driver waits for us
-  await streamToArray(getData());
+  // We need to destroy the data stream or the driver waits for us
+  const data = getData();
+  data.destroy();
+
   // Assert
   expect(column[0]).toEqual({ name: 'id', type: 'number' });
   expect(column[1]).toEqual({ name: 'name', type: 'string' });
   expect(column[2]).toEqual({ name: 'enabled', type: 'boolean' });
 }, 30000);
+
+it('Data source should release connection when readable stream is destroyed', async () => {
+  // Arrange
+  dataSource = new PGDataSource({}, '', [
+    {
+      name: 'profile1',
+      type: 'pg',
+      connection: {
+        host: pg.host,
+        user: pg.user,
+        password: pg.password,
+        database: pg.database,
+        port: pg.port,
+      } as PGOptions,
+      allow: '*',
+    },
+  ]);
+  await dataSource.activate();
+  // Act
+  const { getData } = await dataSource.execute({
+    statement: 'select * from users limit 100',
+    bindParams: new Map(),
+    profileName: 'profile1',
+    operations: {} as any,
+  });
+  const readStream = getData();
+  const rows: any[] = [];
+  let resolve: any;
+  const waitForStream = () => new Promise((res) => (resolve = res));
+  const writeStream = new Writable({
+    write(chunk, _, cb) {
+      rows.push(chunk);
+      // After read 5 records, destroy the upstream
+      if (rows.length === 5) {
+        readStream.destroy();
+        resolve();
+      } else cb();
+    },
+    objectMode: true,
+  });
+  readStream.pipe(writeStream);
+  await waitForStream();
+  // Assert
+  expect(rows.length).toBe(5);
+  // afterEach hook will timeout if any leak occurred.
+}, 30000);

packages/serve/src/lib/response-formatter/jsonFormatter.ts

@@ -71,6 +71,10 @@ export class JsonFormatter extends BaseResponseFormatter {
       })
       .on('end', () => {
         logger.debug('convert to json format stream > done.');
+      })
+      .on('close', () => {
+        console.log('cccc');
+        data.destroy();
       });
 
     return jsonStream;
@@ -83,5 +87,9 @@ export class JsonFormatter extends BaseResponseFormatter {
     // set json stream to response in context ( data is json stream, no need to convert. )
     ctx.response.body = stream;
     ctx.response.set('Content-type', 'application/json');
+    ctx.req.on('close', () => {
+      console.log(88);
+      stream.destroy();
+    });
   }
 }