Added support for compilation overrides in Data Preparations (#1818)

fernst · web-flow · commit 674b1e629b71 · 2024-09-04T10:38:59.000-04:00
* Added support for compilation overrides in Data Preparations

* Added explicit test for unknown properties.

* Fixed tests order

* Fixed whitespace

* Removed logic to ignore schema mismatch

* Add error table to resolution.

* Add error table to resolution.

* Addressed PR Comments
diff --git a/core/actions/data_preparation.ts b/core/actions/data_preparation.ts
@@ -35,7 +35,6 @@ export class DataPreparation extends ActionBuilder<dataform.DataPreparation> {
     config.filename = resolveActionsConfigFilename(config.filename, configPath);
     const dataPreparationAsJson = nativeRequire(config.filename).asJson;
     const dataPreparationDefinition = parseDataPreparationDefinitionJson(dataPreparationAsJson);
-    this.proto.dataPreparationContents = dataform.dataprep.DataPreparation.encode(dataPreparationDefinition).finish();
 
     // Find targets
     const targets = getTargets(dataPreparationDefinition);
@@ -46,6 +45,10 @@ export class DataPreparation extends ActionBuilder<dataform.DataPreparation> {
       this.applySessionToTarget(target, session.canonicalProjectConfig)
     );
 
+    // Resolve all table references with compilation overrides and encode resolved proto instance
+    const resolvedDefinition = applySessionToDataPreparationContents(this, dataPreparationDefinition);
+    this.proto.dataPreparationContents = dataform.dataprep.DataPreparation.encode(resolvedDefinition).finish();
+
     // Set the unique target key as the first target defined.
     // TODO: Remove once multiple targets are supported.
     this.proto.target = this.proto.targets[0];
@@ -128,6 +131,70 @@ function parseDataPreparationDefinitionJson(dataPreparationAsJson: {
   }
 }
 
+function applySessionToDataPreparationContents(
+    actionBuilder: ActionBuilder<dataform.DataPreparation>,
+    definition: dataform.dataprep.DataPreparation
+): dataform.dataprep.DataPreparation {
+  const resolvedDataPreparation = dataform.dataprep.DataPreparation.create(definition);
+
+  // Resolve error table, if set
+  const errorTable = definition.configuration?.errorTable;
+  if (errorTable) {
+    resolvedDataPreparation.configuration.errorTable =
+        applySessionToTableReference(actionBuilder, errorTable);
+  }
+
+  // Loop through all nodes and resolve the compilation overrides for
+  // all source and destination tables.
+  definition.nodes.forEach((node, index) => {
+
+        // Resolve source tables, if set.
+        const sourceTable = node.source.table;
+        if (sourceTable) {
+          resolvedDataPreparation.nodes[index].source.table =
+              applySessionToTableReference(actionBuilder, sourceTable);
+        }
+
+        // Resolve destination tables, if set.
+        const destinationTable = node.destination?.table;
+        if (destinationTable) {
+          resolvedDataPreparation.nodes[index].destination.table =
+              applySessionToTableReference(actionBuilder, destinationTable);
+        }
+      }
+
+    );
+
+  return resolvedDataPreparation;
+}
+
+function applySessionToTableReference(
+    actionBuilder: ActionBuilder<dataform.DataPreparation>,
+    tableReference: dataform.dataprep.ITableReference
+): dataform.dataprep.ITableReference {
+  const target: dataform.ITarget = {
+    database: tableReference.project,
+    schema: tableReference.dataset,
+    name: tableReference.table
+  }
+  const resolvedTarget =
+      actionBuilder.applySessionToTarget(
+          dataform.Target.create(target),
+          actionBuilder.session.projectConfig)
+  // Convert resolved target into a Data Preparation Table Reference
+  const resolvedTableReference = dataform.dataprep.TableReference.create({
+    table: resolvedTarget.name
+  });
+  if (resolvedTarget.database) {
+    resolvedTableReference.project = resolvedTarget.database;
+  }
+  if (resolvedTarget.schema) {
+    resolvedTableReference.dataset = resolvedTarget.schema;
+  }
+  return resolvedTableReference;
+}
+
+
 function getTargets(definition: dataform.dataprep.DataPreparation): dataform.Target[] {
   const targets: dataform.Target[] = [];
 
diff --git a/core/main_test.ts b/core/main_test.ts
@@ -5,7 +5,7 @@ import { dump as dumpYaml, load as loadYaml } from "js-yaml";
 import * as path from "path";
 import { CompilerFunction, NodeVM } from "vm2";
 
-import { decode64, encode64, verifyObjectMatchesProto } from "df/common/protos";
+import { decode64, encode64, verifyObjectMatchesProto, VerifyProtoErrorBehaviour } from "df/common/protos";
 import { compile } from "df/core/compilers";
 import { version } from "df/core/version";
 import { dataform } from "df/protos/ts";
@@ -913,7 +913,7 @@ nodes:
 `
 
       fs.writeFileSync(
-        path.join(projectDir, "definitions/data_preparation.yaml"),
+          path.join(projectDir, "definitions/data_preparation.yaml"),
           dataPreparationYaml
       );
 
@@ -923,7 +923,8 @@ nodes:
           dataform.dataprep.DataPreparation,
           dataPreparationAsObject as {
             [key: string]: any;
-          }
+          },
+          VerifyProtoErrorBehaviour.DEFAULT
       );
       const base64encodedContents = encode64(
           dataform.dataprep.DataPreparation,
@@ -934,37 +935,200 @@ nodes:
 
       expect(result.compile.compiledGraph.graphErrors.compilationErrors).deep.equals([]);
       expect(asPlainObject(result.compile.compiledGraph.dataPreparations)).deep.equals(
-        asPlainObject([
-          {
-            target: {
-              database: "prj",
-              schema: "ds",
-              name: "dest"
-            },
-            canonicalTarget: {
-              database: "prj",
-              schema: "ds",
-              name: "dest"
-            },
-            targets: [
-              {
+          asPlainObject([
+            {
+              target: {
                 database: "prj",
                 schema: "ds",
                 name: "dest"
-              }
-            ],
-            canonicalTargets: [
-              {
+              },
+              canonicalTarget: {
                 database: "prj",
                 schema: "ds",
                 name: "dest"
-              }
-            ],
-            fileName: "definitions/data_preparation.yaml",
-            // Base64 encoded representation of the data preparation definition proto.
-            dataPreparationContents: base64encodedContents
-          }
-        ])
+              },
+              targets: [
+                {
+                  database: "prj",
+                  schema: "ds",
+                  name: "dest"
+                }
+              ],
+              canonicalTargets: [
+                {
+                  database: "prj",
+                  schema: "ds",
+                  name: "dest"
+                }
+              ],
+              fileName: "definitions/data_preparation.yaml",
+              // Base64 encoded representation of the data preparation definition proto.
+              dataPreparationContents: base64encodedContents
+            }
+          ])
+      );
+    });
+
+    test(`data preparations resolves compilation overrides before encoding`, () => {
+      const projectDir = createSimpleDataPreparationProject();
+      const dataPreparationYaml = `
+configuration:
+  errorTable:
+    table: error
+nodes:
+- id: node1
+  source:
+    table:
+      table: src
+  generated:
+    sourceGenerated:
+      sourceSchema:
+        tableSchema:
+          field:
+          - name: a
+            type: STRING
+            mode: NULLABLE
+    outputSchema:
+      field:
+      - name: a
+        type: INT64
+        mode: NULLABLE
+- id: node2
+  source:
+    nodeId: node1
+  destination:
+    table:
+      table: dest
+  generated:
+    sourceGenerated:
+      sourceSchema:
+        nodeSchema:
+          field:
+          - name: a
+            type: STRING
+            mode: NULLABLE
+    outputSchema:
+      field:
+      - name: a
+        type: INT64
+        mode: NULLABLE
+    destinationGenerated:
+      schema:
+        field:
+        - name: a
+          type: STRING
+          mode: NULLABLE
+`
+
+      fs.writeFileSync(
+          path.join(projectDir, "definitions/data_preparation.yaml"),
+          dataPreparationYaml
+      );
+
+      const resolvedYaml = `
+configuration:
+  errorTable:
+    project: defaultProject
+    dataset: defaultDataset
+    table: error
+nodes:
+- id: node1
+  source:
+    table:
+      project: defaultProject
+      dataset: defaultDataset
+      table: src
+  generated:
+    sourceGenerated:
+      sourceSchema:
+        tableSchema:
+          field:
+          - name: a
+            type: STRING
+            mode: NULLABLE
+    outputSchema:
+      field:
+      - name: a
+        type: INT64
+        mode: NULLABLE
+- id: node2
+  source:
+    nodeId: node1
+  destination:
+    table:
+      project: defaultProject
+      dataset: defaultDataset
+      table: dest
+  generated:
+    sourceGenerated:
+      sourceSchema:
+        nodeSchema:
+          field:
+          - name: a
+            type: STRING
+            mode: NULLABLE
+    outputSchema:
+      field:
+      - name: a
+        type: INT64
+        mode: NULLABLE
+    destinationGenerated:
+      schema:
+        field:
+        - name: a
+          type: STRING
+          mode: NULLABLE
+`
+
+      // Generate Base64 encoded representation of the YAML.
+      const dataPreparationAsObject = loadYaml(resolvedYaml);
+      const dataPreparationDefinition = verifyObjectMatchesProto(
+          dataform.dataprep.DataPreparation,
+          dataPreparationAsObject as {
+            [key: string]: any;
+          },
+          VerifyProtoErrorBehaviour.DEFAULT
+      );
+      const base64encodedContents = encode64(
+          dataform.dataprep.DataPreparation,
+          dataPreparationDefinition
+      );
+
+      const result = runMainInVm(coreExecutionRequestFromPath(projectDir));
+
+      expect(result.compile.compiledGraph.graphErrors.compilationErrors).deep.equals([]);
+      expect(asPlainObject(result.compile.compiledGraph.dataPreparations)).deep.equals(
+          asPlainObject([
+            {
+              target: {
+                database: "defaultProject",
+                schema: "defaultDataset",
+                name: "dest"
+              },
+              canonicalTarget: {
+                database: "defaultProject",
+                schema: "defaultDataset",
+                name: "dest"
+              },
+              targets: [
+                {
+                  database: "defaultProject",
+                  schema: "defaultDataset",
+                  name: "dest"
+                }
+              ],
+              canonicalTargets: [
+                {
+                  database: "defaultProject",
+                  schema: "defaultDataset",
+                  name: "dest"
+                }
+              ],
+              fileName: "definitions/data_preparation.yaml",
+              // Base64 encoded representation of the data preparation definition proto.
+              dataPreparationContents: base64encodedContents
+            }
+          ])
       );
     });
   });
