Add repository snapshot destination field to NotebookRuntimeOptions (#2078)

* Add repository snapshot destination field to NotebookRuntimeOptions

* Improve code clarity and add more tests

Author: mfiglus

core/actions/notebook_test.ts

@@ -108,6 +108,8 @@ defaultLocation: US
 defaultNotebookRuntimeOptions:
   outputBucket: gs://some-bucket
   runtimeTemplateName: projects/test-project/locations/us-central1/notebookRuntimeTemplates/test-template
+  repositorySnapshotDestination:
+    repositorySnapshotUri: gs://some-other-bucket
   `);
     fs.writeFileSync(path.join(projectDir, "definitions/notebook.ipynb"), EMPTY_NOTEBOOK_CONTENTS);
 
@@ -120,12 +122,59 @@ defaultNotebookRuntimeOptions:
       defaultNotebookRuntimeOptions: {
         outputBucket: "gs://some-bucket",
         runtimeTemplateName:
-          "projects/test-project/locations/us-central1/notebookRuntimeTemplates/test-template"
+          "projects/test-project/locations/us-central1/notebookRuntimeTemplates/test-template",
+        repositorySnapshotDestination: {
+          repositorySnapshotUri: "gs://some-other-bucket"
+        }
+      },
+      warehouse: "bigquery"
+    });
+  });
+
+  test(`notebook default runtime options snapshot destination defaults to output bucket`, () => {
+    const projectDir = createSimpleNotebookProject(`
+defaultProject: dataform
+defaultLocation: US
+defaultNotebookRuntimeOptions:
+  outputBucket: gs://some-bucket
+  runtimeTemplateName: projects/test-project/locations/us-central1/notebookRuntimeTemplates/test-template
+  repositorySnapshotDestination: {}
+  `);
+    fs.writeFileSync(path.join(projectDir, "definitions/notebook.ipynb"), EMPTY_NOTEBOOK_CONTENTS);
+
+    const result = runMainInVm(coreExecutionRequestFromPath(projectDir));
+
+    expect(result.compile.compiledGraph.graphErrors.compilationErrors).deep.equals([]);
+    expect(asPlainObject(result.compile.compiledGraph.projectConfig)).deep.equals({
+      defaultDatabase: "dataform",
+      defaultLocation: "US",
+      defaultNotebookRuntimeOptions: {
+        outputBucket: "gs://some-bucket",
+        runtimeTemplateName:
+          "projects/test-project/locations/us-central1/notebookRuntimeTemplates/test-template",
+        repositorySnapshotDestination: {
+          repositorySnapshotUri: "gs://some-bucket"
+        }
       },
       warehouse: "bigquery"
     });
   });
 
+  test(`notebook default runtime options throw for snapshot destination with no uri or output bucket`, () => {
+    const projectDir = createSimpleNotebookProject(`
+defaultProject: dataform
+defaultLocation: US
+defaultNotebookRuntimeOptions:
+  runtimeTemplateName: projects/test-project/locations/us-central1/notebookRuntimeTemplates/test-template
+  repositorySnapshotDestination: {}
+  `);
+    fs.writeFileSync(path.join(projectDir, "definitions/notebook.ipynb"), EMPTY_NOTEBOOK_CONTENTS);
+
+    expect(() => runMainInVm(coreExecutionRequestFromPath(projectDir))).to.throw(
+      "Invalid repository_snapshot_destination: either repository_snapshot_uri or output_bucket has to be defined"
+    );
+  });
+
   suite("sqlx and JS API config options", () => {
     test(`for notebooks`, () => {
       const projectDir = tmpDirFixture.createNewTmpDir();

core/workflow_settings.ts

@@ -124,13 +124,27 @@ export function workflowSettingsAsProjectConfig(
   }
   if (workflowSettings.defaultNotebookRuntimeOptions) {
     projectConfig.defaultNotebookRuntimeOptions = {};
-    if (workflowSettings.defaultNotebookRuntimeOptions.outputBucket) {
-      projectConfig.defaultNotebookRuntimeOptions.outputBucket =
-        workflowSettings.defaultNotebookRuntimeOptions.outputBucket;
+    const {outputBucket, runtimeTemplateName, repositorySnapshotDestination} =
+      workflowSettings.defaultNotebookRuntimeOptions;
+    if (outputBucket) {
+      projectConfig.defaultNotebookRuntimeOptions.outputBucket = outputBucket;
     }
-    if (workflowSettings.defaultNotebookRuntimeOptions.runtimeTemplateName) {
-      projectConfig.defaultNotebookRuntimeOptions.runtimeTemplateName =
-        workflowSettings.defaultNotebookRuntimeOptions.runtimeTemplateName;
+    if (runtimeTemplateName) {
+      projectConfig.defaultNotebookRuntimeOptions.runtimeTemplateName = runtimeTemplateName;
+    }
+    if (repositorySnapshotDestination) {
+      projectConfig.defaultNotebookRuntimeOptions.repositorySnapshotDestination = {};
+      if (repositorySnapshotDestination.repositorySnapshotUri) {
+        projectConfig.defaultNotebookRuntimeOptions.repositorySnapshotDestination.repositorySnapshotUri =
+          repositorySnapshotDestination.repositorySnapshotUri;
+      } else if (outputBucket) {
+        projectConfig.defaultNotebookRuntimeOptions.repositorySnapshotDestination.repositorySnapshotUri =
+          outputBucket;
+      } else {
+        throw Error(
+          "Invalid repository_snapshot_destination: either repository_snapshot_uri or output_bucket " +
+            "has to be defined");
+      }
     }
   }
   if(workflowSettings.defaultIcebergConfig) {

examples/extreme_weather_programming/workflow_settings.yaml

@@ -5,3 +5,5 @@ defaultAssertionDataset: dataform_extreme_weather_programming_example_assertions
 defaultNotebookRuntimeOptions:
   outputBucket: gs://some-bucket
   runtimeTemplateName: projects/test-project/locations/us-central1/notebookRuntimeTemplates/test-template
+  repositorySnapshotDestination:
+    repositorySnapshotUri: gs://some-other-bucket

protos/configs.proto

@@ -680,6 +680,11 @@ message ActionConfig {
   }
 }
 
+message RepositorySnapshotDestinationConfig {
+  // Storage URI to upload the repository snapshot to.
+  string repository_snapshot_uri = 1;
+}
+
 message NotebookRuntimeOptionsConfig {
   oneof output_sink {
     // Storage bucket to output notebooks to after their execution.
@@ -689,4 +694,13 @@ message NotebookRuntimeOptionsConfig {
   // Colab runtime template (https://cloud.google.com/colab/docs/runtimes), from
   // which a runtime is created for notebook executions.
   string runtime_template_name = 2;
+
+  // The destination for the snapshot of repository files to be available for
+  // read-only access inside a notebook runtime.
+  oneof repository_snapshot_storage {
+    // Storage URI to upload the snapshot to.
+    // For empty URI it defaults to the provided output_bucket.
+    RepositorySnapshotDestinationConfig repository_snapshot_destination = 3;
+  }
+
 }

protos/core.proto

@@ -310,12 +310,20 @@ message Notebook {
   string notebook_contents = 7;
 }
 
+message RepositorySnapshotDestination {
+  string repository_snapshot_uri = 1;
+}
+
 message NotebookRuntimeOptions {
   oneof output_sink {
     string output_bucket = 1;
   }
 
   string runtime_template_name = 2;
+
+  oneof repository_snapshot_storage {
+    RepositorySnapshotDestination repository_snapshot_destination = 3;
+  }
 }
 
 // Data Preparation Related entries