Merge pull request #90 from juni-b-queer/filter-jetstream-by-did

minor: ability to filter jetstream by DIDs

src/subscriptions/firehose/JetstreamSubscription.ts

@@ -32,10 +32,12 @@ export class JetstreamSubscription extends AbstractSubscription {
      *
      * @param {JetstreamSubscriptionHandlers} handlerControllers - An array of handler controllers.
      * @param {string} wsURL - The WebSocket URL to connect to. Defaults to `wss://bsky.network`.
+     * @param dids
      */
     constructor(
         protected handlerControllers: JetstreamSubscriptionHandlers,
-        protected wsURL: string = 'ws://localhost:6008/subscribe'
+        protected wsURL: string = 'ws://localhost:6008/subscribe',
+        protected wantedDids: string[] = []
     ) {
         super(handlerControllers);
         this.generateWsURL();
@@ -61,6 +63,17 @@ export class JetstreamSubscription extends AbstractSubscription {
         if (queryParams.length > 0) {
             this.setWsURL = `${this.wsURL}?${queryParams.join('&')}`;
         }
+
+        if (this.wantedDids.length > 0) {
+            const dids = this.wantedDids
+                .map((did) => `wantedDids=${did}`)
+                .join('&');
+            if (queryParams.length > 0) {
+                this.setWsURL = `${this.wsURL}&${dids}`;
+            } else {
+                this.setWsURL = `${this.wsURL}?${dids}`;
+            }
+        }
     }
 
     /**

src/subscriptions/firehose/README.md

@@ -65,3 +65,14 @@ const handlers: JetstreamSubscriptionHandlers = {
 ```
 
 This also allows the JetstreamSubscription to automatically subscribe to only events that it's handling.
+
+
+### Wanted DIDs
+If you want to only listen for events from a specific user (or users) you can pass an optional 3rd parameter into the jetstream constructor
+```typescript
+new JetstreamSubscription(
+    handlers,
+    'ws://localhost:6010/subscribe',
+    ['did:plc:123', 'did:plc:124']
+);
+```

tests/firehose/JeststreamSubscription.test.ts

@@ -10,34 +10,35 @@ import WebSocket from 'ws';
 
 describe('JetstreamSubscription', () => {
     let jetSub: JetstreamSubscription;
-    const handlers: JetstreamSubscriptionHandlers = {
-        post: {
-            c: [],
-            d: [],
-        },
-        like: {
-            c: [],
-            d: [],
-        },
-        repost: {
-            c: [],
-            d: [],
-        },
-        follow: {
-            c: [],
-            d: [],
-        },
-        block: {
-            c: [],
-            d: [],
-        },
-    };
+    let handlers: JetstreamSubscriptionHandlers;
     // A dummy message handler for testing
     const dummyHandler: MessageHandler = {
         handle: jest.fn(),
     } as unknown as MessageHandler;
 
     beforeEach(() => {
+        handlers = {
+            post: {
+                c: [],
+                d: [],
+            },
+            like: {
+                c: [],
+                d: [],
+            },
+            repost: {
+                c: [],
+                d: [],
+            },
+            follow: {
+                c: [],
+                d: [],
+            },
+            block: {
+                c: [],
+                d: [],
+            },
+        };
         jetSub = new JetstreamSubscription(handlers);
         (dummyHandler.handle as jest.Mock).mockClear();
     });
@@ -58,6 +59,38 @@ describe('JetstreamSubscription', () => {
         expect((jetSub as any).wsURL).toContain('post');
     });
 
+    test('generateWsURL with did', () => {
+        handlers = {
+            post: {
+                c: [dummyHandler],
+            },
+        };
+        jetSub = new JetstreamSubscription(
+            handlers,
+            'ws://localhost:6010/subscribe',
+            ['did:plc:123']
+        );
+        jetSub.setWsURL = 'ws://localhost:6010/subscribe';
+        jetSub.generateWsURL();
+        expect((jetSub as any).wsURL).toBe(
+            'ws://localhost:6010/subscribe?wantedCollections=app.bsky.feed.post&wantedDids=did:plc:123'
+        );
+    });
+
+    test('generateWsURL with dids', () => {
+        handlers = {};
+        jetSub = new JetstreamSubscription(
+            handlers,
+            'ws://localhost:6010/subscribe',
+            ['did:plc:123', 'did:plc:124']
+        );
+        jetSub.setWsURL = 'ws://localhost:6010/subscribe';
+        jetSub.generateWsURL();
+        expect((jetSub as any).wsURL).toBe(
+            'ws://localhost:6010/subscribe?wantedDids=did:plc:123&wantedDids=did:plc:124'
+        );
+    });
+
     test('handleCreate post', () => {
         // @ts-ignore
         handlers.post.c = [dummyHandler];