Set Dagger shell for empty .dag notebooks (#1949)

stateful · Feb 8, 2025 · cca2d68 · cca2d68
1 parent 29d8529
commit cca2d68
Show file tree

Hide file tree

Showing 2 changed files with 95 additions and 0 deletions.
diff --git a/src/extension/serializer/serializer.ts b/src/extension/serializer/serializer.ts
@@ -134,9 +134,35 @@ export abstract class GrpcSerializer implements ISerializer {
       return
     }
 
+    await this.applyFrontmatterForFileExtension(doc)
+
     this.cacheDocUriMapping.set(cacheId, doc.uri)
   }
 
+  private async applyFrontmatterForFileExtension(notebook: NotebookDocument): Promise<boolean> {
+    // don't tamper with existing docs
+    if (notebook.cellCount > 0 || !!notebook.metadata?.['runme.dev/frontmatterParsed']) {
+      return false
+    }
+
+    const fileExtension = path.extname(notebook.uri.fsPath)
+    if (fileExtension.toLowerCase() === '.dag') {
+      const metadata = {
+        ...notebook.metadata,
+        'runme.dev/frontmatter': '---\n\shell: dagger shell\n---',
+        'runme.dev/frontmatterParsed': { shell: 'dagger shell' },
+      }
+      const notebookEdit = NotebookEdit.updateNotebookMetadata(metadata)
+      const edit = new WorkspaceEdit()
+      edit.set(notebook.uri, [notebookEdit])
+      await workspace.applyEdit(edit)
+      await notebook.save()
+      return true
+    }
+
+    return false
+  }
+
   private async handleSaveNotebookOutputs(doc: NotebookDocument) {
     const cacheId = getDocumentCacheId(doc.metadata)
 

diff --git a/tests/extension/serializer.test.ts b/tests/extension/serializer.test.ts
@@ -386,6 +386,75 @@ describe('ConnectSerializer', () => {
     })
   })
 
+  describe('#applyFrontmatterForFileExtension', () => {
+    const fakeSrcDocUri = { fsPath: '/tmp/fake/source.md' } as any
+
+    it('should not run for documents with cells', async () => {
+      const serializer: any = new ConnectSerializer(context, new Server(), new Kernel())
+
+      const applied = await serializer['applyFrontmatterForFileExtension']({
+        cellCount: 2,
+        uri: fakeSrcDocUri,
+      })
+
+      expect(applied).toBeFalsy()
+    })
+
+    it('should not run documents that are not empty', async () => {
+      const fixture = deepCopyFixture()
+
+      const serializer: any = new ConnectSerializer(context, new Server(), new Kernel())
+
+      const applied = await serializer['applyFrontmatterForFileExtension']({
+        cellCount: 0,
+        uri: fakeSrcDocUri,
+        metadata: fixture.metadata,
+      })
+
+      expect(applied).toBeFalsy()
+    })
+
+    it('should apply dagger shell for .dag file extensions', async () => {
+      vi.mocked(workspace.applyEdit).mockResolvedValue(true)
+
+      const fixture = deepCopyFixture()
+      fakeSrcDocUri.fsPath = '/tmp/fake/source.dag'
+
+      const serializer: any = new ConnectSerializer(context, new Server(), new Kernel())
+
+      delete fixture.metadata['runme.dev/frontmatterParsed']
+      delete fixture.metadata['runme.dev/frontmatter']
+      const save = vi.fn()
+      const applied = await serializer['applyFrontmatterForFileExtension']({
+        save,
+        uri: fakeSrcDocUri,
+        metadata: fixture.metadata,
+      })
+
+      const expectedEdits = new Map()
+      expectedEdits.set(
+        {
+          fsPath: '/tmp/fake/source.dag',
+        },
+        [
+          {
+            metadata: {
+              'runme.dev/cacheId': '01J97S5FVEKBPD9GAH0AZBV0HB',
+              'runme.dev/finalLineBreaks': '1',
+              'runme.dev/frontmatter': '---\nshell: dagger shell\n---',
+              'runme.dev/frontmatterParsed': { shell: 'dagger shell' },
+            },
+            type: 'updateNotebookMetadata',
+          },
+        ],
+      )
+
+      expect(save).toBeCalled()
+      expect(workspace.applyEdit).toBeCalledWith(expectedEdits)
+      expect(applied).toBeTruthy()
+    })
+  })
+
   describe('cell execution summary marshaling', () => {
     it('should not misrepresenting uninitialized values', () => {
       // i.e. undefined is not sucess=false