sqlmesh 0.194.2.dev4__tar.gz → 0.194.2.dev5__tar.gz

{sqlmesh-0.194.2.dev4 → sqlmesh-0.194.2.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.194.2.dev4
+Version: 0.194.2.dev5
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.194.2.dev4 → sqlmesh-0.194.2.dev5}/sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.194.2.dev4'
-__version_tuple__ = version_tuple = (0, 194, 2, 'dev4')
+__version__ = version = '0.194.2.dev5'
+__version_tuple__ = version_tuple = (0, 194, 2, 'dev5')

{sqlmesh-0.194.2.dev4 → sqlmesh-0.194.2.dev5}/sqlmesh.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.194.2.dev4
+Version: 0.194.2.dev5
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.194.2.dev4 → sqlmesh-0.194.2.dev5}/vscode/extension/src/commands/format.ts

@@ -11,12 +11,18 @@ export const format =
   (
     authProvider: AuthenticationProviderTobikoCloud,
     lsp: LSPClient | undefined,
+    restartLSP: () => Promise<void>,
   ) =>
   async (): Promise<void> => {
     traceLog('Calling format')
     const out = await internalFormat(lsp)
     if (isErr(out)) {
-      return handleError(authProvider, out.error, 'Project format failed')
+      return handleError(
+        authProvider,
+        restartLSP,
+        out.error,
+        'Project format failed',
+      )
     }
     vscode.window.showInformationMessage('Project formatted successfully')
   }

sqlmesh-0.194.2.dev5/vscode/extension/src/commands/signin.ts

@@ -0,0 +1,30 @@
+import { AuthenticationProviderTobikoCloud } from '../auth/auth'
+import * as vscode from 'vscode'
+import { isCodespaces } from '../utilities/isCodespaces'
+import { traceInfo } from '../utilities/common/log'
+
+export const signIn =
+  (
+    authenticationProvider: AuthenticationProviderTobikoCloud,
+    onSignInSuccess: () => Promise<void>,
+  ) =>
+  async () => {
+    if (isCodespaces()) {
+      await authenticationProvider.sign_in_device_flow()
+    } else {
+      await authenticationProvider.createSession()
+    }
+
+    // Do not await this, as this will block the thread, you just need to show the message, but not block
+    vscode.window.showInformationMessage('Signed in successfully')
+
+    // Execute callback after successful sign-in
+    if (onSignInSuccess) {
+      traceInfo('Executing post sign-in callback')
+      try {
+        await onSignInSuccess()
+      } catch (error) {
+        traceInfo(`Error in post sign-in callback: ${error}`)
+      }
+    }
+  }

{sqlmesh-0.194.2.dev4 → sqlmesh-0.194.2.dev5}/vscode/extension/src/commands/signinSpecifyFlow.ts

@@ -3,7 +3,11 @@ import { traceInfo } from '../utilities/common/log'
 import { window } from 'vscode'
 
 export const signInSpecifyFlow =
-  (authenticationProvider: AuthenticationProviderTobikoCloud) => async () => {
+  (
+    authenticationProvider: AuthenticationProviderTobikoCloud,
+    onSignInSuccess?: () => Promise<void>,
+  ) =>
+  async () => {
     traceInfo('Sign in specify flow')
     const flowOptions = [
       {
@@ -24,11 +28,31 @@ export const signInSpecifyFlow =
       await authenticationProvider.sign_in_oauth_flow()
       await authenticationProvider.getSessions()
       await window.showInformationMessage('Sign in success')
+
+      // Execute callback after successful sign-in
+      if (onSignInSuccess) {
+        traceInfo('Executing post sign-in callback')
+        try {
+          await onSignInSuccess()
+        } catch (error) {
+          traceInfo(`Error in post sign-in callback: ${error}`)
+        }
+      }
       return
     } else if (selectedFlow.label === 'Device Flow') {
       await authenticationProvider.sign_in_device_flow()
       await authenticationProvider.getSessions()
       await window.showInformationMessage('Sign in success')
+
+      // Execute callback after successful sign-in
+      if (onSignInSuccess) {
+        traceInfo('Executing post sign-in callback')
+        try {
+          await onSignInSuccess()
+        } catch (error) {
+          traceInfo(`Error in post sign-in callback: ${error}`)
+        }
+      }
       return
     } else {
       traceInfo('Invalid flow selected')

{sqlmesh-0.194.2.dev4 → sqlmesh-0.194.2.dev5}/vscode/extension/src/extension.ts

@@ -45,12 +45,21 @@ export async function activate(context: vscode.ExtensionContext) {
   traceInfo('Authentication provider registered')
 
   context.subscriptions.push(
-    vscode.commands.registerCommand('sqlmesh.signin', signIn(authProvider)),
+    vscode.commands.registerCommand(
+      'sqlmesh.signin',
+      signIn(authProvider, async () => {
+        traceInfo('Restarting LSP after sign-in')
+        await restart()
+      }),
+    ),
   )
   context.subscriptions.push(
     vscode.commands.registerCommand(
       'sqlmesh.signinSpecifyFlow',
-      signInSpecifyFlow(authProvider),
+      signInSpecifyFlow(authProvider, async () => {
+        traceInfo('Restarting LSP after sign-in')
+        await restart()
+      }),
     ),
   )
   context.subscriptions.push(
@@ -76,13 +85,6 @@ export async function activate(context: vscode.ExtensionContext) {
     ),
   )
 
-  context.subscriptions.push(
-    vscode.commands.registerCommand(
-      'sqlmesh.format',
-      format(authProvider, lspClient),
-    ),
-  )
-
   // Register the webview
   const lineagePanel = new LineagePanel(context.extensionUri, lspClient)
   context.subscriptions.push(
@@ -118,14 +120,35 @@ export async function activate(context: vscode.ExtensionContext) {
       if (isErr(restartResult)) {
         return handleError(
           authProvider,
+          restart,
           restartResult.error,
           'LSP restart failed',
         )
       }
       context.subscriptions.push(lspClient)
+    } else {
+      lspClient = new LSPClient()
+      const result = await lspClient.start()
+      if (isErr(result)) {
+        return handleError(
+          authProvider,
+          restart,
+          result.error,
+          'Failed to start LSP',
+        )
+      } else {
+        context.subscriptions.push(lspClient)
+      }
     }
   }
 
+  context.subscriptions.push(
+    vscode.commands.registerCommand(
+      'sqlmesh.format',
+      format(authProvider, lspClient, restart),
+    ),
+  )
+
   context.subscriptions.push(
     onDidChangePythonInterpreter(async () => {
       await restart()
@@ -140,7 +163,12 @@ export async function activate(context: vscode.ExtensionContext) {
 
   const result = await lspClient.start()
   if (isErr(result)) {
-    return handleError(authProvider, result.error, 'Failed to start LSP')
+    return handleError(
+      authProvider,
+      restart,
+      result.error,
+      'Failed to start LSP',
+    )
   } else {
     context.subscriptions.push(lspClient)
   }

{sqlmesh-0.194.2.dev4 → sqlmesh-0.194.2.dev5}/vscode/extension/src/utilities/errors.ts

@@ -69,6 +69,7 @@ interface SqlmeshLspDependenciesMissingError {
 
 export async function handleError(
   authProvider: AuthenticationProviderTobikoCloud,
+  restartLsp: () => Promise<void>,
   error: ErrorType,
   genericErrorPrefix?: string,
 ): Promise<void> {
@@ -85,7 +86,7 @@ export async function handleError(
       )
       return
     case 'not_signed_in':
-      return handleNotSignedInError(authProvider)
+      return handleNotSignedInError(authProvider, restartLsp)
     case 'sqlmesh_not_found':
       return handleSqlmeshNotFoundError()
     case 'sqlmesh_lsp_not_found':
@@ -110,6 +111,7 @@ export async function handleError(
  */
 const handleNotSignedInError = async (
   authProvider: AuthenticationProviderTobikoCloud,
+  restartLsp: () => Promise<void>,
 ): Promise<void> => {
   traceInfo('handleNotSginedInError')
   const result = await window.showInformationMessage(
@@ -117,7 +119,7 @@ const handleNotSignedInError = async (
     'Sign In',
   )
   if (result === 'Sign In') {
-    await signIn(authProvider)()
+    await signIn(authProvider, restartLsp)()
   }
 }
 

sqlmesh-0.194.2.dev5/vscode/extension/tests/tcloud.spec.ts

@@ -0,0 +1,423 @@
+import { expect, test } from '@playwright/test'
+import path from 'path'
+import fs from 'fs-extra'
+import os from 'os'
+import {
+  createVirtualEnvironment,
+  pipInstall,
+  REPO_ROOT,
+  startVSCode,
+  SUSHI_SOURCE_PATH,
+} from './utils'
+import { setTcloudVersion, setupAuthenticatedState } from './tcloud_utils'
+
+/**
+ * Helper function to create and set up a Python virtual environment
+ */
+async function setupPythonEnvironment(envDir: string): Promise<string> {
+  // Create virtual environment
+  const pythonDetails = await createVirtualEnvironment(envDir)
+
+  // Install the mock tcloud package
+  const mockTcloudPath = path.join(__dirname, 'tcloud')
+  await pipInstall(pythonDetails, [mockTcloudPath])
+
+  // Install sqlmesh from the local repository with LSP support
+  const customMaterializations = path.join(
+    REPO_ROOT,
+    'examples',
+    'custom_materializations',
+  )
+  const sqlmeshWithExtras = `${REPO_ROOT}[lsp,bigquery]`
+  await pipInstall(pythonDetails, [sqlmeshWithExtras, customMaterializations])
+
+  return pythonDetails.pythonPath
+}
+
+test('not signed in, shows sign in 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-'),
+  )
+  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`,
+    )
+
+    // 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)
+
+    // 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()
+
+    // Wait for the file to open
+    await window.waitForTimeout(2000)
+
+    await window.waitForSelector(
+      'text=Please sign in to Tobiko Cloud to use SQLMesh',
+    )
+
+    // Close VS Code
+    await close()
+  } finally {
+    // Clean up
+    await fs.remove(tempDir)
+  }
+})
+
+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-'),
+  )
+  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)
+
+    // 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()
+
+    await window.waitForSelector('text=Installing enterprise python package')
+    expect(
+      await window.locator('text=Installing enterprise python package'),
+    ).toHaveCount(2)
+
+    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 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)
+  }
+})
+
+// This test is skipped becuase of the way the sign in window is shown is not useable by playwright. It's not solvable
+// but the test is still useful when running it manually.
+test.skip('tcloud not signed in and not installed, shows sign in window and then fact that loaded', 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')
+
+  // 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`,
+  )
+
+  // Set up Python environment with mock tcloud and sqlmesh
+  const pythonPath = await setupPythonEnvironment(pythonEnvDir)
+
+  // 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,
+  })
+
+  // Set tcloud version to 2.10.0
+  await setTcloudVersion(tempDir, '2.10.1')
+
+  // Start VS Code
+  const { window, close } = await startVSCode(tempDir)
+
+  try {
+    // Copy sushi project
+    await fs.copy(SUSHI_SOURCE_PATH, 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 sign in window is shown
+    await window.waitForSelector(
+      'text=Please sign in to Tobiko Cloud to use SQLMesh',
+    )
+
+    // Click on the sign in button
+    await window
+      .getByRole('button', { name: 'Sign in' })
+      .filter({ hasText: 'Sign in' })
+      .click()
+    await window.waitForSelector('text="Signed in successfully"')
+
+    await window.waitForSelector('text=Installing enterprise python package')
+
+    await window.waitForSelector('text=Loaded SQLMesh context')
+  } finally {
+    // Clean up
+    await close()
+    await fs.remove(tempDir)
+  }
+})

sqlmesh-0.194.2.dev4/vscode/extension/src/commands/signin.ts

@@ -1,13 +0,0 @@
-import { AuthenticationProviderTobikoCloud } from '../auth/auth'
-import * as vscode from 'vscode'
-import { isCodespaces } from '../utilities/isCodespaces'
-
-export const signIn =
-  (authenticationProvider: AuthenticationProviderTobikoCloud) => async () => {
-    if (isCodespaces()) {
-      await authenticationProvider.sign_in_device_flow()
-    } else {
-      await authenticationProvider.createSession()
-    }
-    await vscode.window.showInformationMessage('Signed in successfully')
-  }