Handling for Data Preparation actions (#1789)

* Add logic to handle Data preparation definitions and parse data preparation YAML into the proto representation of the operation.

* Addressed PR Comments

* Fixed lint warnings

Author: fernst

core/BUILD

@@ -1,8 +1,8 @@
-load("//tools:ts_library.bzl", "ts_library")
-load("//tools:expand_template.bzl", "expand_template")
 load("//:version.bzl", "DF_VERSION")
 load("//testing:index.bzl", "ts_test_suite")
+load("//tools:expand_template.bzl", "expand_template")
 load("//tools:node_modules.bzl", "node_modules")
+load("//tools:ts_library.bzl", "ts_library")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -19,14 +19,15 @@ ts_library(
     name = "core",
     srcs = [
         "actions/assertion.ts",
+        "actions/data_preparation.ts",
         "actions/declaration.ts",
-        "actions/view.ts",
         "actions/incremental_table.ts",
         "actions/index.ts",
-        "actions/operation.ts",
         "actions/notebook.ts",
+        "actions/operation.ts",
         "actions/table.ts",
         "actions/test.ts",
+        "actions/view.ts",
         "column_descriptors.ts",
         "common.ts",
         "compilers.ts",

core/actions/data_preparation.ts

@@ -0,0 +1,142 @@
+import { verifyObjectMatchesProto, VerifyProtoErrorBehaviour } from "df/common/protos";
+import { ActionBuilder } from "df/core/actions";
+import { Resolvable } from "df/core/common";
+import * as Path from "df/core/path";
+import { Session } from "df/core/session";
+import {
+  addDependenciesToActionDependencyTargets,
+  nativeRequire,
+  resolveActionsConfigFilename
+} from "df/core/utils";
+import { dataform } from "df/protos/ts";
+
+/**
+ * @hidden
+ */
+export class DataPreparation extends ActionBuilder<dataform.DataPreparation> {
+  public session: Session;
+
+  // TODO: make this field private, to enforce proto update logic to happen in this class.
+  public proto: dataform.IDataPreparation = dataform.DataPreparation.create();
+
+  constructor(
+    session?: Session,
+    config?: dataform.ActionConfig.DataPreparationConfig,
+    configPath?: string
+  ) {
+    super(session);
+    this.session = session;
+
+    if (!config.name) {
+      config.name = Path.basename(config.filename);
+    }
+
+    config.filename = resolveActionsConfigFilename(config.filename, configPath);
+    const dataPreparationContents = nativeRequire(config.filename).asJson;
+    const dataPreparationDefinition = parseDataPreparationDefinitionJson(dataPreparationContents);
+    this.proto.dataPreparation = dataPreparationDefinition;
+
+    // Find targets
+    const targets = getTargets(dataPreparationDefinition)
+    this.proto.targets = targets.map((target) => this.applySessionToTarget(
+      target,
+      session.projectConfig,
+      config.filename,
+      true
+    ));
+    this.proto.canonicalTargets = targets.map((target) => this.applySessionToTarget(target, session.canonicalProjectConfig));
+
+    // Set the unique target key as the first target defined.
+    // TODO: Remove once multiple targets are supported.
+    this.proto.target = this.proto.targets[0];
+    this.proto.canonicalTarget = this.proto.canonicalTargets[0];
+
+    this.proto.tags = config.tags;
+    this.dependencies(config.dependencyTargets);
+    this.proto.fileName = config.filename;
+    if (config.disabled) {
+      this.proto.disabled = config.disabled;
+    }
+  }
+
+  /**
+   * @hidden
+   */
+  public config(config: any) {
+    return this;
+  }
+
+  /**
+   * @hidden
+   */
+  public dependencies(value: Resolvable | Resolvable[]) {
+    const newDependencies = Array.isArray(value) ? value : [value];
+    newDependencies.forEach(resolvable =>
+        addDependenciesToActionDependencyTargets(this, resolvable)
+    );
+    return this;
+  }
+
+  /**
+   * @hidden
+   */
+  public getFileName() {
+    return this.proto.fileName;
+  }
+
+  /**
+   * @hidden
+   */
+  public getTarget() {
+    // Return only the first target for now.
+    return dataform.Target.create(this.proto.target);
+  }
+
+  public compile() {
+    return verifyObjectMatchesProto(
+      dataform.DataPreparation,
+      this.proto,
+      VerifyProtoErrorBehaviour.SUGGEST_REPORTING_TO_DATAFORM_TEAM
+    );
+  }
+}
+
+function parseDataPreparationDefinitionJson(
+    dataPreparationAsJson: { [key: string]: unknown }): dataform.DataPreparationDefinition  {
+  try {
+    return dataform.DataPreparationDefinition.create(
+        verifyObjectMatchesProto(
+            dataform.DataPreparationDefinition,
+            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;
+  }
+}
+
+function getTargets(
+  definition: dataform.DataPreparationDefinition
+): dataform.Target[] {
+  const targets: dataform.Target[] = [];
+
+  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;
+}

core/actions/index.ts

@@ -1,4 +1,5 @@
 import { Assertion } from "df/core/actions/assertion";
+import { DataPreparation } from "df/core/actions/data_preparation";
 import { Declaration } from "df/core/actions/declaration";
 import { IncrementalTable } from "df/core/actions/incremental_table";
 import { Notebook } from "df/core/actions/notebook";
@@ -15,7 +16,8 @@ export type Action =
   | Operation
   | Assertion
   | Declaration
-  | Notebook;
+  | Notebook
+  | DataPreparation;
 
 // TODO(ekrekr): In v4, make all method on inheritors of this private, forcing users to use
 // constructors in order to populate actions.

core/main.ts

@@ -5,6 +5,7 @@ import {
   VerifyProtoErrorBehaviour
 } from "df/common/protos";
 import { Assertion } from "df/core/actions/assertion";
+import { DataPreparation } from "df/core/actions/data_preparation";
 import { Declaration } from "df/core/actions/declaration";
 import { IncrementalTable } from "df/core/actions/incremental_table";
 import { Notebook } from "df/core/actions/notebook";
@@ -179,6 +180,14 @@ function loadActionConfigs(session: Session, filePaths: string[]) {
               actionConfigsPath
             )
           );
+        } else if (actionConfig.dataPreparation) {
+          session.actions.push(
+              new DataPreparation(
+                  session,
+                  dataform.ActionConfig.DataPreparationConfig.create(actionConfig.dataPreparation),
+                  actionConfigsPath
+              )
+          );
         } else {
           throw Error("Empty action configs are not permitted.");
         }

core/main_test.ts

@@ -858,6 +858,155 @@ defaultNotebookRuntimeOptions:
     });
   });
 
+  suite("data preparations", () => {
+    const createSimpleDataPreparationProject = (
+        workflowSettingsYaml = VALID_WORKFLOW_SETTINGS_YAML
+    ): string => {
+      const projectDir = tmpDirFixture.createNewTmpDir();
+      fs.writeFileSync(path.join(projectDir, "workflow_settings.yaml"), workflowSettingsYaml);
+      fs.mkdirSync(path.join(projectDir, "definitions"));
+      fs.writeFileSync(
+          path.join(projectDir, "definitions/actions.yaml"),
+          `
+actions:
+- dataPreparation:
+    filename: data_preparation.yaml`
+      );
+      return projectDir;
+    };
+
+    test(`data preparations can be loaded via an actions config file`, () => {
+      const projectDir = createSimpleDataPreparationProject();
+      fs.writeFileSync(
+          path.join(projectDir, "definitions/data_preparation.yaml"),
+          `
+nodes:
+- id: node1
+  source:
+    table:
+      project: prj
+      dataset: ds
+      table: src
+  destination:
+    table:
+      project: prj
+      dataset: ds
+      table: dest
+  generated:
+    outputSchema:
+      field:
+      - name: a
+        type: INT64
+        mode: NULLABLE
+    sourceGenerated:
+      sourceSchema:
+        tableSchema:
+          field:
+          - name: a
+            type: STRING
+            mode: NULLABLE
+    destinationGenerated:
+      schema:
+        field:
+        - name: a
+          type: STRING
+          mode: NULLABLE
+`
+      );
+
+      const result = runMainInVm(coreExecutionRequestFromPath(projectDir));
+
+      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: [
+                {
+                  database: "prj",
+                  schema: "ds",
+                  name: "dest"
+                }
+              ],
+              canonicalTargets: [
+                {
+                  database: "prj",
+                  schema: "ds",
+                  name: "dest"
+                }
+              ],
+              fileName: "definitions/data_preparation.yaml",
+              dataPreparation: {
+                nodes: [
+                  {
+                    id: "node1",
+                    source: {
+                      table: {
+                        dataset: "ds",
+                        project: "prj",
+                        table: "src"
+                      }
+                    },
+                    destination: {
+                      table: {
+                        dataset: "ds",
+                        project: "prj",
+                        table: "dest"
+                      }
+                    },
+                    generated: {
+                      destinationGenerated: {
+                        schema: {
+                          field: [
+                            {
+                              mode: "NULLABLE",
+                              name: "a",
+                              type: "STRING"
+                            }
+                          ]
+                        }
+                      },
+                      outputSchema: {
+                        field: [
+                          {
+                            mode: "NULLABLE",
+                            name: "a",
+                            type: "INT64"
+                          }
+                        ]
+                      },
+                      sourceGenerated: {
+                        sourceSchema: {
+                          tableSchema: {
+                            field: [
+                              {
+                                mode: "NULLABLE",
+                                name: "a",
+                                type: "STRING"
+                              }
+                            ]
+                          }
+                        }
+                      }
+                    }
+                  }
+                ]
+              }
+            }
+          ])
+      );
+    });
+  });
+
   suite("action configs", () => {
     test(`operations can be loaded`, () => {
       const projectDir = tmpDirFixture.createNewTmpDir();