feat: expand API Gateway Management connector and add disconnect/query support

- Rename connector websocket.js → apigatewayclient.js to match SDK client naming convention
- Add GetConnectionCommand and DeleteConnectionCommand to apigatewayclient connector
- Add disconnectConnections sink (with silent 410 handling) to sinks/websocket.js
- Add queryConnection step in new queries/websocket.js
- Add guard clauses to skip uow missing connectionId or messageField across all three steps
- Add tests
- Fix flavors/websocket.test.js assertion for no-connections case
- Update README WebSocket Support section

Author: gmoz22

README.md

@@ -447,7 +447,6 @@ These features are intended for implementing intra-service logic.
 * `Connector` - connector for the S3 SDK
 * `putObjectToS3` - stream steps to put an object in a bucket
 
-
 ## SNS Support
 These features are intended for implementing intra-service logic.
 * `fromSns` - creates a stream from an SNS topic
@@ -464,9 +463,11 @@ These features are intended for implementing intra-service logic. They are frequ
 
 ## WebSocket Support
 These features are intended for broadcasting events to connected WebSocket clients via API Gateway WebSocket API. Useful for lightweight real-time notifications (e.g. job completion, status updates) as an alternative to polling or SQS-based live data.
-* `Connector` - connector for the API Gateway Management API SDK (`PostToConnectionCommand`)
+* `Connector` - connector for the API Gateway Management API SDK (`PostToConnectionCommand`, `GetConnectionCommand`, `DeleteConnectionCommand`)
 * `publishToConnections` - stream steps to post a message to connected WebSocket clients
-* `broadcastToWebSocket` - a flavor that composes the full pipeline: filter on event type → format message via `rule.toMessage` → resolve target connections via `rule.toConnections` → fan out and publish to each connection
+* `disconnectConnections` - stream steps to forcibly disconnect WebSocket clients
+* `queryConnection` - query step to fetch connection metadata (identity, last active time) and adorn it to the `uow` as `getConnectionResponse`
+* `broadcastToWebSocket` - a flavor that composes the full pipeline: filter on event type, format message via `rule.toMessage`, resolve target connections via `rule.toConnections`, fan out and publish to each connection
 
 ## Encryption Support
 * https://github.com/jgilbert01/aws-lambda-stream/issues/20

src/connectors/websocket.js src/connectors/apigatewayclient.jssrc/connectors/websocket.js renamed to src/connectors/apigatewayclient.js

@@ -2,7 +2,12 @@
 
 import Promise from 'bluebird';
 
-import { ApiGatewayManagementApiClient, PostToConnectionCommand } from '@aws-sdk/client-apigatewaymanagementapi';
+import {
+  ApiGatewayManagementApiClient,
+  DeleteConnectionCommand,
+  GetConnectionCommand,
+  PostToConnectionCommand,
+} from '@aws-sdk/client-apigatewaymanagementapi';
 import { NodeHttpHandler } from '@smithy/node-http-handler';
 import { ConfiguredRetryStrategy } from '@smithy/util-retry';
 import { omit, pick } from 'lodash';
@@ -55,6 +60,14 @@ class Connector {
     return this._sendCommand(new PostToConnectionCommand(params), ctx);
   }
 
+  getConnection(connectionId, ctx) {
+    return this._sendCommand(new GetConnectionCommand({ ConnectionId: connectionId }), ctx);
+  }
+
+  deleteConnection(connectionId, ctx) {
+    return this._sendCommand(new DeleteConnectionCommand({ ConnectionId: connectionId }), ctx);
+  }
+
   _sendCommand(command, ctx) {
     this.opt.metrics?.capture(this.client, command, 'ws', this.opt, ctx);
     return Promise.resolve(this.client.send(command))

src/connectors/index.js

@@ -8,4 +8,4 @@ export { default as SecretsMgrConnector } from './secretsmgr';
 export { default as S3Connector } from './s3';
 export { default as SnsConnector } from './sns';
 export { default as SqsConnector } from './sqs';
-export { default as WebSocketConnector } from './websocket';
+export { default as ApiGatewayClientConnector } from './apigatewayclient';

src/queries/index.js

@@ -2,3 +2,4 @@ export * from './claimcheck';
 export * from './dynamodb';
 export * from './secretsmgr';
 export * from './s3';
+export * from './websocket';

src/queries/websocket.js

@@ -0,0 +1,36 @@
+import _ from 'highland';
+
+import Connector from '../connectors/apigatewayclient';
+
+import { rejectWithFault } from '../utils/faults';
+import { debug as d } from '../utils/print';
+import { ratelimit } from '../utils/ratelimit';
+
+export const queryConnection = ({ // eslint-disable-line import/prefer-default-export
+  id: pipelineId,
+  debug = d('ws'),
+  endpoint = process.env.WEBSOCKET_ENDPOINT,
+  getConnectionResponseField = 'getConnectionResponse',
+  parallel = Number(process.env.WS_PARALLEL) || Number(process.env.PARALLEL) || 8,
+  step = 'getConnection',
+  ...opt
+} = {}) => {
+  const connector = new Connector({
+    pipelineId, debug, endpoint, ...opt,
+  });
+
+  const get = (uow) => {
+    if (!uow.connectionId) return _(Promise.resolve(uow));
+
+    const p = () => connector.getConnection(uow.connectionId, uow)
+      .then((getConnectionResponse) => ({ ...uow, [getConnectionResponseField]: getConnectionResponse }))
+      .catch(rejectWithFault(uow));
+
+    return _(uow.metrics?.w(p, step) || /* istanbul ignore next */ p()); // wrap promise in a stream
+  };
+
+  return (s) => s
+    .through(ratelimit(opt))
+    .map(get)
+    .parallel(parallel);
+};

src/sinks/websocket.js

@@ -1,12 +1,12 @@
 import _ from 'highland';
 
-import Connector from '../connectors/websocket';
+import Connector from '../connectors/apigatewayclient';
 
 import { rejectWithFault } from '../utils/faults';
 import { debug as d } from '../utils/print';
 import { ratelimit } from '../utils/ratelimit';
 
-export const publishToConnections = ({ // eslint-disable-line import/prefer-default-export
+export const publishToConnections = ({
   id: pipelineId,
   debug = d('ws'),
   endpoint = process.env.WEBSOCKET_ENDPOINT,
@@ -20,6 +20,8 @@ export const publishToConnections = ({ // eslint-disable-line import/prefer-defa
   });
 
   const post = (uow) => {
+    if (!uow.connectionId || !uow[messageField]) return _(Promise.resolve(uow));
+
     const p = () => connector.postToConnection(uow.connectionId, uow[messageField], uow)
       .then((postResponse) => ({ ...uow, postResponse }))
       .catch((err) => {
@@ -38,3 +40,37 @@ export const publishToConnections = ({ // eslint-disable-line import/prefer-defa
     .map(post)
     .parallel(parallel);
 };
+
+export const disconnectConnections = ({
+  id: pipelineId,
+  debug = d('ws'),
+  endpoint = process.env.WEBSOCKET_ENDPOINT,
+  parallel = Number(process.env.WS_PARALLEL) || Number(process.env.PARALLEL) || 8,
+  step = 'deleteConnection',
+  ...opt
+} = {}) => {
+  const connector = new Connector({
+    pipelineId, debug, endpoint, ...opt,
+  });
+
+  const disconnect = (uow) => {
+    if (!uow.connectionId) return _(Promise.resolve(uow));
+
+    const p = () => connector.deleteConnection(uow.connectionId, uow)
+      .then((deleteResponse) => ({ ...uow, deleteResponse }))
+      .catch((err) => {
+        // 410 = connection already gone, clean up silently
+        if (err.statusCode === 410 || err.$metadata?.httpStatusCode === 410) {
+          return { ...uow, deleteResponse: { statusCode: 410, connectionId: uow.connectionId } };
+        }
+        return rejectWithFault(uow)(err);
+      });
+
+    return _(uow.metrics?.w(p, step) || /* istanbul ignore next */ p()); // wrap promise in a stream
+  };
+
+  return (s) => s
+    .through(ratelimit(opt))
+    .map(disconnect)
+    .parallel(parallel);
+};

test/unit/connectors/websocket.test.js …unit/connectors/apigatewayclient.test.jstest/unit/connectors/websocket.test.js renamed to test/unit/connectors/apigatewayclient.test.js test/unit/connectors/websocket.test.js renamed to test/unit/connectors/apigatewayclient.test.js

@@ -2,9 +2,14 @@ import 'mocha';
 import { expect } from 'chai';
 import sinon from 'sinon';
 import { mockClient } from 'aws-sdk-client-mock';
-import { ApiGatewayManagementApiClient, PostToConnectionCommand } from '@aws-sdk/client-apigatewaymanagementapi';
+import {
+  ApiGatewayManagementApiClient,
+  DeleteConnectionCommand,
+  GetConnectionCommand,
+  PostToConnectionCommand,
+} from '@aws-sdk/client-apigatewaymanagementapi';
 
-import Connector from '../../../src/connectors/websocket';
+import Connector from '../../../src/connectors/apigatewayclient';
 
 import { debug } from '../../../src/utils';
 
@@ -59,4 +64,30 @@ describe('connectors/websocket.js', () => {
     });
     expect(data).to.deep.equal({});
   });
+
+  it('should get connection', async () => {
+    const spy = sinon.spy((_) => ({ ConnectedAt: '2024-01-01T00:00:00Z', LastActiveAt: '2024-01-01T01:00:00Z' }));
+    mockWs.on(GetConnectionCommand).callsFake(spy);
+
+    const data = await new Connector({
+      debug: debug('ws'),
+      endpoint: 'https://test.execute-api.us-west-2.amazonaws.com/dev',
+    }).getConnection('conn-1');
+
+    expect(spy).to.have.been.calledWith({ ConnectionId: 'conn-1' });
+    expect(data).to.deep.equal({ ConnectedAt: '2024-01-01T00:00:00Z', LastActiveAt: '2024-01-01T01:00:00Z' });
+  });
+
+  it('should delete connection', async () => {
+    const spy = sinon.spy((_) => ({}));
+    mockWs.on(DeleteConnectionCommand).callsFake(spy);
+
+    const data = await new Connector({
+      debug: debug('ws'),
+      endpoint: 'https://test.execute-api.us-west-2.amazonaws.com/dev',
+    }).deleteConnection('conn-1');
+
+    expect(spy).to.have.been.calledWith({ ConnectionId: 'conn-1' });
+    expect(data).to.deep.equal({});
+  });
 });

test/unit/flavors/websocket.test.js

@@ -7,13 +7,13 @@ import { initialize, initializeFrom } from '../../../src';
 import { toDynamodbRecords, fromDynamodb } from '../../../src/from/dynamodb';
 
 import { defaultOptions } from '../../../src/utils/opt';
-import { WebSocketConnector } from '../../../src/connectors';
+import { ApiGatewayClientConnector } from '../../../src/connectors';
 
 import { broadcastToWebSocket } from '../../../src/flavors/websocket';
 
 describe('flavors/websocket.js', () => {
   beforeEach(() => {
-    sinon.stub(WebSocketConnector.prototype, 'postToConnection').resolves({});
+    sinon.stub(ApiGatewayClientConnector.prototype, 'postToConnection').resolves({});
   });
 
   afterEach(sinon.restore);
@@ -121,7 +121,7 @@ describe('flavors/websocket.js', () => {
       .tap((collected) => {
         expect(collected.length).to.equal(1);
         expect(collected[0].connectionId).to.be.undefined;
-        expect(collected[0].postResponse).to.deep.equal({});
+        expect(collected[0].postResponse).to.be.undefined;
       })
       .done(done);
   });

test/unit/queries/websocket.test.js

@@ -0,0 +1,76 @@
+import 'mocha';
+import { expect } from 'chai';
+import sinon from 'sinon';
+import _ from 'highland';
+
+import { queryConnection } from '../../../src/queries/websocket';
+
+import Connector from '../../../src/connectors/apigatewayclient';
+
+describe('queries/websocket.js', () => {
+  afterEach(sinon.restore);
+
+  it('should skip uow without connectionId', (done) => {
+    const stub = sinon.stub(Connector.prototype, 'getConnection');
+
+    const uows = [{ someField: 'value' }];
+
+    _(uows)
+      .through(queryConnection())
+      .collect()
+      .tap((collected) => {
+        expect(collected.length).to.equal(1);
+        expect(collected[0].getConnectionResponse).to.be.undefined;
+        expect(stub).to.not.have.been.called;
+      })
+      .done(done);
+  });
+
+  it('should get connection info', (done) => {
+    const connectionInfo = { ConnectedAt: '2024-01-01T00:00:00Z', LastActiveAt: '2024-01-01T01:00:00Z' };
+    sinon.stub(Connector.prototype, 'getConnection').resolves(connectionInfo);
+
+    const uows = [{ connectionId: 'conn-1' }];
+
+    _(uows)
+      .through(queryConnection())
+      .collect()
+      .tap((collected) => {
+        expect(collected.length).to.equal(1);
+        expect(collected[0].connectionId).to.equal('conn-1');
+        expect(collected[0].getConnectionResponse).to.deep.equal(connectionInfo);
+      })
+      .done(done);
+  });
+
+  it('should support custom response field name', (done) => {
+    sinon.stub(Connector.prototype, 'getConnection').resolves({ ConnectedAt: '2024-01-01T00:00:00Z' });
+
+    const uows = [{ connectionId: 'conn-1' }];
+
+    _(uows)
+      .through(queryConnection({ getConnectionResponseField: 'wsInfo' }))
+      .collect()
+      .tap((collected) => {
+        expect(collected.length).to.equal(1);
+        expect(collected[0].wsInfo).to.deep.equal({ ConnectedAt: '2024-01-01T00:00:00Z' });
+      })
+      .done(done);
+  });
+
+  it('should fault on error', (done) => {
+    const err = new Error('Forbidden');
+    sinon.stub(Connector.prototype, 'getConnection').rejects(err);
+
+    const uows = [{ connectionId: 'conn-1' }];
+
+    _(uows)
+      .through(queryConnection())
+      .errors((e) => {
+        expect(e.message).to.equal('Forbidden');
+        expect(e.uow).to.not.be.undefined;
+        done();
+      })
+      .resume();
+  });
+});