Add connect protocol handler to SQL command line service (#7007)

- This already handles commandline which needs same logic - Includes visual confirmation to avoid connecting by default to untrusted server - If successful will connect & expand in Connections tree - If failed, will show dialog with all presets - Shows New Query by default, can add option for New Notebook instead in the future Connection protocol is `azuredatastudio://connect?server=mysqldw.database.windows.net&database=db1&user=myusername&aad=true` Same args as commandline, namely: - server - database - user - aad - integrated I have also added `provider` arg to allow non-MSSQL providers. In the future, we should update logic to prompt for install of correct extension for that provider if it's missing.
2026-01-27 01:25:36 -05:00 · 2019-08-29 17:26:50 -07:00
parent 9b003ad285
commit f51f1989aa
2 changed files with 246 additions and 11 deletions
--- a/src/sql/workbench/parts/commandLine/electron-browser/commandLine.ts
+++ b/src/sql/workbench/parts/commandLine/electron-browser/commandLine.ts
@@ -3,6 +3,7 @@
 *  Licensed under the Source EULA. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

+import * as querystring from 'querystring';
 import * as azdata from 'azdata';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 import { ConnectionProfileGroup } from 'sql/platform/connection/common/connectionProfileGroup';
@@ -26,8 +27,24 @@ import { INotificationService } from 'vs/platform/notification/common/notificati
 import { IWorkbenchContribution } from 'vs/workbench/common/contributions';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { openNewQuery } from 'sql/workbench/parts/query/browser/queryActions';
+import { IURLService, IURLHandler } from 'vs/platform/url/common/url';
+import { getErrorMessage } from 'vs/base/common/errors';
+import { IDialogService } from 'vs/platform/dialogs/common/dialogs';

-export class CommandLineWorkbenchContribution implements IWorkbenchContribution {
+const connectAuthority = 'connect';
+
+interface SqlArgs {
+	_?: string[];
+	aad?: boolean;
+	database?: string;
+	integrated?: boolean;
+	server?: string;
+	user?: string;
+	command?: string;
+	provider?: string;
+}
+
+export class CommandLineWorkbenchContribution implements IWorkbenchContribution, IURLHandler {

 	constructor(
 		@ICapabilitiesService private readonly _capabilitiesService: ICapabilitiesService,
@@ -38,7 +55,9 @@ export class CommandLineWorkbenchContribution implements IWorkbenchContribution
 		@IConfigurationService private readonly _configurationService: IConfigurationService,
 		@INotificationService private readonly _notificationService: INotificationService,
 		@ILogService private readonly logService: ILogService,
-		@IInstantiationService private readonly instantiationService: IInstantiationService
+		@IInstantiationService private readonly instantiationService: IInstantiationService,
+		@IURLService urlService: IURLService,
+		@IDialogService private readonly dialogService: IDialogService
 	) {
 		if (ipc) {
 			ipc.on('ads:processCommandLine', (event: any, args: ParsedArgs) => this.onLaunched(args));
@@ -47,6 +66,9 @@ export class CommandLineWorkbenchContribution implements IWorkbenchContribution
 		if (environmentService) {
 			this.onLaunched(environmentService.args);
 		}
+		if (urlService) {
+			urlService.registerHandler(this);
+		}
 	}

 	private onLaunched(args: ParsedArgs) {
@@ -69,7 +91,7 @@ export class CommandLineWorkbenchContribution implements IWorkbenchContribution
 	// (null, commandName) => Launch the command with a null connection. If the command implementation needs a connection, it will need to create it.
 	// (serverName, null) => Connect object explorer and open a new query editor if no file names are passed. If file names are passed, connect their editors to the server.
 	// (null, null) => Prompt for a connection unless there are registered servers
-	public async processCommandLine(args: ParsedArgs): Promise<void> {
+	public async processCommandLine(args: SqlArgs): Promise<void> {
 		let profile: IConnectionProfile = undefined;
 		let commandName = undefined;
 		if (args) {
@@ -99,7 +121,7 @@ export class CommandLineWorkbenchContribution implements IWorkbenchContribution
 				let updatedProfile = this._connectionManagementService.getConnectionProfileById(profile.id);
 				connectedContext = { connectionProfile: new ConnectionProfile(this._capabilitiesService, updatedProfile).toIConnectionProfile() };
 			} catch (err) {
-				this.logService.warn('Failed to connect due to error' + err.message);
+				this.logService.warn('Failed to connect due to error' + getErrorMessage(err));
 			}
 		}
 		if (commandName) {
@@ -130,6 +152,52 @@ export class CommandLineWorkbenchContribution implements IWorkbenchContribution
 		}
 	}

+	public async handleURL(uri: URI): Promise<boolean> {
+		// Catch file URLs
+		let authority = uri.authority.toLowerCase();
+		if (authority === connectAuthority) {
+			try {
+				let args = this.parseProtocolArgs(uri);
+				if (!args.server) {
+					this._notificationService.warn(localize('warnServerRequired', "Cannot connect as no server information was provided"));
+					return true;
+				}
+				let isOpenOk = await this.confirmConnect(args);
+				if (isOpenOk) {
+					await this.processCommandLine(args);
+				}
+			} catch (err) {
+				this._notificationService.error(localize('errConnectUrl', "Could not open URL due to error {0}", getErrorMessage(err)));
+			}
+			// Handled either way
+			return true;
+		}
+
+		return false;
+	}
+
+	private async confirmConnect(args: SqlArgs): Promise<boolean> {
+		let detail = args && args.server ? localize('connectServerDetail', "This will connect to server {0}", args.server) : '';
+		const result = await this.dialogService.confirm({
+			message: localize('confirmConnect', "Are you sure you want to connect?"),
+			detail: detail,
+			primaryButton: localize('open', "&&Open"),
+			type: 'question'
+		});
+
+		if (result.confirmed) {
+			return true;
+		}
+		return false;
+	}
+
+	private parseProtocolArgs(uri: URI): SqlArgs {
+		let args: SqlArgs = querystring.parse(uri.query);
+		// Clear out command, not supporting arbitrary command via this path
+		args.command = undefined;
+		return args;
+	}
+
 	// If an open and connectable query editor exists for the given URI, attach it to the connection profile
 	private async processFile(uriString: string, profile: IConnectionProfile, warnOnConnectFailure: boolean): Promise<void> {
 		let activeEditor = this._editorService.editors.filter(v => v.getResource().toString() === uriString).pop();
@@ -151,11 +219,11 @@ export class CommandLineWorkbenchContribution implements IWorkbenchContribution
 		}
 	}

-	private readProfileFromArgs(args: ParsedArgs) {
+	private readProfileFromArgs(args: SqlArgs) {
 		let profile = new ConnectionProfile(this._capabilitiesService, null);
 		// We want connection store to use any matching password it finds
 		profile.savePassword = true;
-		profile.providerName = Constants.mssqlProviderName;
+		profile.providerName = args.provider ? args.provider : Constants.mssqlProviderName;
 		profile.serverName = args.server;
 		profile.databaseName = args.database ? args.database : '';
 		profile.userName = args.user ? args.user : '';
--- a/src/sql/workbench/parts/commandLine/test/electron-browser/commandLine.test.ts
+++ b/src/sql/workbench/parts/commandLine/test/electron-browser/commandLine.test.ts
@@ -21,11 +21,14 @@ import { TestCommandService } from 'vs/editor/test/browser/editorTestServices';
 import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { assertThrowsAsync } from 'sqltest/utils/testUtils';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
-import { TestEditorService } from 'vs/workbench/test/workbenchTestServices';
+import { TestEditorService, TestDialogService } from 'vs/workbench/test/workbenchTestServices';
 import { QueryInput, QueryEditorState } from 'sql/workbench/parts/query/common/queryInput';
 import { URI } from 'vs/base/common/uri';
 import { ILogService, NullLogService } from 'vs/platform/log/common/log';
 import { TestConfigurationService } from 'vs/platform/configuration/test/common/testConfigurationService';
+import { IDialogService } from 'vs/platform/dialogs/common/dialogs';
+import { INotificationService } from 'vs/platform/notification/common/notification';
+import { TestNotificationService } from 'vs/platform/notification/test/common/testNotificationService';

 class TestParsedArgs implements ParsedArgs {
 	[arg: string]: any;
@@ -93,7 +96,6 @@ class TestParsedArgs implements ParsedArgs {
 suite('commandLineService tests', () => {

 	let capabilitiesService: TestCapabilitiesService;
-
 	setup(() => {
 		capabilitiesService = new TestCapabilitiesService();
 	});
@@ -104,7 +106,9 @@ suite('commandLineService tests', () => {
 		capabilitiesService?: ICapabilitiesService,
 		commandService?: ICommandService,
 		editorService?: IEditorService,
-		logService?: ILogService
+		logService?: ILogService,
+		dialogService?: IDialogService,
+		notificationService?: INotificationService
 	): CommandLineWorkbenchContribution {
 		return new CommandLineWorkbenchContribution(
 			capabilitiesService,
@@ -113,9 +117,11 @@ suite('commandLineService tests', () => {
 			editorService,
 			commandService,
 			configurationService,
-			undefined,
+			notificationService,
 			logService,
-			undefined
+			undefined,
+			undefined,
+			dialogService
 		);
 	}

@@ -389,4 +395,165 @@ suite('commandLineService tests', () => {
 		queryInput.verifyAll();
 		connectionManagementService.verifyAll();
 	});
+
+	suite('URL Handler', () => {
+
+		let dialogService: TypeMoq.Mock<TestDialogService>;
+
+		setup(() => {
+			dialogService = TypeMoq.Mock.ofType(TestDialogService);
+		});
+
+
+		test('handleUrl ignores non-connect URLs', async () => {
+			// Given a URI pointing to a server
+			let uri: URI = URI.parse('azuredatastudio://file?server=myserver&database=mydatabase&user=myuser');
+
+			const connectionManagementService: TypeMoq.Mock<IConnectionManagementService>
+				= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
+			const configurationService = getConfigurationServiceMock(true);
+			const logService = new NullLogService();
+			let contribution = getCommandLineContribution(connectionManagementService.object, configurationService.object, capabilitiesService, undefined, undefined, logService, dialogService.object);
+
+			// When I call the URL handler and user confirms they should connect
+			dialogService.setup(d => d.confirm(TypeMoq.It.isAny())).returns(() => Promise.resolve({ confirmed: true }));
+			let result = await contribution.handleURL(uri);
+
+			// Then I expect connection management service to have been called
+			assert.equal(result, false, 'Expected URL to be ignored');
+		});
+
+		test('handleUrl opens a new connection if a server name is passed', async () => {
+			// Given a URI pointing to a server
+			let uri: URI = URI.parse('azuredatastudio://connect?server=myserver&database=mydatabase&user=myuser');
+
+			const connectionManagementService: TypeMoq.Mock<IConnectionManagementService>
+				= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
+
+			connectionManagementService.setup((c) => c.showConnectionDialog()).verifiable(TypeMoq.Times.never());
+			connectionManagementService.setup(c => c.hasRegisteredServers()).returns(() => true).verifiable(TypeMoq.Times.atMostOnce());
+			connectionManagementService.setup(c => c.getConnectionGroups(TypeMoq.It.isAny())).returns(() => []);
+			let originalProfile: IConnectionProfile = undefined;
+			connectionManagementService.setup(c => c.connectIfNotConnected(TypeMoq.It.is<ConnectionProfile>(p => p.serverName === 'myserver' && p.authenticationType === Constants.sqlLogin), 'connection', true))
+				.returns((conn) => {
+					originalProfile = conn;
+					return Promise.resolve('unused');
+				})
+				.verifiable(TypeMoq.Times.once());
+			connectionManagementService.setup(c => c.getConnectionProfileById(TypeMoq.It.isAnyString())).returns(() => originalProfile);
+			const configurationService = getConfigurationServiceMock(true);
+			const logService = new NullLogService();
+			let contribution = getCommandLineContribution(connectionManagementService.object, configurationService.object, capabilitiesService, undefined, undefined, logService, dialogService.object);
+
+			// When I call the URL handler and user confirms they should connect
+			dialogService.setup(d => d.confirm(TypeMoq.It.isAny())).returns(() => Promise.resolve({ confirmed: true }));
+			let result = await contribution.handleURL(uri);
+
+			// Then I expect connection management service to have been called
+			assert.equal(result, true, 'Expected URL to be handled');
+			connectionManagementService.verifyAll();
+		});
+
+		test('handleUrl does nothing if a user does not confirm', async () => {
+			// Given a URI pointing to a server
+			let uri: URI = URI.parse('azuredatastudio://connect?server=myserver&database=mydatabase&user=myuser');
+
+			const connectionManagementService: TypeMoq.Mock<IConnectionManagementService>
+				= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
+
+			connectionManagementService.setup((c) => c.showConnectionDialog()).verifiable(TypeMoq.Times.never());
+			connectionManagementService.setup(c => c.hasRegisteredServers()).returns(() => true).verifiable(TypeMoq.Times.atMostOnce());
+			connectionManagementService.setup(c => c.getConnectionGroups(TypeMoq.It.isAny())).returns(() => []);
+			let originalProfile: IConnectionProfile = undefined;
+			connectionManagementService.setup(c => c.connectIfNotConnected(TypeMoq.It.is<ConnectionProfile>(p => p.serverName === 'myserver' && p.authenticationType === Constants.sqlLogin), 'connection', true))
+				.returns((conn) => {
+					originalProfile = conn;
+					return Promise.resolve('unused');
+				})
+				// Note: should not run since we expect to cancel before this
+				.verifiable(TypeMoq.Times.never());
+			connectionManagementService.setup(c => c.getConnectionProfileById(TypeMoq.It.isAnyString())).returns(() => originalProfile);
+			const configurationService = getConfigurationServiceMock(true);
+			const logService = new NullLogService();
+			let contribution = getCommandLineContribution(connectionManagementService.object, configurationService.object, capabilitiesService, undefined, undefined, logService, dialogService.object);
+
+			// When I call the URL handler and user says no on confirmation dialog
+			dialogService.setup(d => d.confirm(TypeMoq.It.isAny())).returns(() => Promise.resolve({ confirmed: false }));
+			let result = await contribution.handleURL(uri);
+
+			// Then I expect no connection, but the URL should still be handled
+			assert.equal(result, true, 'Expected URL to be handled');
+			connectionManagementService.verifyAll();
+		});
+
+		test('handleUrl ignores commands', async () => {
+			// Given I pass a command
+			let uri: URI = URI.parse('azuredatastudio://connect?command=mycommand');
+
+			const connectionManagementService: TypeMoq.Mock<IConnectionManagementService>
+				= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
+			const commandService: TypeMoq.Mock<ICommandService> = TypeMoq.Mock.ofType<ICommandService>(TestCommandService);
+
+			connectionManagementService.setup(c => c.hasRegisteredServers()).returns(() => true);
+			commandService.setup(c => c.executeCommand('mycommand'))
+				.returns(() => Promise.resolve())
+				.verifiable(TypeMoq.Times.never());
+			const configurationService = getConfigurationServiceMock(true);
+
+			const notificationService = TypeMoq.Mock.ofType(TestNotificationService);
+			notificationService.setup(n => n.warn(TypeMoq.It.isAny())).returns(() => undefined)
+				.verifiable(TypeMoq.Times.once());
+			let contribution = getCommandLineContribution(connectionManagementService.object, configurationService.object, capabilitiesService, commandService.object, undefined, new NullLogService(), dialogService.object, notificationService.object);
+
+			// When I handle the command URL
+			let result = await contribution.handleURL(uri);
+
+			// Then command service should not have been called, and instead connection should be handled
+			assert.equal(result, true);
+			commandService.verifyAll();
+			notificationService.verifyAll();
+		});
+
+		test('handleUrl ignores commands and connects', async () => {
+			// Given I pass a command
+			let uri: URI = URI.parse('azuredatastudio://connect?command=mycommand&server=myserver&database=mydatabase&user=myuser');
+
+			const connectionManagementService: TypeMoq.Mock<IConnectionManagementService>
+				= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
+			const commandService: TypeMoq.Mock<ICommandService> = TypeMoq.Mock.ofType<ICommandService>(TestCommandService);
+
+			connectionManagementService.setup((c) => c.showConnectionDialog()).verifiable(TypeMoq.Times.never());
+			connectionManagementService.setup(c => c.hasRegisteredServers()).returns(() => true).verifiable(TypeMoq.Times.atMostOnce());
+			connectionManagementService.setup(c => c.getConnectionGroups(TypeMoq.It.isAny())).returns(() => []);
+			let originalProfile: IConnectionProfile = undefined;
+			connectionManagementService.setup(c => c.connectIfNotConnected(TypeMoq.It.is<ConnectionProfile>(p => p.serverName === 'myserver' && p.authenticationType === Constants.sqlLogin), 'connection', true))
+				.returns((conn) => {
+					originalProfile = conn;
+					return Promise.resolve('unused');
+				})
+				.verifiable(TypeMoq.Times.once());
+
+			commandService.setup(c => c.executeCommand('mycommand'))
+				.returns(() => Promise.resolve())
+				.verifiable(TypeMoq.Times.never());
+			const configurationService = getConfigurationServiceMock(true);
+
+			const notificationService = TypeMoq.Mock.ofType(TestNotificationService);
+			notificationService.setup(n => n.warn(TypeMoq.It.isAny())).returns(() => undefined)
+				.verifiable(TypeMoq.Times.never());
+			let contribution = getCommandLineContribution(connectionManagementService.object, configurationService.object, capabilitiesService, commandService.object, undefined, new NullLogService(), dialogService.object, notificationService.object);
+
+			// When I handle the command URL
+			dialogService.setup(d => d.confirm(TypeMoq.It.isAny())).returns(() => Promise.resolve({ confirmed: true }));
+			let result = await contribution.handleURL(uri);
+
+			// Then command service should not have been called, and instead connection should be handled
+			assert.equal(result, true);
+			commandService.verifyAll();
+			notificationService.verifyAll();
+			connectionManagementService.verifyAll();
+		});
+
+
+	});
 });