From 986030757d5e29ea373e7e14bfac760d14decfea Mon Sep 17 00:00:00 2001
From: Eric Wyne <ecwyne@gmail.com>
Date: Wed, 18 Nov 2020 10:36:10 -0800
Subject: [PATCH] updates to stream types

---
 src/types/Client.d.ts | 14 ++++++++---
 src/types/Stream.d.ts | 54 +++++++++++++++++++++----------------------
 2 files changed, 37 insertions(+), 31 deletions(-)

diff --git a/src/types/Client.d.ts b/src/types/Client.d.ts
index 0fd7423ea..4a4bc78ce 100644
--- a/src/types/Client.d.ts
+++ b/src/types/Client.d.ts
@@ -2,9 +2,9 @@ import Expr from './Expr'
 import { ExprArg } from './query'
 import PageHelper from './PageHelper'
 import RequestResult from './RequestResult'
-import { Subscription } from './Stream'
+import { Subscription, SubscriptionEventHandlers } from './Stream'
 
-type StreamEventFields = ['action', 'document', 'diff', 'prev']
+type StreamEventFields = 'action' | 'document' | 'diff' | 'prev'
 
 export interface ClientConfig {
   secret: string
@@ -28,5 +28,13 @@ export default class Client {
   query<T = object>(expr: ExprArg, options?: QueryOptions): Promise<T>
   paginate(expr: Expr, params?: object, options?: QueryOptions): PageHelper
   ping(scope?: string, timeout?: number): Promise<string>
-  stream(expr: Expr, options?: { fields?: StreamEventFields[] }): Subscription
+  stream: {
+    (expr: Expr, options?: { fields?: StreamEventFields[] }): Subscription<
+      Omit<SubscriptionEventHandlers, 'snapshot'>
+    >
+    document: (
+      expr: Expr,
+      options?: { fields?: StreamEventFields[] }
+    ) => Subscription<SubscriptionEventHandlers>
+  }
 }
diff --git a/src/types/Stream.d.ts b/src/types/Stream.d.ts
index ac376d72b..5cc7915f5 100644
--- a/src/types/Stream.d.ts
+++ b/src/types/Stream.d.ts
@@ -1,43 +1,41 @@
-import Client from './Client'
-import { ExprArg } from './query'
-import { values } from './values'
-
-export interface Subscription {
-  on: <T extends keyof SubscriptionEventHandlers>(
+export interface Subscription<TEventHandlerMap> {
+  on: <T extends keyof TEventHandlerMap>(
     type: T,
-    callback: SubscriptionEventHandlers[T]
+    callback: TEventHandlerMap[T]
   ) => this
   start: () => this
   close: () => void
 }
 
-type SubscriptionEventHandlers = {
-  start: (type: 'start', txn: number, event: number) => void
-  error: (
-    type: 'error',
-    txn: number,
-    event: {
+type Handler<TEventType extends string, TEventData> = (
+  data: TEventData,
+  event: { type: TEventType; txn: number; event: TEventData }
+) => void
+
+export type SubscriptionEventHandlers = {
+  start: Handler<'start', number>
+  error: Handler<
+    'error',
+    {
       code: string
       description: string
     }
-  ) => void
-  version: (
-    type: 'version',
-    txn: number,
-    event: {
+  >
+  version: Handler<
+    'version',
+    {
       action: 'create' | 'update' | 'delete'
-      document: object
-      diff: object
-      prev: object
+      document?: object
+      diff?: object
+      prev?: object
     }
-  ) => void
-  history_rewrite: (
-    type: 'history_rewrite',
-    txn: number,
-    event: {
+  >
+  history_rewrite: Handler<
+    'history_rewrite',
+    {
       action: 'history_rewrite'
       document: object
     }
-  ) => void
-  snapshot: (type: 'snapshot', txn: number, event: object) => void
+  >
+  snapshot: Handler<'snapshot', object>
 }