autotel-cloudflare 2.12.0 → 2.14.0

package/src/bindings/vectorize.test.ts

@@ -0,0 +1,340 @@
+import { describe, it, expect, beforeEach, vi } from 'vitest';
+import { trace, SpanStatusCode, SpanKind } from '@opentelemetry/api';
+import { instrumentVectorize } from './vectorize';
+
+describe('Vectorize Binding Instrumentation', () => {
+  let mockTracer: any;
+  let mockSpan: any;
+  let getTracerSpy: any;
+  let mockVectorize: any;
+
+  beforeEach(() => {
+    mockSpan = {
+      spanContext: () => ({
+        traceId: 'test-trace-id',
+        spanId: 'test-span-id',
+        traceFlags: 1,
+      }),
+      setAttribute: vi.fn(),
+      setAttributes: vi.fn(),
+      setStatus: vi.fn(),
+      recordException: vi.fn(),
+      end: vi.fn(),
+      isRecording: () => true,
+      updateName: vi.fn(),
+      addEvent: vi.fn(),
+    };
+
+    mockTracer = {
+      startActiveSpan: vi.fn((name, options, fn) => {
+        return fn(mockSpan);
+      }),
+    };
+
+    getTracerSpy = vi.spyOn(trace, 'getTracer').mockReturnValue(mockTracer as any);
+
+    mockVectorize = {
+      query: vi.fn(async () => ({
+        matches: [
+          { id: 'vec-1', score: 0.95 },
+          { id: 'vec-2', score: 0.87 },
+          { id: 'vec-3', score: 0.72 },
+        ],
+        count: 3,
+      })),
+      insert: vi.fn(async () => ({ mutationId: 'mut-1', count: 2 })),
+      upsert: vi.fn(async () => ({ mutationId: 'mut-2', count: 3 })),
+      deleteByIds: vi.fn(async () => ({ mutationId: 'mut-3', count: 1 })),
+      getByIds: vi.fn(async () => [{ id: 'vec-1', values: [0.1, 0.2] }]),
+      describe: vi.fn(async () => ({
+        dimensions: 128,
+        vectorCount: 1000,
+        processedUpTo: 'ts-123',
+      })),
+      // Non-instrumented method
+      toString: vi.fn(() => 'VectorizeIndex'),
+    } as unknown as VectorizeIndex;
+  });
+
+  describe('query()', () => {
+    it('should create span with correct attributes', async () => {
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      // The source reads topK from args[0], so we pass a query object as the first argument
+      await instrumented.query({ topK: 5 } as any, {} as any);
+
+      expect(mockTracer.startActiveSpan).toHaveBeenCalledWith(
+        'Vectorize my-index: query',
+        expect.objectContaining({
+          kind: SpanKind.CLIENT,
+          attributes: expect.objectContaining({
+            'db.system': 'cloudflare-vectorize',
+            'db.operation': 'query',
+            'db.collection.name': 'my-index',
+            'db.vectorize.top_k': 5,
+          }),
+        }),
+        expect.any(Function),
+      );
+    });
+
+    it('should record matches_count from result', async () => {
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      await instrumented.query([0.1, 0.2, 0.3] as any, {} as any);
+
+      expect(mockSpan.setAttribute).toHaveBeenCalledWith('db.vectorize.matches_count', 3);
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({ code: SpanStatusCode.OK });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+
+    it('should handle query without topK in first argument', async () => {
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      await instrumented.query([0.1, 0.2, 0.3] as any, {} as any);
+
+      const attributes = mockTracer.startActiveSpan.mock.calls[0][1].attributes;
+      expect(attributes['db.vectorize.top_k']).toBeUndefined();
+    });
+  });
+
+  describe('insert()', () => {
+    it('should set vectors_count from input array length', async () => {
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      const vectors = [
+        { id: 'vec-1', values: [0.1, 0.2] },
+        { id: 'vec-2', values: [0.3, 0.4] },
+      ];
+      await instrumented.insert(vectors as any);
+
+      expect(mockTracer.startActiveSpan).toHaveBeenCalledWith(
+        'Vectorize my-index: insert',
+        expect.objectContaining({
+          kind: SpanKind.CLIENT,
+          attributes: expect.objectContaining({
+            'db.system': 'cloudflare-vectorize',
+            'db.operation': 'insert',
+            'db.collection.name': 'my-index',
+            'db.vectorize.vectors_count': 2,
+          }),
+        }),
+        expect.any(Function),
+      );
+
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({ code: SpanStatusCode.OK });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+  });
+
+  describe('upsert()', () => {
+    it('should set vectors_count from input array length', async () => {
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      const vectors = [
+        { id: 'vec-1', values: [0.1, 0.2] },
+        { id: 'vec-2', values: [0.3, 0.4] },
+        { id: 'vec-3', values: [0.5, 0.6] },
+      ];
+      await instrumented.upsert(vectors as any);
+
+      expect(mockTracer.startActiveSpan).toHaveBeenCalledWith(
+        'Vectorize my-index: upsert',
+        expect.objectContaining({
+          kind: SpanKind.CLIENT,
+          attributes: expect.objectContaining({
+            'db.system': 'cloudflare-vectorize',
+            'db.operation': 'upsert',
+            'db.collection.name': 'my-index',
+            'db.vectorize.vectors_count': 3,
+          }),
+        }),
+        expect.any(Function),
+      );
+
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({ code: SpanStatusCode.OK });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+  });
+
+  describe('deleteByIds()', () => {
+    it('should create correct span', async () => {
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      await instrumented.deleteByIds(['vec-1', 'vec-2']);
+
+      expect(mockTracer.startActiveSpan).toHaveBeenCalledWith(
+        'Vectorize my-index: deleteByIds',
+        expect.objectContaining({
+          kind: SpanKind.CLIENT,
+          attributes: expect.objectContaining({
+            'db.system': 'cloudflare-vectorize',
+            'db.operation': 'deleteByIds',
+            'db.collection.name': 'my-index',
+          }),
+        }),
+        expect.any(Function),
+      );
+
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({ code: SpanStatusCode.OK });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+  });
+
+  describe('describe()', () => {
+    it('should create correct span', async () => {
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      await instrumented.describe();
+
+      expect(mockTracer.startActiveSpan).toHaveBeenCalledWith(
+        'Vectorize my-index: describe',
+        expect.objectContaining({
+          kind: SpanKind.CLIENT,
+          attributes: expect.objectContaining({
+            'db.system': 'cloudflare-vectorize',
+            'db.operation': 'describe',
+            'db.collection.name': 'my-index',
+          }),
+        }),
+        expect.any(Function),
+      );
+
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({ code: SpanStatusCode.OK });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+  });
+
+  describe('Error handling', () => {
+    it('should record exception and set error status on query() failure', async () => {
+      mockVectorize.query = vi.fn(async () => {
+        throw new Error('Vectorize query failed');
+      });
+
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      await expect(instrumented.query([0.1, 0.2] as any, {} as any)).rejects.toThrow('Vectorize query failed');
+
+      expect(mockSpan.recordException).toHaveBeenCalled();
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({
+        code: SpanStatusCode.ERROR,
+        message: 'Vectorize query failed',
+      });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+
+    it('should record exception and set error status on insert() failure', async () => {
+      mockVectorize.insert = vi.fn(async () => {
+        throw new Error('Insert failed');
+      });
+
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      await expect(instrumented.insert([{ id: 'v1', values: [0.1] }] as any)).rejects.toThrow('Insert failed');
+
+      expect(mockSpan.recordException).toHaveBeenCalled();
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({
+        code: SpanStatusCode.ERROR,
+        message: 'Insert failed',
+      });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+
+    it('should record exception and set error status on upsert() failure', async () => {
+      mockVectorize.upsert = vi.fn(async () => {
+        throw new Error('Upsert failed');
+      });
+
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      await expect(instrumented.upsert([{ id: 'v1', values: [0.1] }] as any)).rejects.toThrow('Upsert failed');
+
+      expect(mockSpan.recordException).toHaveBeenCalled();
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({
+        code: SpanStatusCode.ERROR,
+        message: 'Upsert failed',
+      });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+
+    it('should record exception and set error status on deleteByIds() failure', async () => {
+      mockVectorize.deleteByIds = vi.fn(async () => {
+        throw new Error('Delete failed');
+      });
+
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      await expect(instrumented.deleteByIds(['vec-1'])).rejects.toThrow('Delete failed');
+
+      expect(mockSpan.recordException).toHaveBeenCalled();
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({
+        code: SpanStatusCode.ERROR,
+        message: 'Delete failed',
+      });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+
+    it('should record exception and set error status on describe() failure', async () => {
+      mockVectorize.describe = vi.fn(async () => {
+        throw new Error('Describe failed');
+      });
+
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      await expect(instrumented.describe()).rejects.toThrow('Describe failed');
+
+      expect(mockSpan.recordException).toHaveBeenCalled();
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({
+        code: SpanStatusCode.ERROR,
+        message: 'Describe failed',
+      });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+
+    it('should handle non-Error exceptions', async () => {
+      mockVectorize.query = vi.fn(async () => {
+        throw 'string error';
+      });
+
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      await expect(instrumented.query([0.1] as any, {} as any)).rejects.toThrow('string error');
+
+      expect(mockSpan.recordException).toHaveBeenCalled();
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({
+        code: SpanStatusCode.ERROR,
+        message: 'string error',
+      });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+  });
+
+  describe('Non-instrumented methods', () => {
+    it('should pass through non-traced methods without creating spans', () => {
+      const instrumented = instrumentVectorize(mockVectorize, 'my-index');
+
+      const result = instrumented.toString();
+
+      expect(result).toBe('VectorizeIndex');
+      expect(mockTracer.startActiveSpan).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('Default index name', () => {
+    it('should use "vectorize" as default index name when none provided', async () => {
+      const instrumented = instrumentVectorize(mockVectorize);
+
+      await instrumented.query([0.1, 0.2] as any, {} as any);
+
+      expect(mockTracer.startActiveSpan).toHaveBeenCalledWith(
+        'Vectorize vectorize: query',
+        expect.objectContaining({
+          attributes: expect.objectContaining({
+            'db.collection.name': 'vectorize',
+          }),
+        }),
+        expect.any(Function),
+      );
+    });
+  });
+});

package/src/bindings/vectorize.ts

@@ -0,0 +1,86 @@
+/**
+ * Vectorize binding instrumentation
+ */
+
+import {
+  trace,
+  SpanKind,
+  SpanStatusCode,
+} from '@opentelemetry/api';
+import type { WorkerTracer } from 'autotel-edge';
+import { wrap, setAttr } from './common';
+
+const TRACED_METHODS = ['query', 'insert', 'upsert', 'deleteByIds', 'getByIds', 'describe'] as const;
+
+/**
+ * Instrument Vectorize index binding
+ */
+export function instrumentVectorize<T extends VectorizeIndex>(vectorize: T, indexName?: string): T {
+  const name = indexName || 'vectorize';
+
+  const handler: ProxyHandler<T> = {
+    get(target, prop) {
+      const value = Reflect.get(target, prop);
+
+      if (typeof prop === 'string' && TRACED_METHODS.includes(prop as any) && typeof value === 'function') {
+        return new Proxy(value, {
+          apply: (fnTarget, thisArg, args) => {
+            const operation = prop as string;
+            const tracer = trace.getTracer('autotel-edge') as WorkerTracer;
+
+            const attributes: Record<string, string | number> = {
+              'db.system': 'cloudflare-vectorize',
+              'db.operation': operation,
+              'db.collection.name': name,
+            };
+
+            // Per-operation attributes
+            if (operation === 'query') {
+              const queryInput = args[0] as { topK?: number } | undefined;
+              if (queryInput?.topK !== undefined) {
+                attributes['db.vectorize.top_k'] = queryInput.topK;
+              }
+            }
+
+            if ((operation === 'insert' || operation === 'upsert') && Array.isArray(args[0])) {
+              attributes['db.vectorize.vectors_count'] = args[0].length;
+            }
+
+            return tracer.startActiveSpan(
+              `Vectorize ${name}: ${operation}`,
+              {
+                kind: SpanKind.CLIENT,
+                attributes,
+              },
+              async (span) => {
+                try {
+                  const result = await Reflect.apply(fnTarget, thisArg, args);
+
+                  if (operation === 'query' && result?.matches) {
+                    setAttr(span, 'db.vectorize.matches_count', result.matches.length);
+                  }
+
+                  span.setStatus({ code: SpanStatusCode.OK });
+                  return result;
+                } catch (error) {
+                  span.recordException(error as Error);
+                  span.setStatus({
+                    code: SpanStatusCode.ERROR,
+                    message: error instanceof Error ? error.message : String(error),
+                  });
+                  throw error;
+                } finally {
+                  span.end();
+                }
+              },
+            );
+          },
+        });
+      }
+
+      return value;
+    },
+  };
+
+  return wrap(vectorize, handler);
+}