Added collectRawUsage to Hive Client core package (#5116)

graphql-hive · Jul 3, 2024 · f1e43c6 · f1e43c6
1 parent a2406d8
commit f1e43c6
Show file tree

Hide file tree

Showing 5 changed files with 166 additions and 148 deletions.
diff --git a/.changeset/unlucky-singers-shop.md b/.changeset/unlucky-singers-shop.md
@@ -0,0 +1,5 @@
+---
+'@graphql-hive/core': minor
+---
+
+Added `collectRawUsage` to Hive Client core package
diff --git a/packages/libraries/core/src/client/client.ts b/packages/libraries/core/src/client/client.ts
@@ -37,6 +37,10 @@ export function createHive(options: HivePluginOptions): HiveClient {
     return usage.collect();
   }
 
+  function collectRawUsage(...args: Parameters<typeof usage.collectRaw>) {
+    return usage.collectRaw(...args);
+  }
+
   async function dispose() {
     await Promise.all([schemaReporter.dispose(), usage.dispose()]);
   }
@@ -207,6 +211,7 @@ export function createHive(options: HivePluginOptions): HiveClient {
     info,
     reportSchema,
     collectUsage,
+    collectRawUsage,
     dispose,
     collectSubscriptionUsage: usage.collectSubscription,
     createInstrumentedSubscribe,

diff --git a/packages/libraries/core/src/client/types.ts b/packages/libraries/core/src/client/types.ts
@@ -10,6 +10,11 @@ export interface HiveClient {
   reportSchema: SchemaReporter['report'];
   /** Collect usage for Query and Mutation operations */
   collectUsage(): CollectUsageCallback;
+  collectRawUsage(
+    args: ExecutionArgs,
+    result: GraphQLErrorsResult | AbortAction,
+    duration: number,
+  ): void;
   /** Collect usage for Subscription operations */
   collectSubscriptionUsage(args: { args: ExecutionArgs }): void;
   createInstrumentedExecute(executeImpl: any): any;

diff --git a/packages/libraries/core/src/client/usage.ts b/packages/libraries/core/src/client/usage.ts
@@ -16,13 +16,19 @@ import type {
   AbortAction,
   ClientInfo,
   CollectUsageCallback,
+  GraphQLErrorsResult,
   HivePluginOptions,
   HiveUsagePluginOptions,
 } from './types.js';
 import { cache, cacheDocumentKey, logIf, measureDuration, memo } from './utils.js';
 
 interface UsageCollector {
   collect(): CollectUsageCallback;
+  collectRaw(
+    args: ExecutionArgs,
+    result: GraphQLErrorsResult | AbortAction,
+    durationMs: number,
+  ): void;
   collectSubscription(args: { args: ExecutionArgs }): void;
   dispose(): Promise<void>;
 }
@@ -37,6 +43,7 @@ export function createUsage(pluginOptions: HivePluginOptions): UsageCollector {
       collect() {
         return async () => {};
       },
+      collectRaw() {},
       async dispose() {},
       collectSubscription() {},
     };
@@ -149,87 +156,93 @@ export function createUsage(pluginOptions: HivePluginOptions): UsageCollector {
       ? dynamicSampling(options.sampler)
       : randomSampling(options.sampleRate ?? 1.0);
 
+  const collectRaw: UsageCollector['collectRaw'] = (args, result, duration) => {
+    let providedOperationName: string | undefined = undefined;
+    try {
+      if (isAbortAction(result)) {
+        if (result.logging) {
+          logger.info(result.reason);
+        }
+        return;
+      }
+
+      const document = args.document;
+      const rootOperation = document.definitions.find(
+        o => o.kind === Kind.OPERATION_DEFINITION,
+      ) as OperationDefinitionNode;
+      providedOperationName = args.operationName || rootOperation.name?.value;
+      const operationName = providedOperationName || 'anonymous';
+      // Check if operationName is a match with any string or regex in excludeSet
+      const isMatch = Array.from(excludeSet).some(excludingValue =>
+        excludingValue instanceof RegExp
+          ? excludingValue.test(operationName)
+          : operationName === excludingValue,
+      );
+      if (
+        !isMatch &&
+        shouldInclude({
+          operationName,
+          document,
+          variableValues: args.variableValues,
+          contextValue: args.contextValue,
+        })
+      ) {
+        const errors =
+          result.errors?.map(error => ({
+            message: error.message,
+            path: error.path?.join('.'),
+          })) ?? [];
+        const collect = collector({
+          schema: args.schema,
+          max: options.max ?? 1000,
+          ttl: options.ttl,
+          processVariables: options.processVariables ?? false,
+        });
+
+        agent.capture(
+          collect(document, args.variableValues ?? null).then(({ key, value: info }) => {
+            return {
+              type: 'request',
+              data: {
+                key,
+                timestamp: Date.now(),
+                operationName,
+                operation: info.document,
+                fields: info.fields,
+                execution: {
+                  ok: errors.length === 0,
+                  duration,
+                  errorsTotal: errors.length,
+                  errors,
+                },
+                // TODO: operationHash is ready to accept hashes of persisted operations
+                client: pickClientInfoProperties(
+                  typeof args.contextValue !== 'undefined' &&
+                    typeof options.clientInfo !== 'undefined'
+                    ? options.clientInfo(args.contextValue)
+                    : createDefaultClientInfo()(args.contextValue),
+                ),
+              },
+            };
+          }),
+        );
+      }
+    } catch (error) {
+      const details = providedOperationName ? ` (name: "${providedOperationName}")` : '';
+      logger.error(`Failed to collect operation${details}`, error);
+    }
+  };
+
   return {
     dispose: agent.dispose,
+    collectRaw,
     collect() {
       const finish = measureDuration();
 
       return async function complete(args, result) {
         const duration = finish();
-        let providedOperationName: string | undefined = undefined;
-        try {
-          if (isAbortAction(result)) {
-            if (result.logging) {
-              logger.info(result.reason);
-            }
-            return;
-          }
 
-          const document = args.document;
-          const rootOperation = document.definitions.find(
-            o => o.kind === Kind.OPERATION_DEFINITION,
-          ) as OperationDefinitionNode;
-          providedOperationName = args.operationName || rootOperation.name?.value;
-          const operationName = providedOperationName || 'anonymous';
-          // Check if operationName is a match with any string or regex in excludeSet
-          const isMatch = Array.from(excludeSet).some(excludingValue =>
-            excludingValue instanceof RegExp
-              ? excludingValue.test(operationName)
-              : operationName === excludingValue,
-          );
-          if (
-            !isMatch &&
-            shouldInclude({
-              operationName,
-              document,
-              variableValues: args.variableValues,
-              contextValue: args.contextValue,
-            })
-          ) {
-            const errors =
-              result.errors?.map(error => ({
-                message: error.message,
-                path: error.path?.join('.'),
-              })) ?? [];
-            const collect = collector({
-              schema: args.schema,
-              max: options.max ?? 1000,
-              ttl: options.ttl,
-              processVariables: options.processVariables ?? false,
-            });
-
-            agent.capture(
-              collect(document, args.variableValues ?? null).then(({ key, value: info }) => {
-                return {
-                  type: 'request',
-                  data: {
-                    key,
-                    timestamp: Date.now(),
-                    operationName,
-                    operation: info.document,
-                    fields: info.fields,
-                    execution: {
-                      ok: errors.length === 0,
-                      duration,
-                      errorsTotal: errors.length,
-                      errors,
-                    },
-                    // TODO: operationHash is ready to accept hashes of persisted operations
-                    client: pickClientInfoProperties(
-                      typeof args.contextValue !== 'undefined' &&
-                        typeof options.clientInfo !== 'undefined'
-                        ? options.clientInfo(args.contextValue)
-                        : createDefaultClientInfo()(args.contextValue),
-                    ),
-                  },
-                };
-              }),
-            );
-          }
-        } catch (error) {
-          const details = providedOperationName ? ` (name: "${providedOperationName}")` : '';
-          logger.error(`Failed to collect operation${details}`, error);
-        }
+        return collectRaw(args, result, duration);
       };
     },
     async collectSubscription({ args }) {