add uri transforms to query and connection (#6586)

src/sql/workbench/api/common/extHostDataProtocol.ts

@@ -10,6 +10,8 @@ import { IMainContext } from 'vs/workbench/api/common/extHost.protocol';
 import { Disposable } from 'vs/workbench/api/common/extHostTypes';
 import { SqlMainContext, MainThreadDataProtocolShape, ExtHostDataProtocolShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
 import { DataProviderType } from 'sql/workbench/api/common/sqlExtHostTypes';
+import { IURITransformer } from 'vs/base/common/uriIpc';
+import { URI } from 'vs/base/common/uri';
 
 export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 
@@ -24,7 +26,8 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 	private _providersByType = new Map<azdata.DataProviderType, azdata.DataProvider[]>();
 
 	constructor(
-		mainContext: IMainContext
+		mainContext: IMainContext,
+		private uriTransformer: IURITransformer
 	) {
 		super();
 		this._proxy = mainContext.getProxy(SqlMainContext.MainThreadDataProtocol);
@@ -185,7 +188,8 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 
 	// Connection Management handlers
 	$connect(handle: number, connectionUri: string, connection: azdata.ConnectionInfo): Thenable<boolean> {
-		return this._resolveProvider<azdata.ConnectionProvider>(handle).connect(connectionUri, connection);
+		const transformedUri = URI.from(this.uriTransformer.transformIncoming(URI.parse(connectionUri))).toString();
+		return this._resolveProvider<azdata.ConnectionProvider>(handle).connect(transformedUri, connection);
 	}
 
 	$disconnect(handle: number, connectionUri: string): Thenable<boolean> {
@@ -222,6 +226,7 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 	}
 
 	$onConnectComplete(handle: number, connectionInfoSummary: azdata.ConnectionInfoSummary): void {
+		connectionInfoSummary.ownerUri = URI.from(this.uriTransformer.transformOutgoing(URI.parse(connectionInfoSummary.ownerUri))).toString();
 		this._proxy.$onConnectionComplete(handle, connectionInfoSummary);
 	}
 
@@ -245,7 +250,8 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 	}
 
 	$runQuery(handle: number, ownerUri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Thenable<void> {
-		return this._resolveProvider<azdata.QueryProvider>(handle).runQuery(ownerUri, selection, runOptions);
+		const transformedUri = URI.from(this.uriTransformer.transformIncoming(URI.parse(ownerUri))).toString();
+		return this._resolveProvider<azdata.QueryProvider>(handle).runQuery(transformedUri, selection, runOptions);
 	}
 
 	$runQueryStatement(handle: number, ownerUri: string, line: number, column: number): Thenable<void> {
@@ -273,29 +279,37 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 	}
 
 	$getQueryRows(handle: number, rowData: azdata.QueryExecuteSubsetParams): Thenable<azdata.QueryExecuteSubsetResult> {
+		rowData.ownerUri = URI.from(this.uriTransformer.transformIncoming(URI.parse(rowData.ownerUri))).toString();
 		return this._resolveProvider<azdata.QueryProvider>(handle).getQueryRows(rowData);
 	}
 
 	$disposeQuery(handle: number, ownerUri: string): Thenable<void> {
-		return this._resolveProvider<azdata.QueryProvider>(handle).disposeQuery(ownerUri);
+		const transformedUri = URI.from(this.uriTransformer.transformOutgoing(URI.parse(ownerUri))).toString();
+		return this._resolveProvider<azdata.QueryProvider>(handle).disposeQuery(transformedUri);
 	}
 
 	$onQueryComplete(handle: number, result: azdata.QueryExecuteCompleteNotificationResult): void {
+		result.ownerUri = URI.from(this.uriTransformer.transformOutgoing(URI.parse(result.ownerUri))).toString();
 		this._proxy.$onQueryComplete(handle, result);
 	}
 	$onBatchStart(handle: number, batchInfo: azdata.QueryExecuteBatchNotificationParams): void {
+		batchInfo.ownerUri = URI.from(this.uriTransformer.transformOutgoing(URI.parse(batchInfo.ownerUri))).toString();
 		this._proxy.$onBatchStart(handle, batchInfo);
 	}
 	$onBatchComplete(handle: number, batchInfo: azdata.QueryExecuteBatchNotificationParams): void {
+		batchInfo.ownerUri = URI.from(this.uriTransformer.transformOutgoing(URI.parse(batchInfo.ownerUri))).toString();
 		this._proxy.$onBatchComplete(handle, batchInfo);
 	}
 	$onResultSetAvailable(handle: number, resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void {
+		resultSetInfo.ownerUri = URI.from(this.uriTransformer.transformOutgoing(URI.parse(resultSetInfo.ownerUri))).toString();
 		this._proxy.$onResultSetAvailable(handle, resultSetInfo);
 	}
 	$onResultSetUpdated(handle: number, resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void {
+		resultSetInfo.ownerUri = URI.from(this.uriTransformer.transformOutgoing(URI.parse(resultSetInfo.ownerUri))).toString();
 		this._proxy.$onResultSetUpdated(handle, resultSetInfo);
 	}
 	$onQueryMessage(handle: number, message: azdata.QueryExecuteMessageParams): void {
+		message.ownerUri = URI.from(this.uriTransformer.transformOutgoing(URI.parse(message.ownerUri))).toString();
 		this._proxy.$onQueryMessage(handle, message);
 	}
 

src/sql/workbench/api/electron-browser/extensionHost.contribution.ts

@@ -3,10 +3,6 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { IWorkbenchContribution, IWorkbenchContributionsRegistry, Extensions as WorkbenchExtensions } from 'vs/workbench/common/contributions';
-import { Registry } from 'vs/platform/registry/common/platform';
-import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
-
 // --- SQL contributions
 import 'sql/workbench/api/browser/mainThreadConnectionManagement';
 import 'sql/workbench/api/browser/mainThreadCredentialManagement';
@@ -25,22 +21,3 @@ import 'sql/workbench/api/browser/mainThreadNotebook';
 import 'sql/workbench/api/browser/mainThreadNotebookDocumentsAndEditors';
 import 'sql/workbench/api/browser/mainThreadAccountManagement';
 import 'sql/workbench/api/browser/mainThreadExtensionManagement';
-import { LifecyclePhase } from 'vs/platform/lifecycle/common/lifecycle';
-
-export class SqlExtHostContribution implements IWorkbenchContribution {
-
-	constructor(
-		@IInstantiationService private instantiationService: IInstantiationService
-	) {
-	}
-
-	public getId(): string {
-		return 'sql.api.sqlExtHost';
-	}
-}
-
-// Register File Tracker
-Registry.as<IWorkbenchContributionsRegistry>(WorkbenchExtensions.Workbench).registerWorkbenchContribution(
-	SqlExtHostContribution,
-	LifecyclePhase.Restored
-);

src/sql/workbench/api/node/sqlExtHost.api.impl.ts

@@ -69,7 +69,7 @@ export function createApiFactory(
 	const extHostAccountManagement = rpcProtocol.set(SqlExtHostContext.ExtHostAccountManagement, new ExtHostAccountManagement(rpcProtocol));
 	const extHostConnectionManagement = rpcProtocol.set(SqlExtHostContext.ExtHostConnectionManagement, new ExtHostConnectionManagement(rpcProtocol));
 	const extHostCredentialManagement = rpcProtocol.set(SqlExtHostContext.ExtHostCredentialManagement, new ExtHostCredentialManagement(rpcProtocol));
-	const extHostDataProvider = rpcProtocol.set(SqlExtHostContext.ExtHostDataProtocol, new ExtHostDataProtocol(rpcProtocol));
+	const extHostDataProvider = rpcProtocol.set(SqlExtHostContext.ExtHostDataProtocol, new ExtHostDataProtocol(rpcProtocol, uriTransformer));
 	const extHostObjectExplorer = rpcProtocol.set(SqlExtHostContext.ExtHostObjectExplorer, new ExtHostObjectExplorer(rpcProtocol));
 	const extHostSerializationProvider = rpcProtocol.set(SqlExtHostContext.ExtHostSerializationProvider, new ExtHostSerializationProvider(rpcProtocol));
 	const extHostResourceProvider = rpcProtocol.set(SqlExtHostContext.ExtHostResourceProvider, new ExtHostResourceProvider(rpcProtocol));

src/sql/workbench/test/electron-browser/api/extHostDataProtocol.test.ts

@@ -21,7 +21,7 @@ suite('ExtHostDataProtocol', () => {
 					$registerConnectionProvider: (providerId, handle) => Promise.resolve()
 				} as any;
 			}
-		} as any);
+		} as any, undefined);
 	});
 
 	test('Providers are exposed to other extensions', () => {

src/vs/workbench/workbench.main.ts

@@ -392,9 +392,8 @@ import 'vs/workbench/contrib/relauncher/electron-browser/relauncher.contribution
 // registerSingleton(ITaskService, TaskService, true);
 
 // Remote
-// {{SQL CARBON EDIT}} @anthonydresser comment our remote
+import 'vs/workbench/contrib/remote/common/remote.contribution';
-// import 'vs/workbench/contrib/remote/common/remote.contribution';
+import 'vs/workbench/contrib/remote/electron-browser/remote.contribution';
-// import 'vs/workbench/contrib/remote/electron-browser/remote.contribution';
 
 // Emmet
 // import 'vs/workbench/contrib/emmet/browser/emmet.contribution'; {{SQL CARBON EDIT}} @anthonydresser comment our emmet