From 8950c2e279771afc43682ca8f60fe858955eff3d Mon Sep 17 00:00:00 2001
From: Anirudh Singh <anirudh@quizizz.com>
Date: Tue, 29 Nov 2022 13:44:56 +0530
Subject: [PATCH] Add an API for registering watch listeners for relationship
 updates

---
 src/lib/authzed.ts             | 52 +++++++++++++++++++++++++++++++---
 tests/lib/authzed.spec.ts      | 52 ++++++++++++++++++++++++++++++++--
 tests/setup/docker-compose.yml |  2 +-
 3 files changed, 98 insertions(+), 8 deletions(-)

diff --git a/src/lib/authzed.ts b/src/lib/authzed.ts
index 409f4e1..acf602f 100644
--- a/src/lib/authzed.ts
+++ b/src/lib/authzed.ts
@@ -1,8 +1,10 @@
-import { ConsoleLogger, ILogger } from '../logger';
+import { Readable } from 'stream';
 import { v1 } from '@authzed/authzed-node';
 import { ClientSecurity as AZClientSecurity } from '@authzed/authzed-node/dist/src/util';
 import { RelationshipUpdate_Operation } from '@authzed/authzed-node/dist/src/v1';
-import { Readable } from 'stream';
+import { EventEmitter } from 'node:events';
+
+import { ConsoleLogger, ILogger } from '../logger';
 
 type AuthZedClientParams = {
   host: string;
@@ -11,7 +13,13 @@ type AuthZedClientParams = {
 };
 
 type ZedToken = v1.ZedToken;
-export { AZClientSecurity as ClientSecurity, ZedToken };
+type RelationshipUpdate = v1.RelationshipUpdate;
+export {
+  AZClientSecurity as ClientSecurity,
+  ZedToken,
+  RelationshipUpdate,
+  RelationshipUpdate_Operation as RelationshipUpdateOperation,
+};
 
 export declare type PartialMessage<T extends object> = {
   [K in keyof T]?: PartialField<T[K]>;
@@ -117,9 +125,16 @@ type ListAccessorsForResourceResponse = {
   zedToken?: string;
 }[];
 
+type RegisterWatchEventListenerParams = {
+  emitter: EventEmitter;
+  watchFromToken?: ZedToken;
+  objectTypes?: string[];
+};
+
 export class AuthZed {
   private _client: ReturnType<typeof v1.NewClient>;
   private logger: ILogger;
+  private watchEventListeners: EventEmitter[];
 
   constructor(
     params: AuthZedClientParams,
@@ -378,7 +393,7 @@ export class AuthZed {
     return response;
   }
 
-  async listAccesorsForResource(
+  async listAccessorsForResource(
     params: ListAccessorsForResourceParams,
   ): Promise<ListAccessorsForResourceResponse> {
     const lookupSubjectsRequest = v1.LookupSubjectsRequest.create({
@@ -404,4 +419,33 @@ export class AuthZed {
 
     return accessors;
   }
+
+  registerWatchEventListener(params: RegisterWatchEventListenerParams): void {
+    const watchStream = this._client.watch({
+      optionalStartCursor: params.watchFromToken,
+      optionalObjectTypes: params.objectTypes ?? [],
+    });
+
+    this.logger.debugj({
+      msg: 'Registered watch listener',
+      params,
+    });
+
+    const emitter = params.emitter;
+
+    watchStream.on('data', (watchEvent: v1.WatchResponse) => {
+      this.logger.debugj({
+        msg: 'Got watch data',
+        watchEvent,
+      });
+      emitter.emit('data', {
+        zedToken: watchEvent.changesThrough,
+        updates: watchEvent.updates,
+      });
+    });
+
+    watchStream.on('close', () => emitter.emit('close'));
+    watchStream.on('end', () => emitter.emit('end'));
+    watchStream.on('error', (err) => emitter.emit('error', err));
+  }
 }
diff --git a/tests/lib/authzed.spec.ts b/tests/lib/authzed.spec.ts
index 3f60f5c..975e4c3 100644
--- a/tests/lib/authzed.spec.ts
+++ b/tests/lib/authzed.spec.ts
@@ -1,5 +1,14 @@
 import { ClientSecurity } from '@authzed/authzed-node/dist/src/util';
-import { AuthZed } from '../../src/lib/authzed';
+import {
+  RelationshipUpdate,
+  RelationshipUpdate_Operation,
+} from '@authzed/authzed-node/dist/src/v1';
+import { EventEmitter } from 'node:stream';
+import {
+  AuthZed,
+  RelationshipUpdateOperation,
+  ZedToken,
+} from '../../src/lib/authzed';
 
 const schema = `
 definition quizizz/quiz {
@@ -16,7 +25,7 @@ describe('AuthZed Wrapper', () => {
   const client = new AuthZed(
     {
       token: 'quizizz',
-      host: '127.0.0.1:50052',
+      host: '127.0.0.1:50051',
       security: ClientSecurity.INSECURE_PLAINTEXT_CREDENTIALS,
     },
     {},
@@ -123,7 +132,7 @@ describe('AuthZed Wrapper', () => {
   });
 
   it('lists all users that have a permission to a particular resource', async () => {
-    const accessors = await client.listAccesorsForResource({
+    const accessors = await client.listAccessorsForResource({
       resource: {
         id: 'quiz_2',
         type: 'quizizz/quiz',
@@ -137,4 +146,41 @@ describe('AuthZed Wrapper', () => {
 
     expect(accessors.length).toBeGreaterThan(0);
   });
+
+  it.only('receives watch events correctly', async () => {
+    const emitter = new EventEmitter({
+      captureRejections: true,
+    });
+
+    client.registerWatchEventListener({
+      emitter,
+    });
+
+    const updates: { updates: RelationshipUpdate[]; zedToken: ZedToken }[] = [];
+
+    emitter.on('data', (event) => {
+      updates.push(event);
+    });
+
+    await client.addRelations({
+      relations: [
+        {
+          relation: 'owner',
+          resource: {
+            id: 'quiz_1',
+            type: 'quizizz/quiz',
+          },
+          subject: {
+            id: 'user_1',
+            type: 'quizizz/user',
+          },
+        },
+      ],
+    });
+
+    expect(updates.length).toBeGreaterThan(0);
+    expect(updates[0].updates[0]).toMatchObject({
+      operation: RelationshipUpdateOperation.TOUCH,
+    });
+  });
 });
diff --git a/tests/setup/docker-compose.yml b/tests/setup/docker-compose.yml
index db93ce3..b8f00e6 100644
--- a/tests/setup/docker-compose.yml
+++ b/tests/setup/docker-compose.yml
@@ -4,6 +4,6 @@ services:
   spicedb:
     image: authzed/spicedb
     ports:
-      - "50052:50051"
+      - "50051:50051"
     command: serve --grpc-preshared-key "quizizz"
   
\ No newline at end of file