Skip to content

Commit f11f8f3

Browse files
SuchodolskiEdvinSuchodolskiEdvin
committed
Support incremental_strategy parameter and new insert_overwrite strategy
- updated proto with new parameters - added new tests - added validation for chosen incremental_strategies - added new insert_overwrite strategy logic
1 parent 3b7eb1f commit f11f8f3

11 files changed

Lines changed: 545 additions & 19 deletions

‎cli/api/dbadapters/execution_sql.ts‎

Lines changed: 74 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -154,24 +154,8 @@ from (${query}) as insertions`;
154154
this.buildIncrementalSchemaChangeTasks(tasks, table);
155155
// Fall through to run the static DML after the procedure alters the schema
156156
case dataform.OnSchemaChange.IGNORE:
157-
default:
158-
tasks.add(
159-
Task.statement(
160-
table.uniqueKey && table.uniqueKey.length > 0
161-
? this.mergeInto(
162-
table.target,
163-
tableMetadata?.fields.map(f => f.name),
164-
this.getIncrementalQuery(table),
165-
table.uniqueKey,
166-
table.bigquery && table.bigquery.updatePartitionFilter
167-
)
168-
: this.insertInto(
169-
table.target,
170-
tableMetadata?.fields.map(f => f.name).map(column => `\`${column}\``),
171-
this.getIncrementalQuery(table)
172-
)
173-
)
174-
);
157+
const columns = tableMetadata?.fields.map(f => f.name) || [];
158+
tasks.add(Task.statement(this.getIncrementalDmlStatement(table, columns)));
175159
break;
176160
}
177161
}
@@ -451,7 +435,7 @@ DROP TABLE IF EXISTS ${emptyTempTableName};
451435
create or replace view ${this.resolveTarget(target)} as ${query}`;
452436
}
453437

454-
private mergeInto(
438+
private mergeInto(
455439
target: dataform.ITarget,
456440
columns: string[],
457441
query: string,
@@ -470,6 +454,77 @@ when matched then
470454
when not matched then
471455
insert (${backtickedColumns.join(",")}) values (${backtickedColumns.join(",")})`;
472456
}
457+
458+
private insertOverwrite(
459+
target: dataform.ITarget,
460+
columns: string[],
461+
query: string,
462+
partitionBy: string,
463+
updatePartitionFilter: string
464+
): string {
465+
const uniqueId = this.uniqueIdGenerator();
466+
const stagingTableUnqualified = `staging_table_temp_${uniqueId}`;
467+
const backtickedColumns = columns.map(column => `\`${column}\``);
468+
const resolveTargetTable = this.resolveTarget(target);
469+
470+
return `CREATE OR REPLACE TEMP TABLE \`${stagingTableUnqualified}\` AS (
471+
${query}
472+
);
473+
474+
BEGIN
475+
DECLARE partitions_for_replacement DEFAULT (
476+
ARRAY(
477+
SELECT DISTINCT ${partitionBy}
478+
FROM \`${stagingTableUnqualified}\`
479+
WHERE ${partitionBy} IS NOT NULL
480+
)
481+
);
482+
483+
MERGE ${resolveTargetTable} T
484+
USING \`${stagingTableUnqualified}\` S
485+
ON FALSE
486+
WHEN NOT MATCHED BY SOURCE AND ${partitionBy} IN UNNEST(partitions_for_replacement) ${updatePartitionFilter ? `and T.${updatePartitionFilter}` : ""} THEN
487+
DELETE
488+
WHEN NOT MATCHED BY TARGET THEN
489+
INSERT (${backtickedColumns.join(",")}) VALUES (${backtickedColumns.join(",")});
490+
END;
491+
492+
DROP TABLE IF EXISTS \`${stagingTableUnqualified}\`;`;
493+
}
494+
495+
private getIncrementalDmlStatement(
496+
table: dataform.ITable,
497+
columns: string[]
498+
): string {
499+
const incrementalQuery = this.getIncrementalQuery(table);
500+
501+
switch (table.incrementalStrategy) {
502+
case dataform.IncrementalStrategy.INSERT_OVERWRITE:
503+
return this.insertOverwrite(
504+
table.target,
505+
columns,
506+
incrementalQuery,
507+
table.bigquery && table.bigquery.partitionBy,
508+
table.bigquery && table.bigquery.updatePartitionFilter
509+
);
510+
case dataform.IncrementalStrategy.MERGE:
511+
default:
512+
if (table.uniqueKey && table.uniqueKey.length > 0) {
513+
return this.mergeInto(
514+
table.target,
515+
columns,
516+
incrementalQuery,
517+
table.uniqueKey,
518+
table.bigquery && table.bigquery.updatePartitionFilter
519+
);
520+
}
521+
return this.insertInto(
522+
table.target,
523+
columns.map(column => `\`${column}\``),
524+
incrementalQuery
525+
);
526+
}
527+
}
473528
}
474529

475530
export function collectEvaluationQueries(

‎cli/api/execution_sql_test.ts‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,4 +86,34 @@ suite("ExecutionSql with 'onSchemaChange'", () => {
8686
const expectedSql = fs.readFileSync("cli/api/goldens/on_schema_change_ignore.sql", "utf8");
8787
expect(procedureSql).to.equal(expectedSql.trim());
8888
});
89+
90+
test("generates INSERT_OVERWRITE script for IGNORE strategy", () => {
91+
const table = {
92+
...baseTable,
93+
incrementalStrategy: dataform.IncrementalStrategy.INSERT_OVERWRITE,
94+
bigquery: {
95+
partitionBy: "DATE(ts)",
96+
updatePartitionFilter: "ts >= '2024-01-01'"
97+
}
98+
};
99+
const tasks = executionSql.publishTasks(table, { fullRefresh: false }, tableMetadata);
100+
const sql = tasks.build().map(t => t.statement).join("\n;\n");
101+
const expectedSql = fs.readFileSync("cli/api/goldens/insert_overwrite_ignore.sql", "utf8");
102+
expect(sql).to.equal(expectedSql.trim());
103+
});
104+
105+
test("generates INSERT_OVERWRITE script for EXTEND strategy", () => {
106+
const table = {
107+
...baseTable,
108+
incrementalStrategy: dataform.IncrementalStrategy.INSERT_OVERWRITE,
109+
onSchemaChange: dataform.OnSchemaChange.EXTEND,
110+
bigquery: {
111+
partitionBy: "DATE(ts)"
112+
}
113+
};
114+
const tasks = executionSql.publishTasks(table, { fullRefresh: false }, tableMetadata);
115+
const sql = tasks.build().map(t => t.statement).join("\n;\n");
116+
const expectedSql = fs.readFileSync("cli/api/goldens/insert_overwrite_extend.sql", "utf8");
117+
expect(sql).to.equal(expectedSql.trim());
118+
});
89119
});

‎cli/api/goldens/insert_overwrite_extend.sql‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
CREATE OR REPLACE PROCEDURE `project-id.dataset-id.df_osc_test_uuid`()
2+
OPTIONS(strict_mode=false)
3+
BEGIN
4+
5+
-- Create empty table to extract schema of new query.
6+
CREATE OR REPLACE TABLE `project-id.dataset-id.incremental_on_schema_change_df_temp_test_uuid_empty` AS (
7+
SELECT * FROM (select 1 as id, 'a' as field1, 'new' as field2) AS insertions LIMIT 0
8+
);
9+
10+
11+
-- Compare schemas
12+
DECLARE dataform_columns ARRAY<STRING>;
13+
DECLARE temp_table_columns ARRAY<STRUCT<column_name STRING, data_type STRING>>;
14+
DECLARE columns_added ARRAY<STRUCT<column_name STRING, data_type STRING>>;
15+
DECLARE columns_removed ARRAY<STRING>;
16+
17+
SET dataform_columns = (
18+
SELECT IFNULL(ARRAY_AGG(DISTINCT column_name), [])
19+
FROM `project-id.dataset-id.INFORMATION_SCHEMA.COLUMNS`
20+
WHERE table_name = 'incremental_on_schema_change'
21+
);
22+
23+
SET temp_table_columns = (
24+
SELECT IFNULL(ARRAY_AGG(STRUCT(column_name, data_type)), [])
25+
FROM `project-id.dataset-id.INFORMATION_SCHEMA.COLUMNS`
26+
WHERE table_name = 'incremental_on_schema_change_df_temp_test_uuid_empty'
27+
);
28+
29+
SET columns_added = (
30+
SELECT IFNULL(ARRAY_AGG(column_info), [])
31+
FROM UNNEST(temp_table_columns) AS column_info
32+
WHERE column_info.column_name NOT IN UNNEST(dataform_columns)
33+
);
34+
SET columns_removed = (
35+
SELECT IFNULL(ARRAY_AGG(column_name), [])
36+
FROM UNNEST(dataform_columns) AS column_name
37+
WHERE column_name NOT IN (SELECT col.column_name FROM UNNEST(temp_table_columns) AS col)
38+
);
39+
40+
41+
-- Apply schema change strategy (EXTEND).
42+
IF ARRAY_LENGTH(columns_removed) > 0 THEN
43+
RAISE USING MESSAGE = FORMAT(
44+
"Column removals are not allowed when on_schema_change = 'EXTEND'. Removed columns: %T",
45+
columns_removed
46+
);
47+
END IF;
48+
49+
IF ARRAY_LENGTH(columns_added) > 0 THEN
50+
EXECUTE IMMEDIATE (
51+
"ALTER TABLE `project-id.dataset-id.incremental_on_schema_change` " ||
52+
(
53+
SELECT STRING_AGG(FORMAT("ADD COLUMN IF NOT EXISTS %s %s", column_info.column_name, column_info.data_type), ", ")
54+
FROM UNNEST(columns_added) AS column_info
55+
)
56+
);
57+
END IF;
58+
59+
60+
61+
-- Cleanup temporary tables.
62+
DROP TABLE IF EXISTS `project-id.dataset-id.incremental_on_schema_change_df_temp_test_uuid_empty`;
63+
64+
END
65+
;
66+
BEGIN
67+
CALL `project-id.dataset-id.df_osc_test_uuid`();
68+
EXCEPTION WHEN ERROR THEN
69+
DROP TABLE IF EXISTS `project-id.dataset-id.incremental_on_schema_change_df_temp_test_uuid_empty`;
70+
DROP PROCEDURE IF EXISTS `project-id.dataset-id.df_osc_test_uuid`;
71+
RAISE;
72+
END;
73+
DROP PROCEDURE IF EXISTS `project-id.dataset-id.df_osc_test_uuid`
74+
;
75+
CREATE OR REPLACE TEMP TABLE `staging_table_temp_test_uuid` AS (
76+
select 1 as id, 'a' as field1, 'new' as field2
77+
);
78+
79+
BEGIN
80+
DECLARE partitions_for_replacement DEFAULT (
81+
ARRAY(
82+
SELECT DISTINCT DATE(ts)
83+
FROM `staging_table_temp_test_uuid`
84+
WHERE DATE(ts) IS NOT NULL
85+
)
86+
);
87+
88+
MERGE `project-id.dataset-id.incremental_on_schema_change` T
89+
USING `staging_table_temp_test_uuid` S
90+
ON FALSE
91+
WHEN NOT MATCHED BY SOURCE AND DATE(ts) IN UNNEST(partitions_for_replacement) THEN
92+
DELETE
93+
WHEN NOT MATCHED BY TARGET THEN
94+
INSERT (`id`,`field1`) VALUES (`id`,`field1`);
95+
END;
96+
97+
DROP TABLE IF EXISTS `staging_table_temp_test_uuid`

‎cli/api/goldens/insert_overwrite_ignore.sql‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
CREATE OR REPLACE TEMP TABLE `staging_table_temp_test_uuid` AS (
2+
select 1 as id, 'a' as field1, 'new' as field2
3+
);
4+
5+
BEGIN
6+
DECLARE partitions_for_replacement DEFAULT (
7+
ARRAY(
8+
SELECT DISTINCT DATE(ts)
9+
FROM `staging_table_temp_test_uuid`
10+
WHERE DATE(ts) IS NOT NULL
11+
)
12+
);
13+
14+
MERGE `project-id.dataset-id.incremental_on_schema_change` T
15+
USING `staging_table_temp_test_uuid` S
16+
ON FALSE
17+
WHEN NOT MATCHED BY SOURCE AND DATE(ts) IN UNNEST(partitions_for_replacement) and T.ts >= '2024-01-01' THEN
18+
DELETE
19+
WHEN NOT MATCHED BY TARGET THEN
20+
INSERT (`id`,`field1`) VALUES (`id`,`field1`);
21+
END;
22+
23+
DROP TABLE IF EXISTS `staging_table_temp_test_uuid`

‎core/actions/incremental_table.ts‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,9 @@ export class IncrementalTable extends ActionBuilder<dataform.Table> {
217217
} : {}),
218218
});
219219
this.proto.onSchemaChange = this.mapOnSchemaChange(config.onSchemaChange);
220+
this.proto.incrementalStrategy = this.mapIncrementalStrategy(config.incrementalStrategy);
221+
222+
this.checkIncrementalStrategyRequirements(config);
220223

221224
if (config.reservation) {
222225
if (!this.proto.actionDescriptor) {
@@ -736,6 +739,63 @@ export class IncrementalTable extends ActionBuilder<dataform.Table> {
736739
throw new Error(`OnSchemaChange value "${onSchemaChange}" is not supported`);
737740
}
738741
}
742+
743+
private mapIncrementalStrategy(
744+
incrementalStrategy?: string | number
745+
): dataform.IncrementalStrategy {
746+
if (!incrementalStrategy) {
747+
return dataform.IncrementalStrategy.INCREMENTAL_STRATEGY_UNSPECIFIED;
748+
}
749+
750+
if (typeof incrementalStrategy === "number") {
751+
switch (incrementalStrategy) {
752+
case dataform.ActionConfig.IncrementalStrategy.INCREMENTAL_STRATEGY_UNSPECIFIED:
753+
return dataform.IncrementalStrategy.INCREMENTAL_STRATEGY_UNSPECIFIED;
754+
case dataform.ActionConfig.IncrementalStrategy.MERGE:
755+
return dataform.IncrementalStrategy.MERGE;
756+
case dataform.ActionConfig.IncrementalStrategy.INSERT_OVERWRITE:
757+
return dataform.IncrementalStrategy.INSERT_OVERWRITE;
758+
default:
759+
throw new Error(`IncrementalStrategy value "${incrementalStrategy}" is not supported`);
760+
}
761+
}
762+
763+
switch (incrementalStrategy.toString().toUpperCase()) {
764+
case "INCREMENTAL_STRATEGY_UNSPECIFIED":
765+
return dataform.IncrementalStrategy.INCREMENTAL_STRATEGY_UNSPECIFIED;
766+
case "MERGE":
767+
return dataform.IncrementalStrategy.MERGE;
768+
case "INSERT_OVERWRITE":
769+
return dataform.IncrementalStrategy.INSERT_OVERWRITE;
770+
default:
771+
throw new Error(`IncrementalStrategy value "${incrementalStrategy}" is not supported`);
772+
}
773+
}
774+
775+
private checkIncrementalStrategyRequirements(config: dataform.ActionConfig.IIncrementalTableConfig) {
776+
switch (this.proto.incrementalStrategy) {
777+
case dataform.IncrementalStrategy.INSERT_OVERWRITE:
778+
if (!this.proto.bigquery || !this.proto.bigquery.partitionBy) {
779+
this.session.compileError(
780+
new Error("incrementalStrategy 'insert_overwrite' requires 'partitionBy' to be set."),
781+
config.filename,
782+
this.proto.target
783+
);
784+
}
785+
break;
786+
case dataform.IncrementalStrategy.MERGE:
787+
if (!this.proto.uniqueKey || this.proto.uniqueKey.length === 0) {
788+
this.session.compileError(
789+
new Error("incrementalStrategy 'merge' requires 'uniqueKey' to be set."),
790+
config.filename,
791+
this.proto.target
792+
);
793+
}
794+
break;
795+
default:
796+
break;
797+
}
798+
}
739799
}
740800

741801
/**

0 commit comments

Comments
 (0)