feat(drivers): sql_preamble data source option - #11464
Conversation
…estination safety
…he credential namespace
github-actions
Bot
added
driver:mysql
|
Claude finished @igorlukanin's task in 7m 3s —— View job Review complete — solid feature, 2 medium correctness gaps + missing docs (2 medium, 1 medium-doc, 4 low)Full review
GeneralThis is unusually careful work. The rationale-in-code is genuinely useful rather than decorative (why the pre-agg preamble inherits while credentials don't, why the version is resolved from the env rather than a driver, why Findings
Also, minor and not worth an inline: Security / performanceNo injection surface added that didn't already exist — the preamble is operator-supplied config, not tenant input, and nothing interpolates request data into it. Worth noting for the record that it is not safe to expose via TestingCoverage is strong — env resolution end-to-end through real constructors, delivery asserted per driver on both paths, real DuckDB execution including the stream-replay case. Gaps map to the findings above: no test for the ambiguous-split path reaching the BigQuery destination guard (#1), none for a reused JDBC connection (#2), none for a trailing comment-only segment (#4), and none for backslash-escaped quotes under a MySQL-style dialect (#5). The new |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| protected static isScriptSafePreamble(preamble: string): boolean { | ||
| const statements = splitSqlPreamble(preamble); | ||
|
|
||
| return statements.length > 0 && statements.every( | ||
| statement => /^create\s+(or\s+replace\s+)?temp(orary)?\s+function\b/i | ||
| .test(BigQueryDriver.withoutLeadingComments(statement)) | ||
| ); | ||
| } |
There was a problem hiding this comment.
The destination guard inherits the splitter's ambiguous fallback, so the exact failure it exists to prevent can still slip through.
splitSqlPreamble returns [normalized] — the whole blob as one statement — whenever it meets something it cannot parse confidently (nested block comment, bare $$, unterminated construct). isScriptSafePreamble then sees statements.length === 1 and only checks that this single "statement" starts with CREATE TEMP FUNCTION, so a blob that in fact contains further statements is judged script-safe.
Concretely, with
CREATE TEMP FUNCTION f(x INT64) AS (x); /* uses /* nested */ SET @@dataset_id = 'analytics'the splitter goes ambiguous (it nests block comments, BigQuery does not) → one statement → passes the regex → prepended onto a destination job. BigQuery parses three statements, runs the request as a script, ignores destinationTable, and the pre-aggregation build reports success having written nothing — the silent-empty-table case the guard was added for.
Ambiguity should be fail-closed here: have splitSqlPreamble signal that it gave up (a second return value / a sibling trySplitSqlPreamble) and treat that as not script-safe, or at minimum reject when the single returned statement still contains a ; that is not its terminator.
Reviewing
|
| // A streamed query has to run in the preamble's context too; this path | ||
| // used to skip the connection queries the query path replays. | ||
| for (const statementSql of this.prepareConnectionQueries()) { | ||
| await this.executeStatement(conn, statementSql); | ||
| } |
There was a problem hiding this comment.
JDBC is the one pooled driver that doesn't get the already-applied tolerance.
The PR fixes exactly this for Postgres and MySQL via applySqlPreambleStatements (pool reuses connections → a CREATE … preamble fails on the second query). JDBC pools connections the same way (this.pool.acquire() here and in queryPromised), but both paths execute the statements raw through executeStatement, so an "already exists" on a reused connection propagates as a query failure.
This was safe before the feature because the built-in prepareConnectionQueries are only SETs; it stops being safe now that user SQL — whose headline use case is creating a function — flows through the same list. Suggest running the preamble portion through applySqlPreambleStatements while keeping the built-ins raw:
for (const statementSql of builtInQueries) {
await this.executeStatement(conn, statementSql);
}
await applySqlPreambleStatements(preamble, statement => this.executeStatement(conn, statement));
| } else if (char === ';') { | ||
| const statement = normalizeSqlPreamble(current); | ||
| if (statement) { | ||
| statements.push(statement); | ||
| } | ||
| current = ''; | ||
| index += 1; | ||
| end = -2; |
There was a problem hiding this comment.
A trailing comment becomes its own statement.
Comments are appended into current, so 'SET a = 1;\n-- note' splits into ['SET a = 1', '-- note'] — the second entry has no executable token. Every driver that executes statements one at a time then sends a comment-only statement to the engine: Snowflake rejects an empty statement outright, JDBC's executeQuery on a comment-only string is driver-dependent, and on BigQuery withoutLeadingComments reduces it to '' so isScriptSafePreamble refuses a perfectly legitimate
CREATE TEMP FUNCTION median(x FLOAT64) AS (x);
-- keep in sync with the data modelon pre-aggregation builds, with a message telling the user to do what they already did.
Dropping segments that contain no non-comment token (or tracking whether anything but whitespace/comments was consumed since the last ;) fixes all three at once, and the 'SET a = 1;\n-- note' / trailing-comment case is worth a unit test either way.
| export function getPreAggregationSqlPreamble(preAggregation): string | undefined { | ||
| let fromEnv: string | undefined; | ||
|
|
||
| try { | ||
| fromEnv = getEnv('dbSqlPreamble', { | ||
| dataSource: preAggregation.dataSource || 'default', | ||
| preAggregations: true, | ||
| }); | ||
| } catch (e) { | ||
| // `assertDataSource` throws for a data source missing from | ||
| // CUBEJS_DATASOURCES. This runs inside version computation and inside the | ||
| // builder for the "no partitions were built" message, so throwing here would | ||
| // replace an actionable error with a confusing one. A key without the | ||
| // preamble is the pre-feature behavior, which is the safe fallback. | ||
| return undefined; | ||
| } | ||
|
|
||
| // Normalized the same way the drivers normalize it, so the version reflects | ||
| // the preamble that actually runs: without this, re-indenting a preamble — or | ||
| // setting a whitespace-only one, which the drivers treat as none at all — | ||
| // would change the key and rebuild every pre-aggregation for no behavior | ||
| // difference. | ||
| return normalizeSqlPreamble(fromEnv); |
There was a problem hiding this comment.
The version key participates for every driver, but only eight drivers implement the preamble.
This reads the env unconditionally, so on a deployment whose data source is ClickHouse, MS SQL, Athena, Firebolt, Druid, Elasticsearch, … setting CUBEJS_DB_SQL_PREAMBLE re-keys and rebuilds every pre-aggregation while changing nothing about how queries run — the driver silently ignores the option. That's the worst combination: cost with no effect and no signal.
Two mitigations worth considering, ideally both:
- Warn once at driver construction (or in
BaseDriver) when a preamble is configured for a driver that doesn't implement it, so the no-op is visible. - Or gate version participation on driver support rather than on the env var merely being present.
Separately, the catch swallows any failure, not just assertDataSource. If getEnv ever throws for another reason the key silently loses the preamble and stale tables get served — narrowing the catch (or logging it) would keep the fallback intentional.
| // Finds the end of a quoted string or identifier, or -1 when it never closes. | ||
| // | ||
| // A doubled quote is an escaped quote in every dialect here. A backslash is | ||
| // NOT treated as an escape: Postgres, DuckDB and Snowflake follow the standard | ||
| // (`standard_conforming_strings`), where a trailing backslash is a literal | ||
| // character, so consuming the following quote would swallow the terminator and | ||
| // merge two statements. | ||
| const endOfQuoted = (start: number, quote: string): number => { | ||
| let cursor = start + 1; | ||
|
|
||
| while (cursor < normalized.length) { | ||
| if (normalized[cursor] !== quote) { | ||
| cursor += 1; | ||
| } else if (normalized[cursor + 1] === quote) { | ||
| cursor += 2; | ||
| } else { | ||
| return cursor + 1; | ||
| } | ||
| } | ||
|
|
||
| return -1; | ||
| }; |
There was a problem hiding this comment.
The no-backslash-escape rule is right for Postgres/DuckDB/Snowflake and wrong for MySQL and BigQuery — and those two are also drivers that execute the split statements individually.
MySQL escapes with backslashes by default (NO_BACKSLASH_ESCAPES off), and BigQuery always does. So SET a = 'it\'s; fine'; SET b = 2 in a MySQL preamble terminates the literal at the escaped quote here, splits on the ; inside it, and hands MySQL two syntactically broken statements. The comment explains the choice as universal, but it's dialect-dependent — either take an optional dialect hint from the caller (each driver knows its own), or treat a backslash-immediately-before-quote as ambiguous and fall back to the unsplit blob, which is the conservative posture the rest of this function already takes.
Also dialect-specific: # is unconditionally a line comment (line 127), but in Postgres # is a valid operator (bitwise XOR, geometric ops), so SET a = 5 # 3; SET b = 1 swallows the separator and merges the two statements. Lower practical risk, same root cause.
| for (const statement of splitSqlPreamble(preamble ?? legacy)) { | ||
| try { | ||
| await execAsync(statement); | ||
| } catch (e) { | ||
| if (legacy) { | ||
| if (this.logger) { | ||
| console.error('DuckDB - error on init sql (skipping)', { e }); | ||
| } | ||
| } else if (!/already exists/i.test((e as Error)?.message ?? '')) { | ||
| throw e; | ||
| } | ||
| } |
There was a problem hiding this comment.
This re-implements isAlreadyAppliedPreambleError with a narrower regex (/already exists/i only). The shared helper also covers duplicate function / is already defined, and any future spelling added there won't reach here. Since the loop is otherwise identical to applySqlPreambleStatements, the legacy branch aside, importing the predicate keeps the two definitions of "already applied" from drifting.
| for (const statement of splitSqlPreamble(preamble ?? legacy)) { | |
| try { | |
| await execAsync(statement); | |
| } catch (e) { | |
| if (legacy) { | |
| if (this.logger) { | |
| console.error('DuckDB - error on init sql (skipping)', { e }); | |
| } | |
| } else if (!/already exists/i.test((e as Error)?.message ?? '')) { | |
| throw e; | |
| } | |
| } | |
| for (const statement of splitSqlPreamble(preamble ?? legacy)) { | |
| try { | |
| await execAsync(statement); | |
| } catch (e) { | |
| if (legacy) { | |
| if (this.logger) { | |
| console.error('DuckDB - error on init sql (skipping)', { e }); | |
| } | |
| } else if (!isAlreadyAppliedPreambleError(e)) { | |
| throw e; | |
| } |
|
Claude finished @igorlukanin's task in 3m 16s —— View job Re-review after 3 new commits — docs gap closed, 5 prior findings still stand, 1 new (doc, low)Full review
What changed since the last pass
Still applies (re-affirmed, not reposted)Verified against the current tree; all five are unchanged in the code:
The DuckDB nit ( New this pass1 · Low · Minor, no inline: the new TestingThe new Security / performanceNothing new. The preamble remains operator-supplied config with no request data interpolated; |
|
|
||
| </Warning> | ||
|
|
||
| When the preamble runs depends on how the data source handles connections. On |
There was a problem hiding this comment.
This paragraph describes when the preamble runs but never says which drivers implement it — and that omission has a cost, because the pre-aggregation cache key participates unconditionally.
Eight drivers wire it up (BigQuery, Snowflake, Postgres, Crate, Materialize, MySQL, DuckDB, JDBC-based). On any other data source — ClickHouse, MS SQL, Athena, Firebolt, Druid, Elasticsearch, … — setting this variable silently does nothing to how queries run, while getPreAggregationSqlPreamble still folds it into the structure and content versions and rebuilds every pre-aggregation. A user following this page on ClickHouse pays a full rebuild for a no-op, with the reference page implying it should work.
Worth an explicit supported-drivers list here (the "On Snowflake … On Postgres, MySQL and the JDBC-based drivers … on BigQuery" sentence already omits Crate, Materialize and DuckDB), plus a <Warning> that the variable is ignored elsewhere. Pairs with the unresolved thread on PreAggregations.ts suggesting a construction-time warning or gating version participation on driver support.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #11464 +/- ##
===========================================
- Coverage 79.43% 59.34% -20.09%
===========================================
Files 480 228 -252
Lines 98778 18226 -80552
Branches 3636 3715 +79
===========================================
- Hits 78460 10817 -67643
+ Misses 19800 6883 -12917
- Partials 518 526 +8
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
1 participant
Summary
Adds a
sql_preambledata source option: SQL that runs in the context of every query Cube sends to a data source. Its main use is defining temporary UDFs the data model's SQL then calls, which on a stateless data source like BigQuery is the only way to get a user-defined function in scope. Configurable asCUBEJS_DB_SQL_PREAMBLE(with the usual per-data-source and pre-aggregation variants) or as thesqlPreambledriver option.One semantic, per-driver mechanisms. What holds everywhere is that the primary query executes in the preamble's context; how that is achieved differs because the execution models do:
ALTER SESSIONso a user can override those defaults deliberately.applySqlPreamblehook so the two drivers that replaceprepareConnectionstill apply it.downloadQueryResults.init()prepares, plus a replay on the separate connectionstream()opens.Every driver's
stream()is a separate hook from its query path, so each is wired and covered individually.Pre-aggregations. A pre-aggregation-specific preamble follows the existing pre-aggregation credentials convention, but with one deliberate difference: it inherits the regular preamble when unset. A preamble usually defines a function the model depends on, and a build that silently ran without it would fail later on the missing function; credentials do not inherit because a half-inherited connection target is worse than none. Setting it to an empty value opts builds out entirely.
The preamble also participates in the pre-aggregation structure and content versions, so changing it rebuilds the pre-aggregations on that data source. That is the point — a preamble can change what identical SQL returns, so serving a table built under a different one would be wrong. A deployment that sets no preamble computes a byte-identical version key, so upgrading rebuilds nothing.
Two subtleties worth calling out, both of which were bugs before they were fixed here:
driver_factorytherefore does not participate in the key; that is documented.No breaking changes.
initSql(DuckDB) andprepareConnectionQueries(JDBC) keep working as deprecated aliases with their existing behavior —initSqlstill swallows failures wherever it runs,prepareConnectionQueriesstill replaces the built-ins rather than appending, and both warn once per driver instance.sql_preambleitself fails loudly, so a preamble meant to define a function does not resurface as a baffling "function does not exist".Also fixed, because the feature surfaced them:
CREATE TEMP FUNCTION-then-query shape BigQuery exempts is now refused with an actionable message.standard_conforming_strings),#comments, nested block comments, and bare$$. It is now conservative: input it cannot parse confidently is passed to the engine whole rather than split on a guess.JDBCDriver's connection queries never running on the query path is fixed separately in #11456; this branch will rebase onto it.Test plan
unitscripts added where credential-free tests existed but never ran in CI