[tracing] Propagate TraceID to the backend. Create a new setupEnv span. Fix span hierarchy

packages/client/src/api/fetcher.ts

@@ -97,7 +97,7 @@ export async function fetch<
   url: path,
   method,
   body,
-  headers,
+  headers: extraHeaders,
   pathParams,
   queryParams,
   fetchImpl,
@@ -107,8 +107,8 @@ export async function fetch<
   trace
 }: FetcherOptions<TBody, THeaders, TQueryParams, TPathParams> & FetcherExtraProps): Promise<TData> {
   return trace(
-    `${method.toUpperCase()} ${path}`,
-    async ({ setAttributes }) => {
+    `[HTTP] ${method.toUpperCase()} ${path}`,
+    async ({ setAttributes, setHeaders }) => {
       const baseUrl = buildBaseUrl({ path, workspacesApiUrl, pathParams, apiUrl });
       const fullUrl = resolveUrl(baseUrl, queryParams, pathParams);
 
@@ -120,16 +120,20 @@ export async function fetch<
         [TraceAttributes.HTTP_TARGET]: resolveUrl(path, queryParams, pathParams)
       });
 
+      const headers = {
+        'Content-Type': 'application/json',
+        'User-Agent': `Xata client-ts/${VERSION}`,
+        ...extraHeaders,
+        ...hostHeader(fullUrl),
+        Authorization: `Bearer ${apiKey}`
+      };
+
+      setHeaders(headers);
+
       const response = await fetchImpl(url, {
         method: method.toUpperCase(),
         body: body ? JSON.stringify(body) : undefined,
-        headers: {
-          'Content-Type': 'application/json',
-          'User-Agent': `Xata client-ts/${VERSION}`,
-          ...headers,
-          ...hostHeader(fullUrl),
-          Authorization: `Bearer ${apiKey}`
-        }
+        headers
       });
 
       // No content

packages/client/src/schema/index.ts

@@ -13,6 +13,7 @@ export { isIdentifiable, isXataRecord } from './record';
 export type { BaseData, EditableData, Identifiable, Link, XataRecord } from './record';
 export { Repository, RestRepository } from './repository';
 export * from './selection';
+export * from './tracing';
 
 export type SchemaDefinition = {
   table: string;

packages/client/src/schema/repository.ts

@@ -27,7 +27,7 @@ import { Query } from './query';
 import { EditableData, Identifiable, isIdentifiable, XataRecord } from './record';
 import { SelectableColumn, SelectedPick } from './selection';
 import { buildSortFilter } from './sorting';
-import { AttributeDictionary, defaultTrace, TraceAttributes, TraceFunction } from './tracing';
+import { AttributeDictionary, defaultTrace, TraceAttributes, TraceFunction, TraceFunctionCallback } from './tracing';
 
 /**
  * Common interface for performing operations on a table.
@@ -427,12 +427,8 @@ export class RestRepository<Record extends XataRecord>
     this.#schemaTables = options.schemaTables;
 
     const trace = options.pluginOptions.trace ?? defaultTrace;
-    this.#trace = async <T>(
-      name: string,
-      fn: (options: { setAttributes: (attrs: AttributeDictionary) => void }) => T,
-      options: AttributeDictionary = {}
-    ) => {
-      return trace<T>(name, fn, {
+    this.#trace = <T>(name: string, fn: TraceFunctionCallback<T>, options: AttributeDictionary = {}) => {
+      return trace<T>(`[SDK] ${name}`, fn, {
         ...options,
         [TraceAttributes.TABLE]: this.#table,
         [TraceAttributes.KIND]: 'sdk-operation',

packages/client/src/schema/tracing.ts

@@ -1,19 +1,32 @@
 export type AttributeDictionary = Record<string, string | number | boolean | undefined>;
 
+export type SetAttributesFn = (
+  attrs: AttributeDictionary,
+  extra: { headers: Record<string, string | undefined> }
+) => void;
+
+export type TraceFunctionCallback<T> = (options: {
+  setAttributes: (attrs: AttributeDictionary) => void;
+  setHeaders: (headers: AttributeDictionary) => void;
+}) => T;
+
 export type TraceFunction = <T>(
   name: string,
-  fn: (options: { setAttributes: (attrs: AttributeDictionary) => void }) => T,
-  options?: AttributeDictionary
+  fn: TraceFunctionCallback<T>,
+  attributes?: AttributeDictionary
 ) => Promise<T>;
 
 export const defaultTrace: TraceFunction = async <T>(
   _name: string,
-  fn: (options: { setAttributes: (attrs: Record<string, string | number | boolean | undefined>) => void }) => T,
-  _options?: Record<string, any>
+  fn: TraceFunctionCallback<T>,
+  _attributes?: Record<string, any>
 ): Promise<T> => {
   return await fn({
     setAttributes: () => {
       return;
+    },
+    setHeaders: () => {
+      return;
     }
   });
 };

packages/plugin-client-opentelemetry/src/index.ts

@@ -1,38 +1,45 @@
-import { SpanStatusCode, Tracer } from '@opentelemetry/api';
+import { SpanStatusCode, propagation, Tracer, context, trace } from '@opentelemetry/api';
+import { TraceFunctionCallback, AttributeDictionary } from '@xata.io/client';
 
 export const buildTraceFunction =
   (tracer: Tracer) =>
-  async <T>(
-    name: string,
-    fn: (options: { setAttributes: (attrs: Record<string, any>) => void }) => T,
-    attributes: Record<string, string | number | boolean | undefined> = {}
-  ): Promise<T> => {
-    return await tracer.startActiveSpan(name, { attributes }, async (span) => {
-      try {
-        const setAttributes = (attrs: Record<string, any>) => {
-          for (const [key, value] of Object.entries(attrs)) {
-            span.setAttribute(key, value);
-          }
-        };
-
-        return await fn({ setAttributes });
-      } catch (error: any) {
-        const message = error.message ?? error.toString();
-
-        // We only log errors 500, as they are the only ones that are not expected
-        if (isHttpError(error) && !String(error.status).startsWith('5')) {
-          span.setStatus({ code: SpanStatusCode.UNSET, message });
-        } else {
-          span.setStatus({ code: SpanStatusCode.ERROR, message });
-        }
-
-        span.recordException(message);
-
-        throw error;
-      } finally {
-        span.end();
+  async <T>(name: string, fn: TraceFunctionCallback<T>, attributes: AttributeDictionary = {}): Promise<T> => {
+    const span = tracer.startSpan(name, { attributes });
+
+    const setAttributes = (attrs: AttributeDictionary) => {
+      for (const [key, value] of Object.entries(attrs)) {
+        if (value) span.setAttribute(key, value);
       }
-    });
+    };
+
+    const setHeaders = (headers: AttributeDictionary) => {
+      propagation.inject(context.active(), headers);
+    };
+
+    try {
+      const spanCtx = trace.setSpan(context.active(), span);
+
+      const result = await context.with(spanCtx, async () => {
+        return await fn({ setAttributes, setHeaders });
+      });
+
+      return result;
+    } catch (error: any) {
+      const message = error.message ?? error.toString();
+
+      // We only log errors 500, as they are the only ones that are not expected
+      if (isHttpError(error) && !String(error.status).startsWith('5')) {
+        span.setStatus({ code: SpanStatusCode.UNSET, message });
+      } else {
+        span.setStatus({ code: SpanStatusCode.ERROR, message });
+      }
+
+      span.recordException(message);
+
+      throw error;
+    } finally {
+      span.end();
+    }
   };
 
 function isHttpError(error: any): error is Error & { status?: number } {

test/integration/branches.test.ts

@@ -1,9 +1,10 @@
 import { execSync } from 'child_process';
-import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'vitest';
+import { afterAll, afterEach, beforeAll, beforeEach, expect, describe } from 'vitest';
 import { XataApiClient } from '../../packages/client/src';
 import { getCurrentBranchName } from '../../packages/client/src/util/config';
 import { XataClient } from '../../packages/codegen/example/xata';
 import { setUpTestEnvironment, TestEnvironmentResult } from '../utils/setup';
+import { test } from '../utils/tracing';
 
 let api: XataApiClient;
 let workspace: string;

test/integration/cache.test.ts

@@ -1,7 +1,8 @@
-import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'vitest';
+import { afterAll, afterEach, beforeAll, beforeEach, expect, describe } from 'vitest';
 import { BaseClientOptions, SimpleCache } from '../../packages/client/src';
 import { XataClient } from '../../packages/codegen/example/xata';
 import { setUpTestEnvironment, TestEnvironmentResult } from '../utils/setup';
+import { test } from '../utils/tracing';
 
 const cache = new SimpleCache();
 

test/integration/create.test.ts

@@ -1,6 +1,7 @@
-import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'vitest';
+import { afterAll, afterEach, beforeAll, beforeEach, expect, describe } from 'vitest';
 import { XataClient } from '../../packages/codegen/example/xata';
 import { setUpTestEnvironment, TestEnvironmentResult } from '../utils/setup';
+import { test } from '../utils/tracing';
 
 let xata: XataClient;
 let hooks: TestEnvironmentResult['hooks'];

test/integration/createOrUpdate.test.ts

@@ -1,6 +1,7 @@
-import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'vitest';
+import { afterAll, afterEach, beforeAll, beforeEach, expect, describe } from 'vitest';
 import { XataClient } from '../../packages/codegen/example/xata';
 import { setUpTestEnvironment, TestEnvironmentResult } from '../utils/setup';
+import { test } from '../utils/tracing';
 
 let xata: XataClient;
 let hooks: TestEnvironmentResult['hooks'];

test/integration/dates.test.ts

@@ -1,7 +1,8 @@
-import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'vitest';
+import { afterAll, afterEach, beforeAll, beforeEach, expect, describe } from 'vitest';
 import { gte, is, lt } from '../../packages/client/src';
 import { XataClient } from '../../packages/codegen/example/xata';
 import { setUpTestEnvironment, TestEnvironmentResult } from '../utils/setup';
+import { test } from '../utils/tracing';
 
 let xata: XataClient;
 let hooks: TestEnvironmentResult['hooks'];

test/integration/delete.test.ts

@@ -1,6 +1,7 @@
-import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'vitest';
+import { afterAll, afterEach, beforeAll, beforeEach, expect, describe } from 'vitest';
 import { XataClient } from '../../packages/codegen/example/xata';
 import { setUpTestEnvironment, TestEnvironmentResult } from '../utils/setup';
+import { test } from '../utils/tracing';
 
 let xata: XataClient;
 let hooks: TestEnvironmentResult['hooks'];

test/integration/query.test.ts

@@ -1,4 +1,4 @@
-import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'vitest';
+import { afterAll, afterEach, beforeAll, beforeEach, expect, describe } from 'vitest';
 import {
   BaseClient,
   contains,
@@ -19,6 +19,7 @@ import {
 import { UsersRecord, XataClient } from '../../packages/codegen/example/xata';
 import { animalUsers, fruitUsers, mockUsers, ownerAnimals, ownerFruits } from '../mock_data';
 import { setUpTestEnvironment, TestEnvironmentResult } from '../utils/setup';
+import { test } from '../utils/tracing';
 
 let xata: XataClient;
 let api: XataApiClient;

test/integration/read.test.ts

@@ -1,6 +1,7 @@
-import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'vitest';
+import { afterAll, afterEach, beforeAll, beforeEach, expect, describe } from 'vitest';
 import { XataClient } from '../../packages/codegen/example/xata';
 import { setUpTestEnvironment, TestEnvironmentResult } from '../utils/setup';
+import { test } from '../utils/tracing';
 
 let xata: XataClient;
 let hooks: TestEnvironmentResult['hooks'];

test/integration/search.test.ts

@@ -1,7 +1,8 @@
-import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'vitest';
+import { afterAll, afterEach, beforeAll, beforeEach, expect, describe } from 'vitest';
 import { XataClient } from '../../packages/codegen/example/xata';
 import { mockUsers } from '../mock_data';
 import { setUpTestEnvironment, TestEnvironmentResult } from '../utils/setup';
+import { test } from '../utils/tracing';
 
 let xata: XataClient;
 let hooks: TestEnvironmentResult['hooks'];

test/integration/update.test.ts

@@ -1,6 +1,7 @@
-import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'vitest';
+import { afterAll, afterEach, beforeAll, beforeEach, expect, describe } from 'vitest';
 import { XataClient } from '../../packages/codegen/example/xata';
 import { setUpTestEnvironment, TestEnvironmentResult } from '../utils/setup';
+import { test } from '../utils/tracing';
 
 let xata: XataClient;
 let hooks: TestEnvironmentResult['hooks'];