Watch for dvc.yaml changes for manually added stages (#3365)

* Trigger a new webview message on dvc yaml change

* Correctly update the webview if there is a stage or not and if the yaml is valid

* Add test

* Remove useless test id

---------

Co-authored-by: Matt Seddon <[email protected]>

Author: sroy3

extension/src/experiments/index.ts

@@ -156,6 +156,7 @@ export class Experiments extends BaseRepository<TableData> {
         if (hadCheckpoints !== this.hasCheckpoints()) {
           this.checkpointsChanged.fire()
         }
+        void this.webviewMessages.changeHasConfig(true)
       })
     )
 

extension/src/experiments/webview/contract.ts

@@ -100,6 +100,7 @@ export type TableData = {
   hasColumns: boolean
   hasConfig: boolean
   hasRunningExperiment: boolean
+  hasValidDvcYaml: boolean
   rows: Row[]
   sorts: SortDefinition[]
   filteredCounts: FilteredCounts

extension/src/experiments/webview/messages.ts

@@ -43,6 +43,7 @@ export class WebviewMessages {
   private readonly hasStages: () => Promise<string>
 
   private hasConfig = false
+  private hasValidDvcYaml = true
 
   private readonly addStage: () => Promise<boolean>
 
@@ -74,8 +75,11 @@ export class WebviewMessages {
     this.addStage = addStage
   }
 
-  public async changeHasConfig() {
-    this.hasConfig = !!(await this.hasStages())
+  public async changeHasConfig(update?: boolean) {
+    const stages = await this.hasStages()
+    this.hasValidDvcYaml = stages !== undefined
+    this.hasConfig = !!stages
+    update && this.sendWebviewMessage()
   }
 
   public sendWebviewMessage() {
@@ -218,14 +222,14 @@ export class WebviewMessages {
       hasColumns: this.columns.hasNonDefaultColumns(),
       hasConfig: this.hasConfig,
       hasRunningExperiment: this.experiments.hasRunningExperiment(),
+      hasValidDvcYaml: this.hasValidDvcYaml,
       rows: this.experiments.getRowData(),
       sorts: this.experiments.getSorts()
     }
   }
 
   private async addConfiguration() {
     await this.addStage()
-    await this.changeHasConfig()
   }
 
   private async setMaxTableHeadDepth() {

extension/src/test/fixtures/expShow/base/tableData.ts

@@ -10,6 +10,7 @@ const tableDataFixture: TableData = {
   hasCheckpoints: true,
   hasConfig: true,
   hasRunningExperiment: true,
+  hasValidDvcYaml: true,
   hasColumns: true,
   sorts: [],
   changes: [],

extension/src/test/fixtures/expShow/dataTypes/tableData.ts

@@ -11,6 +11,7 @@ export const data: TableData = {
   hasCheckpoints: false,
   hasConfig: true,
   hasRunningExperiment: false,
+  hasValidDvcYaml: true,
   sorts: [],
   columns,
   hasColumns: true,

extension/src/test/fixtures/expShow/deeplyNested/tableData.ts

@@ -14,6 +14,7 @@ const data: TableData = {
   hasCheckpoints: false,
   hasConfig: true,
   hasRunningExperiment: false,
+  hasValidDvcYaml: true,
   sorts: [
     {
       path: 'params:params.yaml:nested1.doubled',

extension/src/test/fixtures/expShow/survival/tableData.ts

@@ -11,6 +11,7 @@ const data: TableData = {
   hasConfig: true,
   hasRunningExperiment: true,
   hasColumns: true,
+  hasValidDvcYaml: true,
   sorts: [],
   changes: [],
   columnOrder: [],

extension/src/test/suite/experiments/index.test.ts

@@ -153,6 +153,7 @@ suite('Experiments Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: rowsFixture,
         sorts: []
       }
@@ -208,6 +209,7 @@ suite('Experiments Test Suite', () => {
         hasColumns: true,
         hasConfig: false,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: rowsFixture,
         sorts: []
       }
@@ -236,6 +238,7 @@ suite('Experiments Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: rowsFixture,
         sorts: []
       }
@@ -968,6 +971,7 @@ suite('Experiments Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: rowsFixture,
         sorts: []
       }

extension/src/test/suite/experiments/model/filterBy/tree.test.ts

@@ -127,6 +127,7 @@ suite('Experiments Filter By Tree Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: filteredRows,
         sorts: []
       }
@@ -159,6 +160,7 @@ suite('Experiments Filter By Tree Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: [workspace, main],
         sorts: []
       }
@@ -426,6 +428,7 @@ suite('Experiments Filter By Tree Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: filteredRows,
         sorts: []
       }

webview/src/experiments/components/AddStage.tsx

@@ -5,16 +5,28 @@ import { IconButton } from '../../shared/components/button/IconButton'
 import { Add } from '../../shared/components/icons'
 import { sendMessage } from '../../shared/vscode'
 
-export const AddStage: React.FC = () => (
+interface AddStageProps {
+  hasValidDvcYaml: boolean
+}
+
+export const AddStage: React.FC<AddStageProps> = ({ hasValidDvcYaml }) => (
   <div className={styles.addConfigButton}>
     <p>Easily and efficiently reproduce your experiments </p>
     <IconButton
       icon={Add}
       onClick={() =>
+        hasValidDvcYaml &&
         sendMessage({ type: MessageFromWebviewType.ADD_CONFIGURATION })
       }
       text="Add a Pipeline Stage"
+      disabled={!hasValidDvcYaml}
     />
+    {!hasValidDvcYaml && (
+      <p className={styles.error}>
+        Your dvc.yaml file should contain valid yaml before adding any pipeline
+        stages.
+      </p>
+    )}
     <p>
       <a href="https://dvc.org/doc/user-guide/project-structure/dvcyaml-files#stages">
         Learn more