feat(prestodb-driver): Tag Presto/Trino queries with requestId trace token (#12080)

PrestoDriver (and so TrinoDriver) now sends the Cube request UUID
(extractRequestUUID(requestId)) as X-Trino-Trace-Token / X-Presto-Trace-Token
on the initial POST /v1/statement for query(), stream(), downloadQueryResults()
and the export-bucket unload path, so Trino/Presto queries can be traced back
to the Cube request, like BigQuery job labels and ClickHouse query_id. The
token is skipped for ids that fail isValidRequestId, since
scheduledRefreshContexts ids are only warned about and an invalid header value
would make Node fail the query.

BaseDriver.downloadQueryResults now forwards requestId to query(), which also
fixes BigQuery: its non-streamImport downloads were missing the
cube_request_id label.

The Trino and Presto docs pages get a "Query attribution" section.

Fixes #12072

Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: ovr

docs-mintlify/admin/connect-to-data/data-sources/presto.mdx

@@ -121,6 +121,13 @@ To enable SSL-encrypted connections between Cube and Presto, set the
 configure custom certificates, please check out [Enable SSL Connections to the
 Database][ref-recipe-enable-ssl].
 
+## Query attribution
+
+Cube sends the identifier shown for the query in Query History as the
+`X-Presto-Trace-Token` header of every statement it runs, so all statements of one Cube
+query share it. Presto records it as the session's trace token, returned as
+`session.traceToken` by the query info endpoint (`GET /v1/query/<query-id>`).
+
 ## Custom headers
 
 The Presto driver supports forwarding custom HTTP headers on every request to

docs-mintlify/admin/connect-to-data/data-sources/trino.mdx

@@ -121,6 +121,15 @@ To enable SSL-encrypted connections between Cube and Trino, set the
 configure custom certificates, please check out [Enable SSL Connections to the
 Database][ref-recipe-enable-ssl].
 
+## Query attribution
+
+Cube sends the identifier shown for the query in Query History as the
+`X-Trino-Trace-Token` header of every statement it runs, so all statements of one Cube
+query share it. Trino records it as the session's trace token, returned as
+`session.traceToken` by the query info endpoint (`GET /v1/query/<query-id>`).
+It is also passed to [event listeners][trino-docs-event-listener] as the query context's
+`traceToken`.
+
 ## Custom headers
 
 The Trino driver supports forwarding custom HTTP headers (e.g., `X-Trino-Source`,
@@ -209,4 +218,5 @@ module.exports = {
 [trino-docs-approx-agg-fns]:
   https://trino.io/docs/current/functions/aggregate.html#approximate-aggregate-functions
 [ref-recipe-enable-ssl]: /recipes/configuration/using-ssl-connections-to-data-source
-[ref-schema-ref-types-formats-countdistinctapprox]: /reference/data-modeling/measures#type
+[ref-schema-ref-types-formats-countdistinctapprox]: /reference/data-modeling/measures#type
+[trino-docs-event-listener]: https://trino.io/docs/current/develop/event-listener.html

packages/cubejs-base-driver/src/BaseDriver.ts

@@ -296,8 +296,8 @@ export abstract class BaseDriver implements DriverInterface {
     throw new TypeError('Driver\'s .streamQuery() method is not implemented yet.');
   }
 
-  public async downloadQueryResults(query: string, values: unknown[], _options: DownloadQueryResultsOptions): Promise<DownloadQueryResultsResult> {
-    const rows = await this.query<Row>(query, values);
+  public async downloadQueryResults(query: string, values: unknown[], options: DownloadQueryResultsOptions): Promise<DownloadQueryResultsResult> {
+    const rows = await this.query<Row>(query, values, { requestId: options?.requestId });
     const types = detectTypesFromTabular(rows);
 
     return {

packages/cubejs-prestodb-driver/src/PrestoDriver.ts

@@ -7,6 +7,7 @@
 import {
   DownloadQueryResultsOptions, DownloadQueryResultsResult,
   DriverCapabilities, DriverInterface,
+  QueryOptions,
   StreamOptions,
   StreamTableData,
   TableStructure,
@@ -17,6 +18,8 @@ import {
   getEnv,
   assertDataSource,
   formatAnsi,
+  extractRequestUUID,
+  isValidRequestId,
 } from '@cubejs-backend/shared';
 
 import { Transform, TransformCallback } from 'stream';
@@ -82,6 +85,8 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
 
   protected useSelectTestConnection: boolean;
 
+  private readonly traceTokenHeader: string;
+
   /**
    * Class constructor.
    */
@@ -127,6 +132,7 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
       ...config
     };
     this.catalog = this.config.catalog;
+    this.traceTokenHeader = this.config.engine === 'trino' ? 'X-Trino-Trace-Token' : 'X-Presto-Trace-Token';
     this.client = new presto.Client({
       timeout: this.config.queryTimeout,
       engine: 'presto',
@@ -212,15 +218,19 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
     await this.queryPromised('SELECT 1', false);
   }
 
-  public query(query: string, values: unknown[]): Promise<any[]> {
-    return <Promise<any[]>> this.queryPromised(this.prepareQueryWithParams(query, values), false);
+  public query(query: string, values: unknown[], options?: QueryOptions): Promise<any[]> {
+    return <Promise<any[]>> this.queryPromised(this.prepareQueryWithParams(query, values), false, options?.requestId);
   }
 
   protected prepareQueryWithParams(query: string, values: unknown[]) {
     return formatAnsi(query, values || []);
   }
 
-  public queryPromised(query: string, streaming: boolean): Promise<any[] | StreamTableData> {
+  public queryPromised(query: string, streaming: boolean, requestId?: string): Promise<any[] | StreamTableData> {
+    const traceToken = requestId && extractRequestUUID(requestId);
+    const headers = traceToken && isValidRequestId(traceToken)
+      ? { ...this.config.headers, [this.traceTokenHeader]: traceToken }
+      : this.config.headers;
     const toError = (error: any) => new Error(error.error ? `${error.message}\n${error.error}` : error.message);
     if (streaming) {
       const rowStream = new Transform({
@@ -236,7 +246,7 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
         this.client.execute({
           query,
           schema: this.config.schema || 'default',
-          headers: this.config.headers,
+          headers,
           session: this.config.queryTimeout ? `query_max_run_time=${this.config.queryTimeout}s` : undefined,
           columns: (error: any, columns: TableStructure) => {
             resolve({
@@ -266,7 +276,7 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
         this.client.execute({
           query,
           schema: this.config.schema || 'default',
-          headers: this.config.headers,
+          headers,
           data: (error: any, data: any[], columns: TableStructure) => {
             const normalData = this.normalizeResultOverColumns(data, columns);
             fullData = concat(normalData, fullData);
@@ -347,10 +357,10 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
     return map(arrayToObject, data || []);
   }
 
-  public stream(query: string, values: unknown[], _options: StreamOptions): Promise<StreamTableData> {
+  public stream(query: string, values: unknown[], options: StreamOptions): Promise<StreamTableData> {
     const queryWithParams = this.prepareQueryWithParams(query, values);
 
-    return <Promise<StreamTableData>> this.queryPromised(queryWithParams, true);
+    return <Promise<StreamTableData>> this.queryPromised(queryWithParams, true, options?.requestId);
   }
 
   public capabilities(): DriverCapabilities {
@@ -381,8 +391,8 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
     }
 
     const types = options.query
-      ? await this.unloadWithSql(tableName, options.query.sql, options.query.params)
-      : await this.unloadWithTable(tableName);
+      ? await this.unloadWithSql(tableName, options.query.sql, options.query.params, options.requestId)
+      : await this.unloadWithTable(tableName, options.requestId);
 
     const csvFile = await this.getCsvFiles(tableName);
 
@@ -403,33 +413,36 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
     return types.map((c) => `CAST(${c.name} AS varchar) ${c.name}`).join(', ');
   }
 
-  private async unloadWithSql(tableFullName: string, sql: string, params: any[]) {
+  private async unloadWithSql(tableFullName: string, sql: string, params: any[], requestId?: string) {
     return this.unloadGeneric({
       tableFullName,
       typeSql: sql,
       typeParams: params,
       fromSql: sql,
-      fromParams: params
+      fromParams: params,
+      requestId,
     });
   }
 
-  private async unloadWithTable(tableFullName: string) {
+  private async unloadWithTable(tableFullName: string, requestId?: string) {
     return this.unloadGeneric({
       tableFullName,
       typeSql: `SELECT * FROM ${tableFullName}`,
       typeParams: [],
       fromSql: tableFullName,
-      fromParams: []
+      fromParams: [],
+      requestId,
     });
   }
 
-  private async unloadGeneric(params: { tableFullName: string, typeSql: string, typeParams: any[], fromSql: string, fromParams: any[] }) {
+  private async unloadGeneric(params: { tableFullName: string, typeSql: string, typeParams: any[], fromSql: string, fromParams: any[], requestId?: string }) {
     if (!this.config.exportBucket) {
       throw new Error('Export bucket is not configured.');
     }
 
     const { bucketType, exportBucket } = this.config;
-    const types = await this.queryColumnTypes(params.typeSql, params.typeParams);
+    const { requestId } = params;
+    const types = await this.queryColumnTypes(params.typeSql, params.typeParams, { requestId });
 
     const { schema, tableName } = this.splitTableFullName(params.tableFullName);
     const tableWithCatalogAndSchema = `${this.config.catalog}.${schema}.${tableName}`;
@@ -448,16 +461,17 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
       await this.query(
         createTableQuery,
         params.fromParams,
+        { requestId },
       );
     } finally {
-      await this.query(`DROP TABLE IF EXISTS ${tableWithCatalogAndSchema}`, []);
+      await this.query(`DROP TABLE IF EXISTS ${tableWithCatalogAndSchema}`, [], { requestId });
     }
 
     return types;
   }
 
-  public async queryColumnTypes(sql: string, params: unknown[]): Promise<{ name: string; type: string; }[]> {
-    const response = await this.stream(`${sql} LIMIT 0`, params || [], { highWaterMark: 1 });
+  public async queryColumnTypes(sql: string, params: unknown[], options?: QueryOptions): Promise<{ name: string; type: string; }[]> {
+    const response = await this.stream(`${sql} LIMIT 0`, params || [], { highWaterMark: 1, requestId: options?.requestId });
     const result = [];
 
     for (const column of response.types || []) {

packages/cubejs-prestodb-driver/test/unit/headers.test.ts

@@ -129,4 +129,118 @@ describe('PrestoDriver custom headers', () => {
     expect(poll!.headers['X-Custom-Header']).toBe('custom-value');
     expect(poll!.headers['Proxy-Authorization']).toBe('Basic dGVzdA==');
   });
+
+  describe('requestId trace token', () => {
+    const createDriver = (config: Record<string, unknown> = {}) => new PrestoDriver({
+      host: 'coordinator.local',
+      port: '8080',
+      catalog: 'test',
+      schema: 'default',
+      dataSource: 'default',
+      checkInterval: 1,
+      ...config,
+    } as any);
+
+    it('tags the initial POST of query() with the request UUID', async () => {
+      const driver = createDriver({ headers: { 'X-Custom-Header': 'custom-value' } });
+
+      await driver.query('SELECT 1', [], { requestId: 'abc-123-span-2' });
+
+      const post = mockRecorded.find((r) => r.method === 'POST');
+      const poll = mockRecorded.find((r) => r.method === 'GET');
+
+      expect(post!.headers['X-Presto-Trace-Token']).toBe('abc-123');
+      expect(post!.headers['X-Custom-Header']).toBe('custom-value');
+      expect(poll!.headers['X-Presto-Trace-Token']).toBeUndefined();
+      expect(poll!.headers['X-Custom-Header']).toBe('custom-value');
+    });
+
+    it('tags the initial POST of stream()', async () => {
+      const driver = createDriver();
+
+      const { rowStream } = await driver.stream('SELECT 1', [], { highWaterMark: 1, requestId: 'stream-req-span-1' });
+
+      for await (const _row of rowStream) {
+        // drain
+      }
+
+      const post = mockRecorded.find((r) => r.method === 'POST');
+      expect(post!.headers['X-Presto-Trace-Token']).toBe('stream-req');
+    });
+
+    it('uses the Trino header name for the trino engine', async () => {
+      const driver = createDriver({ engine: 'trino' });
+
+      await driver.query('SELECT 1', [], { requestId: 'trino-req' });
+
+      const post = mockRecorded.find((r) => r.method === 'POST');
+      expect(post!.headers['X-Trino-Trace-Token']).toBe('trino-req');
+      expect(post!.headers['X-Presto-Trace-Token']).toBeUndefined();
+    });
+
+    it('skips the trace token for an invalid requestId and still runs the query', async () => {
+      const driver = createDriver();
+
+      for (const requestId of ['bad\u0001id', 'refresh-Ж', 'a'.repeat(129)]) {
+        mockRecorded.length = 0;
+        const rows = await driver.query('SELECT 1', [], { requestId });
+        expect(rows).toEqual([{ one: 1 }]);
+
+        const post = mockRecorded.find((r) => r.method === 'POST');
+        expect(post!.headers['X-Presto-Trace-Token']).toBeUndefined();
+      }
+    });
+
+    it('tags downloadQueryResults() without streamImport', async () => {
+      const driver = createDriver();
+
+      await driver.downloadQueryResults('SELECT 1', [], { highWaterMark: 1, requestId: 'download-req-span-1' });
+
+      const post = mockRecorded.find((r) => r.method === 'POST');
+      expect(post!.headers['X-Presto-Trace-Token']).toBe('download-req');
+    });
+
+    it('tags every statement of an export bucket unload', async () => {
+      class TestPrestoDriver extends PrestoDriver {
+        protected override async extractUnloadedFilesFromS3(): Promise<string[]> {
+          return [];
+        }
+      }
+
+      const driver = new TestPrestoDriver({
+        host: 'coordinator.local',
+        port: '8080',
+        catalog: 'test',
+        schema: 'default',
+        dataSource: 'default',
+        checkInterval: 1,
+        bucketType: 's3',
+        exportBucket: 'bucket',
+      } as any);
+
+      await driver.unload('stb.orders', {
+        maxFileSize: 64,
+        query: { sql: 'SELECT 1 AS one', params: [] },
+        requestId: 'unload-req-span-1',
+      });
+
+      const posts = mockRecorded.filter((r) => r.method === 'POST');
+      // column type probe, CREATE TABLE ... AS, DROP TABLE
+      expect(posts).toHaveLength(3);
+
+      for (const post of posts) {
+        expect(post.headers['X-Presto-Trace-Token']).toBe('unload-req');
+      }
+    });
+
+    it('sends no trace token without a requestId', async () => {
+      const driver = createDriver();
+
+      await driver.query('SELECT 1', []);
+
+      const post = mockRecorded.find((r) => r.method === 'POST');
+      expect(post!.headers['X-Presto-Trace-Token']).toBeUndefined();
+      expect(post!.headers['X-Trino-Trace-Token']).toBeUndefined();
+    });
+  });
 });

packages/cubejs-trino-driver/test/unit/headers.test.ts

@@ -85,4 +85,20 @@ describe('TrinoDriver headers', () => {
       'X-Trino-Routing-Group': 'etl',
     });
   });
+
+  it('adds X-Trino-Trace-Token from requestId on query()', async () => {
+    const driver = new TrinoDriver({
+      host: 'trino.local',
+      port: '8080',
+      headers: { 'X-Trino-Source': 'cube' },
+    });
+
+    await driver.query('SELECT 1', [], { requestId: 'req-42-span-3' });
+
+    const [executeOpts] = mockExecute.mock.calls[0];
+    expect(executeOpts.headers).toEqual({
+      'X-Trino-Source': 'cube',
+      'X-Trino-Trace-Token': 'req-42',
+    });
+  });
 });