feat: use new tcloud sqlmesh_lsp call (#4762)

Author: benfdking

vscode/extension/src/utilities/semver.test.ts

@@ -0,0 +1,37 @@
+import { describe, it, expect } from 'vitest'
+import { isSemVerGreaterThanOrEqual } from './semver'
+
+describe('isSemVerGreaterThanOrEqual', () => {
+  it('should return true when major version is greater', () => {
+    expect(isSemVerGreaterThanOrEqual([2, 0, 0], [1, 0, 0])).toBe(true)
+    expect(isSemVerGreaterThanOrEqual([3, 0, 0], [2, 5, 10])).toBe(true)
+  })
+
+  it('should return false when major version is less', () => {
+    expect(isSemVerGreaterThanOrEqual([1, 0, 0], [2, 0, 0])).toBe(false)
+    expect(isSemVerGreaterThanOrEqual([0, 10, 20], [1, 0, 0])).toBe(false)
+  })
+
+  it('should compare minor version when major versions are equal', () => {
+    expect(isSemVerGreaterThanOrEqual([1, 2, 0], [1, 1, 0])).toBe(true)
+    expect(isSemVerGreaterThanOrEqual([1, 1, 0], [1, 2, 0])).toBe(false)
+    expect(isSemVerGreaterThanOrEqual([2, 5, 0], [2, 3, 10])).toBe(true)
+  })
+
+  it('should compare patch version when major and minor versions are equal', () => {
+    expect(isSemVerGreaterThanOrEqual([1, 1, 2], [1, 1, 1])).toBe(true)
+    expect(isSemVerGreaterThanOrEqual([1, 1, 1], [1, 1, 2])).toBe(false)
+    expect(isSemVerGreaterThanOrEqual([2, 3, 10], [2, 3, 5])).toBe(true)
+  })
+
+  it('should return true when versions are equal', () => {
+    expect(isSemVerGreaterThanOrEqual([1, 0, 0], [1, 0, 0])).toBe(true)
+    expect(isSemVerGreaterThanOrEqual([2, 5, 10], [2, 5, 10])).toBe(true)
+  })
+
+  it('should handle zero versions correctly', () => {
+    expect(isSemVerGreaterThanOrEqual([0, 0, 1], [0, 0, 0])).toBe(true)
+    expect(isSemVerGreaterThanOrEqual([0, 1, 0], [0, 0, 10])).toBe(true)
+    expect(isSemVerGreaterThanOrEqual([0, 0, 0], [0, 0, 0])).toBe(true)
+  })
+})

vscode/extension/src/utilities/semver.ts

@@ -0,0 +1,24 @@
+type SemVer = [number, number, number]
+
+/**
+ * Check if a is greater than or equal to b.
+ *
+ * @param a - The first version.
+ * @param b - The second version.
+ * @returns True if a is greater than b, false otherwise.
+ */
+export function isSemVerGreaterThanOrEqual(a: SemVer, b: SemVer): boolean {
+  if (a[0] > b[0]) {
+    return true
+  }
+  if (a[0] < b[0]) {
+    return false
+  }
+  if (a[1] > b[1]) {
+    return true
+  }
+  if (a[1] < b[1]) {
+    return false
+  }
+  return a[2] >= b[2]
+}

vscode/extension/src/utilities/sqlmesh/sqlmesh.ts

@@ -12,6 +12,7 @@ import z from 'zod'
 import { ProgressLocation, window } from 'vscode'
 import { IS_WINDOWS } from '../isWindows'
 import { resolveProjectPath } from '../config'
+import { isSemVerGreaterThanOrEqual } from '../semver'
 
 export interface SqlmeshExecInfo {
   workspacePath: string
@@ -394,6 +395,23 @@ export const sqlmeshLspExec = async (): Promise<
       if (isErr(ensured)) {
         return ensured
       }
+      const tcloudBinVersion = await getTcloudBinVersion()
+      if (isErr(tcloudBinVersion)) {
+        return tcloudBinVersion
+      }
+      // TODO: Remove this once we have a stable version of tcloud that supports sqlmesh_lsp.
+      if (isSemVerGreaterThanOrEqual(tcloudBinVersion.value, [2, 10, 1])) {
+        return ok ({
+          bin: tcloudBin.value,
+          workspacePath,
+          env: {
+            PYTHONPATH: interpreterDetails.path?.[0],
+            VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
+            PATH: interpreterDetails.binPath!,
+          },
+          args: ['sqlmesh_lsp'],
+        })
+      }
     }
     const binPath = path.join(interpreterDetails.binPath!, sqlmeshLSP)
     traceLog(`Bin path: ${binPath}`)
@@ -445,4 +463,31 @@ async function doesExecutableExist(executable: string): Promise<boolean> {
     traceLog(`Checked if ${executable} exists with ${command}, errored, returning false`)
     return false
   }
+}
+
+/**
+ * Get the version of the tcloud bin.
+ *
+ * @returns The version of the tcloud bin.
+ */
+async function getTcloudBinVersion(): Promise<Result<[number, number, number], ErrorType>> {
+  const tcloudBin = await getTcloudBin()
+  if (isErr(tcloudBin)) {
+    return tcloudBin
+  }
+  const called = await execAsync(tcloudBin.value, ['--version'])
+  if (called.exitCode !== 0) {
+    return err({
+      type: 'generic',
+      message: `Failed to get tcloud bin version: ${called.stderr}`,
+    })
+  }
+  const version = called.stdout.split('.').map(Number)
+  if (version.length !== 3) {
+    return err({
+      type: 'generic',
+      message: `Failed to get tcloud bin version: ${called.stdout}`,
+    })
+  }
+  return ok(version as [number, number, number])
 }

vscode/extension/tests/tcloud.spec.ts

@@ -54,6 +54,17 @@ async function setupAuthenticatedState(tempDir: string): Promise<void> {
   await fs.writeJson(authStateFile, authState)
 }
 
+/**
+ * Helper function to set the tcloud version for testing
+ */
+async function setTcloudVersion(
+  tempDir: string,
+  version: string,
+): Promise<void> {
+  const versionStateFile = path.join(tempDir, '.tcloud_version_state.json')
+  await fs.writeJson(versionStateFile, { version })
+}
+
 test.describe('Tcloud', () => {
   test('not signed in, shows sign in window', async ({}, testInfo) => {
     testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation
@@ -77,6 +88,9 @@ test.describe('Tcloud', () => {
         `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`,
       )
 
+      // Set tcloud version to 2.10.0
+      await setTcloudVersion(tempDir, '2.10.0')
+
       // Set up Python environment with mock tcloud and sqlmesh
       const pythonPath = await setupPythonEnvironment(pythonEnvDir)
 
@@ -126,7 +140,7 @@ test.describe('Tcloud', () => {
     }
   })
 
-  test('signed in and not installed shows installation window and can see output', async ({}, testInfo) => {
+  test('signed in and not installed shows installation window', async ({}, testInfo) => {
     testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation
     const tempDir = await fs.mkdtemp(
       path.join(os.tmpdir(), 'vscode-test-tcloud-'),
@@ -151,6 +165,9 @@ test.describe('Tcloud', () => {
       // Write mock ".tcloud_auth_state.json" file
       await setupAuthenticatedState(tempDir)
 
+      // Set tcloud version to 2.10.0
+      await setTcloudVersion(tempDir, '2.10.0')
+
       // Set up Python environment with mock tcloud and sqlmesh
       const pythonPath = await setupPythonEnvironment(pythonEnvDir)
 
@@ -199,4 +216,160 @@ test.describe('Tcloud', () => {
       await fs.remove(tempDir)
     }
   })
+
+  test('tcloud sqlmesh_lsp command starts the sqlmesh_lsp in old version when ready', async ({}, testInfo) => {
+    testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation
+    const tempDir = await fs.mkdtemp(
+      path.join(os.tmpdir(), 'vscode-test-tcloud-'),
+    )
+    const pythonEnvDir = path.join(tempDir, '.venv')
+
+    try {
+      // Copy sushi project
+      await fs.copy(SUSHI_SOURCE_PATH, tempDir)
+
+      // Create a tcloud.yaml to mark this as a tcloud project
+      const tcloudConfig = {
+        url: 'https://mock.tobikodata.com',
+        org: 'test-org',
+        project: 'test-project',
+      }
+      await fs.writeFile(
+        path.join(tempDir, 'tcloud.yaml'),
+        `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`,
+      )
+
+      // Write mock ".tcloud_auth_state.json" file
+      await setupAuthenticatedState(tempDir)
+
+      // Set tcloud version to 2.10.0
+      await setTcloudVersion(tempDir, '2.10.0')
+
+      // Set up Python environment with mock tcloud and sqlmesh
+      const pythonPath = await setupPythonEnvironment(pythonEnvDir)
+
+      // Mark sqlmesh as installed
+      const binDir = path.dirname(pythonPath)
+      const installStateFile = path.join(binDir, '.sqlmesh_installed')
+      await fs.writeFile(installStateFile, '')
+
+      // Configure VS Code settings to use our Python environment
+      const settings = {
+        'python.defaultInterpreterPath': pythonPath,
+        'sqlmesh.environmentPath': pythonEnvDir,
+      }
+      await fs.ensureDir(path.join(tempDir, '.vscode'))
+      await fs.writeJson(
+        path.join(tempDir, '.vscode', 'settings.json'),
+        settings,
+        { spaces: 2 },
+      )
+
+      // Start VS Code
+      const { window, close } = await startVSCode(tempDir)
+
+      // Open a SQL file to trigger SQLMesh activation
+      // Wait for the models folder to be visible
+      await window.waitForSelector('text=models')
+
+      // Click on the models folder
+      await window
+        .getByRole('treeitem', { name: 'models', exact: true })
+        .locator('a')
+        .click()
+
+      // Open the top_waiters model
+      await window
+        .getByRole('treeitem', { name: 'customers.sql', exact: true })
+        .locator('a')
+        .click()
+
+      // Verify the context loads successfully
+      await window.waitForSelector('text=Loaded SQLMesh context')
+
+      // Close VS Code
+      await close()
+    } finally {
+      // Clean up
+      await fs.remove(tempDir)
+    }
+  })
+
+  test('tcloud sqlmesh_lsp command starts the sqlmesh_lsp in new version when ready', async ({}, testInfo) => {
+    testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation
+    const tempDir = await fs.mkdtemp(
+      path.join(os.tmpdir(), 'vscode-test-tcloud-'),
+    )
+    const pythonEnvDir = path.join(tempDir, '.venv')
+
+    try {
+      // Copy sushi project
+      await fs.copy(SUSHI_SOURCE_PATH, tempDir)
+
+      // Create a tcloud.yaml to mark this as a tcloud project
+      const tcloudConfig = {
+        url: 'https://mock.tobikodata.com',
+        org: 'test-org',
+        project: 'test-project',
+      }
+      await fs.writeFile(
+        path.join(tempDir, 'tcloud.yaml'),
+        `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`,
+      )
+
+      // Write mock ".tcloud_auth_state.json" file
+      await setupAuthenticatedState(tempDir)
+
+      // Set tcloud version to 2.10.0
+      await setTcloudVersion(tempDir, '2.10.1')
+
+      // Set up Python environment with mock tcloud and sqlmesh
+      const pythonPath = await setupPythonEnvironment(pythonEnvDir)
+
+      // Mark sqlmesh as installed
+      const binDir = path.dirname(pythonPath)
+      const installStateFile = path.join(binDir, '.sqlmesh_installed')
+      await fs.writeFile(installStateFile, '')
+
+      // Configure VS Code settings to use our Python environment
+      const settings = {
+        'python.defaultInterpreterPath': pythonPath,
+        'sqlmesh.environmentPath': pythonEnvDir,
+      }
+      await fs.ensureDir(path.join(tempDir, '.vscode'))
+      await fs.writeJson(
+        path.join(tempDir, '.vscode', 'settings.json'),
+        settings,
+        { spaces: 2 },
+      )
+
+      // Start VS Code
+      const { window, close } = await startVSCode(tempDir)
+
+      // Open a SQL file to trigger SQLMesh activation
+      // Wait for the models folder to be visible
+      await window.waitForSelector('text=models')
+
+      // Click on the models folder
+      await window
+        .getByRole('treeitem', { name: 'models', exact: true })
+        .locator('a')
+        .click()
+
+      // Open the top_waiters model
+      await window
+        .getByRole('treeitem', { name: 'customers.sql', exact: true })
+        .locator('a')
+        .click()
+
+      // Verify the context loads successfully
+      await window.waitForSelector('text=Loaded SQLMesh context')
+
+      // Close VS Code
+      await close()
+    } finally {
+      // Clean up
+      await fs.remove(tempDir)
+    }
+  })
 })