fix: handle multiple cache keys and more complex calls

Author: logaretm

packages/nuxt/src/runtime/plugins/storage.server.ts

@@ -1,5 +1,6 @@
 import {
   type SpanAttributes,
+  type StartSpanOptions,
   captureException,
   debug,
   flushIfServerless,
@@ -26,20 +27,6 @@ type MaybeInstrumentedDriver = MaybeInstrumented<Driver>;
 
 type DriverMethod = keyof Driver;
 
-/**
- * Methods that should have a key argument.
- */
-const KEYED_METHODS = new Set<DriverMethod>([
-  'hasItem',
-  'getItem',
-  'getItemRaw',
-  'getItems',
-  'setItem',
-  'setItemRaw',
-  'setItems',
-  'removeItem',
-]);
-
 /**
  * Methods that should have a attribute to indicate a cache hit.
  */
@@ -127,45 +114,35 @@ function createMethodWrapper(
 ): (...args: unknown[]) => unknown {
   return new Proxy(original, {
     async apply(target, thisArg, args) {
-      const attributes = getSpanAttributes(methodName, driver, mountBase, args);
+      const options = createSpanStartOptions(methodName, driver, mountBase, args);
 
       DEBUG_BUILD && debug.log(`[storage] Running method: "${methodName}" on driver: "${driver.name ?? 'unknown'}"`);
 
-      const spanName = KEYED_METHODS.has(methodName)
-        ? `${mountBase}${args?.[0]}`
-        : `storage.${normalizeMethodName(methodName)}`;
-
-      return startSpan(
-        {
-          name: spanName,
-          attributes,
-        },
-        async span => {
-          try {
-            const result = await target.apply(thisArg, args);
-            span.setStatus({ code: SPAN_STATUS_OK });
-
-            if (CACHE_HIT_METHODS.has(methodName)) {
-              span.setAttribute(SEMANTIC_ATTRIBUTE_CACHE_HIT, !isEmptyValue(result));
-            }
-
-            return result;
-          } catch (error) {
-            span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
-            captureException(error, {
-              mechanism: {
-                handled: false,
-                type: attributes[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN],
-              },
-            });
-
-            // Re-throw the error to be handled by the caller
-            throw error;
-          } finally {
-            await flushIfServerless();
+      return startSpan(options, async span => {
+        try {
+          const result = await target.apply(thisArg, args);
+          span.setStatus({ code: SPAN_STATUS_OK });
+
+          if (CACHE_HIT_METHODS.has(methodName)) {
+            span.setAttribute(SEMANTIC_ATTRIBUTE_CACHE_HIT, !isEmptyValue(result));
           }
-        },
-      );
+
+          return result;
+        } catch (error) {
+          span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
+          captureException(error, {
+            mechanism: {
+              handled: false,
+              type: options.attributes?.[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN],
+            },
+          });
+
+          // Re-throw the error to be handled by the caller
+          throw error;
+        } finally {
+          await flushIfServerless();
+        }
+      });
     },
   });
 }
@@ -191,11 +168,30 @@ function wrapStorageMount(storage: Storage): Storage['mount'] {
 
   return mountWithInstrumentation;
 }
+/**
+ * Normalizes the method name to snake_case to be used in span names or op.
+ */
+function normalizeMethodName(methodName: string): string {
+  return methodName.replace(/[A-Z]/g, letter => `_${letter.toLowerCase()}`);
+}
 
 /**
- * Gets the span attributes for the storage method.
+ * Checks if the value is empty, used for cache hit detection.
  */
-function getSpanAttributes(methodName: string, driver: Driver, mountBase: string, args: unknown[]): SpanAttributes {
+function isEmptyValue(value: unknown): boolean {
+  return value === null || value === undefined;
+}
+
+/**
+ * Creates the span start options for the storage method.
+ */
+function createSpanStartOptions(
+  methodName: keyof Driver,
+  driver: Driver,
+  mountBase: string,
+  args: unknown[],
+): StartSpanOptions {
+  const key = normalizeKey(args?.[0]);
   const attributes: SpanAttributes = {
     [SEMANTIC_ATTRIBUTE_SENTRY_OP]: `cache.${normalizeMethodName(methodName)}`,
     [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.cache.nuxt',
@@ -204,24 +200,38 @@ function getSpanAttributes(methodName: string, driver: Driver, mountBase: string
     'db.system.name': driver.name ?? 'unknown',
   };
 
-  // Add the key if it's a get/set/del call
-  if (args?.[0] && typeof args[0] === 'string') {
-    attributes[SEMANTIC_ATTRIBUTE_CACHE_KEY] = `${mountBase}${args[0]}`;
+  if (key) {
+    attributes[SEMANTIC_ATTRIBUTE_CACHE_KEY] = `${mountBase}${key}`;
   }
 
-  return attributes;
+  return {
+    name: `${normalizeMethodName(methodName)} ${mountBase}${key}`,
+    attributes,
+  };
 }
 
 /**
- * Normalizes the method name to snake_case to be used in span names or op.
+ * Normalizes the key to a string for display purposes.
+ * @param key The key to normalize.
  */
-function normalizeMethodName(methodName: string): string {
-  return methodName.replace(/[A-Z]/g, letter => `_${letter.toLowerCase()}`);
-}
+function normalizeKey(key: unknown): string {
+  if (isEmptyValue(key)) {
+    return '';
+  }
 
-/**
- * Checks if the value is empty, used for cache hit detection.
- */
-function isEmptyValue(value: unknown): boolean {
-  return value === null || value === undefined;
+  if (typeof key === 'string') {
+    return key;
+  }
+
+  // Handles an object with a key property
+  if (typeof key === 'object' && key !== null && 'key' in key) {
+    return `${key.key}`;
+  }
+
+  // Handles an array of keys
+  if (Array.isArray(key)) {
+    return key.map(k => normalizeKey(k)).join(', ');
+  }
+
+  return String(key);
 }