Add upgrade dvc button to setup when incompatible (#3904)

extension/src/setup/autoInstall.ts

@@ -16,6 +16,31 @@ export const findPythonBinForInstall = async (): Promise<
   )
 }
 
+const showUpgradeProgress = (
+  root: string,
+  pythonBinPath: string
+): Thenable<unknown> =>
+  Toast.showProgress('Upgrading DVC', async progress => {
+    progress.report({ increment: 0 })
+
+    progress.report({ increment: 25, message: 'Updating packages...' })
+
+    try {
+      await Toast.runCommandAndIncrementProgress(
+        async () => {
+          await installPackages(root, pythonBinPath, 'dvc')
+          return 'Upgraded successfully'
+        },
+        progress,
+        75
+      )
+
+      return Toast.delayProgressClosing()
+    } catch (error: unknown) {
+      return Toast.reportProgressError(error, progress)
+    }
+  })
+
 const showInstallProgress = (
   root: string,
   pythonBinPath: string
@@ -52,7 +77,9 @@ const showInstallProgress = (
     }
   })
 
-export const autoInstallDvc = async (): Promise<unknown> => {
+const getArgsAndRunCommand = async (
+  command: (root: string, pythonBinPath: string) => Thenable<unknown>
+): Promise<unknown> => {
   const pythonBinPath = await findPythonBinForInstall()
   const root = getFirstWorkspaceFolder()
 
@@ -68,5 +95,13 @@ export const autoInstallDvc = async (): Promise<unknown> => {
     )
   }
 
-  return showInstallProgress(root, pythonBinPath)
+  return command(root, pythonBinPath)
+}
+
+export const autoInstallDvc = (): Promise<unknown> => {
+  return getArgsAndRunCommand(showInstallProgress)
+}
+
+export const autoUpgradeDvc = (): Promise<unknown> => {
+  return getArgsAndRunCommand(showUpgradeProgress)
 }

extension/src/setup/webview/messages.ts

@@ -9,7 +9,7 @@ import { BaseWebview } from '../../webview'
 import { sendTelemetryEvent } from '../../telemetry'
 import { EventName } from '../../telemetry/constants'
 import { selectPythonInterpreter } from '../../extensions/python'
-import { autoInstallDvc } from '../autoInstall'
+import { autoInstallDvc, autoUpgradeDvc } from '../autoInstall'
 import {
   RegisteredCliCommands,
   RegisteredCommands
@@ -79,6 +79,8 @@ export class WebviewMessages {
         return this.selectPythonInterpreter()
       case MessageFromWebviewType.INSTALL_DVC:
         return this.installDvc()
+      case MessageFromWebviewType.UPGRADE_DVC:
+        return this.upgradeDvc()
       case MessageFromWebviewType.SETUP_WORKSPACE:
         return commands.executeCommand(
           RegisteredCommands.EXTENSION_SETUP_WORKSPACE
@@ -131,6 +133,12 @@ export class WebviewMessages {
     return selectPythonInterpreter()
   }
 
+  private upgradeDvc() {
+    sendTelemetryEvent(EventName.VIEWS_SETUP_UPGRADE_DVC, undefined, undefined)
+
+    return autoUpgradeDvc()
+  }
+
   private installDvc() {
     sendTelemetryEvent(EventName.VIEWS_SETUP_INSTALL_DVC, undefined, undefined)
 

extension/src/telemetry/constants.ts

@@ -83,6 +83,7 @@ export const EventName = Object.assign(
     VIEWS_SETUP_SELECT_PYTHON_INTERPRETER:
       'views.setup.selectPythonInterpreter',
     VIEWS_SETUP_SHOW_SCM_FOR_COMMIT: 'views.setup.showScmForCommit',
+    VIEWS_SETUP_UPGRADE_DVC: 'view.setup.upgradeDvc',
 
     VIEWS_TERMINAL_CLOSED: 'views.terminal.closed',
     VIEWS_TERMINAL_CREATED: 'views.terminal.created',
@@ -269,6 +270,7 @@ export interface IEventNamePropertyMapping {
   [EventName.VIEWS_SETUP_SHOW_SCM_FOR_COMMIT]: undefined
   [EventName.VIEWS_SETUP_INIT_GIT]: undefined
   [EventName.VIEWS_SETUP_INSTALL_DVC]: undefined
+  [EventName.VIEWS_SETUP_UPGRADE_DVC]: undefined
 
   [EventName.SETUP_SHOW]: undefined
   [EventName.SETUP_SHOW_EXPERIMENTS]: undefined

extension/src/test/suite/setup/autoInstall.test.ts

@@ -5,7 +5,7 @@ import { window } from 'vscode'
 import { Disposable } from '../../../extension'
 import * as PythonExtension from '../../../extensions/python'
 import * as Python from '../../../python'
-import { autoInstallDvc } from '../../../setup/autoInstall'
+import { autoInstallDvc, autoUpgradeDvc } from '../../../setup/autoInstall'
 import * as WorkspaceFolders from '../../../vscode/workspaceFolders'
 import { bypassProgressCloseDelay } from '../util'
 import { Toast } from '../../../vscode/toast'
@@ -23,9 +23,98 @@ suite('Auto Install Test Suite', () => {
     disposable.dispose()
   })
 
-  describe('autoInstallDvc', () => {
-    const defaultPython = getDefaultPython()
+  const defaultPython = getDefaultPython()
+
+  describe('autoUpgradeDvc', () => {
+    it('should return early if no Python interpreter is found', async () => {
+      stub(PythonExtension, 'getPythonExecutionDetails').resolves(undefined)
+      stub(Python, 'findPythonBin').resolves(undefined)
+      const mockInstallPackages = stub(Python, 'installPackages').resolves(
+        undefined
+      )
+
+      const showProgressSpy = spy(window, 'withProgress')
+      const showErrorSpy = spy(window, 'showErrorMessage')
+
+      await autoUpgradeDvc()
+
+      expect(showProgressSpy).not.to.be.called
+      expect(showErrorSpy).to.be.called
+      expect(mockInstallPackages).not.to.be.called
+    })
+
+    it('should return early if there is no workspace folder open', async () => {
+      stub(PythonExtension, 'getPythonExecutionDetails').resolves(undefined)
+      stub(Python, 'findPythonBin').resolves(defaultPython)
+      const mockInstallPackages = stub(Python, 'installPackages').resolves(
+        undefined
+      )
+      stub(WorkspaceFolders, 'getFirstWorkspaceFolder').returns(undefined)
 
+      const showProgressSpy = spy(window, 'withProgress')
+      const showErrorSpy = spy(window, 'showErrorMessage')
+
+      await autoUpgradeDvc()
+
+      expect(showProgressSpy).not.to.be.called
+      expect(showErrorSpy).to.be.called
+      expect(mockInstallPackages).not.to.be.called
+    })
+
+    it('should install DVC if a Python interpreter is found', async () => {
+      bypassProgressCloseDelay()
+      const cwd = __dirname
+      stub(PythonExtension, 'getPythonExecutionDetails').resolves(undefined)
+      stub(Python, 'findPythonBin').resolves(defaultPython)
+      const mockInstallPackages = stub(Python, 'installPackages').resolves(
+        undefined
+      )
+      stub(WorkspaceFolders, 'getFirstWorkspaceFolder').returns(cwd)
+
+      const showProgressSpy = spy(window, 'withProgress')
+      const showErrorSpy = spy(window, 'showErrorMessage')
+
+      await autoUpgradeDvc()
+
+      expect(showProgressSpy).to.be.called
+      expect(showErrorSpy).not.to.be.called
+      expect(mockInstallPackages).to.be.called
+      expect(mockInstallPackages).to.be.calledWithExactly(
+        cwd,
+        defaultPython,
+        'dvc'
+      )
+    })
+
+    it('should show an error message if DVC fails to install', async () => {
+      bypassProgressCloseDelay()
+      const cwd = __dirname
+      stub(PythonExtension, 'getPythonExecutionDetails').resolves(undefined)
+      stub(Python, 'findPythonBin').resolves(defaultPython)
+      const mockInstallPackages = stub(Python, 'installPackages')
+        .onFirstCall()
+        .rejects(new Error('Network error'))
+      stub(WorkspaceFolders, 'getFirstWorkspaceFolder').returns(cwd)
+
+      const showProgressSpy = spy(window, 'withProgress')
+      const showErrorSpy = spy(window, 'showErrorMessage')
+      const reportProgressErrorSpy = spy(Toast, 'reportProgressError')
+
+      await autoUpgradeDvc()
+
+      expect(showProgressSpy).to.be.called
+      expect(showErrorSpy).not.to.be.called
+      expect(reportProgressErrorSpy).to.be.calledOnce
+      expect(mockInstallPackages).to.be.called
+      expect(mockInstallPackages).to.be.calledWithExactly(
+        cwd,
+        defaultPython,
+        'dvc'
+      )
+    })
+  })
+
+  describe('autoInstallDvc', () => {
     it('should return early if no Python interpreter is found', async () => {
       stub(PythonExtension, 'getPythonExecutionDetails').resolves(undefined)
       stub(Python, 'findPythonBin').resolves(undefined)

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

@@ -138,6 +138,22 @@ suite('Setup Test Suite', () => {
       expect(mockAutoInstallDvc).to.be.calledOnce
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
+    it('should handle an auto upgrade dvc message from the webview', async () => {
+      const { messageSpy, setup, mockAutoUpgradeDvc } = buildSetup(disposable)
+
+      const webview = await setup.showWebview()
+      await webview.isReady()
+
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+
+      messageSpy.resetHistory()
+      mockMessageReceived.fire({
+        type: MessageFromWebviewType.UPGRADE_DVC
+      })
+
+      expect(mockAutoUpgradeDvc).to.be.calledOnce
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
     it('should handle a select Python interpreter message from the webview', async () => {
       const { messageSpy, mockExecuteCommand, setup } = buildSetup(disposable)
       const setInterpreterCommand = 'python.setInterpreter'

extension/src/test/suite/setup/util.ts

@@ -62,6 +62,7 @@ export const buildSetup = (
   )
 
   const mockAutoInstallDvc = stub(AutoInstall, 'autoInstallDvc')
+  const mockAutoUpgradeDvc = stub(AutoInstall, 'autoUpgradeDvc')
   stub(AutoInstall, 'findPythonBinForInstall').resolves(undefined)
 
   const mockShowWebview = stub(WorkspaceExperiments.prototype, 'showWebview')
@@ -105,6 +106,7 @@ export const buildSetup = (
     internalCommands,
     messageSpy,
     mockAutoInstallDvc,
+    mockAutoUpgradeDvc,
     mockExecuteCommand,
     mockGetGitRepositoryRoot,
     mockGlobalVersion,

extension/src/webview/contract.ts

@@ -62,6 +62,7 @@ export enum MessageFromWebviewType {
   INITIALIZE_GIT = 'initialize-git',
   SHOW_SCM_PANEL = 'show-scm-panel',
   INSTALL_DVC = 'install-dvc',
+  UPGRADE_DVC = 'upgrade-dvc',
   SETUP_WORKSPACE = 'setup-workspace',
   ZOOM_PLOT = 'zoom-plot',
   SHOW_MORE_COMMITS = 'show-more-commits',
@@ -214,6 +215,7 @@ export type MessageFromWebview =
   | { type: MessageFromWebviewType.INITIALIZE_GIT }
   | { type: MessageFromWebviewType.SHOW_SCM_PANEL }
   | { type: MessageFromWebviewType.INSTALL_DVC }
+  | { type: MessageFromWebviewType.UPGRADE_DVC }
   | { type: MessageFromWebviewType.SETUP_WORKSPACE }
   | { type: MessageFromWebviewType.OPEN_STUDIO }
   | { type: MessageFromWebviewType.OPEN_STUDIO_PROFILE }

webview/src/setup/components/App.test.tsx

@@ -97,6 +97,9 @@ describe('App', () => {
       })
 
       expect(screen.getByText('DVC is incompatible')).toBeInTheDocument()
+      expect(
+        screen.getByText('Please update your install and try again.')
+      ).toBeInTheDocument()
 
       const button = screen.getByText('Check Compatibility')
       expect(button).toBeInTheDocument()
@@ -107,18 +110,27 @@ describe('App', () => {
       })
     })
 
-    it('should show a screen saying that DVC is not installed if the cli is unavailable', () => {
+    it('should tell the user than they can auto upgrade DVC if DVC is incompatible and python is available', () => {
       renderApp({
-        cliCompatible: undefined,
+        cliCompatible: false,
         dvcCliDetails: {
           command: 'dvc',
-          version: undefined
-        }
+          version: '1.0.0'
+        },
+        pythonBinPath: 'python'
       })
 
-      expect(screen.getAllByText('DVC is currently unavailable')).toHaveLength(
-        3
-      )
+      expect(screen.getByText('DVC is incompatible')).toBeInTheDocument()
+
+      const compatibityButton = screen.getByText('Check Compatibility')
+      expect(compatibityButton).toBeInTheDocument()
+      const upgradeButton = screen.getByText('Upgrade (pip)')
+      expect(upgradeButton).toBeInTheDocument()
+
+      fireEvent.click(upgradeButton)
+      expect(mockPostMessage).toHaveBeenCalledWith({
+        type: MessageFromWebviewType.UPGRADE_DVC
+      })
     })
 
     it('should tell the user they cannot install DVC without a Python interpreter', () => {

webview/src/setup/components/dvc/CliIncompatible.tsx

@@ -1,15 +1,38 @@
 import React, { PropsWithChildren } from 'react'
+import { useSelector } from 'react-redux'
+import styles from './styles.module.scss'
 import { EmptyState } from '../../../shared/components/emptyState/EmptyState'
 import { Button } from '../../../shared/components/button/Button'
-import { checkCompatibility } from '../messages'
+import { checkCompatibility, upgradeDvc } from '../messages'
+import { SetupState } from '../../store'
 
-export const CliIncompatible: React.FC<PropsWithChildren> = ({ children }) => (
-  <EmptyState isFullScreen={false}>
-    <div>
-      <h1>DVC is incompatible</h1>
-      {children}
+export const CliIncompatible: React.FC<PropsWithChildren> = ({ children }) => {
+  const pythonBinPath = useSelector(
+    (state: SetupState) => state.dvc.pythonBinPath
+  )
+  const canUpgrade = !!pythonBinPath
+
+  const conditionalContents = canUpgrade ? (
+    <>
+      <div className={styles.sideBySideButtons}>
+        <Button onClick={upgradeDvc} text="Upgrade (pip)" />
+        <Button text="Check Compatibility" onClick={checkCompatibility} />
+      </div>
+    </>
+  ) : (
+    <>
       <p>Please update your install and try again.</p>
       <Button text="Check Compatibility" onClick={checkCompatibility} />
-    </div>
-  </EmptyState>
-)
+    </>
+  )
+
+  return (
+    <EmptyState isFullScreen={false}>
+      <div>
+        <h1>DVC is incompatible</h1>
+        {children}
+        {conditionalContents}
+      </div>
+    </EmptyState>
+  )
+}

webview/src/setup/components/messages.ts

@@ -25,6 +25,10 @@ export const installDvc = () => {
   sendMessage({ type: MessageFromWebviewType.INSTALL_DVC })
 }
 
+export const upgradeDvc = () => {
+  sendMessage({ type: MessageFromWebviewType.UPGRADE_DVC })
+}
+
 export const selectPythonInterpreter = () => {
   sendMessage({ type: MessageFromWebviewType.SELECT_PYTHON_INTERPRETER })
 }