Add STS root folder override (#16927)

* Add STS root folder override

* Display message to user

* Show once for any service

extensions/mssql/src/credentialstore/credentialstore.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { SqlOpsDataClient, ClientOptions, SqlOpsFeature } from 'dataprotocol-client';
-import { IConfig, ServerProvider } from '@microsoft/ads-service-downloader';
+import { IConfig } from '@microsoft/ads-service-downloader';
 import { ServerOptions, RPCMessageType, ClientCapabilities, ServerCapabilities, TransportKind } from 'vscode-languageclient';
 import { Disposable } from 'vscode';
 import * as UUID from 'vscode-languageclient/lib/utils/uuid';
@@ -77,17 +77,15 @@ export class CredentialStore {
 		}
 	}
 
-	public start() {
+	public async start(): Promise<void> {
-		let serverdownloader = new ServerProvider(this._config);
 		let clientOptions: ClientOptions = {
 			providerId: Constants.providerId,
 			features: [CredentialsFeature]
 		};
-		return serverdownloader.getOrDownloadServer().then(e => {
+		const serverPath = await Utils.getOrDownloadServer(this._config);
-			let serverOptions = this.generateServerOptions(e);
+		const serverOptions = this.generateServerOptions(serverPath);
 		this._client = new SqlOpsDataClient(Constants.serviceName, serverOptions, clientOptions);
 		this._client.start();
-		});
 	}
 
 	dispose() {

extensions/mssql/src/resourceProvider/resourceProvider.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as azdata from 'azdata';
-import { IConfig, ServerProvider } from '@microsoft/ads-service-downloader';
+import { IConfig } from '@microsoft/ads-service-downloader';
 import { SqlOpsDataClient, SqlOpsFeature, ClientOptions } from 'dataprotocol-client';
 import { ServerCapabilities, ClientCapabilities, RPCMessageType, ServerOptions, TransportKind } from 'vscode-languageclient';
 import * as UUID from 'vscode-languageclient/lib/utils/uuid';
@@ -81,17 +81,15 @@ export class AzureResourceProvider {
 		}
 	}
 
-	public start() {
+	public async start(): Promise<void> {
-		let serverdownloader = new ServerProvider(this._config);
 		let clientOptions: ClientOptions = {
 			providerId: Constants.providerId,
 			features: [FireWallFeature]
 		};
-		return serverdownloader.getOrDownloadServer().then(e => {
+		const serverPath = await Utils.getOrDownloadServer(this._config);
-			let serverOptions = this.generateServerOptions(e);
+		let serverOptions = this.generateServerOptions(serverPath);
 		this._client = new SqlOpsDataClient(Constants.serviceName, serverOptions, clientOptions);
 		this._client.start();
-		});
 	}
 
 	public dispose() {

extensions/mssql/src/sqlToolsServer.ts

@@ -3,12 +3,12 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ServerProvider, IConfig, Events } from '@microsoft/ads-service-downloader';
+import { IConfig, Events } from '@microsoft/ads-service-downloader';
 import { ServerOptions, TransportKind } from 'vscode-languageclient';
 import * as Constants from './constants';
 import * as vscode from 'vscode';
 import * as path from 'path';
-import { getCommonLaunchArgsAndCleanupOldLogFiles } from './utils';
+import { getCommonLaunchArgsAndCleanupOldLogFiles, getOrDownloadServer } from './utils';
 import { Telemetry, LanguageClientErrorHandler } from './telemetry';
 import { SqlOpsDataClient, ClientOptions } from 'dataprotocol-client';
 import { TelemetryFeature, AgentServicesFeature, SerializationFeature, AccountFeature, SqlAssessmentServicesFeature, ProfilerFeature } from './features';
@@ -82,10 +82,7 @@ export class SqlToolsServer {
 		this.config.installDirectory = path.join(configDir, this.config.installDirectory);
 		this.config.proxy = vscode.workspace.getConfiguration('http').get('proxy');
 		this.config.strictSSL = vscode.workspace.getConfiguration('http').get('proxyStrictSSL') || true;
-
+		return getOrDownloadServer(this.config, handleServerProviderEvent);
-		const serverdownloader = new ServerProvider(this.config);
-		serverdownloader.eventEmitter.onAny(generateHandleServerProviderEvent());
-		return serverdownloader.getOrDownloadServer();
 	}
 
 	private activateFeatures(context: AppContext): Promise<void> {
@@ -109,9 +106,8 @@ function generateServerOptions(logPath: string, executablePath: string): ServerO
 	return { command: executablePath, args: launchArgs, transport: TransportKind.stdio };
 }
 
-function generateHandleServerProviderEvent() {
+function handleServerProviderEvent(e: string, ...args: any[]): void {
 	let dots = 0;
-	return (e: string, ...args: any[]) => {
 	switch (e) {
 		case Events.INSTALL_START:
 			outputChannel.show(true);
@@ -141,7 +137,6 @@ function generateHandleServerProviderEvent() {
 			outputChannel.appendLine(localize('entryExtractedChannelMsg', "Extracted {0} ({1}/{2})", args[0], args[1], args[2]));
 			break;
 	}
-	};
 }
 
 function getClientOptions(context: AppContext): ClientOptions {

extensions/mssql/src/utils.ts

@@ -12,6 +12,8 @@ import * as os from 'os';
 import * as findRemoveSync from 'find-remove';
 import * as constants from './constants';
 import { promises as fs } from 'fs';
+import { IConfig, ServerProvider } from '@microsoft/ads-service-downloader';
+import { env } from 'process';
 
 const configTracingLevel = 'tracingLevel';
 const configLogRetentionMinutes = 'logRetentionMinutes';
@@ -304,3 +306,45 @@ export async function exists(path: string): Promise<boolean> {
 		return false;
 	}
 }
+
+const STS_OVERRIDE_ENV_VAR = 'ADS_SQLTOOLSSERVICE';
+let overrideMessageDisplayed = false;
+/**
+ * Gets the full path to the EXE for the specified tools service, downloading it in the process if necessary. The location
+ * for this can be overridden with an environment variable for debugging or other purposes.
+ * @param config The configuration values of the server to get/download
+ * @param handleServerEvent A callback for handling events from the server downloader
+ * @returns The path to the server exe
+ */
+export async function getOrDownloadServer(config: IConfig, handleServerEvent?: (e: string, ...args: any[]) => void): Promise<string> {
+	// This env var is used to override the base install location of STS - primarily to be used for debugging scenarios.
+	try {
+		const stsRootPath = env[STS_OVERRIDE_ENV_VAR];
+		if (stsRootPath) {
+			for (const exeFile of config.executableFiles) {
+				const serverFullPath = path.join(stsRootPath, exeFile);
+				if (await exists(serverFullPath)) {
+					const overrideMessage = `Using ${exeFile} from ${stsRootPath}`;
+					// Display message to the user so they know the override is active, but only once so we don't show too many
+					if (!overrideMessageDisplayed) {
+						overrideMessageDisplayed = true;
+						vscode.window.showInformationMessage(overrideMessage);
+					}
+					console.log(overrideMessage);
+					return serverFullPath;
+				}
+			}
+			console.warn(`Could not find valid SQL Tools Service EXE from ${JSON.stringify(config.executableFiles)} at ${stsRootPath}, falling back to config`);
+		}
+	} catch (err) {
+		console.warn('Unexpected error getting override path for SQL Tools Service client ', err);
+		// Fall back to config if something unexpected happens here
+	}
+
+	const serverdownloader = new ServerProvider(config);
+	if (handleServerEvent) {
+		serverdownloader.eventEmitter.onAny(handleServerEvent);
+	}
+
+	return serverdownloader.getOrDownloadServer();
+}