feat: add WebSocket support via API Gateway Management API

- Add WebSocketConnector wrapping PostToConnectionCommand with retry and timeout support
- Add publishToConnections sink with silent 410 (stale connection) handling
- Add broadcastToWebSocket flavor: filter by event type → format message → fan out to connections
- Add @aws-sdk/client-apigatewaymanagementapi dev dependency
- Unit tests
- Document WebSocket support in README

Author: gmoz22

README.md

@@ -462,6 +462,12 @@ These features are intended for implementing intra-service logic. They are frequ
 * `Connector` - connector for the SQS SDK
 * `sendToSqs` - stream steps to send a message to a queue
 
+## 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`)
+* `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
+
 ## Encryption Support
 * https://github.com/jgilbert01/aws-lambda-stream/issues/20
 

package-lock.json

@@ -1,6 +1,6 @@
 {
   "name": "aws-lambda-stream",
-  "version": "1.1.20",
+  "version": "1.1.21",
   "description": "Create stream processors with AWS Lambda functions.",
   "keywords": [
     "aws",
@@ -64,6 +64,7 @@
     "@aws-sdk/client-kinesis": "^3.450.0",
     "@aws-sdk/client-kms": "^3.450.0",
     "@aws-sdk/client-lambda": "^3.450.0",
+    "@aws-sdk/client-apigatewaymanagementapi": "^3.450.0",
     "@aws-sdk/client-s3": "^3.450.0",
     "@aws-sdk/client-secrets-manager": "^3.450.0",
     "@aws-sdk/client-scheduler": "^3.450.0",

package.json

@@ -8,3 +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';

src/connectors/index.js

@@ -0,0 +1,66 @@
+/* eslint import/no-extraneous-dependencies: ["error", {"devDependencies": true}] */
+
+import Promise from 'bluebird';
+
+import { ApiGatewayManagementApiClient, PostToConnectionCommand } from '@aws-sdk/client-apigatewaymanagementapi';
+import { NodeHttpHandler } from '@smithy/node-http-handler';
+import { ConfiguredRetryStrategy } from '@smithy/util-retry';
+import { omit, pick } from 'lodash';
+import { defaultBackoffDelay } from '../utils/retry';
+import { defaultDebugLogger } from '../utils/log';
+
+class Connector {
+  constructor({
+    debug,
+    pipelineId,
+    endpoint = process.env.WEBSOCKET_ENDPOINT,
+    timeout = Number(process.env.WS_TIMEOUT) || Number(process.env.TIMEOUT) || 1000,
+    additionalClientOpts = {},
+    ...opt
+  }) {
+    this.debug = (msg) => debug('%j', msg);
+    this.endpoint = endpoint || 'undefined';
+    this.client = Connector.getClient(pipelineId, debug, timeout, endpoint, additionalClientOpts);
+    this.opt = opt;
+  }
+
+  static clients = {};
+
+  static getClient(pipelineId, debug, timeout, endpoint, additionalClientOpts) {
+    const addlRequestHandlerOpts = pick(additionalClientOpts, ['requestHandler']);
+    const addlClientOpts = omit(additionalClientOpts, ['requestHandler']);
+
+    if (!this.clients[pipelineId]) {
+      this.clients[pipelineId] = new ApiGatewayManagementApiClient({
+        endpoint,
+        requestHandler: new NodeHttpHandler({
+          requestTimeout: timeout,
+          connectionTimeout: timeout,
+          ...addlRequestHandlerOpts,
+        }),
+        retryStrategy: new ConfiguredRetryStrategy(11, defaultBackoffDelay),
+        logger: defaultDebugLogger(debug),
+        ...addlClientOpts,
+      });
+    }
+    return this.clients[pipelineId];
+  }
+
+  postToConnection(connectionId, data, ctx) {
+    const params = {
+      ConnectionId: connectionId,
+      Data: typeof data === 'string' ? data : JSON.stringify(data),
+    };
+
+    return this._sendCommand(new PostToConnectionCommand(params), ctx);
+  }
+
+  _sendCommand(command, ctx) {
+    this.opt.metrics?.capture(this.client, command, 'ws', this.opt, ctx);
+    return Promise.resolve(this.client.send(command))
+      .tap(this.debug)
+      .tapCatch(this.debug);
+  }
+}
+
+export default Connector;

src/connectors/websocket.js

@@ -9,3 +9,4 @@ export * from './materialize';
 export * from './materializeS3';
 export * from './sendMessages';
 export * from './update';
+export * from './websocket';

src/flavors/index.js

@@ -0,0 +1,44 @@
+import _ from 'highland';
+
+import {
+  printStartPipeline, printEndPipeline,
+  faulty, faultyAsyncStream, faultify,
+} from '../utils';
+
+import { publishToConnections } from '../sinks/websocket';
+import { filterOnEventType, filterOnContent } from '../filters';
+
+export const broadcastToWebSocket = (rule) => (s) => s // eslint-disable-line import/prefer-default-export
+  .filter(onEventType(rule))
+  .tap(printStartPipeline)
+
+  .filter(onContent(rule))
+
+  .map(toMessage(rule))
+  .parallel(rule.parallel || Number(process.env.PARALLEL) || 4)
+
+  .flatMap(toConnections(rule))
+
+  .through(publishToConnections(rule))
+
+  .tap(printEndPipeline);
+
+const onEventType = (rule) => faulty((uow) => filterOnEventType(rule, uow));
+const onContent = (rule) => faulty((uow) => filterOnContent(rule, uow));
+
+const toMessage = (rule) => faultyAsyncStream(async (uow) => ({
+  ...uow,
+  message: await faultify(rule.toMessage)(uow, rule),
+}));
+
+// rule.toConnections returns a promise resolving to an array of { connectionId }
+// flatMap fans out the uow into one per connection
+const toConnections = (rule) => (uow) => {
+  if (!rule.toConnections) return _([uow]);
+  const p = faultify(rule.toConnections)(uow, rule)
+    .then((connections) => connections.map((conn) => ({
+      ...uow,
+      connectionId: conn.connectionId,
+    })));
+  return _(p).flatten();
+};

src/flavors/websocket.js

@@ -7,3 +7,4 @@ export * from './eventbridge';
 export * from './s3';
 export * from './sns';
 export * from './sqs';
+export * from './websocket';

src/sinks/index.js

@@ -0,0 +1,40 @@
+import _ from 'highland';
+
+import Connector from '../connectors/websocket';
+
+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
+  id: pipelineId,
+  debug = d('ws'),
+  endpoint = process.env.WEBSOCKET_ENDPOINT,
+  messageField = 'message',
+  parallel = Number(process.env.WS_PARALLEL) || Number(process.env.PARALLEL) || 8,
+  step = 'postToConnection',
+  ...opt
+} = {}) => {
+  const connector = new Connector({
+    pipelineId, debug, endpoint, ...opt,
+  });
+
+  const post = (uow) => {
+    const p = () => connector.postToConnection(uow.connectionId, uow[messageField], uow)
+      .then((postResponse) => ({ ...uow, postResponse }))
+      .catch((err) => {
+        // 410 = connection is gone, clean up silently
+        if (err.statusCode === 410 || err.$metadata?.httpStatusCode === 410) {
+          return { ...uow, postResponse: { 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(post)
+    .parallel(parallel);
+};

src/sinks/websocket.js

@@ -0,0 +1,62 @@
+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 Connector from '../../../src/connectors/websocket';
+
+import { debug } from '../../../src/utils';
+
+describe('connectors/websocket.js', () => {
+  let mockWs = mockClient(ApiGatewayManagementApiClient);
+
+  beforeEach(() => {
+    mockWs = mockClient(ApiGatewayManagementApiClient);
+  });
+
+  afterEach(() => {
+    mockWs.restore();
+  });
+
+  it('should reuse client per pipeline', () => {
+    const client1 = Connector.getClient('ws-test1', debug('test'), 1000, 'https://test');
+    const client2 = Connector.getClient('ws-test1', debug('test'), 1000, 'https://test');
+    const client3 = Connector.getClient('ws-test2', debug('test'), 1000, 'https://test');
+
+    expect(client1).to.eq(client2);
+    expect(client2).to.not.eq(client3);
+  });
+
+  it('should post to connection with object data', async () => {
+    const spy = sinon.spy((_) => ({}));
+    mockWs.on(PostToConnectionCommand).callsFake(spy);
+
+    const data = await new Connector({
+      debug: debug('ws'),
+      endpoint: 'https://test.execute-api.us-west-2.amazonaws.com/dev',
+    }).postToConnection('conn-1', { type: 'thing-updated', data: { id: '1' } });
+
+    expect(spy).to.have.been.calledWith({
+      ConnectionId: 'conn-1',
+      Data: JSON.stringify({ type: 'thing-updated', data: { id: '1' } }),
+    });
+    expect(data).to.deep.equal({});
+  });
+
+  it('should post to connection with string data', async () => {
+    const spy = sinon.spy((_) => ({}));
+    mockWs.on(PostToConnectionCommand).callsFake(spy);
+
+    const data = await new Connector({
+      debug: debug('ws'),
+      endpoint: 'https://test.execute-api.us-west-2.amazonaws.com/dev',
+    }).postToConnection('conn-1', 'raw string');
+
+    expect(spy).to.have.been.calledWith({
+      ConnectionId: 'conn-1',
+      Data: 'raw string',
+    });
+    expect(data).to.deep.equal({});
+  });
+});