Merge branch 'main' into add-show-options

extension/package.json

@@ -1634,7 +1634,7 @@
     "@types/sinon-chai": "3.2.9",
     "@types/uuid": "9.0.1",
     "@types/vscode": "1.64.0",
-    "@vscode/test-electron": "2.2.3",
+    "@vscode/test-electron": "2.3.0",
     "@vscode/vsce": "2.18.0",
     "@wdio/cli": "8.4.0",
     "@wdio/local-runner": "8.4.0",

extension/src/cli/dvc/executor.test.ts

@@ -5,7 +5,7 @@ import { Flag, GcPreserveFlag } from './constants'
 import { DvcExecutor } from './executor'
 import { CliResult, CliStarted } from '..'
 import { createProcess } from '../../process/execution'
-import { getMockedProcess } from '../../test/util/jest'
+import { flushPromises, getMockedProcess } from '../../test/util/jest'
 import { getProcessEnv } from '../../env'
 import { Config } from '../../config'
 import { ContextKey, setContextValue } from '../../vscode/context'
@@ -28,6 +28,9 @@ const mockedEnv = {
 
 const mockedSetContextValue = jest.mocked(setContextValue)
 
+const mockedGetStudioLiveShareToken = jest.fn()
+const mockedGetRepoUrl = jest.fn()
+
 beforeEach(() => {
   jest.resetAllMocks()
   mockedGetProcessEnv.mockReturnValueOnce(mockedEnv)
@@ -48,6 +51,8 @@ describe('CliExecutor', () => {
       getCliPath: () => undefined,
       getPythonBinPath: () => undefined
     } as unknown as Config,
+    mockedGetStudioLiveShareToken,
+    mockedGetRepoUrl,
     {
       processCompleted: {
         event: jest.fn(),
@@ -602,6 +607,39 @@ describe('CliExecutor', () => {
         executable: 'dvc'
       })
     })
+
+    it("should call createProcess with the correct parameters to start the experiment's queue and send live updates to Studio", async () => {
+      const cwd = __dirname
+      const jobs = '91231324'
+      const mockedToken = 'isat_notarealtoken'
+      const mockedUrl = '[email protected]:iterative/vscode-dvc-demo.git'
+
+      mockedGetStudioLiveShareToken.mockReturnValueOnce(mockedToken)
+      mockedGetRepoUrl.mockResolvedValueOnce(
+        '[email protected]:iterative/vscode-dvc-demo.git'
+      )
+
+      const stdout = `Started '${jobs}' new experiments task queue workers.`
+
+      mockedCreateProcess.mockReturnValueOnce(getMockedProcess(stdout))
+
+      void dvcExecutor.queueStart(cwd, jobs)
+      await flushPromises()
+
+      expect(mockedGetRepoUrl).toHaveBeenCalledWith(cwd)
+
+      expect(mockedCreateProcess).toHaveBeenCalledWith({
+        args: ['queue', 'start', '-j', jobs],
+        cwd,
+        detached: true,
+        env: {
+          ...mockedEnv,
+          STUDIO_REPO_URL: mockedUrl,
+          STUDIO_TOKEN: mockedToken
+        },
+        executable: 'dvc'
+      })
+    })
   })
 
   describe('queueStop', () => {

extension/src/cli/dvc/executor.ts

@@ -1,3 +1,4 @@
+import { EventEmitter } from 'vscode'
 import { DvcCli } from '.'
 import {
   Args,
@@ -8,8 +9,10 @@ import {
   GcPreserveFlag,
   QueueSubCommand
 } from './constants'
-import { typeCheckCommands } from '..'
+import { addStudioAccessToken } from './options'
+import { CliResult, CliStarted, typeCheckCommands } from '..'
 import { ContextKey, setContextValue } from '../../vscode/context'
+import { Config } from '../../config'
 
 export const autoRegisteredCommands = {
   ADD: 'add',
@@ -38,8 +41,24 @@ export class DvcExecutor extends DvcCli {
     this
   )
 
+  private readonly getStudioLiveShareToken: () => string | undefined
+  private readonly getRepoUrl: (cwd: string) => Promise<string>
   private scmCommandRunning = false
 
+  constructor(
+    config: Config,
+    getStudioLiveShareToken: () => string | undefined,
+    getRepoUrl: (cwd: string) => Promise<string>,
+    emitters?: {
+      processStarted: EventEmitter<CliStarted>
+      processCompleted: EventEmitter<CliResult>
+    }
+  ) {
+    super(config, emitters)
+    this.getStudioLiveShareToken = getStudioLiveShareToken
+    this.getRepoUrl = getRepoUrl
+  }
+
   public add(cwd: string, target: string) {
     return this.blockAndExecuteProcess(cwd, Command.ADD, target)
   }
@@ -139,14 +158,20 @@ export class DvcExecutor extends DvcCli {
     )
   }
 
-  public queueStart(cwd: string, jobs: string) {
-    return this.createBackgroundDvcProcess(
+  public async queueStart(cwd: string, jobs: string) {
+    const options = this.getOptions(
       cwd,
       Command.QUEUE,
       QueueSubCommand.START,
       Flag.JOBS,
       jobs
     )
+    const studioAccessToken = this.getStudioLiveShareToken()
+    const repoUrl = studioAccessToken ? await this.getRepoUrl(cwd) : undefined
+
+    return this.createBackgroundProcess(
+      addStudioAccessToken(options, studioAccessToken, repoUrl)
+    )
   }
 
   public queueStop(cwd: string, ...args: Args) {

extension/src/cli/dvc/index.ts

@@ -34,7 +34,7 @@ export class DvcCli extends Cli {
     return this.createBackgroundProcess(options)
   }
 
-  private getOptions(cwd: string, ...args: Args) {
+  protected getOptions(cwd: string, ...args: Args) {
     return getOptions(
       this.config.getPythonBinPath(),
       this.config.getCliPath(),

extension/src/cli/dvc/options.ts

@@ -49,3 +49,29 @@ export const getOptions = (
     executable
   }
 }
+
+export const addStudioAccessToken = (
+  options: ExecutionOptions,
+  studioAccessToken: string | undefined,
+  repoUrl?: string
+): ExecutionOptions => {
+  if (!studioAccessToken) {
+    return options
+  }
+
+  if (!repoUrl) {
+    return {
+      ...options,
+      env: { ...options.env, STUDIO_TOKEN: studioAccessToken }
+    }
+  }
+
+  return {
+    ...options,
+    env: {
+      ...options.env,
+      STUDIO_REPO_URL: repoUrl,
+      STUDIO_TOKEN: studioAccessToken
+    }
+  }
+}

extension/src/cli/dvc/runner.ts

@@ -5,7 +5,7 @@ import {
   ExperimentFlag,
   ExperimentSubCommand
 } from './constants'
-import { getOptions } from './options'
+import { addStudioAccessToken, getOptions } from './options'
 import { CliResult, CliStarted, ICli, typeCheckCommands } from '..'
 import { getCommandString } from '../command'
 import { Config } from '../../config'
@@ -183,15 +183,7 @@ export class DvcRunner extends Disposable implements ICli {
     )
 
     const studioAccessToken = this.getStudioLiveShareToken()
-
-    if (!studioAccessToken) {
-      return options
-    }
-
-    return {
-      ...options,
-      env: { ...options.env, STUDIO_TOKEN: studioAccessToken }
-    }
+    return addStudioAccessToken(options, studioAccessToken)
   }
 
   private startProcess(cwd: string, args: Args) {

extension/src/extension.ts

@@ -94,16 +94,23 @@ export class Extension extends Disposable {
       new Connect(context, this.resourceLocator.dvcIcon)
     )
 
-    this.dvcExecutor = this.dispose.track(new DvcExecutor(config))
+    this.gitExecutor = this.dispose.track(new GitExecutor())
+    this.gitReader = this.dispose.track(new GitReader())
+
+    const getStudioLiveShareToken = () => this.connect.getStudioLiveShareToken()
+
+    this.dvcExecutor = this.dispose.track(
+      new DvcExecutor(config, getStudioLiveShareToken, cwd =>
+        this.gitReader.getRemoteUrl(cwd)
+      )
+    )
+
     this.dvcReader = this.dispose.track(new DvcReader(config))
     this.dvcRunner = this.dispose.track(
-      new DvcRunner(config, () => this.connect.getStudioLiveShareToken())
+      new DvcRunner(config, getStudioLiveShareToken)
     )
     this.dvcViewer = this.dispose.track(new DvcViewer(config))
 
-    this.gitExecutor = this.dispose.track(new GitExecutor())
-    this.gitReader = this.dispose.track(new GitReader())
-
     const clis = [
       this.dvcExecutor,
       this.dvcReader,

extension/src/test/cli/util.ts

@@ -15,7 +15,11 @@ const config = {
 } as Config
 
 export const dvcReader = new DvcReader(config)
-export const dvcExecutor = new DvcExecutor(config)
+export const dvcExecutor = new DvcExecutor(
+  config,
+  () => undefined,
+  () => Promise.resolve('')
+)
 const gitExecutor = new GitExecutor()
 
 let demoInitialized: Promise<string>

extension/src/test/suite/util.ts

@@ -142,7 +142,13 @@ export const buildInternalCommands = (disposer: Disposer) => {
   const config = disposer.track(new Config())
   const dvcReader = disposer.track(new DvcReader(config))
   const dvcRunner = disposer.track(new DvcRunner(config, () => undefined))
-  const dvcExecutor = disposer.track(new DvcExecutor(config))
+  const dvcExecutor = disposer.track(
+    new DvcExecutor(
+      config,
+      () => undefined,
+      () => Promise.resolve('')
+    )
+  )
   const dvcViewer = disposer.track(new DvcViewer(config))
   const gitReader = disposer.track(new GitReader())
 

package.json

@@ -76,9 +76,7 @@
     "loader-utils": "2.0.4",
     "terser": "5.16.5",
     "trim-newlines": "3.0.1",
-    "trim": "1.0.1",
-    "vega-functions": "5.13.0",
-    "vega-scale": "7.2.0"
+    "trim": "1.0.1"
   },
   "packageManager": "[email protected]"
 }

yarn.lock

@@ -4793,15 +4793,15 @@
     "@microsoft/applicationinsights-web-basic" "^2.8.9"
     applicationinsights "2.4.1"
 
-"@vscode/test-electron@2.2.3":
-  version "2.2.3"
-  resolved "https://registry.yarnpkg.com/@vscode/test-electron/-/test-electron-2.2.3.tgz#bf6a77542970b5d34561d0671259900632e5eb94"
-  integrity sha512-7DmdGYQTqRNaLHKG3j56buc9DkstriY4aV0S3Zj32u0U9/T0L8vwWAC9QGCh1meu1VXDEla1ze27TkqysHGP0Q==
+"@vscode/test-electron@2.3.0":
+  version "2.3.0"
+  resolved "https://registry.yarnpkg.com/@vscode/test-electron/-/test-electron-2.3.0.tgz#de0ba2f5d36546a83cd481b458cbdbb7cc0f7049"
+  integrity sha512-fwzA9RtazH1GT/sckYlbxu6t5e4VaMXwCVtyLv4UAG0hP6NTfnMaaG25XCfWqlVwFhBMcQXHBCy5dmz2eLUnkw==
   dependencies:
     http-proxy-agent "^4.0.1"
     https-proxy-agent "^5.0.0"
-    rimraf "^3.0.2"
-    unzipper "^0.10.11"
+    jszip "^3.10.1"
+    semver "^7.3.8"
 
 "@vscode/test-electron@^2.2.1":
   version "2.2.1"
@@ -11467,6 +11467,11 @@ ignore@^5.1.1, ignore@^5.1.9, ignore@^5.2.0:
   resolved "https://registry.yarnpkg.com/ignore/-/ignore-5.2.0.tgz#6d3bac8fa7fe0d45d9f9be7bac2fc279577e345a"
   integrity sha512-CmxgYGiEPCLhfLnpPp1MoRmifwEIOgjcHXxOBjv7mY96c+eWScsOP9c112ZyLdWHi0FxHjI+4uVhKYp/gcdRmQ==
 
+immediate@~3.0.5:
+  version "3.0.6"
+  resolved "https://registry.yarnpkg.com/immediate/-/immediate-3.0.6.tgz#9db1dbd0faf8de6fbe0f5dd5e56bb606280de69b"
+  integrity sha512-XXOFtyqDjNDAQxVfYxuF7g9Il/IbWmmlQg2MYKOH8ExIT1qg6xc4zyS3HaEEATgs1btfzxq15ciUiY7gjSXRGQ==
+
 immer@^9.0.7:
   version "9.0.14"
   resolved "https://registry.yarnpkg.com/immer/-/immer-9.0.14.tgz#e05b83b63999d26382bb71676c9d827831248a48"
@@ -13066,6 +13071,16 @@ jsx-ast-utils@^3.3.3:
     array-includes "^3.1.5"
     object.assign "^4.1.3"
 
+jszip@^3.10.1:
+  version "3.10.1"
+  resolved "https://registry.yarnpkg.com/jszip/-/jszip-3.10.1.tgz#34aee70eb18ea1faec2f589208a157d1feb091c2"
+  integrity sha512-xXDvecyTpGLrqFrvkrUSoxxfJI5AH7U8zxxtVclpsUtMCq4JQ290LY8AW5c7Ggnr/Y/oK+bQMbqK2qmtk3pN4g==
+  dependencies:
+    lie "~3.3.0"
+    pako "~1.0.2"
+    readable-stream "~2.3.6"
+    setimmediate "^1.0.5"
+
 junk@^3.1.0:
   version "3.1.0"
   resolved "https://registry.yarnpkg.com/junk/-/junk-3.1.0.tgz#31499098d902b7e98c5d9b9c80f43457a88abfa1"
@@ -13200,6 +13215,13 @@ levn@~0.3.0:
     prelude-ls "~1.1.2"
     type-check "~0.3.2"
 
+lie@~3.3.0:
+  version "3.3.0"
+  resolved "https://registry.yarnpkg.com/lie/-/lie-3.3.0.tgz#dcf82dee545f46074daf200c7c1c5a08e0f40f6a"
+  integrity sha512-UaiMJzeWRlEujzAuw5LokY1L5ecNQYZKfmyZ9L7wDHb/p5etKaxXhohBcrw0EYby+G/NA52vRSN4N39dxHAIwQ==
+  dependencies:
+    immediate "~3.0.5"
+
 light-my-request@^4.2.0:
   version "4.12.0"
   resolved "https://registry.yarnpkg.com/light-my-request/-/light-my-request-4.12.0.tgz#fd59329a7b4f794842103c7bef69e12252478831"
@@ -14992,7 +15014,7 @@ package-hash@^4.0.0:
     lodash.flattendeep "^4.4.0"
     release-zalgo "^1.0.0"
 
-pako@~1.0.5:
+pako@~1.0.2, pako@~1.0.5:
   version "1.0.11"
   resolved "https://registry.yarnpkg.com/pako/-/pako-1.0.11.tgz#6c9599d340d54dfd3946380252a35705a6b992bf"
   integrity sha512-4hLB8Py4zZce5s4yd9XzopqwVv/yGNhV1Bl8NTmCq1763HeK2+EwVTv+leGeL13Dnh2wfbqowVPXCIO0z4taYw==
@@ -17042,7 +17064,7 @@ set-value@^2.0.0, set-value@^2.0.1:
     is-plain-object "^2.0.3"
     split-string "^3.0.1"
 
-setimmediate@^1.0.4, setimmediate@~1.0.4:
+setimmediate@^1.0.4, setimmediate@^1.0.5, setimmediate@~1.0.4:
   version "1.0.5"
   resolved "https://registry.yarnpkg.com/setimmediate/-/setimmediate-1.0.5.tgz#290cbb232e306942d7d7ea9b83732ab7856f8285"
   integrity sha1-KQy7Iy4waULX1+qbg3Mqt4VvgoU=
@@ -18972,7 +18994,7 @@ vega-format@^1.1.0, vega-format@~1.1.0:
     vega-time "^2.0.3"
     vega-util "^1.15.2"
 
-vega-functions@5.13.0, vega-functions@^5.12.1, vega-functions@^5.13.0, vega-functions@~5.13.0:
+vega-functions@^5.12.1, vega-functions@^5.13.0, vega-functions@~5.13.0:
   version "5.13.0"
   resolved "https://registry.yarnpkg.com/vega-functions/-/vega-functions-5.13.0.tgz#c9ab8c6eedbf39f75b424cca6776b1d0b8c74b32"
   integrity sha512-Mf53zNyx+c9fFqagEI0T8zc9nMlx0zozOngr8oOpG1tZDKOgwOnUgN99zQKbLHjyv+UzWrq3LYTnSLyVe0ZmhQ==
@@ -19110,7 +19132,7 @@ vega-runtime@^6.1.3, vega-runtime@~6.1.3:
     vega-dataflow "^5.7.3"
     vega-util "^1.15.2"
 
-vega-scale@7.2.0, vega-scale@^7.0.3, vega-scale@^7.1.1, vega-scale@^7.2.0, vega-scale@~7.2.0:
+vega-scale@^7.0.3, vega-scale@^7.1.1, vega-scale@^7.2.0, vega-scale@~7.2.0:
   version "7.2.0"
   resolved "https://registry.yarnpkg.com/vega-scale/-/vega-scale-7.2.0.tgz#9e298cc02ad340498cb56847436b19439911f0fc"
   integrity sha512-QYltO/otrZHLrCGGf06Y99XtPtqWXITr6rw7rO9oL+l3d9o5RFl9sjHrVxiM7v+vGoZVWbBd5IPbFhPsXZ6+TA==