diff --git a/index.js b/index.js index 3630144..54342b9 100644 --- a/index.js +++ b/index.js @@ -275,46 +275,90 @@ function applyReservationToAction(action, actionToReservation) { action._queriesPatched = true } - // Prefer modifying data structure directly if we know it's a safe type - // This handles both Builders (via .proto) and Compiled Objects (direct) + // For table/view/incremental builders, preOps and the main query are supplied + // AFTER the builder is created, via chained .preOps()/.query() calls (e.g. + // `publish(...).preOps('DECLARE ...').query(...)`). Injecting the reservation + // eagerly at creation time puts it BEFORE a later DECLARE, which BigQuery + // rejects (a DECLARE must be the first statement in a script). Mirror the + // operations `.queries()` patch above: intercept both methods and inject the + // reservation only when there is no outer DECLARE — never inject-then-strip. + // Every valid table/view builder exposes both .preOps() and .query(); when it + // does, we defer to this lazy path and skip the direct injection below. + if (hasType && typeof action.preOps === 'function' && typeof action.query === 'function' && !action._reservationPatched) { + const originalPreOpsFn = action.preOps + const originalQueryFn = action.query + + action.preOps = function (ops) { + if (hasOuterDeclare(ops)) { + // User supplies their own leading DECLARE — leave it untouched. + action._reservationHandled = true + return originalPreOpsFn.apply(this, arguments) + } + // No DECLARE: prepend the reservation to the user's preOps so it runs first. + const withReservation = typeof ops === 'function' + ? (ctx) => prependStatement(ops(ctx), statement) + : prependStatement(ops, statement) + action._reservationHandled = true + return originalPreOpsFn.apply(this, [withReservation]) + } - // 1. Try contextablePreOps (Tables/Views Builders before resolution) - if (action.contextablePreOps) { - if (!hasOuterDeclare(action.contextablePreOps)) { - action.contextablePreOps = prependStatement(action.contextablePreOps, statement) + action.query = function () { + // If the builder had no preOps at all, there is no DECLARE to worry about, + // so inject the reservation now (equivalent to the operations queries patch). + if (!action._reservationHandled) { + originalPreOpsFn.call(action, statement) + action._reservationHandled = true + } + return originalQueryFn.apply(this, arguments) } + + action._reservationPatched = true } - // 2. Try contextableQueries (Operations Builders before resolution) - else if (action.contextableQueries) { - // Skip if there is an outer DECLARE - if (!hasOuterDeclare(action.contextableQueries)) { - action.contextableQueries = prependStatement(action.contextableQueries, statement) + + // Direct injection for compiled objects and already-resolved builders that do + // not expose the chained .preOps()/.query() API handled above. + // Skip entirely when the lazy table/view patch is installed to avoid double + // injection (the reservation is added inside .preOps()/.query() instead). + if (!action._reservationPatched) { + + // 1. Try contextablePreOps (Tables/Views Builders before resolution) + if (action.contextablePreOps) { + if (!hasOuterDeclare(action.contextablePreOps)) { + action.contextablePreOps = prependStatement(action.contextablePreOps, statement) + } } - } - // 3. Try proto.preOps (Compiled Tables/Views or Resolved Builders) - else if (hasType) { - if (!hasOuterDeclare(proto.preOps || [])) { - if (!proto.preOps) { - proto.preOps = [] + // 2. Try contextableQueries (Operations Builders before resolution) + else if (action.contextableQueries) { + // Skip if there is an outer DECLARE + if (!hasOuterDeclare(action.contextableQueries)) { + action.contextableQueries = prependStatement(action.contextableQueries, statement) } - - if (isArrayOrString(proto.preOps)) { - proto.preOps = prependStatement(proto.preOps, statement) - } else if (hasPreOpsFn) { - action.preOps(statement) + } + // 3. Try proto.preOps (Compiled Tables/Views or Resolved Builders) + else if (hasType) { + if (!hasOuterDeclare(proto.preOps || [])) { + if (!proto.preOps) { + proto.preOps = [] + } + + if (isArrayOrString(proto.preOps)) { + proto.preOps = prependStatement(proto.preOps, statement) + } else if (hasPreOpsFn) { + action.preOps(statement) + } } } - } - // 4. Try proto.queries (Compiled Operations or Resolved Builders) - else if (proto.queries) { + // 4. Try proto.queries (Compiled Operations or Resolved Builders) + else if (proto.queries) { // Skip if there is an outer DECLARE - if (!hasOuterDeclare(proto.queries)) { - proto.queries = prependStatement(proto.queries, statement) + if (!hasOuterDeclare(proto.queries)) { + proto.queries = prependStatement(proto.queries, statement) + } + } + // 5. Fallback to function API (likely Tables/Views) + else if (hasPreOpsFn) { + action.preOps(statement) } - } - // 5. Fallback to function API (likely Tables/Views) - else if (hasPreOpsFn) { - action.preOps(statement) } } } diff --git a/test/index.test.js b/test/index.test.js index 1610b1b..70aca30 100644 --- a/test/index.test.js +++ b/test/index.test.js @@ -755,6 +755,103 @@ describe('Dataform package', () => { const action = global.dataform.actions[0] expect(action.proto.queries).not.toContain('SET @@reservation=\'projects/test/locations/US/reservations/prod\';') }) + + // Table/view builders supply their preOps and query AFTER creation via + // chained .preOps()/.query() calls. The reservation is injected at + // creation time by intercepting those methods (mirroring the operations + // .queries() patch), so it is never emitted before a later DECLARE and + // nothing is stripped afterward. These tests exercise that path with a + // realistic builder that writes through to proto.preOps / proto.query. + describe('chained .preOps()/.query() builder interception', () => { + let originalPublishBuilder + + // A builder that mimics a real Dataform publish() action: preOps and + // query are chained on after creation and stored on the proto. + const makeBuilder = (name) => { + const action = { + proto: { + type: 'table', + target: { name, database: 'test-project', schema: 'test-schema' } + }, + preOps: function (ops) { + this.proto.preOps = Array.isArray(ops) ? ops : [ops] + return this + }, + query: function (q) { + this.proto.query = q + return this + } + } + global.dataform.actions.push(action) + return action + } + + beforeEach(() => { + originalPublishBuilder = global.publish + global.publish = jest.fn((name) => makeBuilder(name)) + }) + + afterEach(() => { + global.publish = originalPublishBuilder + }) + + const reservationStatement = 'SET @@reservation=\'projects/test/locations/US/reservations/prod\';' + const builderConfig = [ + { + tag: 'test', + reservation: 'projects/test/locations/US/reservations/prod', + actions: [ + 'test-project.test-schema.plain_table', + 'test-project.test-schema.declare_table', + 'test-project.test-schema.normal_preops_table' + ] + } + ] + + test('should inject reservation via .query() when the builder has no preOps', () => { + autoAssignActions(builderConfig) + const builder = global.publish('plain_table') + builder.query('SELECT 1') + + const action = global.dataform.actions[0] + expect(action.proto.preOps).toEqual([reservationStatement]) + }) + + test('should NOT inject reservation when preOps has an outer DECLARE', () => { + autoAssignActions(builderConfig) + const builder = global.publish('declare_table') + builder.preOps('DECLARE x INT64 DEFAULT 1;') + builder.query('SELECT x') + + const action = global.dataform.actions[0] + // DECLARE must stay the first statement; reservation is never added. + expect(action.proto.preOps).toEqual(['DECLARE x INT64 DEFAULT 1;']) + expect(action.proto.preOps).not.toContain(reservationStatement) + }) + + test('should prepend reservation before non-DECLARE preOps', () => { + autoAssignActions(builderConfig) + const builder = global.publish('normal_preops_table') + builder.preOps('CREATE TEMP FUNCTION f() AS (1);') + builder.query('SELECT f()') + + const action = global.dataform.actions[0] + expect(action.proto.preOps).toEqual([ + reservationStatement, + 'CREATE TEMP FUNCTION f() AS (1);' + ]) + }) + + test('should not touch preOps for an unconfigured builder', () => { + autoAssignActions(builderConfig) + const builder = global.publish('unconfigured_table') + builder.preOps('DECLARE y INT64 DEFAULT 2;') + builder.query('SELECT y') + + const action = global.dataform.actions[0] + expect(action.proto.preOps).toEqual(['DECLARE y INT64 DEFAULT 2;']) + }) + }) } }) })