Add command to add remote (#3929)

* add add remote command

* add tests

* refactor

* apply review feedback

.eslintrc.js

@@ -3,6 +3,7 @@ module.exports = {
   env: {
     'jest/globals': true
   },
+
   extends: [
     'prettier-standard/prettier-file',
     'plugin:@typescript-eslint/eslint-recommended',
@@ -22,7 +23,8 @@ module.exports = {
     '**/*.js',
     '*.d.ts',
     'tsconfig.json',
-    'webpack.config.ts'
+    'webpack.config.ts',
+    'scripts/virtualenv-install.ts'
   ],
   overrides: [
     {

extension/package.json

@@ -100,6 +100,12 @@
         "category": "DVC",
         "icon": "$(add)"
       },
+      {
+        "title": "Add Remote",
+        "command": "dvc.addRemote",
+        "category": "DVC",
+        "icon": "$(add)"
+      },
       {
         "title": "Filter Experiments Table to Starred",
         "command": "dvc.addStarredExperimentsTableFilter",
@@ -654,6 +660,10 @@
           "command": "dvc.addExperimentsTableSort",
           "when": "dvc.commands.available && dvc.project.available"
         },
+        {
+          "command": "dvc.addRemote",
+          "when": "dvc.commands.available && dvc.project.available && !dvc.experiment.running.workspace"
+        },
         {
           "command": "dvc.addStarredExperimentsTableFilter",
           "when": "dvc.commands.available && dvc.project.available"

extension/src/cli/dvc/constants.ts

@@ -38,6 +38,7 @@ export enum Command {
 }
 
 export enum SubCommand {
+  ADD = 'add',
   DIFF = 'diff',
   LIST = 'list',
   STATUS = 'status',
@@ -47,13 +48,15 @@ export enum SubCommand {
 export enum Flag {
   ALL_COMMITS = '-A',
   FOLLOW = '-f',
+  DEFAULT = '-d',
   FORCE = '-f',
   GLOBAL = '--global',
   GRANULAR = '--granular',
-  LOCAL = '--local',
   JOBS = '-j',
   JSON = '--json',
   KILL = '--kill',
+  LOCAL = '--local',
+  PROJECT = '--project',
   NUM_COMMIT = '-n',
   OUTPUT_PATH = '-o',
   SUBDIRECTORY = '--subdir',

extension/src/cli/dvc/executor.ts

@@ -7,8 +7,7 @@ import {
   ExperimentSubCommand,
   Flag,
   GcPreserveFlag,
-  QueueSubCommand,
-  SubCommand
+  QueueSubCommand
 } from './constants'
 import { addStudioAccessToken } from './options'
 import { CliResult, CliStarted, typeCheckCommands } from '..'
@@ -198,8 +197,8 @@ export class DvcExecutor extends DvcCli {
     return this.blockAndExecuteProcess(cwd, Command.REMOVE, ...args)
   }
 
-  public remote(cwd: string, arg: typeof SubCommand.LIST) {
-    return this.executeDvcProcess(cwd, Command.REMOTE, arg)
+  public remote(cwd: string, ...args: Args) {
+    return this.executeDvcProcess(cwd, Command.REMOTE, ...args)
   }
 
   private executeExperimentProcess(cwd: string, ...args: Args) {

extension/src/commands/external.ts

@@ -41,6 +41,8 @@ export enum RegisteredCliCommands {
   REMOVE_TARGET = 'dvc.removeTarget',
   RENAME_TARGET = 'dvc.renameTarget',
 
+  REMOTE_ADD = 'dvc.addRemote',
+
   GIT_STAGE_ALL = 'dvc.gitStageAll',
   GIT_UNSTAGE_ALL = 'dvc.gitUnstageAll'
 }

extension/src/extension.ts

@@ -47,7 +47,7 @@ import { Flag } from './cli/dvc/constants'
 import { LanguageClient } from './languageClient'
 import { collectRunningExperimentPids } from './experiments/processExecution/collect'
 import { DvcViewer } from './cli/dvc/viewer'
-import { registerSetupCommands } from './setup/register'
+import { registerSetupCommands } from './setup/commands/register'
 import { Status } from './status'
 import { registerPersistenceCommands } from './persistence/register'
 

extension/src/setup/commands/index.ts

@@ -0,0 +1,86 @@
+import { Setup } from '..'
+import { Flag, SubCommand } from '../../cli/dvc/constants'
+import { AvailableCommands, InternalCommands } from '../../commands/internal'
+import { definedAndNonEmpty } from '../../util/array'
+import { getInput } from '../../vscode/inputBox'
+import { quickPickYesOrNo } from '../../vscode/quickPick'
+import { Title } from '../../vscode/title'
+import { Toast } from '../../vscode/toast'
+import { getOnlyOrPickProject } from '../../workspace/util'
+
+const noExistingOrUserConfirms = async (
+  internalCommands: InternalCommands,
+  dvcRoot: string
+): Promise<boolean | undefined> => {
+  const remoteList = await internalCommands.executeCommand(
+    AvailableCommands.REMOTE,
+    dvcRoot,
+    SubCommand.LIST
+  )
+
+  if (!remoteList) {
+    return true
+  }
+
+  return await quickPickYesOrNo(
+    'make this new remote the default',
+    'keep the current default',
+    {
+      placeHolder: 'Would you like to set this new remote as the default?',
+      title: Title.SET_REMOTE_AS_DEFAULT
+    }
+  )
+}
+
+const addRemoteToProject = async (
+  internalCommands: InternalCommands,
+  dvcRoot: string
+): Promise<void> => {
+  const name = await getInput(Title.ENTER_REMOTE_NAME)
+  if (!name) {
+    return
+  }
+
+  const url = await getInput(Title.ENTER_REMOTE_URL)
+  if (!url) {
+    return
+  }
+
+  const args = [Flag.PROJECT, name, url]
+
+  const shouldSetAsDefault = await noExistingOrUserConfirms(
+    internalCommands,
+    dvcRoot
+  )
+  if (shouldSetAsDefault === undefined) {
+    return
+  }
+
+  if (shouldSetAsDefault) {
+    args.unshift(Flag.DEFAULT)
+  }
+
+  return await Toast.showOutput(
+    internalCommands.executeCommand(
+      AvailableCommands.REMOTE,
+      dvcRoot,
+      SubCommand.ADD,
+      ...args
+    )
+  )
+}
+
+export const getAddRemoteCommand =
+  (setup: Setup, internalCommands: InternalCommands) =>
+  async (): Promise<void> => {
+    const dvcRoots = setup.getRoots()
+    if (!definedAndNonEmpty(dvcRoots)) {
+      return Toast.showError('Cannot add a remote without a DVC project')
+    }
+    const dvcRoot = await getOnlyOrPickProject(dvcRoots)
+
+    if (!dvcRoot) {
+      return
+    }
+    return addRemoteToProject(internalCommands, dvcRoot)
+  }

extension/src/setup/register.ts → extension/src/setup/commands/register.ts

@@ -1,10 +1,14 @@
 import { commands } from 'vscode'
-import { Setup } from '.'
-import { run } from './runner'
-import { SetupSection } from './webview/contract'
-import { AvailableCommands, InternalCommands } from '../commands/internal'
-import { RegisteredCliCommands, RegisteredCommands } from '../commands/external'
-import { getFirstWorkspaceFolder } from '../vscode/workspaceFolders'
+import { getAddRemoteCommand } from '.'
+import { Setup } from '..'
+import { run } from '../runner'
+import { SetupSection } from '../webview/contract'
+import { AvailableCommands, InternalCommands } from '../../commands/internal'
+import {
+  RegisteredCliCommands,
+  RegisteredCommands
+} from '../../commands/external'
+import { getFirstWorkspaceFolder } from '../../vscode/workspaceFolders'
 
 const registerSetupConfigCommands = (
   setup: Setup,
@@ -100,6 +104,11 @@ export const registerSetupCommands = (
     }
   )
 
+  internalCommands.registerExternalCliCommand(
+    RegisteredCliCommands.REMOTE_ADD,
+    getAddRemoteCommand(setup, internalCommands)
+  )
+
   registerSetupConfigCommands(setup, internalCommands)
   registerSetupShowCommands(setup, internalCommands)
   registerSetupStudioCommands(setup, internalCommands)

extension/src/setup/webview/messages.ts

@@ -104,6 +104,8 @@ export class WebviewMessages {
         )
       case MessageFromWebviewType.OPEN_EXPERIMENTS_WEBVIEW:
         return commands.executeCommand(RegisteredCommands.EXPERIMENT_SHOW)
+      case MessageFromWebviewType.REMOTE_ADD:
+        return commands.executeCommand(RegisteredCliCommands.REMOTE_ADD)
 
       default:
         Logger.error(`Unexpected message: ${JSON.stringify(message)}`)

extension/src/telemetry/constants.ts

@@ -187,6 +187,8 @@ export interface IEventNamePropertyMapping {
   [EventName.REMOVE_TARGET]: undefined
   [EventName.RENAME_TARGET]: undefined
 
+  [EventName.REMOTE_ADD]: undefined
+
   [EventName.GIT_STAGE_ALL]: undefined
   [EventName.GIT_UNSTAGE_ALL]: undefined
 

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

@@ -69,6 +69,7 @@ suite('Setup Test Suite', () => {
     ])
   })
 
+  // eslint-disable-next-line sonarjs/cognitive-complexity
   describe('Setup', () => {
     it('should handle an initialize git message from the webview', async () => {
       const { messageSpy, setup, mockInitializeGit } = buildSetup(disposable)
@@ -843,6 +844,103 @@ suite('Setup Test Suite', () => {
       expect(mockShowWebview).to.be.calledOnce
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
+    it('should handle a message to add a remote', async () => {
+      const { messageSpy, setup, mockExecuteCommand } = buildSetup(disposable)
+
+      const webview = await setup.showWebview()
+      await webview.isReady()
+      mockExecuteCommand.restore()
+
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+
+      const mockRemote = stub(DvcExecutor.prototype, 'remote')
+
+      const remoteAdded = new Promise(resolve =>
+        mockRemote.callsFake((cwd, ...args) => {
+          if (args.includes('add')) {
+            resolve(undefined)
+          }
+          return Promise.resolve('')
+        })
+      )
+
+      const mockShowInputBox = stub(window, 'showInputBox')
+        .onFirstCall()
+        .resolves('storage')
+        .onSecondCall()
+        .resolves('s3://my-bucket')
+
+      messageSpy.resetHistory()
+      mockMessageReceived.fire({
+        type: MessageFromWebviewType.REMOTE_ADD
+      })
+
+      await remoteAdded
+
+      expect(mockShowInputBox).to.be.calledTwice
+      expect(
+        mockRemote,
+        'new remote is set as the default'
+      ).to.be.calledWithExactly(
+        dvcDemoPath,
+        'add',
+        '-d',
+        '--project',
+        'storage',
+        's3://my-bucket'
+      )
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
+    it('should be able to add a remote', async () => {
+      const mockRemote = stub(DvcExecutor.prototype, 'remote')
+
+      const remoteAdded = new Promise(resolve =>
+        mockRemote.callsFake((cwd, ...args) => {
+          if (args.includes('list')) {
+            return Promise.resolve('storage s3://my-bucket')
+          }
+
+          if (args.includes('add')) {
+            resolve(undefined)
+          }
+          return Promise.resolve('')
+        })
+      )
+
+      const mockShowInputBox = stub(window, 'showInputBox')
+        .onFirstCall()
+        .resolves('backup')
+        .onSecondCall()
+        .resolves('s3://my-backup-bucket')
+
+      const mockShowQuickPick = (
+        stub(window, 'showQuickPick') as SinonStub<
+          [items: readonly QuickPickItem[], options: QuickPickOptionsWithTitle],
+          Thenable<QuickPickItemWithValue<boolean> | undefined>
+        >
+      ).resolves({
+        label: 'No',
+        value: false
+      })
+
+      await commands.executeCommand(RegisteredCliCommands.REMOTE_ADD)
+
+      await remoteAdded
+
+      expect(mockShowInputBox).to.be.calledTwice
+      expect(mockShowQuickPick).to.be.calledOnce
+      expect(
+        mockRemote,
+        'should not set a remote as the default unless the user explicitly chooses to'
+      ).to.be.calledWithExactly(
+        dvcDemoPath,
+        'add',
+        '--project',
+        'backup',
+        's3://my-backup-bucket'
+      )
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
     it('should send the appropriate messages to the webview to focus different sections', async () => {
       const { setup, messageSpy } = buildSetup(disposable)
       messageSpy.restore()

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

@@ -44,7 +44,7 @@ export const buildSetup = (
 
   const mockEmitter = disposer.track(new EventEmitter())
   stub(dvcReader, 'root').resolves(mockDvcRoot)
-  stub(dvcExecutor, 'remote').resolves('')
+  const mockRemote = stub(dvcExecutor, 'remote').resolves('')
   const mockVersion = stub(dvcReader, 'version').resolves(MIN_CLI_VERSION)
   const mockGlobalVersion = stub(dvcReader, 'globalVersion').resolves(
     MIN_CLI_VERSION
@@ -112,6 +112,7 @@ export const buildSetup = (
     mockGlobalVersion,
     mockInitializeGit,
     mockOpenExternal,
+    mockRemote,
     mockRunSetup,
     mockShowWebview,
     mockVersion,

extension/src/vscode/title.ts

@@ -6,6 +6,8 @@ export enum Title {
   ENTER_EXPERIMENT_WORKER_COUNT = 'Enter the Number of Queue Workers',
   ENTER_FILTER_VALUE = 'Enter a Filter Value',
   ENTER_RELATIVE_DESTINATION = 'Enter a Destination Relative to the Root',
+  ENTER_REMOTE_NAME = 'Enter a Name for the Remote',
+  ENTER_REMOTE_URL = 'Enter the URL for the Remote',
   ENTER_PATH_OR_CHOOSE_FILE = 'Enter the path to your training script or select it',
   ENTER_STUDIO_USERNAME = 'Enter your Studio username',
   ENTER_STUDIO_TOKEN = 'Enter your Studio access token',
@@ -36,7 +38,8 @@ export enum Title {
   SELECT_SORTS_TO_REMOVE = 'Select Sort(s) to Remove',
   SELECT_TRAINING_SCRIPT = 'Select your training script',
   SETUP_WORKSPACE = 'Setup the Workspace',
-  SET_EXPERIMENTS_HEADER_HEIGHT = 'Set Maximum Experiment Table Header Height'
+  SET_EXPERIMENTS_HEADER_HEIGHT = 'Set Maximum Experiment Table Header Height',
+  SET_REMOTE_AS_DEFAULT = 'Set Default Remote'
 }
 
 export const getEnterValueTitle = (path: string): Title =>