Update Data Preparation processor to remove proto dependency. (#1839)

fernst · web-flow · commit 86254ebf9f57 · 2024-09-11T15:34:52.000-04:00
* Update Data Preparation processor to remove proto dependency. The data preparation in the compiled graph now contains the resolved YAML after compiling.

* Fixed compilation issue
diff --git a/core/actions/data_preparation.ts b/core/actions/data_preparation.ts
@@ -1,3 +1,5 @@
+import { dump as dumpYaml } from "js-yaml";
+
 import { verifyObjectMatchesProto, VerifyProtoErrorBehaviour } from "df/common/protos";
 import { ActionBuilder } from "df/core/actions";
 import { Resolvable } from "df/core/common";
@@ -34,10 +36,12 @@ export class DataPreparation extends ActionBuilder<dataform.DataPreparation> {
 
     config.filename = resolveActionsConfigFilename(config.filename, configPath);
     const dataPreparationAsJson = nativeRequire(config.filename).asJson;
-    const dataPreparationDefinition = parseDataPreparationDefinitionJson(dataPreparationAsJson);
+    const dataPreparationDefinition = this.parseDataPreparationDefinitionJson(dataPreparationAsJson);
 
     // Find targets
-    const targets = getTargets(dataPreparationDefinition);
+    const targets = this.getTargets(dataPreparationAsJson as {
+      [key: string]: any;
+    });
     this.proto.targets = targets.map(target =>
       this.applySessionToTarget(target, session.projectConfig, config.filename, true)
     );
@@ -46,8 +50,8 @@ export class DataPreparation extends ActionBuilder<dataform.DataPreparation> {
     );
 
     // 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();
+    const resolvedDefinition = this.applySessionToDataPreparationContents(dataPreparationAsJson);
+    this.proto.dataPreparationYaml = dumpYaml(resolvedDefinition)
 
     // Set the unique target key as the first target defined.
     // TODO: Remove once multiple targets are supported.
@@ -108,107 +112,112 @@ export class DataPreparation extends ActionBuilder<dataform.DataPreparation> {
       VerifyProtoErrorBehaviour.SUGGEST_REPORTING_TO_DATAFORM_TEAM
     );
   }
-}
 
-function parseDataPreparationDefinitionJson(dataPreparationAsJson: {
-  [key: string]: unknown;
-}): dataform.dataprep.DataPreparation {
-  try {
-    return dataform.dataprep.DataPreparation.create(
-      verifyObjectMatchesProto(
-        dataform.dataprep.DataPreparation,
-        dataPreparationAsJson as {
-          [key: string]: any;
-        },
-        VerifyProtoErrorBehaviour.SHOW_DOCS_LINK
-      )
-    );
-  } catch (e) {
-    if (e instanceof ReferenceError) {
-      throw ReferenceError(`Data Preparation parsing error: ${e.message}`);
+  private applySessionToDataPreparationContents(
+      definition: {[key: string]: any}
+  ): {[key: string]: any} {
+    // Resolve error table, if set
+    // @ts-ignore
+    const errorTable = definition.configuration?.errorTable;
+    if (errorTable) {
+      definition.configuration.errorTable =
+          this.applySessionToTableReference(errorTable as {[key: string]: string});
     }
-    throw e;
-  }
-}
-
-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) => {
+    // Loop through all nodes and resolve the compilation overrides for
+    // all source and destination tables.
+    if (definition.nodes) {
+      (definition.nodes as Array<{ [key: string]: any }>).forEach((node, index) => {
 
         // Resolve source tables, if set.
-        const sourceTable = node.source.table;
+        const sourceTable = node.source?.table;
         if (sourceTable) {
-          resolvedDataPreparation.nodes[index].source.table =
-              applySessionToTableReference(actionBuilder, sourceTable);
+          definition.nodes[index].source.table =
+              this.applySessionToTableReference(sourceTable as {[key: string]: string});
         }
 
         // Resolve destination tables, if set.
         const destinationTable = node.destination?.table;
         if (destinationTable) {
-          resolvedDataPreparation.nodes[index].destination.table =
-              applySessionToTableReference(actionBuilder, destinationTable);
+          definition.nodes[index].destination.table =
+              this.applySessionToTableReference(destinationTable as {[key: string]: string});
         }
-      }
+      });
+    }
 
-    );
+    return definition;
+  }
 
-  return resolvedDataPreparation;
-}
+  private applySessionToTableReference(
+      tableReference: {[key: string]: string}
+  ): object {
+    const target: dataform.ITarget = {
+      database: tableReference.project,
+      schema: tableReference.dataset,
+      name: tableReference.table
+    }
+    const resolvedTarget =
+        this.applySessionToTarget(
+            dataform.Target.create(target),
+            this.session.projectConfig)
+    // Convert resolved target into a Data Preparation Table Reference
+    let resolvedTableReference : {[key: string]: string} = {
+      table: resolvedTarget.name,
+    }
 
-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;
+    // Ensure project and dataset field are added in order
+    if (resolvedTarget.schema) {
+      resolvedTableReference = { dataset: resolvedTarget.schema, ...resolvedTableReference }
+    }
+    if (resolvedTarget.database) {
+      resolvedTableReference = { project: resolvedTarget.database, ...resolvedTableReference }
+    }
+    return resolvedTableReference;
   }
-  if (resolvedTarget.schema) {
-    resolvedTableReference.dataset = resolvedTarget.schema;
+
+  private parseDataPreparationDefinitionJson(dataPreparationAsJson: {
+    [key: string]: unknown;
+  }): dataform.dataprep.DataPreparation {
+    try {
+      return dataform.dataprep.DataPreparation.create(
+          verifyObjectMatchesProto(
+              dataform.dataprep.DataPreparation,
+              dataPreparationAsJson as {
+                [key: string]: any;
+              },
+              VerifyProtoErrorBehaviour.SHOW_DOCS_LINK
+          )
+      );
+    } catch (e) {
+      if (e instanceof ReferenceError) {
+        throw ReferenceError(`Data Preparation parsing error: ${e.message}`);
+      }
+      throw e;
+    }
   }
-  return resolvedTableReference;
-}
 
 
-function getTargets(definition: dataform.dataprep.DataPreparation): dataform.Target[] {
-  const targets: dataform.Target[] = [];
+  private getTargets(definition: {
+    [key: string]: any;
+  }): dataform.Target[] {
+    const targets: dataform.Target[] = [];
+
+    if (definition.nodes) {
+      (definition.nodes as Array<{ [key: string]: any }>).forEach(node => {
+        const table = node.destination?.table;
+        if (table) {
+          const compiledGraphTarget: dataform.ITarget = {
+            database: table.project,
+            schema: table.dataset,
+            name: table.table
+          };
+          targets.push(dataform.Target.create(compiledGraphTarget));
 
-  definition.nodes.forEach(node => {
-    const table = node.destination?.table;
-    if (table) {
-      const compiledGraphTarget: dataform.ITarget = {
-        database: table.project,
-        schema: table.dataset,
-        name: table.table
-      };
-      targets.push(dataform.Target.create(compiledGraphTarget));
+        }
+      });
     }
-  });
 
-  return targets;
+    return targets;
+  }
 }
+
diff --git a/core/main_test.ts b/core/main_test.ts
@@ -917,20 +917,6 @@ nodes:
           dataPreparationYaml
       );
 
-      // Generate Base64 encoded representation of the YAML.
-      const dataPreparationAsObject = loadYaml(dataPreparationYaml);
-      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([]);
@@ -962,8 +948,7 @@ nodes:
                 }
               ],
               fileName: "definitions/data_preparation.yaml",
-              // Base64 encoded representation of the data preparation definition proto.
-              dataPreparationContents: base64encodedContents
+              dataPreparationYaml: dumpYaml(loadYaml(dataPreparationYaml))
             }
           ])
       );
@@ -1080,20 +1065,6 @@ nodes:
           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([]);
@@ -1125,8 +1096,7 @@ nodes:
                 }
               ],
               fileName: "definitions/data_preparation.yaml",
-              // Base64 encoded representation of the data preparation definition proto.
-              dataPreparationContents: base64encodedContents
+              dataPreparationYaml: dumpYaml(loadYaml(resolvedYaml))
             }
           ])
       );
diff --git a/protos/core.proto b/protos/core.proto
@@ -267,10 +267,10 @@ message DataPreparation {
 
   bool disabled = 6;
 
-  // Binary encoded contents of the Data preparation proto
-  bytes data_preparation_contents = 10;
+  // YAML file containing the data preparation YAML with resolved dependencies.
+  string data_preparation_yaml = 11;
 
-  reserved 7;
+  reserved 7, 10;
 }
 
 message CompiledGraph {
