perf(trino-driver): native REST polling to cut coordinator wait/drain latency

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

@@ -1,6 +1,6 @@
 ---
 title: Trino
-description: Map Cube to Trino with host credentials plus catalog, schema, SSL, and optional bearer-token variables on the shared Presto driver.
+description: Map Cube to Trino with host credentials plus catalog, schema, SSL, optional bearer-token variables, and configurable query polling.
 ---
 
 ## Prerequisites
@@ -34,6 +34,11 @@ CUBEJS_DB_SCHEMA=my_trino_schema
 | [`CUBEJS_DB_PASS`](/reference/configuration/environment-variables#cubejs_db_pass)           | The password used to connect to the database                                        | A valid database password                     |    ✅    |
 | [`CUBEJS_DB_PRESTO_CATALOG`](/reference/configuration/environment-variables#cubejs_db_presto_catalog) | The catalog within Presto to connect to                                             | A valid catalog name within a Presto database |    ✅    |
 | [`CUBEJS_DB_PRESTO_AUTH_TOKEN`](/reference/configuration/environment-variables#cubejs_db_presto_auth_token) | The authentication token to use when connecting to Presto/Trino. It will be sent in the `Authorization` header. | A valid authentication token |    ❌    |
+| [`CUBEJS_DB_TRINO_SOURCE`](/reference/configuration/environment-variables#cubejs_db_trino_source) | Value sent as the `X-Trino-Source` header | A string | ❌ |
+| [`CUBEJS_DB_TRINO_POLL_INITIAL_INTERVAL`](/reference/configuration/environment-variables#cubejs_db_trino_poll_initial_interval) | First wait (ms) between status polls while a query is queued with no rows | A positive integer | ❌ |
+| [`CUBEJS_DB_TRINO_POLL_INCREMENT_STEP`](/reference/configuration/environment-variables#cubejs_db_trino_poll_increment_step) | Added to the wait-phase poll interval after enough empty polls | A positive integer | ❌ |
+| [`CUBEJS_DB_TRINO_POLL_MAX_INTERVAL`](/reference/configuration/environment-variables#cubejs_db_trino_poll_max_interval) | Cap on wait-phase poll interval (ms) | A positive integer | ❌ |
+| [`CUBEJS_DB_TRINO_POLL_TRIES_BEFORE_INCREMENT`](/reference/configuration/environment-variables#cubejs_db_trino_poll_tries_before_increment) | Empty wait-phase polls at the current interval before increasing it | A positive integer | ❌ |
 | [`CUBEJS_DB_SCHEMA`](/reference/configuration/environment-variables#cubejs_db_schema)         | The schema within the database to connect to                                        | A valid schema name within a Presto database  |    ✅    |
 | [`CUBEJS_DB_SSL`](/reference/configuration/environment-variables#cubejs_db_ssl)            | If `true`, enables SSL encryption for database connections from Cube                | `true`, `false`                               |    ❌    |
 | [`CUBEJS_DB_MAX_POOL`](/reference/configuration/environment-variables#cubejs_db_max_pool)       | The maximum number of concurrent database connections to pool. Default is `8`       | A valid number                                |    ❌    |
@@ -130,7 +135,8 @@ protocol][trino-docs-client-protocol] for the list of headers accepted by Trino.
 
 Custom headers can't be configured via environment variables. Instead, use the
 [`driver_factory`](/reference/configuration/config#driver_factory) configuration
-option to pass a `headers` object to the driver:
+option to pass a `headers` object to the driver. Headers are sent on the initial
+statement POST and on every follow-up `nextUri` poll.
 
 <CodeGroup>
 
@@ -199,6 +205,65 @@ module.exports = {
 
 </CodeGroup>
 
+## Query polling
+
+The Trino driver talks to Trino's HTTP protocol directly. While a query is
+queued or planning (no rows yet), it polls `nextUri` with a short progressive
+backoff: 50ms for the first ten empty polls, then 100ms, up to 500ms. After the
+first rows arrive, it drains remaining pages with no extra delay.
+
+Tune this with environment variables, or pass a `pollBackoff` object (and
+optional `drainInterval`, in milliseconds) from
+[`driver_factory`](/reference/configuration/config#driver_factory):
+
+<CodeGroup>
+
+```python title="Python"
+from cube import config
+
+@config('driver_factory')
+def driver_factory(ctx: dict) -> dict:
+  return {
+    'type': 'trino',
+    'pollBackoff': {
+      'initialInterval': 50,
+      'incrementStep': 50,
+      'maxInterval': 500,
+      'triesBeforeIncrement': 10
+    },
+    'drainInterval': 0
+  }
+```
+
+```javascript title="JavaScript"
+module.exports = {
+  driverFactory: () => ({
+    type: "trino",
+    pollBackoff: {
+      initialInterval: 50,
+      incrementStep: 50,
+      maxInterval: 500,
+      triesBeforeIncrement: 10
+    },
+    drainInterval: 0
+  })
+};
+```
+
+</CodeGroup>
+
+To restore constant-interval polling (the previous `presto-client` default of
+800ms between every poll, including while draining rows), set `checkInterval`:
+
+```javascript
+module.exports = {
+  driverFactory: () => ({
+    type: "trino",
+    checkInterval: 800
+  })
+};
+```
+
 [trino-docs-client-protocol]: https://trino.io/docs/current/develop/client-protocol.html
 
 [aws-s3]: https://aws.amazon.com/s3/

docs-mintlify/reference/configuration/environment-variables.mdx

@@ -1042,6 +1042,53 @@ for TLS Secure Contexts [in the Node.js documentation][nodejs-docs-tls-options].
 | ----------------- | ---------------------- | --------------------- |
 | A valid host name | N/A                    | N/A                   |
 
+## `CUBEJS_DB_TRINO_SOURCE`
+
+Value sent as the `X-Trino-Source` header so Trino query history identifies the
+client. Used by the Trino driver. Defaults to `nodejs-client`, matching the
+previous `presto-client` default.
+
+| Possible Values | Default in Development | Default in Production |
+| --------------- | ---------------------- | --------------------- |
+| A string        | `nodejs-client`        | `nodejs-client`       |
+
+## `CUBEJS_DB_TRINO_POLL_INITIAL_INTERVAL`
+
+First wait, in milliseconds, before the next `nextUri` poll while a Trino query
+is still queued or planning with no rows.
+
+| Possible Values     | Default in Development | Default in Production |
+| ------------------- | ---------------------- | --------------------- |
+| A positive integer  | `50`                   | `50`                  |
+
+## `CUBEJS_DB_TRINO_POLL_INCREMENT_STEP`
+
+Added to the wait-phase poll interval after
+[`CUBEJS_DB_TRINO_POLL_TRIES_BEFORE_INCREMENT`](#cubejs_db_trino_poll_tries_before_increment)
+empty polls.
+
+| Possible Values     | Default in Development | Default in Production |
+| ------------------- | ---------------------- | --------------------- |
+| A positive integer  | `50`                   | `50`                  |
+
+## `CUBEJS_DB_TRINO_POLL_MAX_INTERVAL`
+
+Cap on the wait-phase poll interval, in milliseconds.
+
+| Possible Values     | Default in Development | Default in Production |
+| ------------------- | ---------------------- | --------------------- |
+| A positive integer  | `500`                  | `500`                 |
+
+## `CUBEJS_DB_TRINO_POLL_TRIES_BEFORE_INCREMENT`
+
+Number of empty wait-phase polls at the current interval before increasing it.
+With the default initial interval of 50ms, the default of 10 keeps about 500ms
+of fine-grained polling.
+
+| Possible Values     | Default in Development | Default in Production |
+| ------------------- | ---------------------- | --------------------- |
+| A positive integer  | `10`                   | `10`                  |
+
 ## `CUBEJS_DB_TYPE`
 
 A database from the list of [supported databases][ref-config-db].

packages/cubejs-backend-shared/src/env.ts

@@ -71,6 +71,35 @@ export function asPortNumber(input: number, envName: string) {
   return input;
 }
 
+/**
+ * Parse an env var as a base-10 integer. Rejects empty strings, floats,
+ * scientific notation (`1e2`), and unit suffixes (`50ms`) that `parseInt`
+ * would otherwise silently truncate.
+ */
+export function parseNonNegativeIntEnv(
+  envName: string,
+  defaultValue: string,
+  min: number = 0,
+): number {
+  const raw = get(envName).default(defaultValue).asString();
+  if (!/^\d+$/.test(raw)) {
+    throw new InvalidConfiguration(
+      envName,
+      raw,
+      `Must be an integer >= ${min}.`
+    );
+  }
+  const value = parseInt(raw, 10);
+  if (value < min) {
+    throw new InvalidConfiguration(
+      envName,
+      raw,
+      `Must be an integer >= ${min}.`
+    );
+  }
+  return value;
+}
+
 /**
  * Determines whether multiple data sources were declared or not.
  */
@@ -1804,6 +1833,73 @@ const variables: Record<string, (...args: any) => any> = {
     get(keyByDataSource('CUBEJS_DB_PRESTO_AUTH_TOKEN', dataSource, preAggregations)).asString()
   ),
 
+  /**
+   * Value sent as the `X-Trino-Source` header. Identifies the client in
+   * Trino's query history. Default: `nodejs-client` (same as `presto-client`).
+   */
+  trinoSource: ({
+    dataSource,
+    preAggregations,
+  }: DataSourceOpts) => (
+    get(keyByDataSource('CUBEJS_DB_TRINO_SOURCE', dataSource, preAggregations)).asString()
+  ),
+
+  /**
+   * First wait (ms) before the next `nextUri` poll while the query is still
+   * queued/planning with no rows. Default: 50.
+   */
+  trinoPollInitialInterval: ({
+    dataSource,
+    preAggregations,
+  }: DataSourceOpts) => (
+    parseNonNegativeIntEnv(
+      keyByDataSource('CUBEJS_DB_TRINO_POLL_INITIAL_INTERVAL', dataSource, preAggregations),
+      '50',
+    )
+  ),
+
+  /**
+   * Added to the wait-phase poll interval after `trinoPollTriesBeforeIncrement`
+   * empty polls. Default: 50.
+   */
+  trinoPollIncrementStep: ({
+    dataSource,
+    preAggregations,
+  }: DataSourceOpts) => (
+    parseNonNegativeIntEnv(
+      keyByDataSource('CUBEJS_DB_TRINO_POLL_INCREMENT_STEP', dataSource, preAggregations),
+      '50',
+    )
+  ),
+
+  /**
+   * Cap on wait-phase poll interval (ms). Default: 500.
+   */
+  trinoPollMaxInterval: ({
+    dataSource,
+    preAggregations,
+  }: DataSourceOpts) => (
+    parseNonNegativeIntEnv(
+      keyByDataSource('CUBEJS_DB_TRINO_POLL_MAX_INTERVAL', dataSource, preAggregations),
+      '500',
+    )
+  ),
+
+  /**
+   * Empty wait-phase polls at the current interval before increasing it.
+   * Default: 10 (so 50ms × 10 ≈ 500ms of fine-grained polling).
+   */
+  trinoPollTriesBeforeIncrement: ({
+    dataSource,
+    preAggregations,
+  }: DataSourceOpts) => (
+    parseNonNegativeIntEnv(
+      keyByDataSource('CUBEJS_DB_TRINO_POLL_TRIES_BEFORE_INCREMENT', dataSource, preAggregations),
+      '10',
+      1,
+    )
+  ),
+
   /** ***************************************************************
    * Pinot Driver                                                  *
    **************************************************************** */

packages/cubejs-backend-shared/test/env.test.ts

@@ -178,3 +178,85 @@ describe('getEnv(apiSecret / apiSecrets)', () => {
     expect(getEnv('apiSecrets')).toEqual(['only']);
   });
 });
+
+const TRINO_POLL_ENV = {
+  trinoPollInitialInterval: 'CUBEJS_DB_TRINO_POLL_INITIAL_INTERVAL',
+  trinoPollIncrementStep: 'CUBEJS_DB_TRINO_POLL_INCREMENT_STEP',
+  trinoPollMaxInterval: 'CUBEJS_DB_TRINO_POLL_MAX_INTERVAL',
+  trinoPollTriesBeforeIncrement: 'CUBEJS_DB_TRINO_POLL_TRIES_BEFORE_INCREMENT',
+} as const;
+
+describe('getEnv(trino poll / source)', () => {
+  const opts = { dataSource: 'default' as const };
+
+  afterEach(() => {
+    delete process.env.CUBEJS_DB_TRINO_SOURCE;
+    Object.values(TRINO_POLL_ENV).forEach((key) => {
+      delete process.env[key];
+    });
+  });
+
+  test('defaults', () => {
+    expect(getEnv('trinoSource', opts)).toBeUndefined();
+    expect(getEnv('trinoPollInitialInterval', opts)).toBe(50);
+    expect(getEnv('trinoPollIncrementStep', opts)).toBe(50);
+    expect(getEnv('trinoPollMaxInterval', opts)).toBe(500);
+    expect(getEnv('trinoPollTriesBeforeIncrement', opts)).toBe(10);
+  });
+
+  test('valid overrides', () => {
+    process.env.CUBEJS_DB_TRINO_SOURCE = 'my-app';
+    process.env.CUBEJS_DB_TRINO_POLL_INITIAL_INTERVAL = '25';
+    process.env.CUBEJS_DB_TRINO_POLL_INCREMENT_STEP = '0';
+    process.env.CUBEJS_DB_TRINO_POLL_MAX_INTERVAL = '1000';
+    process.env.CUBEJS_DB_TRINO_POLL_TRIES_BEFORE_INCREMENT = '3';
+
+    expect(getEnv('trinoSource', opts)).toBe('my-app');
+    expect(getEnv('trinoPollInitialInterval', opts)).toBe(25);
+    expect(getEnv('trinoPollIncrementStep', opts)).toBe(0);
+    expect(getEnv('trinoPollMaxInterval', opts)).toBe(1000);
+    expect(getEnv('trinoPollTriesBeforeIncrement', opts)).toBe(3);
+  });
+
+  test('empty source is empty string (driver falls back to nodejs-client)', () => {
+    process.env.CUBEJS_DB_TRINO_SOURCE = '';
+    expect(getEnv('trinoSource', opts)).toBe('');
+  });
+
+  test.each(Object.entries(TRINO_POLL_ENV))(
+    '%s rejects improper values',
+    (getter, envKey) => {
+      const read = () => getEnv(getter as keyof typeof TRINO_POLL_ENV, opts);
+
+      process.env[envKey] = '-1';
+      expect(read).toThrow(/Must be an integer >= /);
+
+      process.env[envKey] = 'abc';
+      expect(read).toThrow(/Must be an integer >= /);
+
+      process.env[envKey] = '';
+      expect(read).toThrow(/Must be an integer >= /);
+
+      process.env[envKey] = '1.5';
+      expect(read).toThrow(/Must be an integer >= /);
+
+      process.env[envKey] = '50ms';
+      expect(read).toThrow(/Must be an integer >= /);
+
+      process.env[envKey] = '1e2';
+      expect(read).toThrow(/Must be an integer >= /);
+    }
+  );
+
+  test('triesBeforeIncrement rejects 0', () => {
+    process.env.CUBEJS_DB_TRINO_POLL_TRIES_BEFORE_INCREMENT = '0';
+    expect(() => getEnv('trinoPollTriesBeforeIncrement', opts)).toThrow(
+      /Must be an integer >= 1/
+    );
+  });
+
+  test('interval 0 is allowed (poll immediately)', () => {
+    process.env.CUBEJS_DB_TRINO_POLL_INITIAL_INTERVAL = '0';
+    expect(getEnv('trinoPollInitialInterval', opts)).toBe(0);
+  });
+});

packages/cubejs-schema-compiler/src/adapter/index.ts

@@ -17,6 +17,7 @@ export * from './MysqlQuery';
 export * from './PostgresQuery';
 export * from './MssqlQuery';
 export * from './PrestodbQuery';
+export * from './TrinoQuery';
 
 export { PreAggregationReferences } from '../compiler/CubeEvaluator';
 

packages/cubejs-trino-driver/jest.config.js

@@ -0,0 +1,7 @@
+const base = require('../../jest.base.config');
+
+/** @type {import('jest').Config} */
+module.exports = {
+  ...base,
+  rootDir: '.',
+};

packages/cubejs-trino-driver/package.json

@@ -29,11 +29,9 @@
   },
   "dependencies": {
     "@cubejs-backend/base-driver": "1.7.19",
-    "@cubejs-backend/prestodb-driver": "1.7.19",
     "@cubejs-backend/schema-compiler": "1.7.19",
     "@cubejs-backend/shared": "1.7.19",
-    "node-fetch": "^2.6.1",
-    "presto-client": "^1.2.0"
+    "node-fetch": "^2.6.1"
   },
   "license": "Apache-2.0",
   "publishConfig": {
@@ -42,6 +40,7 @@
   "devDependencies": {
     "@cubejs-backend/linter": "1.7.19",
     "@types/jest": "^29",
+    "@types/node-fetch": "^2.5.8",
     "jest": "^29",
     "testcontainers": "^10.28.0",
     "typescript": "~5.2.2"