Mark asyncIter-originated fetchPage calls

palantir · Jul 12, 2024 · a1a7a69 · a1a7a69
1 parent 5d4c562
commit a1a7a69
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 4 deletions.
diff --git a/packages/client.api/src/objectSet/ObjectSet.ts b/packages/client.api/src/objectSet/ObjectSet.ts
@@ -51,6 +51,7 @@ export interface MinimalObjectSet<Q extends ObjectOrInterfaceDefinition>
     S extends NullabilityAdherence = NullabilityAdherenceDefault,
   >(
     args?: FetchPageArgs<Q, L, R, A, S>,
+    calledByAsyncIter?: boolean,
   ) => Promise<FetchPageResult<Q, L, R, S>>;
 
   fetchPageWithErrors: <

diff --git a/packages/client/src/object/fetchPage.ts b/packages/client/src/object/fetchPage.ts
@@ -41,6 +41,7 @@ import type {
 import { OntologiesV2 } from "@osdk/internal.foundry";
 import type { MinimalClient } from "../MinimalClientContext.js";
 import { addUserAgentAndRequestContextHeaders } from "../util/addUserAgentAndRequestContextHeaders.js";
+import { augmentRequestContext } from "../util/augmentRequestContext.js";
 import { convertWireToOsdkObjects } from "./convertWireToOsdkObjects.js";
 
 export function augment<
@@ -129,11 +130,25 @@ export async function fetchPageInternal<
   objectType: Q,
   objectSet: ObjectSet,
   args: FetchPageArgs<Q, L, R, A, S> = {},
+  calledByAsyncIter: boolean = false,
 ): Promise<FetchPageResult<Q, L, R, S>> {
+  const modifiedClient = calledByAsyncIter
+    ? augmentRequestContext(client, _ => ({ finalMethodCall: "asyncIter" }))
+    : client;
   if (objectType.type === "interface") {
-    return await fetchInterfacePage(client, objectType, args, objectSet) as any; // fixme
+    return await fetchInterfacePage(
+      modifiedClient,
+      objectType,
+      args,
+      objectSet,
+    ) as any; // fixme
   } else {
-    return await fetchObjectPage(client, objectType, args, objectSet) as any; // fixme
+    return await fetchObjectPage(
+      modifiedClient,
+      objectType,
+      args,
+      objectSet,
+    ) as any; // fixme
   }
 }
 

diff --git a/packages/client/src/objectSet/createObjectSet.ts b/packages/client/src/objectSet/createObjectSet.ts
@@ -145,8 +145,7 @@ export function createObjectSet<Q extends ObjectOrInterfaceDefinition>(
     asyncIter: async function*(): AsyncIterableIterator<Osdk<Q, "$all">> {
       let $nextPageToken: string | undefined = undefined;
       do {
-        // TODO: figure out how to mark asyncIter-originated fetchPage requests
-        const result = await base.fetchPage({ $nextPageToken });
+        const result = await base.fetchPage({ $nextPageToken }, true);
 
         for (const obj of result.data) {
           yield obj as Osdk<Q, "$all">;