@milaboratories/pl-client 3.4.1 → 3.5.0

package/src/core/ll_transaction.ts

@@ -24,37 +24,125 @@ export type OneOfKind<T extends { oneofKind: unknown }, Kind extends T["oneofKin
   { oneofKind: Kind }
 >;
 
-interface SingleResponseHandler<Kind extends ServerMessageResponse["oneofKind"]> {
-  kind: Kind;
-  expectMultiResponse: false;
-  resolve: (v: OneOfKind<ServerMessageResponse, Kind>) => void;
+interface SingleResponseHandler {
+  mode: "single";
+  kind: ServerMessageResponse["oneofKind"];
+  resolve: (v: ServerMessageResponse) => void;
   reject: (e: Error) => void;
 }
 
-interface MultiResponseHandler<Kind extends ServerMessageResponse["oneofKind"]> {
-  kind: Kind;
-  expectMultiResponse: true;
-  resolve: (v: OneOfKind<ServerMessageResponse, Kind>[]) => void;
+interface MultiResponseHandler {
+  mode: "multiBuffered";
+  kind: ServerMessageResponse["oneofKind"];
+  resolve: (v: ServerMessageResponse[]) => void;
   reject: (e: Error) => void;
 }
 
-type AnySingleResponseHandler = SingleResponseHandler<ServerMessageResponse["oneofKind"]>;
+interface MultiStreamResponseHandler {
+  mode: "multiStream";
+  kind: ServerMessageResponse["oneofKind"];
+  stream: AsyncMessageStream<ServerMessageResponse>;
+}
 
-type AnyMultiResponseHandler = MultiResponseHandler<ServerMessageResponse["oneofKind"]>;
+type AnySingleResponseHandler = SingleResponseHandler;
+type AnyMultiStreamResponseHandler = MultiStreamResponseHandler;
 
 type AnyResponseHandler =
-  | SingleResponseHandler<ServerMessageResponse["oneofKind"]>
-  | MultiResponseHandler<ServerMessageResponse["oneofKind"]>;
-
-function createResponseHandler<Kind extends ServerMessageResponse["oneofKind"]>(
-  kind: Kind,
-  expectMultiResponse: boolean,
-  resolve:
-    | ((v: OneOfKind<ServerMessageResponse, Kind>) => void)
-    | ((v: OneOfKind<ServerMessageResponse, Kind>[]) => void),
-  reject: (e: Error) => void,
-): AnyResponseHandler {
-  return { kind, expectMultiResponse, resolve, reject } as AnyResponseHandler;
+  | AnySingleResponseHandler
+  | MultiResponseHandler
+  | AnyMultiStreamResponseHandler;
+
+// Implements both AsyncIterable and AsyncIterator: `[Symbol.asyncIterator]()` returns
+// `this`, so multiple `for await` loops over the same stream share state instead of
+// creating independent iterators that would race on `waitingResolve`/`waitingReject`.
+// Only one active consumer is supported; concurrent `next()` calls are not safe.
+class AsyncMessageStream<T> implements AsyncIterable<T>, AsyncIterator<T> {
+  private readonly queue = new Denque<T>();
+  private done = false;
+  private cancelled = false;
+  private failure?: Error;
+  private waitingResolve?: (value: IteratorResult<T>) => void;
+  private waitingReject?: (reason?: unknown) => void;
+
+  public push(value: T): void {
+    if (this.done || this.cancelled) return;
+    if (this.waitingResolve) {
+      const resolve = this.waitingResolve;
+      this.waitingResolve = undefined;
+      this.waitingReject = undefined;
+      resolve({ value, done: false });
+      return;
+    }
+    this.queue.push(value);
+  }
+
+  public end(): void {
+    if (this.done) return;
+    this.done = true;
+    if (this.waitingResolve) {
+      const resolve = this.waitingResolve;
+      this.waitingResolve = undefined;
+      this.waitingReject = undefined;
+      resolve({ value: undefined, done: true });
+    }
+  }
+
+  // Fail-fast: any frames still buffered in `this.queue` are intentionally dropped.
+  // `next()` checks `this.failure` before draining the queue, so consumers see the
+  // error immediately. This differs from `end()`, which drains buffered frames first.
+  // Rationale: when the upstream stream errors, the in-flight response is treated as
+  // corrupt and partial frames must not be surfaced.
+  public fail(error: Error): void {
+    if (this.done) return;
+    this.failure = error;
+    this.done = true;
+    if (this.waitingReject) {
+      const reject = this.waitingReject;
+      this.waitingResolve = undefined;
+      this.waitingReject = undefined;
+      reject(error);
+    }
+  }
+
+  public cancel(): void {
+    this.cancelled = true;
+    this.done = true;
+    this.queue.clear();
+    if (this.waitingResolve) {
+      const resolve = this.waitingResolve;
+      this.waitingResolve = undefined;
+      this.waitingReject = undefined;
+      resolve({ value: undefined, done: true });
+    }
+  }
+
+  public async next(): Promise<IteratorResult<T>> {
+    if (this.failure) throw this.failure;
+
+    const value = this.queue.shift();
+    if (value !== undefined) return { value, done: false };
+
+    if (this.done) return { value: undefined, done: true };
+
+    return await new Promise<IteratorResult<T>>((resolve, reject) => {
+      this.waitingResolve = resolve;
+      this.waitingReject = reject;
+    });
+  }
+
+  public async return(): Promise<IteratorResult<T>> {
+    this.cancel();
+    return { value: undefined, done: true };
+  }
+
+  public async throw(error?: unknown): Promise<IteratorResult<T>> {
+    this.cancel();
+    throw error;
+  }
+
+  [Symbol.asyncIterator](): this {
+    return this;
+  }
 }
 
 function isRecoverable(status: Status): boolean {
@@ -127,6 +215,7 @@ export class LLPlTransaction {
     // to the specific request on which it happened
     let currentHandler: AnyResponseHandler | undefined = undefined;
     let responseAggregator: ServerMessageResponse[] | undefined = undefined;
+    let currentMultiIdx = 0;
     try {
       for await (const message of this.stream.responses) {
         if (currentHandler === undefined) {
@@ -139,9 +228,8 @@ export class LLPlTransaction {
             break;
           }
 
-          // allocating response aggregator array
-          if (currentHandler.expectMultiResponse) responseAggregator = [];
-
+          if (currentHandler.mode === "multiBuffered") responseAggregator = [];
+          currentMultiIdx = 0;
           expectedId++;
         }
 
@@ -157,11 +245,14 @@ export class LLPlTransaction {
           const status = message.error;
 
           if (isRecoverable(status)) {
-            currentHandler.reject(
-              new RethrowError(() => {
-                throw new RecoverablePlError(status);
-              }),
-            );
+            const recoverableError = new RethrowError(() => {
+              throw new RecoverablePlError(status);
+            });
+            if (currentHandler.mode === "single" || currentHandler.mode === "multiBuffered") {
+              currentHandler.reject(recoverableError);
+            } else {
+              currentHandler.stream.fail(recoverableError);
+            }
             currentHandler = undefined;
 
             if (message.multiMessage !== undefined && !message.multiMessage.isLast) {
@@ -174,9 +265,17 @@ export class LLPlTransaction {
             // We can continue to work after recoverable errors
             continue;
           } else {
-            this.assignErrorFactoryIfNotSet(() => {
-              throw new UnrecoverablePlError(status);
-            }, currentHandler.reject);
+            this.assignErrorFactoryIfNotSet(
+              () => {
+                throw new UnrecoverablePlError(status);
+              },
+              currentHandler.mode === "single" || currentHandler.mode === "multiBuffered"
+                ? currentHandler.reject
+                : undefined,
+            );
+            if (currentHandler.mode === "multiStream") {
+              currentHandler.stream.fail(new UnrecoverablePlError(status));
+            }
             currentHandler = undefined;
 
             // In case of unrecoverable errors we close the transaction
@@ -190,20 +289,38 @@ export class LLPlTransaction {
         ) {
           const errorMessage = `inconsistent request response types: ${currentHandler.kind} !== ${message.response.oneofKind}`;
 
-          this.assignErrorFactoryIfNotSet(() => {
-            throw new Error(errorMessage);
-          }, currentHandler.reject);
+          this.assignErrorFactoryIfNotSet(
+            () => {
+              throw new Error(errorMessage);
+            },
+            currentHandler.mode === "single" || currentHandler.mode === "multiBuffered"
+              ? currentHandler.reject
+              : undefined,
+          );
+          if (currentHandler.mode === "multiStream") {
+            currentHandler.stream.fail(new Error(errorMessage));
+          }
           currentHandler = undefined;
 
           break;
         }
 
-        if (currentHandler.expectMultiResponse !== (message.multiMessage !== undefined)) {
-          const errorMessage = `inconsistent multi state: ${currentHandler.expectMultiResponse} !== ${message.multiMessage !== undefined}`;
-
-          this.assignErrorFactoryIfNotSet(() => {
-            throw new Error(errorMessage);
-          }, currentHandler.reject);
+        const expectMultiResponse =
+          currentHandler.mode === "multiBuffered" || currentHandler.mode === "multiStream";
+        if (expectMultiResponse !== (message.multiMessage !== undefined)) {
+          const errorMessage = `inconsistent multi state: ${expectMultiResponse} !== ${message.multiMessage !== undefined}`;
+
+          this.assignErrorFactoryIfNotSet(
+            () => {
+              throw new Error(errorMessage);
+            },
+            currentHandler.mode === "single" || currentHandler.mode === "multiBuffered"
+              ? currentHandler.reject
+              : undefined,
+          );
+          if (currentHandler.mode === "multiStream") {
+            currentHandler.stream.fail(new Error(errorMessage));
+          }
           currentHandler = undefined;
 
           break;
@@ -213,23 +330,38 @@ export class LLPlTransaction {
 
         if (message.multiMessage !== undefined) {
           if (!message.multiMessage.isEmpty) {
-            if (message.multiMessage.id !== responseAggregator!.length + 1) {
-              const errorMessage = `inconsistent multi id: ${message.multiMessage.id} !== ${responseAggregator!.length + 1}`;
-
-              this.assignErrorFactoryIfNotSet(() => {
-                throw new Error(errorMessage);
-              }, currentHandler.reject);
+            if (message.multiMessage.id !== currentMultiIdx + 1) {
+              const errorMessage = `inconsistent multi id: ${message.multiMessage.id} !== ${currentMultiIdx + 1}`;
+
+              this.assignErrorFactoryIfNotSet(
+                () => {
+                  throw new Error(errorMessage);
+                },
+                currentHandler.mode === "multiBuffered" ? currentHandler.reject : undefined,
+              );
+              if (currentHandler.mode === "multiStream") {
+                currentHandler.stream.fail(new Error(errorMessage));
+              }
               currentHandler = undefined;
 
               break;
             }
 
-            responseAggregator!.push(message.response);
+            currentMultiIdx++;
+            if (currentHandler.mode === "multiBuffered") {
+              responseAggregator!.push(message.response);
+            } else if (currentHandler.mode === "multiStream") {
+              currentHandler.stream.push(message.response);
+            }
           }
 
           if (message.multiMessage.isLast) {
-            (currentHandler as AnyMultiResponseHandler).resolve(responseAggregator!);
-            responseAggregator = undefined;
+            if (currentHandler.mode === "multiBuffered") {
+              currentHandler.resolve(responseAggregator!);
+              responseAggregator = undefined;
+            } else if (currentHandler.mode === "multiStream") {
+              currentHandler.stream.end();
+            }
             currentHandler = undefined;
           }
         } else {
@@ -245,9 +377,15 @@ export class LLPlTransaction {
         }
       }
     } catch (e: any) {
-      return this.assignErrorFactoryIfNotSet(() => {
-        rethrowMeaningfulError(e, true);
-      }, currentHandler?.reject);
+      return this.assignErrorFactoryIfNotSet(
+        () => {
+          rethrowMeaningfulError(e, true);
+        },
+        currentHandler &&
+          (currentHandler.mode === "single" || currentHandler.mode === "multiBuffered")
+          ? currentHandler.reject
+          : undefined,
+      );
     } finally {
       await this.close();
     }
@@ -265,8 +403,14 @@ export class LLPlTransaction {
     while (true) {
       const handler = this.responseHandlerQueue.shift();
       if (!handler) break;
-      if (this.errorFactory) handler.reject(new RethrowError(this.errorFactory));
-      else handler.reject(new Error("no reply"));
+      const noReplyError = this.errorFactory
+        ? new RethrowError(this.errorFactory)
+        : new Error("no reply");
+      if (handler.mode === "single" || handler.mode === "multiBuffered") {
+        handler.reject(noReplyError);
+      } else {
+        handler.stream.fail(noReplyError);
+      }
     }
 
     // closing outgoing stream
@@ -311,10 +455,24 @@ export class LLPlTransaction {
 
     // Note: Promise synchronously executes a callback passed to a constructor
     const result = StatefulPromise.fromDeferredReject(
-      new Promise<OneOfKind<ServerMessageResponse, Kind>>((resolve, reject) => {
-        this.responseHandlerQueue.push(
-          createResponseHandler(r.oneofKind, expectMultiResponse, resolve, reject),
-        );
+      new Promise<
+        OneOfKind<ServerMessageResponse, Kind> | OneOfKind<ServerMessageResponse, Kind>[]
+      >((resolve, reject) => {
+        if (expectMultiResponse) {
+          this.responseHandlerQueue.push({
+            mode: "multiBuffered",
+            kind: r.oneofKind,
+            resolve: resolve as (v: ServerMessageResponse[]) => void,
+            reject,
+          });
+        } else {
+          this.responseHandlerQueue.push({
+            mode: "single",
+            kind: r.oneofKind,
+            resolve: resolve as (v: ServerMessageResponse) => void,
+            reject,
+          });
+        }
       }),
     );
 
@@ -333,6 +491,25 @@ export class LLPlTransaction {
     }
   }
 
+  public sendStream<Kind extends ClientMessageRequest["oneofKind"]>(
+    r: OneOfKind<ClientMessageRequest, Kind>,
+  ): AsyncIterable<OneOfKind<ServerMessageResponse, Kind>> {
+    if (this.errorFactory) throw new RethrowError(this.errorFactory);
+    if (this.closed) throw new Error("Transaction already closed");
+
+    const stream = new AsyncMessageStream<ServerMessageResponse>();
+    this.responseHandlerQueue.push({ mode: "multiStream", kind: r.oneofKind, stream });
+
+    void this.stream.requests
+      .send({ requestId: this.requestIdxCounter++, request: r })
+      .catch((e: unknown) => {
+        if (e instanceof Error) stream.fail(e);
+        else stream.fail(new Error("Error while sending request", { cause: e }));
+      });
+
+    return stream as AsyncIterable<OneOfKind<ServerMessageResponse, Kind>>;
+  }
+
   private _completed = false;
 
   /** Safe to call multiple times */

package/src/core/transaction.test.ts

@@ -177,3 +177,41 @@ test("handle KV storage 2", async () => {
     });
   });
 });
+
+test("resourceTree accepts mixed seed types and returns kv entries", async () => {
+  await withTempRoot(async (pl) => {
+    const [rootId, childId] = await pl.withWriteTx("resourceTreeSetup", async (tx) => {
+      const root = tx.createStruct(StructTestResource);
+      const child = tx.createStruct(StructTestResource);
+
+      tx.createField(field(root, "child"), "Dynamic", child);
+      tx.createField(field(tx.clientRoot, "treeRoot"), "Dynamic", root);
+      tx.setKValue(root, "root-k", "root-v");
+
+      await tx.commit();
+      return [await root.globalId, await child.globalId];
+    });
+
+    await pl.withReadTx("resourceTreeRead", async (tx) => {
+      const rootData = await tx.getResourceData(rootId, true);
+      const visitedIds = new Set<string>();
+      const iter = tx.resourceTree([rootData, childId], { includeKv: true });
+
+      for await (const item of iter) {
+        visitedIds.add(item.id);
+        expect(typeof item.traverseWasStopped).toBe("boolean");
+      }
+
+      expect(visitedIds.has(rootId)).toBe(true);
+      expect(visitedIds.has(childId)).toBe(true);
+    });
+  });
+});
+
+test("resourceTree empty seeds fails", async () => {
+  await withTempRoot(async (pl) => {
+    await pl.withReadTx("resourceTreeEmptySeeds", async (tx) => {
+      expect(() => tx.resourceTree([])).toThrow("resourceTree: at least one seed must be provided");
+    });
+  });
+});

package/src/core/transaction.ts

@@ -30,10 +30,19 @@ import type {
   OneOfKind,
   ServerMessageResponse,
 } from "./ll_transaction";
+import type {
+  ResourceAPI_Tree_Filter,
+  ResourceAPI_Tree_SeedResource,
+} from "../proto-grpc/github.com/milaboratory/pl/plapi/plapiproto/api";
 import { TxAPI_Open_Request_WritableTx } from "../proto-grpc/github.com/milaboratory/pl/plapi/plapiproto/api";
 import type { NonUndefined } from "utility-types";
 import { toBytes } from "../util/util";
-import { fieldTypeToProto, protoToField, protoToResource } from "./type_conversion";
+import {
+  fieldTypeToProto,
+  protoToField,
+  protoToResource,
+  resourceIsDeleted,
+} from "./type_conversion";
 import {
   canonicalJsonBytes,
   canonicalJsonGzBytes,
@@ -74,6 +83,19 @@ export interface KeyValueString {
   value: string;
 }
 
+export type ResourceTreeItem = ResourceData & { kv: KeyValue[]; traverseWasStopped: boolean };
+
+/** A single frame from the resourceTree() stream.
+ *
+ * When the server advertises `treeStopMarker:v1`, a stop-matched node is emitted
+ * as a lightweight marker instead of a full resource payload. Callers must
+ * narrow on `frameKind` before accessing type-specific fields.
+ * Note: `frameKind` is distinct from `ResourceData.kind` (which is the resource kind).
+ */
+export type ResourceTreeFrame =
+  | (ResourceTreeItem & { frameKind: "resource" })
+  | { frameKind: "stopMarker"; id: SignedResourceId; traverseWasStopped: true };
+
 interface _FieldId<RId> {
   /** Parent resource id */
   resourceId: RId;
@@ -1057,6 +1079,119 @@ export class PlTransaction {
     await this.ll.await();
   }
 
+  //
+  // Tree streaming
+  //
+
+  /**
+   * Perform a single `ResourceTree` walk and yield each visited resource.
+   *
+   * Items are yielded incrementally as stream frames arrive from server.
+   *
+   * The low-level transaction stream is still ordered, so each active request
+   * must be consumed to completion (or closed via iterator `return()`/`throw()`)
+   * before later responses can be dispatched.
+   *
+   * @param seeds - Roots for the traversal. May include both {@link ResourceData}
+   *   and {@link SignedResourceId} values in the same array.
+   * @param opts.fieldFilter       - Per-field predicate; matching field edges are not descended.
+   *   This wire field accepts only one filter. Compose multiple rules via
+   *   `treeFilter.and(...)` / `treeFilter.or(...)`.
+   * @param opts.traverseStopRules - Per-resource predicate; matching resources are returned
+   *   with `traverseWasStopped = true` and their children are not visited.
+   *   This wire field accepts only one filter. Compose multiple rules via
+   *   `treeFilter.and(...)` / `treeFilter.or(...)`.
+   * @param opts.includeKv         - When true, each yielded item includes KV entries.
+   * @param opts.maxDepth          - Optional depth cap.
+   */
+  public resourceTree(
+    seeds: (ResourceData | SignedResourceId)[],
+    opts?: {
+      fieldFilter?: ResourceAPI_Tree_Filter;
+      traverseStopRules?: ResourceAPI_Tree_Filter;
+      includeKv?: boolean;
+      maxDepth?: number;
+    },
+  ): AsyncIterable<ResourceTreeFrame> {
+    if (seeds.length === 0) {
+      throw new Error("resourceTree: at least one seed must be provided");
+    }
+
+    const seedProtos = seeds.map((seed): ResourceAPI_Tree_SeedResource => {
+      const signedId = typeof seed === "string" ? seed : seed.id;
+      const { globalId, signature } = parseSignedResourceId(signedId);
+      return { resourceId: globalId, resourceSignature: signature };
+    });
+
+    const firstSeed = seedProtos[0];
+
+    const responseStream = this.ll.sendStream({
+      oneofKind: "resourceTree",
+      resourceTree: {
+        // Keep compatibility with request shape that requires non-optional fields.
+        resourceId: firstSeed?.resourceId ?? 0n,
+        resourceSignature: firstSeed?.resourceSignature ?? new Uint8Array(0),
+        seeds: seedProtos,
+        fieldFilter: opts?.fieldFilter,
+        traverseStopRules: opts?.traverseStopRules,
+        includeKv: opts?.includeKv ?? false,
+        maxDepth: opts?.maxDepth,
+      },
+    });
+
+    return {
+      [Symbol.asyncIterator](): AsyncIterator<ResourceTreeFrame> {
+        const iterator = responseStream[Symbol.asyncIterator]();
+
+        return {
+          async next(): Promise<IteratorResult<ResourceTreeFrame>> {
+            while (true) {
+              const item = await iterator.next();
+              if (item.done) return { value: undefined, done: true };
+
+              const frame = item.value.resourceTree;
+              if (frame === undefined) continue;
+
+              if (frame.resource && resourceIsDeleted(frame.resource)) {
+                continue;
+              }
+
+              if (!frame.resource) {
+                // Stop-marker frame: no resource body; server advertised treeStopMarker:v1.
+                const id = createSignedResourceId(
+                  frame.resourceId,
+                  toResourceSignature(frame.resourceSignature),
+                );
+                return {
+                  value: { frameKind: "stopMarker", id, traverseWasStopped: true },
+                  done: false,
+                };
+              }
+
+              return {
+                value: {
+                  frameKind: "resource",
+                  ...protoToResource(frame.resource),
+                  kv: frame.kv,
+                  traverseWasStopped: frame.traverseWasStopped,
+                },
+                done: false,
+              };
+            }
+          },
+          async return(): Promise<IteratorResult<ResourceTreeFrame>> {
+            if (iterator.return) await iterator.return();
+            return { value: undefined, done: true };
+          },
+          async throw(error?: unknown): Promise<IteratorResult<ResourceTreeFrame>> {
+            if (iterator.throw) await iterator.throw(error);
+            throw error;
+          },
+        };
+      },
+    };
+  }
+
   //
   // Helpers
   //