URI Handler (#6)

* URI Handler * use withProgress to notify user things are happening * PR feedback
dbos-inc · Feb 26, 2024 · 9d91336 · 9d91336
1 parent e742ddb
commit 9d91336
Show file tree

Hide file tree

Showing 8 changed files with 130 additions and 42 deletions.
diff --git a/src/DebugProxy.ts b/src/DebugProxy.ts
@@ -121,6 +121,11 @@ export class DebugProxy {
             this._outChannel.error(e, { database });
         });
 
+        proxyProcess.on('close', (code, signal) => {
+            this._proxyProcesses.delete(configHash);
+            this._outChannel.info(`Debug Proxy closed with exit code ${code}`, { database });
+        });
+
         proxyProcess.on("exit", (code, _signal) => {
             this._proxyProcesses.delete(configHash);
             this._outChannel.info(`Debug Proxy exited with exit code ${code}`, { database });

diff --git a/src/ProvenanceDatabase.ts b/src/ProvenanceDatabase.ts
@@ -15,6 +15,8 @@ export interface workflow_status {
     authenticated_roles: string; // Serialized list of roles.
     request: string; // Serialized HTTPRequest
     executor_id: string; // Set to "local" for local deployment, set to microVM ID for cloud deployment.
+    created_at: string;
+    updated_at: string;
 }
 
 export class ProvenanceDatabase {
@@ -44,4 +46,11 @@ export class ProvenanceDatabase {
         const results = await db.query<workflow_status>('SELECT * FROM dbos.workflow_status WHERE name = $1 LIMIT 10', [wfName]);
         return results.rows;
     }
+
+    async getWorkflowStatus(clientConfig: ClientConfig, wfid: string): Promise<workflow_status | undefined> {
+        const db = await this.connect(clientConfig);
+        const results = await db.query<workflow_status>('SELECT * FROM dbos.workflow_status WHERE workflow_uuid = $1 LIMIT 10', [wfid]);
+        if (results.rows.length > 1) { throw new Error(`Multiple workflow status records found for workflow ID ${wfid}`); }
+        return results.rows.length === 1 ? results.rows[0] : undefined;
+    }
 }
diff --git a/src/codeLensProvider.ts b/src/codeLensProvider.ts
@@ -1,6 +1,6 @@
 import * as vscode from 'vscode';
 import ts from 'typescript';
-import { startDebuggingCommandName } from './commands';
+import { startDebuggingCodeLensCommandName } from './commands';
 import { logger } from './extension';
 import { getDbosMethodType, parse } from './sourceParser';
 
@@ -32,7 +32,7 @@ export class TTDbgCodeLensProvider implements vscode.CodeLensProvider {
                     return new vscode.CodeLens(range, {
                         title: '⏳ Time Travel Debug',
                         tooltip: `Debug ${methodInfo.name} with the DBOS Time Travel Debugger`,
-                        command: startDebuggingCommandName,
+                        command: startDebuggingCodeLensCommandName,
                         arguments: [folder, methodInfo.name, methodType]
                     });
                 })

diff --git a/src/commands.ts b/src/commands.ts
@@ -1,50 +1,81 @@
 import * as vscode from 'vscode';
 import { logger, config, provDB, debugProxy } from './extension';
 import { DbosMethodType } from "./sourceParser";
-import { stringify } from './utils';
+import { getWorkspaceFolder, stringify } from './utils';
 import { dbos_cloud_login } from './configuration';
+import { ClientConfig } from 'pg';
 
 export const cloudLoginCommandName = "dbos-ttdbg.cloud-login";
-export const startDebuggingCommandName = "dbos-ttdbg.startDebugging";
+export const startDebuggingCodeLensCommandName = "dbos-ttdbg.start-debugging-code-lens";
+export const startDebuggingUriCommandName = "dbos-ttdbg.start-debugging-uri";
 export const shutdownDebugProxyCommandName = "dbos-ttdbg.shutdown-debug-proxy";
 export const deleteProvDbPasswordsCommandName = "dbos-ttdbg.delete-prov-db-passwords";
 
-export async function startDebugging(folder: vscode.WorkspaceFolder, name: string, $type: DbosMethodType) {
+async function startDebugging(folder: vscode.WorkspaceFolder, getWorkflowID: (clientConfig: ClientConfig) => Promise<string | undefined>) {
     try {
         const clientConfig = await config.getProvDbConfig(folder);
         if (!clientConfig) { return; }
+
+        const debuggerStarted = await vscode.window.withProgress(
+            { 
+                location: vscode.ProgressLocation.Window,
+                title: "Launching DBOS Time Travel Debugger",
+            },
+            async () => {
+                await debugProxy.launch(clientConfig);
+                const workflowID = await getWorkflowID(clientConfig);
+                if (!workflowID) { return; }
 
-        const statuses = await vscode.window.withProgress(
-            { location: vscode.ProgressLocation.Window },
-            () => { return provDB.getWorkflowStatuses(clientConfig, name, $type); });
+                const proxyURL = `http://localhost:${config.proxyPort ?? 2345}`;
+                logger.info(`startDebugging`, { folder: folder.uri.fsPath, database: clientConfig.database, workflowID });
+                return await vscode.debug.startDebugging(
+                    folder,
+                    {
+                        name: `Time-Travel Debug ${workflowID}`,
+                        type: 'node-terminal',
+                        request: 'launch',
+                        command: `npx dbos-sdk debug -x ${proxyURL} -u ${workflowID}`
+                    }
+                );
+            }
+        );
 
-        await debugProxy.launch(clientConfig);
+        if (!debuggerStarted) {
+            throw new Error("vscode.debug.startDebugging returned false");
+        }
 
+    } catch (e) {
+        logger.error("startDebugging", e);
+        vscode.window.showErrorMessage(`Failed to start debugging`);
+    }
+}
+
+export async function startDebuggingFromCodeLens(folder: vscode.WorkspaceFolder, name: string, $type: DbosMethodType) {
+    logger.info(`startDebuggingFromCodeLens`, { folder: folder.uri.fsPath, name, type: $type });
+    await startDebugging(folder, async (clientConfig) => {
         // TODO: eventually, we'll need a better UI than "list all workflow IDs and let the user pick one"
-        const wfID = await vscode.window.showQuickPick(statuses.map(s => s.workflow_uuid), {
+        const statuses = await provDB.getWorkflowStatuses(clientConfig, name, $type);
+        return await vscode.window.showQuickPick(statuses.map(s => s.workflow_uuid), {
             placeHolder: `Select a ${name} workflow ID to debug`,
             canPickMany: false,
         });
+    });
+}
 
-        if (!wfID) { return; }
-
-        logger.info(`Starting debugging for ${name} workflow ${wfID}`);
+export async function startDebuggingFromUri(wfid: string) {
+    const folder = await getWorkspaceFolder();
+    if (!folder) { return; }
 
-        const proxyURL = `http://localhost:${config.proxyPort ?? 2345}`;
-        await vscode.debug.startDebugging(
-            vscode.workspace.getWorkspaceFolder(vscode.window.activeTextEditor!.document.uri),
-            {
-                name: `Debug ${wfID}`,
-                type: 'node-terminal',
-                request: 'launch',
-                command: `npx dbos-sdk debug -x ${proxyURL} -u ${wfID}`
-            }
-        );
-    } catch (e) {
-        const reason = stringify(e);
-        logger.error("startDebugging", e);
-        vscode.window.showErrorMessage(`Failed to start debugging\n${reason}`);
-    }
+    logger.info(`startDebuggingFromUri`, { folder: folder.uri.fsPath, wfid });
+    await startDebugging(folder, async (clientConfig) => {
+        const wfStatus = await provDB.getWorkflowStatus(clientConfig, wfid);
+        if (!wfStatus) {
+            vscode.window.showErrorMessage(`Workflow ID ${wfid} not found in provenance database`);
+            return undefined;
+        } else {
+            return wfid;
+        }
+    });
 }
 
 export function shutdownDebugProxy() {

diff --git a/src/configuration.ts b/src/configuration.ts
@@ -204,21 +204,19 @@ export class Configuration {
                 return { host, port, database, user };
             });
 
-        if (!dbConfig.host || !dbConfig.database || !dbConfig.user) {
+        if (dbConfig.host && dbConfig.database && dbConfig.user) {
+            return {
+                host: dbConfig.host,
+                port: dbConfig.port,
+                database: dbConfig.database,
+                user: dbConfig.user,
+                password: () => this.#getPassword(folder),
+                ssl: { rejectUnauthorized: false }
+            };
+        } else {
             startInvalidCredentialsFlow(folder).catch(e => logger.error("startInvalidCredentialsFlow", e));
             return undefined;
         }
-
-        return {
-            host: dbConfig.host,
-            port: dbConfig.port,
-            database: dbConfig.database,
-            user: dbConfig.user,
-            password: () => this.#getPassword(folder),
-            ssl: {
-                rejectUnauthorized: false,
-            }
-        };
     }
 
     get proxyPort(): number {

diff --git a/src/extension.ts b/src/extension.ts
@@ -1,11 +1,12 @@
 import * as vscode from 'vscode';
 import { S3CloudStorage } from './CloudStorage';
 import { TTDbgCodeLensProvider } from './codeLensProvider';
-import { deleteProvenanceDatabasePasswords, deleteProvDbPasswordsCommandName, startDebuggingCommandName, startDebugging, shutdownDebugProxyCommandName, shutdownDebugProxy, cloudLoginCommandName, cloudLogin } from './commands';
+import { deleteProvenanceDatabasePasswords, deleteProvDbPasswordsCommandName, shutdownDebugProxyCommandName, shutdownDebugProxy, cloudLoginCommandName, cloudLogin, startDebuggingCodeLensCommandName, startDebuggingFromCodeLens, startDebuggingFromUri, startDebuggingUriCommandName } from './commands';
 import { Configuration } from './configuration';
 import { DebugProxy, } from './DebugProxy';
 import { LogOutputChannelTransport, Logger, createLogger } from './logger';
 import { ProvenanceDatabase } from './ProvenanceDatabase';
+import { TTDbgUriHandler } from './uriHandler';
 
 export let logger: Logger;
 export let config: Configuration;
@@ -30,7 +31,9 @@ export async function activate(context: vscode.ExtensionContext) {
   context.subscriptions.push(debugProxy);
 
   context.subscriptions.push(
-    vscode.commands.registerCommand(startDebuggingCommandName, startDebugging));
+    vscode.commands.registerCommand(startDebuggingCodeLensCommandName, startDebuggingFromCodeLens));
+  context.subscriptions.push(
+    vscode.commands.registerCommand(startDebuggingUriCommandName, startDebuggingFromUri));
   context.subscriptions.push(
     vscode.commands.registerCommand(shutdownDebugProxyCommandName, shutdownDebugProxy));
   context.subscriptions.push(
@@ -43,6 +46,9 @@ export async function activate(context: vscode.ExtensionContext) {
       { scheme: 'file', language: 'typescript' },
       new TTDbgCodeLensProvider()));
 
+  context.subscriptions.push(
+    vscode.window.registerUriHandler(new TTDbgUriHandler()));
+
   await debugProxy.update().catch(e => {
     logger.error("Debug Proxy Update Failed", e);
     vscode.window.showErrorMessage(`Debug Proxy Update Failed`);

diff --git a/src/uriHandler.ts b/src/uriHandler.ts
@@ -0,0 +1,27 @@
+import * as vscode from 'vscode';
+import { logger } from './extension';
+import { startDebuggingUriCommandName } from './commands';
+
+export class TTDbgUriHandler implements vscode.UriHandler {
+	async handleUri(uri: vscode.Uri): Promise<void> {
+		logger.debug(`TTDbgUriHandler.handleUri`, { uri: uri.toString() });
+		switch (uri.path) {
+			case '/start-debugging':
+				const searchParams = new URLSearchParams(uri.query);
+				const wfid = searchParams.get('wfid');
+				if (wfid) {
+					vscode.window.withProgress({
+						location: vscode.ProgressLocation.Notification,
+						title: `Starting DBOS Time Travel Debugger for workflow ID ${wfid}`,
+					}, async () => {
+						await vscode.commands.executeCommand(startDebuggingUriCommandName, wfid);
+					});
+				} else {
+					vscode.window.showErrorMessage(`Invalid start-debugging uri: ${uri}`);
+				}
+				break;
+			default:
+				vscode.window.showErrorMessage(`Unsupported uri: ${uri.path}}`);
+		}
+	}
+}
diff --git a/src/utils.ts b/src/utils.ts
@@ -49,3 +49,15 @@ export function hashClientConfig(clientConfig: ClientConfig) {
         : undefined;
 }
 
+export async function getWorkspaceFolder() {
+    const folders = vscode.workspace.workspaceFolders ?? [];
+    if (folders.length === 1) { return folders[0]; }
+
+	if (vscode.window.activeTextEditor) {
+		const folder = vscode.workspace.getWorkspaceFolder(vscode.window.activeTextEditor.document.uri);
+		if (folder) {
+			return folder;
+		}
+	}
+	return await vscode.window.showWorkspaceFolderPick();
+}