add env selector for the visualizer

vscode-extension/src/extension.ts

@@ -15,26 +15,27 @@ import { outputChannel } from "./utils";
 import {
   OpenStudioCommand,
   OpenStudio,
-  SelectEnvAndOpenStudio,
-  SelectEnvAndOpenStudioCommand,
 } from "./modules/studio/open-studio/command";
 import {
   StopStudioCommand,
   StopStudio,
 } from "./modules/studio/stop-studio/command";
 import { OpenStudioCodeLens } from "./modules/studio/open-studio/codelens";
 import { stopStudio } from "./modules/studio/server";
+import {
+  SelectEnv,
+  SelectEnvCommand,
+} from "./modules/internal/select-env.command";
 
 export function activate(context: vscode.ExtensionContext) {
   checkNodeVersion();
 
   context.subscriptions.push(
+    /* Internal */
+    vscode.commands.registerCommand(SelectEnvCommand, SelectEnv),
+
     /* Studio */
     vscode.commands.registerCommand(OpenStudioCommand, OpenStudio),
-    vscode.commands.registerCommand(
-      SelectEnvAndOpenStudioCommand,
-      SelectEnvAndOpenStudio,
-    ),
     vscode.commands.registerCommand(StopStudioCommand, StopStudio),
     vscode.languages.registerCodeLensProvider(
       { pattern: "**/*{drizzle,config}.ts", language: "typescript" },

vscode-extension/src/modules/internal/select-env.command.ts

@@ -0,0 +1,53 @@
+import * as vscode from "vscode";
+
+export const SelectEnvCommand = "drizzle:select_env";
+
+/* Local state */
+let $lastSelectedEnv: string | undefined;
+
+export async function SelectEnv(configPath: string, executeCommand: string) {
+  // Find .env files in the workspace
+  const envFiles = await vscode.workspace.findFiles(
+    "**/.env*",
+    "**/node_modules/**",
+  );
+
+  // Create quick pick items
+  const items = envFiles.map((file) => ({
+    label: vscode.workspace.asRelativePath(file),
+    path: file.fsPath,
+  }));
+
+  // Add option to not use an env file
+  items.unshift({ label: "None", path: "Don't use an env file" });
+
+  // Move last selected item to top if it exists
+  if ($lastSelectedEnv) {
+    const lastSelectedIndex = items.findIndex(
+      (item) => item.path === $lastSelectedEnv,
+    );
+    if (lastSelectedIndex > -1) {
+      const [lastItem] = items.splice(lastSelectedIndex, 1);
+      items.unshift(lastItem);
+    }
+  }
+
+  // Show quick pick
+  const selected = await vscode.window.showQuickPick(items, {
+    placeHolder: "Select an environment file",
+  });
+
+  if (selected) {
+    // Save the selection to workspace configuration (except for "None")
+    if (selected.label !== "None") {
+      $lastSelectedEnv = selected.path;
+    }
+
+    // Call open studio command with the selected env file
+    await vscode.commands.executeCommand(
+      executeCommand,
+      configPath,
+      selected.label === "None" ? undefined : selected.path,
+    );
+  }
+}

vscode-extension/src/modules/studio/open-studio/codelens.ts

@@ -2,7 +2,8 @@ import * as vscode from "vscode";
 
 import { findDrizzleConfigLines } from "../../../utils";
 import { findDrizzleKitPath } from "../server";
-import { SelectEnvAndOpenStudioCommand } from "./command";
+import { SelectEnvCommand } from "../../internal/select-env.command";
+import { OpenStudioCommand } from "./command";
 
 export class OpenStudioCodeLens implements vscode.CodeLensProvider {
   async provideCodeLenses(
@@ -23,9 +24,9 @@ export class OpenStudioCodeLens implements vscode.CodeLensProvider {
 
         return new vscode.CodeLens(range, {
           title: "🌧️ Open Drizzle Studio",
-          command: SelectEnvAndOpenStudioCommand,
+          command: SelectEnvCommand,
           tooltip: "Open Drizzle Studio",
-          arguments: [document.uri.fsPath],
+          arguments: [document.uri.fsPath, OpenStudioCommand],
         });
       },
     );

vscode-extension/src/modules/studio/open-studio/command.ts

@@ -46,56 +46,3 @@ export async function OpenStudio(...args: any[]) {
     return;
   }
 }
-
-export const SelectEnvAndOpenStudioCommand =
-  "drizzle.studio:select_env_and_open";
-
-/* Local state */
-let $lastSelectedEnv: string | undefined;
-
-export async function SelectEnvAndOpenStudio(configPath: string) {
-  // Find .env files in the workspace
-  const envFiles = await vscode.workspace.findFiles(
-    "**/.env*",
-    "**/node_modules/**",
-  );
-
-  // Create quick pick items
-  const items = envFiles.map((file) => ({
-    label: vscode.workspace.asRelativePath(file),
-    path: file.fsPath,
-  }));
-
-  // Add option to not use an env file
-  items.unshift({ label: "None", path: "Don't use an env file" });
-
-  // Move last selected item to top if it exists
-  if ($lastSelectedEnv) {
-    const lastSelectedIndex = items.findIndex(
-      (item) => item.path === $lastSelectedEnv,
-    );
-    if (lastSelectedIndex > -1) {
-      const [lastItem] = items.splice(lastSelectedIndex, 1);
-      items.unshift(lastItem);
-    }
-  }
-
-  // Show quick pick
-  const selected = await vscode.window.showQuickPick(items, {
-    placeHolder: "Select an environment file",
-  });
-
-  if (selected) {
-    // Save the selection to workspace configuration (except for "None")
-    if (selected.label !== "None") {
-      $lastSelectedEnv = selected.path;
-    }
-
-    // Call open studio command with the selected env file
-    await vscode.commands.executeCommand(
-      OpenStudioCommand,
-      configPath,
-      selected.label === "None" ? undefined : selected.path,
-    );
-  }
-}

vscode-extension/src/modules/visualizer/open-visualizer/codelens.ts

@@ -2,6 +2,7 @@ import * as vscode from "vscode";
 
 import { findDrizzleConfigLines } from "../../../utils";
 import { OpenVisualizerCommand } from "./command";
+import { SelectEnvCommand } from "../../internal/select-env.command";
 
 export class OpenVisualizerCodeLens implements vscode.CodeLensProvider {
   async provideCodeLenses(
@@ -15,9 +16,9 @@ export class OpenVisualizerCodeLens implements vscode.CodeLensProvider {
 
       return new vscode.CodeLens(range, {
         title: "🌧️ Open Drizzle Visualizer",
-        command: OpenVisualizerCommand,
+        command: SelectEnvCommand,
         tooltip: "Open Drizzle Schema Visualizer",
-        arguments: [document.uri.fsPath],
+        arguments: [document.uri.fsPath, OpenVisualizerCommand],
       });
     });
   }

vscode-extension/src/modules/visualizer/open-visualizer/command.ts

@@ -7,14 +7,15 @@ export const OpenVisualizerCommand = "drizzle.visualizer:open";
 export async function OpenVisualizer(...args: any[]) {
   const OutputKey = `[${OpenVisualizerCommand}]`;
   const configPath = args[0];
+  const envFile = args[1];
 
   if (!configPath || typeof configPath !== "string") {
     toastError(`${OutputKey} Expected config path to be a string`);
     return;
   }
 
   try {
-    const { port } = await startVisualizer(configPath);
+    const { port } = await startVisualizer(configPath, envFile);
     const panel = createDrizzleVisualizerPanel();
 
     panel.webview.html = render(`

vscode-extension/src/modules/visualizer/server.ts

@@ -20,7 +20,7 @@ interface ServerStartResult {
   port: string;
 }
 
-export async function startVisualizer(configPath: string) {
+export async function startVisualizer(configPath: string, envFile?: string) {
   outputChannel.appendLine(`${OutputKey} extension cwd: ${extensionCwd}`);
   outputChannel.appendLine(`${OutputKey} apps cwd: ${appsCwd}`);
   outputChannel.appendLine(
@@ -63,7 +63,12 @@ export async function startVisualizer(configPath: string) {
       TS_CONFIG_PATH: path.join(cwd, "tsconfig.json"),
     };
 
-    $app = spawn(process.execPath, [binPath], {
+    const envFileArg = envFile ? ["--env-file", envFile] : [];
+    outputChannel.appendLine(
+      `${OutputKey} Using env file: ${envFile ? envFile : "none"}`,
+    );
+
+    $app = spawn(process.execPath, [...envFileArg, binPath], {
       stdio: "pipe",
       detached: true,
       shell: false,