Adding completion extension loading and execution logic (#833)

* Adding ICompletionExtension interface

* Adding extension loading and execution logic

* Fixing compilation error in VS 2017

* Using MEF for completion extension discovery

* using await on GetCompletionItems

* Adding an integration test for completion extension and update the completion extension interface

* Update the completion extension test

* Fix issues based on review comments

* Remove try/cache based on review comments, fix a integration test.

* More changes based on review comments

* Fixing SendResult logic for completion extension loading

* Only load completion extension from the assembly passed in, add more comments in the test

* Adding right assert messages in the test.

* More fixes based on review comments

* Dropping ICompletionExtensionProvider, load assembly only if it's loaded at the first time or updated since last load.

* Fix based on the latest review comments

* Adding missing TSQL functions in default completion list

* Update jsonrpc documentation for completion/extLoad

docs/guide/jsonrpc_protocol.md

@@ -66,6 +66,10 @@ The SQL Tools Service implements the following portion Language Service Protocol
 * :leftwards_arrow_with_hook: [textDocument/references](https://github.com/Microsoft/language-server-protocol/blob/master/protocol.md#textDocument_references) 
 * :leftwards_arrow_with_hook: [textDocument/definition](https://github.com/Microsoft/language-server-protocol/blob/master/protocol.md#textDocument_definition) 
 
+### Language Service Protocol Extensions
+
+* :leftwards_arrow_with_hook: [completion/extLoad](#completion_extLoad)
+
 # Message Protocol
 
 A message consists of two parts: a header section and the message body.  For now, there is
@@ -744,3 +748,37 @@ Save a resultset as JSON to a file.
         public string Messages { get; set; }
     }
 ```
+
+## Language Service Protocol Extensions
+
+### <a name="completion_extload"></a>`completion/extLoad`
+
+Load a completion extension.
+
+#### Request
+
+```typescript
+    public class CompletionExtensionParams
+    {
+        /// <summary>
+        /// Absolute path for the assembly containing the completion extension
+        /// </summary>
+        public string AssemblyPath { get; set; }
+
+        /// <summary>
+        /// The type name for the completion extension
+        /// </summary>
+        public string TypeName { get; set; }
+
+        /// <summary>
+        /// Property bag for initializing the completion extension
+        /// </summary>
+        public Dictionary<string, object> Properties { get; set; }
+    }
+```
+
+#### Response
+
+```typescript
+    bool
+```

sqltoolsservice.sln

@@ -95,6 +95,8 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.SqlTools.ManagedB
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.SqlTools.ManagedBatchParser.IntegrationTests", "test\Microsoft.SqlTools.ManagedBatchParser.IntegrationTests\Microsoft.SqlTools.ManagedBatchParser.IntegrationTests.csproj", "{D3696EFA-FB1E-4848-A726-FF7B168AFB96}"
 EndProject
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.SqlTools.Test.CompletionExtension", "test\CompletionExtSample\Microsoft.SqlTools.Test.CompletionExtension.csproj", "{0EC2B30C-0652-49AE-9594-85B3C3E9CA21}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -228,6 +230,12 @@ Global
 		{D3696EFA-FB1E-4848-A726-FF7B168AFB96}.Integration|Any CPU.Build.0 = Debug|Any CPU
 		{D3696EFA-FB1E-4848-A726-FF7B168AFB96}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{D3696EFA-FB1E-4848-A726-FF7B168AFB96}.Release|Any CPU.Build.0 = Release|Any CPU
+		{0EC2B30C-0652-49AE-9594-85B3C3E9CA21}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{0EC2B30C-0652-49AE-9594-85B3C3E9CA21}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{0EC2B30C-0652-49AE-9594-85B3C3E9CA21}.Integration|Any CPU.ActiveCfg = Debug|Any CPU
+		{0EC2B30C-0652-49AE-9594-85B3C3E9CA21}.Integration|Any CPU.Build.0 = Debug|Any CPU
+		{0EC2B30C-0652-49AE-9594-85B3C3E9CA21}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{0EC2B30C-0652-49AE-9594-85B3C3E9CA21}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -256,6 +264,7 @@ Global
 		{EF02F89F-417E-4A40-B7E6-B102EE2DF24D} = {2BBD7364-054F-4693-97CD-1C395E3E84A9}
 		{3F82F298-700A-48DF-8A69-D048DFBA782C} = {2BBD7364-054F-4693-97CD-1C395E3E84A9}
 		{D3696EFA-FB1E-4848-A726-FF7B168AFB96} = {AB9CA2B8-6F70-431C-8A1D-67479D8A7BE4}
+		{0EC2B30C-0652-49AE-9594-85B3C3E9CA21} = {AB9CA2B8-6F70-431C-8A1D-67479D8A7BE4}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {B31CDF4B-2851-45E5-8C5F-BE97125D9DD8}

src/Microsoft.SqlTools.Hosting.v2/Extensibility/ExtensionServiceProvider.cs

@@ -5,6 +5,7 @@
 
 using System;
 using System.Collections;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Composition.Convention;
 using System.Composition.Hosting;
@@ -29,7 +30,7 @@ namespace Microsoft.SqlTools.Hosting.Extensibility
     /// </summary>
     public class ExtensionServiceProvider : RegisteredServiceProvider
     {
-        private readonly Func<ConventionBuilder, ContainerConfiguration> config;
+        private Func<ConventionBuilder, ContainerConfiguration> config;
 
         public ExtensionServiceProvider(Func<ConventionBuilder, ContainerConfiguration> config)
         {
@@ -106,6 +107,21 @@ namespace Microsoft.SqlTools.Hosting.Extensibility
                 Register(() => store.GetExports<T>());
             }
         }
+
+        /// <summary>
+        /// Merges in new assemblies to the existing container configuration.
+        /// </summary>
+        public void AddAssembliesToConfiguration(IEnumerable<Assembly> assemblies)
+        {
+            Validate.IsNotNull(nameof(assemblies), assemblies);
+            var previousConfig = config;
+            this.config = conventions => {
+                // Chain in the existing configuration function's result, then include additional
+                // assemblies
+                ContainerConfiguration containerConfig = previousConfig(conventions);
+                return containerConfig.WithAssemblies(assemblies, conventions);
+            };
+        }
     }
 
     /// <summary>

src/Microsoft.SqlTools.Hosting/Extensibility/ExtensionServiceProvider.cs

@@ -5,6 +5,7 @@
 
 using System;
 using System.Collections;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Composition.Convention;
 using System.Composition.Hosting;
@@ -114,6 +115,22 @@ namespace Microsoft.SqlTools.Extensibility
                 base.Register<T>(() => store.GetExports<T>());
             }
         }
+
+        /// <summary>
+        /// Merges in new assemblies to the existing container configuration.
+        /// </summary>
+        public void AddAssembliesToConfiguration(IEnumerable<Assembly> assemblies)
+        {
+            Validate.IsNotNull(nameof(assemblies), assemblies);
+            var previousConfig = config;
+            this.config = conventions => {
+                // Chain in the existing configuration function's result, then include additional
+                // assemblies
+                ContainerConfiguration containerConfig = previousConfig(conventions);
+                return containerConfig.WithAssemblies(assemblies, conventions);
+            };
+        }
+
     }
 
     /// <summary>

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/AutoCompleteHelper.cs

@@ -30,13 +30,21 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
         private static readonly string[] DefaultCompletionText = new string[]
         {
+            "abs",
+            "acos",
             "all",
             "alter",
             "and",
             "apply",
+            "approx_count_distinct",
             "as",
             "asc",
+            "ascii",
+            "asin",
             "at",
+            "atan",
+            "atn2",
+            "avg",
             "backup",
             "begin",
             "between",
@@ -48,41 +56,64 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "call",
             "cascade",
             "case",
+            "cast",
             "catch",
+            "ceiling",
             "char",
             "character",
+            "charindex",
             "check",
             "checkpoint",
+            "checksum_agg",
             "close",
             "clustered",
+            "coalesce",
             "column",
             "columnstore",
             "commit",
+            "concat",
+            "concat_ws",
             "connect",
             "constraint",
             "continue",
+            "convert",
+            "cos",
+            "cot",
+            "count",
+            "count_big",
             "create",
             "cross",
             "current_date",
+            "current_timestamp",
+            "current_user",
             "cursor",
             "cursor_close_on_commit",
             "cursor_default",
             "data",
             "data_compression",
             "database",
+            "datalength",
             "date",
+            "dateadd",
+            "datediff",
+            "datefromparts",
+            "datename",
+            "datepart",
             "datetime",
             "datetime2",
+            "day",
             "days",
             "dbcc",
             "dec",
             "decimal",
             "declare",
             "default",
+            "degrees",
             "delete",
             "deny",
             "desc",
             "description",
+            "difference",
             "disabled",
             "disk",
             "distinct",
@@ -101,6 +132,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "execute",
             "exists",
             "exit",
+            "exp",
             "external",
             "fast_forward",
             "fetch",
@@ -111,18 +143,24 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "filter",
             "first",
             "float",
+            "floor",
             "for",
             "foreign",
+            "format",
             "from",
             "full",
             "function",
             "geography",
             "get",
+            "getdate",
+            "getutcdate",
             "global",
             "go",
             "goto",
             "grant",
             "group",
+            "grouping",
+            "grouping_id",
             "hash",
             "hashed",
             "having",
@@ -133,6 +171,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "identity",
             "identitycol",
             "if",
+            "iif",
             "image",
             "immediate",
             "include",
@@ -144,6 +183,9 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "integer",
             "intersect",
             "into",
+            "isdate",
+            "isnull",
+            "isnumeric",
             "isolation",
             "join",
             "json",
@@ -151,18 +193,26 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "language",
             "last",
             "left",
+            "len",
             "level",
             "lineno",
             "load",
             "local",
             "locate",
             "location",
+            "log",
+            "log10",
             "login",
+            "lower",
+            "ltrim",
             "masked",
+            "max",
             "maxdop",
             "merge",
             "message",
+            "min",
             "modify",
+            "month",
             "move",
             "namespace",
             "native_compilation",
@@ -177,6 +227,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "not",
             "now",
             "null",
+            "nullif",
             "numeric",
             "nvarchar",
             "object",
@@ -200,12 +251,15 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "partition",
             "password",
             "path",
+            "patindex",
             "percent",
             "percentage",
             "period",
             "persisted",
+            "pi",
             "plan",
             "policy",
+            "power",
             "precision",
             "predicate",
             "primary",
@@ -216,7 +270,10 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "public",
             "query_store",
             "quoted_identifier",
+            "quotename",
+            "radians",
             "raiserror",
+            "rand",
             "range",
             "raw",
             "read",
@@ -236,6 +293,8 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "relative",
             "remove",
             "reorganize",
+            "replace",
+            "replicate",
             "required",
             "restart",
             "restore",
@@ -243,14 +302,18 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "resume",
             "return",
             "returns",
+            "reverse",
             "revert",
             "revoke",
+            "right",
             "rollback",
             "rollup",
+            "round",
             "row",
             "rowcount",
             "rowguidcol",
             "rows",
+            "rtrim",
             "rule",
             "sample",
             "save",
@@ -266,14 +329,22 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "sequence",
             "server",
             "session",
+            "session_user",
+            "sessionproperty",
             "set",
             "sets",
             "setuser",
+            "sign",
             "simple",
+            "sin",
             "smallint",
             "smallmoney",
             "snapshot",
+            "soundex",
+            "space",
             "sql",
+            "sqrt",
+            "square",
             "standard",
             "start",
             "started",
@@ -283,12 +354,22 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "statistics",
             "statistics_norecompute",
             "status",
+            "stdev",
+            "stdevp",
             "stopped",
+            "str",
+            "string_agg",
+            "stuff",
+            "substring",
+            "sum",
+            "sysdatetime",
             "sysname",
             "system",
             "system_time",
+            "system_user",
             "table",
             "take",
+            "tan",
             "target",
             "then",
             "throw",
@@ -299,23 +380,30 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "top",
             "tran",
             "transaction",
+            "translate",
             "trigger",
+            "trim",
             "truncate",
             "try",
             "tsql",
             "type",
             "uncommitted",
+            "unicode",
             "union",
             "unique",
             "uniqueidentifier",
             "update",
             "updatetext",
+            "upper",
             "use",
             "user",
+            "user_name",
             "using",
             "value",
             "values",
+            "var",
             "varchar",
+            "varp",
             "version",
             "view",
             "waitfor",
@@ -328,6 +416,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             "writetext",
             "xact_abort",
             "xml",
+            "year",
         };
 
         /// <summary>

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/Completion/Extension/CompletionExtensionParams.cs

@@ -0,0 +1,32 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System;
+using System.Collections.Generic;
+using System.Text;
+
+namespace Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion.Extension
+{
+
+    [Serializable]
+    public class CompletionExtensionParams
+    {
+        /// <summary>
+        /// Absolute path for the assembly containing the completion extension
+        /// </summary>
+        public string AssemblyPath { get; set; }
+
+        /// <summary>
+        /// The type name for the completion extension
+        /// </summary>
+        public string TypeName { get; set; }
+
+        /// <summary>
+        /// Property bag for initializing the completion extension
+        /// </summary>
+        public Dictionary<string, object> Properties { get; set; }
+    }
+
+}

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/Completion/Extension/ICompletionExtension.cs

@@ -0,0 +1,40 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts;
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion.Extension
+{
+    public interface ICompletionExtension : IDisposable
+    {
+        /// <summary>
+        /// Unique name for the extension
+        /// </summary>
+        string Name { get; }
+
+        /// <summary>
+        /// Method for initializing the extension, this is called once when the extension is loaded
+        /// </summary>
+        /// <param name="properties">Parameters needed by the extension</param>
+        /// <param name="cancelToken">Cancellation token used to indicate that the initialization should be cancelled</param>
+        /// <returns></returns>
+        Task Initialize(IReadOnlyDictionary<string, object> properties, CancellationToken token);
+
+        /// <summary>
+        /// Handles the completion request, returning the modified CompletionItemList if used
+        /// </summary>
+        /// <param name="connInfo">Connection information for the completion session</param>
+        /// <param name="scriptDocumentInfo">Script parsing information</param>
+        /// <param name="completions">Current completion list</param>
+        /// <param name="cancelToken">Token used to indicate that the completion request should be cancelled</param>
+        /// <returns></returns>
+        Task<CompletionItem[]> HandleCompletionAsync(ConnectionInfo connInfo, ScriptDocumentInfo scriptDocumentInfo, CompletionItem[] completions, CancellationToken cancelToken);
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/Contracts/Completion.cs

@@ -5,6 +5,7 @@
 
 using System.Diagnostics;
 using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+using Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion.Extension;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts
@@ -23,6 +24,13 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts
             RequestType<CompletionItem, CompletionItem>.Create("completionItem/resolve");
     }
 
+    public class CompletionExtLoadRequest
+    {
+        public static readonly
+            RequestType<CompletionExtensionParams, bool> Type =
+            RequestType<CompletionExtensionParams, bool>.Create("completion/extLoad");
+    }
+
     public enum CompletionItemKind
     {
         	Text = 1,
@@ -45,6 +53,19 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts
             Reference = 18
     }
 
+    public class Command
+    {
+        /// <summary>
+        /// The identifier of the actual command handler, like `vsintellicode.completionItemSelected`.
+        /// </summary>
+        public string CommandStr { get; set; }
+
+        /// <summary>
+        /// Arguments that the command handler should be invoked with.
+        /// </summary>
+        public object[] Arguments { get; set; }
+    }
+
     [DebuggerDisplay("Kind = {Kind.ToString()}, Label = {Label}, Detail = {Detail}")]
     public class CompletionItem
     {
@@ -74,5 +95,15 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts
         /// resolve request.
         /// </summary>
         public object Data { get; set; }
+
+        /// <summary>
+        /// Exposing a command field for a completion item for passing telemetry
+        /// </summary>
+        public Command Command { get; set; }
+
+        /// <summary>
+        /// Whether this completion item is preselected or not
+        /// </summary>
+        public bool? Preselect { get; set; }
     }
 }

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/LanguageService.cs

@@ -7,7 +7,10 @@ using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.IO;
 using System.Linq;
+using System.Reflection;
+using System.Runtime.Loader;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.SqlServer.Management.Common;
@@ -17,13 +20,14 @@ using Microsoft.SqlServer.Management.SqlParser.Common;
 using Microsoft.SqlServer.Management.SqlParser.Intellisense;
 using Microsoft.SqlServer.Management.SqlParser.Parser;
 using Microsoft.SqlServer.Management.SqlParser.SqlCodeDom;
+using Microsoft.SqlTools.Extensibility;
 using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Hosting;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion;
+using Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion.Extension;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts;
-using Microsoft.SqlTools.ServiceLayer.QueryExecution;
 using Microsoft.SqlTools.ServiceLayer.Scripting;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Utility;
@@ -72,6 +76,10 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
         internal const int PeekDefinitionTimeout = 10 * OneSecond;
 
+        internal const int ExtensionLoadingTimeout = 10 * OneSecond;
+
+        internal const int CompletionExtTimeout = 200;
+
         private ConnectionService connectionService = null;
 
         private WorkspaceService<SqlToolsSettings> workspaceServiceInstance;
@@ -95,6 +103,9 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         private Lazy<Dictionary<string, ScriptParseInfo>> scriptParseInfoMap
             = new Lazy<Dictionary<string, ScriptParseInfo>>(() => new Dictionary<string, ScriptParseInfo>());
 
+        private readonly ConcurrentDictionary<string, ICompletionExtension> completionExtensions = new ConcurrentDictionary<string, ICompletionExtension>();
+        private readonly ConcurrentDictionary<string, DateTime> extAssemblyLastUpdateTime = new ConcurrentDictionary<string, DateTime>();
+
         /// <summary>
         /// Gets a mapping dictionary for SQL file URIs to ScriptParseInfo objects
         /// </summary>
@@ -245,6 +256,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             serviceHost.SetRequestHandler(CompletionRequest.Type, HandleCompletionRequest);
             serviceHost.SetRequestHandler(DefinitionRequest.Type, HandleDefinitionRequest);
             serviceHost.SetRequestHandler(SyntaxParseRequest.Type, HandleSyntaxParseRequest);
+            serviceHost.SetRequestHandler(CompletionExtLoadRequest.Type, HandleCompletionExtLoadRequest);
             serviceHost.SetEventHandler(RebuildIntelliSenseNotification.Type, HandleRebuildIntelliSenseNotification);
             serviceHost.SetEventHandler(LanguageFlavorChangeNotification.Type, HandleDidChangeLanguageFlavorNotification);
 
@@ -283,6 +295,90 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
         #region Request Handlers
 
+        /// <summary>
+        /// Completion extension load request callback
+        /// </summary>
+        /// <param name="param"></param>
+        /// <param name="requestContext"></param>
+        /// <returns></returns>
+        internal async Task HandleCompletionExtLoadRequest(CompletionExtensionParams param, RequestContext<bool> requestContext)
+        {
+            try
+            {
+                //register the new assembly
+                var serviceProvider = (ExtensionServiceProvider)ServiceHostInstance.ServiceProvider;
+                var assembly = AssemblyLoadContext.Default.LoadFromAssemblyPath(param.AssemblyPath);
+                var assemblies = new Assembly[] { assembly };
+                serviceProvider.AddAssembliesToConfiguration(assemblies);
+                foreach (var ext in serviceProvider.GetServices<ICompletionExtension>())
+                {
+                    var cancellationTokenSource = new CancellationTokenSource(ExtensionLoadingTimeout);
+                    var cancellationToken = cancellationTokenSource.Token;
+                    string extName = ext.Name;
+                    string extTypeName = ext.GetType().FullName;
+                    if (extTypeName != param.TypeName)
+                    {
+                        continue;
+                    }
+
+                    if (!CheckIfAssemblyShouldBeLoaded(param.AssemblyPath, extTypeName))
+                    {
+                        await requestContext.SendError(string.Format("Skip loading {0} because it's already loaded", param.AssemblyPath));
+                        return;
+                    }
+
+                    await ext.Initialize(param.Properties, cancellationToken).WithTimeout(ExtensionLoadingTimeout);
+                    cancellationTokenSource.Dispose();
+                    if (!string.IsNullOrEmpty(extName))
+                    {
+                        completionExtensions[extName] = ext;
+                        await requestContext.SendResult(true);
+                        return;
+                    }
+                    else
+                    {
+                        await requestContext.SendError(string.Format("Skip loading an unnamed completion extension from {0}", param.AssemblyPath));
+                        return;
+                    }
+                }
+            }
+            catch (Exception ex)
+            {
+                await requestContext.SendError(ex.Message);
+                return;
+            }
+
+            await requestContext.SendError(string.Format("Couldn't discover completion extension with type {0} in {1}", param.TypeName, param.AssemblyPath));
+        }
+
+        /// <summary>
+        /// Check whether a particular assembly should be reloaded based on
+        /// whether it's been updated since it was last loaded.
+        /// </summary>
+        /// <param name="assemblyPath">The assembly path</param>
+        /// <param name="extTypeName">The type loading from the assembly</param>
+        /// <returns></returns>
+        private bool CheckIfAssemblyShouldBeLoaded(string assemblyPath, string extTypeName)
+        {
+            var lastModified = File.GetLastWriteTime(assemblyPath);
+            if (extAssemblyLastUpdateTime.ContainsKey(extTypeName))
+            {
+                if (lastModified > extAssemblyLastUpdateTime[extTypeName])
+                {
+                    extAssemblyLastUpdateTime[extTypeName] = lastModified;
+                    return true;
+                }
+            }
+            else
+            {
+                extAssemblyLastUpdateTime[extTypeName] = lastModified;
+                return true;
+            }
+
+            return false;
+
+        }
+
         /// <summary>
         /// T-SQL syntax parse request callback
         /// </summary>
@@ -352,7 +448,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                         scriptFile.ClientFilePath,
                         out connInfo);
 
-                    var completionItems = GetCompletionItems(
+                    var completionItems = await GetCompletionItems(
                         textDocumentPosition, scriptFile, connInfo);
 
                     await requestContext.SendResult(completionItems);
@@ -1466,7 +1562,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// This method does not await cache builds since it expects to return quickly
         /// </summary>
         /// <param name="textDocumentPosition"></param>
-        public CompletionItem[] GetCompletionItems(
+        public async Task<CompletionItem[]> GetCompletionItems(
             TextDocumentPosition textDocumentPosition,
             ScriptFile scriptFile,
             ConnectionInfo connInfo)
@@ -1482,7 +1578,11 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
             if (scriptParseInfo == null)
             {
-                return AutoCompleteHelper.GetDefaultCompletionItems(ScriptDocumentInfo.CreateDefaultDocumentInfo(textDocumentPosition, scriptFile), useLowerCaseSuggestions);
+                var scriptDocInfo = ScriptDocumentInfo.CreateDefaultDocumentInfo(textDocumentPosition, scriptFile);
+                resultCompletionItems = AutoCompleteHelper.GetDefaultCompletionItems(scriptDocInfo, useLowerCaseSuggestions);
+                //call completion extensions only for default completion list
+                resultCompletionItems = await ApplyCompletionExtensions(connInfo, resultCompletionItems, scriptDocInfo);
+                return resultCompletionItems;
             }
 
             ScriptDocumentInfo scriptDocumentInfo = new ScriptDocumentInfo(textDocumentPosition, scriptFile, scriptParseInfo);
@@ -1496,7 +1596,10 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             // if the parse failed then return the default list
             if (scriptParseInfo.ParseResult == null)
             {
-                return AutoCompleteHelper.GetDefaultCompletionItems(scriptDocumentInfo, useLowerCaseSuggestions);
+                resultCompletionItems = AutoCompleteHelper.GetDefaultCompletionItems(scriptDocumentInfo, useLowerCaseSuggestions);
+                //call completion extensions only for default completion list
+                resultCompletionItems = await ApplyCompletionExtensions(connInfo, resultCompletionItems, scriptDocumentInfo);
+                return resultCompletionItems;
             }
             AutoCompletionResult result = completionService.CreateCompletions(connInfo, scriptDocumentInfo, useLowerCaseSuggestions);
             // cache the current script parse info object to resolve completions later
@@ -1507,6 +1610,38 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             if (resultCompletionItems == null)
             {
                 resultCompletionItems = AutoCompleteHelper.GetDefaultCompletionItems(scriptDocumentInfo, useLowerCaseSuggestions);
+                //call completion extensions only for default completion list
+                resultCompletionItems = await ApplyCompletionExtensions(connInfo, resultCompletionItems, scriptDocumentInfo);
+            }
+
+            return resultCompletionItems;
+        }
+
+        /// <summary>
+        /// Run all completion extensions
+        /// </summary>
+        /// <param name="connInfo"></param>
+        /// <param name="resultCompletionItems"></param>
+        /// <param name="scriptDocumentInfo"></param>
+        /// <returns></returns>
+        private async Task<CompletionItem[]> ApplyCompletionExtensions(ConnectionInfo connInfo, CompletionItem[] resultCompletionItems, ScriptDocumentInfo scriptDocumentInfo)
+        {
+            //invoke the completion extensions
+            foreach (var completionExt in completionExtensions.Values)
+            {
+                var cancellationTokenSource = new CancellationTokenSource();
+                cancellationTokenSource.CancelAfter(CompletionExtTimeout);
+                var cancellationToken = cancellationTokenSource.Token;
+                try
+                {
+                    resultCompletionItems = await completionExt.HandleCompletionAsync(connInfo, scriptDocumentInfo, resultCompletionItems, cancellationToken).WithTimeout(CompletionExtTimeout);
+                }
+                catch (Exception e)
+                {
+                    Logger.Write(TraceEventType.Error, string.Format("Exception in calling completion extension {0}:\n{1}", completionExt.Name, e.ToString()));
+                }
+
+                cancellationTokenSource.Dispose();
             }
 
             return resultCompletionItems;

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/ScriptDocumentInfo.cs

@@ -14,7 +14,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion
     /// <summary>
     /// A class to calculate the numbers used by SQL parser using the text positions and content
     /// </summary>
-    internal class ScriptDocumentInfo
+    public class ScriptDocumentInfo
     {
         /// <summary>
         /// Create new instance

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/ScriptParseInfo.cs

@@ -12,7 +12,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
     /// <summary>
     /// Class for storing cached metadata regarding a parsed SQL file
     /// </summary>
-    internal class ScriptParseInfo
+    public class ScriptParseInfo
     {
         private object buildingMetadataLock = new object();
 

src/Microsoft.SqlTools.ServiceLayer/Utility/TaskExtensions.cs

@@ -6,6 +6,7 @@
 using System;
 using System.Diagnostics;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.Utility;
 
@@ -97,5 +98,38 @@ namespace Microsoft.SqlTools.ServiceLayer.Utility
             }
             Logger.Write(TraceEventType.Error, sb.ToString());
         }
+
+        /// <summary>
+        /// This will enforce time out to run an async task with returning result
+        /// </summary>
+        /// <typeparam name="TResult"></typeparam>
+        /// <param name="task">The async task to run</param>
+        /// <param name="timeout">Time out in milliseconds</param>
+        /// <returns></returns>
+        public static async Task<TResult> WithTimeout<TResult>(this Task<TResult> task, int timeout)
+        {
+            if (task == await Task.WhenAny(task, Task.Delay(timeout)))
+            {
+                return await task;
+            }
+            throw new TimeoutException();
+        }
+
+        /// <summary>
+        /// This will enforce time out to run an async task without returning result
+        /// </summary>
+        /// <typeparam name="TResult"></typeparam>
+        /// <param name="task">The async task to run</param>
+        /// <param name="timeout">Time out in milliseconds</param>
+        /// <returns></returns>
+        public static async Task WithTimeout(this Task task, int timeout)
+        {
+            if (task == await Task.WhenAny(task, Task.Delay(timeout)))
+            {
+                await task;
+                return;
+            }
+            throw new TimeoutException();
+        }
     }
 }

test/CompletionExtSample/CompletionExt.cs

@@ -0,0 +1,81 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System;
+using System.Collections.Generic;
+using System.ComponentModel.Composition;
+using System.Threading;
+using System.Threading.Tasks;
+using System.Linq;
+using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion;
+using Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion.Extension;
+using Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts;
+
+namespace Microsoft.SqlTools.Test.CompletionExtension
+{
+
+    [Export(typeof(ICompletionExtension))]
+    public class CompletionExt : ICompletionExtension
+    {
+        public string Name => "CompletionExt";
+
+        private string modelPath;
+
+        public CompletionExt()
+        {
+        }
+
+        void IDisposable.Dispose()
+        {
+        }
+
+        async Task<CompletionItem[]> ICompletionExtension.HandleCompletionAsync(ConnectionInfo connInfo, ScriptDocumentInfo scriptDocumentInfo, CompletionItem[] completions, CancellationToken token)
+        {
+            if (completions == null || completions == null || completions.Length == 0)
+            {
+                return completions;
+            }
+            return await Run(completions, token);
+        }
+
+        async Task ICompletionExtension.Initialize(IReadOnlyDictionary<string, object> properties, CancellationToken token)
+        {
+            modelPath = (string)properties["modelPath"];
+            await LoadModel(token).ConfigureAwait(false);
+            return;
+        }
+
+        private async Task LoadModel(CancellationToken token)
+        {
+            //loading model logic here
+            await Task.Delay(2000).ConfigureAwait(false); //for testing
+            token.ThrowIfCancellationRequested();
+            Console.WriteLine("Model loaded from: " + modelPath);
+        }
+
+        private async Task<CompletionItem[]> Run(CompletionItem[] completions, CancellationToken token)
+        {
+            Console.WriteLine("Enter ExecuteAsync");
+
+            var sortedItems = completions.OrderBy(item => item.SortText);
+            sortedItems.First().Preselect = true;
+            foreach(var item in sortedItems)
+            {
+                item.Command = new Command
+                {
+                    CommandStr = "vsintellicode.completionItemSelected",
+                    Arguments = new object[] { new Dictionary<string, string> { { "IsCommit", "True" } } }
+                };
+            }
+
+            //code to augment the default completion list
+            await Task.Delay(20); // for testing
+            token.ThrowIfCancellationRequested();
+            Console.WriteLine("Exit ExecuteAsync");
+            return sortedItems.ToArray();
+        }
+    }
+}

test/CompletionExtSample/Microsoft.SqlTools.Test.CompletionExtension.csproj

@@ -0,0 +1,11 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <Import Project="../../Common.props" />
+  <PropertyGroup>
+    <TargetFramework>netcoreapp2.2</TargetFramework>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="..\..\src\Microsoft.SqlTools.ServiceLayer\Microsoft.SqlTools.ServiceLayer.csproj" />
+  </ItemGroup>
+
+</Project>

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/LanguageServer/LanguageServiceTests.cs

@@ -3,14 +3,21 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Reflection;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.IntegrationTests.Utility;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices;
+using Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion.Extension;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Test.Common;
+using Microsoft.SqlTools.ServiceLayer.UnitTests.ServiceHost;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
+using Moq;
 using Xunit;
 
 namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.LanguageServer
@@ -75,9 +82,9 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.LanguageServer
             }
         }
 
-        // This test currently requires a live database connection to initialize
+        /// <summary>
-        // SMO connected metadata provider.  Since we don't want a live DB dependency
+        /// This test tests auto completion
-        // in the CI unit tests this scenario is currently disabled.
+        /// </summary>
         [Fact]
         public void AutoCompleteFindCompletions()
         {
@@ -90,11 +97,116 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.LanguageServer
             var completions = autoCompleteService.GetCompletionItems(
                 result.TextDocumentPosition,
                 result.ScriptFile,
-                result.ConnectionInfo);
+                result.ConnectionInfo).Result;
 
             Assert.True(completions.Length > 0);
         }
 
+        public static string AssemblyDirectory
+        {
+            get
+            {
+                string codeBase = Assembly.GetExecutingAssembly().CodeBase;
+                UriBuilder uri = new UriBuilder(codeBase);
+                string path = Uri.UnescapeDataString(uri.Path);
+                return Path.GetDirectoryName(path);
+            }
+        }
+
+        /// <summary>
+        /// This test tests completion extension interface in following aspects
+        /// 1. Loading a sample completion extension assembly
+        /// 2. Initializing a completion extension implementation
+        /// 3. Excuting an auto completion with extension enabled
+        /// </summary>
+        [Fact]
+        public async void AutoCompleteWithExtension()
+        {
+            var result = GetLiveAutoCompleteTestObjects();
+
+            result.TextDocumentPosition.Position.Character = 10;
+            result.ScriptFile = ScriptFileTests.GetTestScriptFile("select * f");
+            result.TextDocumentPosition.TextDocument.Uri = result.ScriptFile.FilePath;
+
+            var autoCompleteService = LanguageService.Instance;
+            var requestContext = new Mock<SqlTools.Hosting.Protocol.RequestContext<bool>>();
+            requestContext.Setup(x => x.SendResult(It.IsAny<bool>()))
+                .Returns(Task.FromResult(true));
+            requestContext.Setup(x => x.SendError(It.IsAny<string>(), 0))
+                .Returns(Task.FromResult(true));
+
+            //Create completion extension parameters
+            var extensionParams = new CompletionExtensionParams()
+            {
+                AssemblyPath = Path.Combine(AssemblyDirectory, "Microsoft.SqlTools.Test.CompletionExtension.dll"),
+                TypeName = "Microsoft.SqlTools.Test.CompletionExtension.CompletionExt",
+                Properties = new Dictionary<string, object> { { "modelPath", "testModel" } }
+            };
+
+            //load and initialize completion extension, expect a success
+            await autoCompleteService.HandleCompletionExtLoadRequest(extensionParams, requestContext.Object);
+
+            requestContext.Verify(x => x.SendResult(It.IsAny<bool>()), Times.Once);
+            requestContext.Verify(x => x.SendError(It.IsAny<string>(), 0), Times.Never);
+
+            //Try to load the same completion extension second time, expect an error sent
+            await autoCompleteService.HandleCompletionExtLoadRequest(extensionParams, requestContext.Object);
+
+            requestContext.Verify(x => x.SendResult(It.IsAny<bool>()), Times.Once);
+            requestContext.Verify(x => x.SendError(It.IsAny<string>(), 0), Times.Once);
+
+            //Try to load the completion extension with new modified timestamp, expect a success
+            var assemblyCopyPath = CopyFileWithNewModifiedTime(extensionParams.AssemblyPath);
+            extensionParams = new CompletionExtensionParams()
+            {
+                AssemblyPath = assemblyCopyPath,
+                TypeName = "Microsoft.SqlTools.Test.CompletionExtension.CompletionExt",
+                Properties = new Dictionary<string, object> { { "modelPath", "testModel" } }
+            };
+            //load and initialize completion extension
+            await autoCompleteService.HandleCompletionExtLoadRequest(extensionParams, requestContext.Object);
+
+            requestContext.Verify(x => x.SendResult(It.IsAny<bool>()), Times.Exactly(2));
+            requestContext.Verify(x => x.SendError(It.IsAny<string>(), 0), Times.Once);
+
+            ScriptParseInfo scriptInfo = new ScriptParseInfo { IsConnected = true };
+            autoCompleteService.ParseAndBind(result.ScriptFile, result.ConnectionInfo);
+            scriptInfo.ConnectionKey = autoCompleteService.BindingQueue.AddConnectionContext(result.ConnectionInfo);
+
+            //Invoke auto completion with extension enabled
+            var completions = autoCompleteService.GetCompletionItems(
+                result.TextDocumentPosition,
+                result.ScriptFile,
+                result.ConnectionInfo).Result;
+
+            //Validate completion list is not empty
+            Assert.True(completions != null && completions.Length > 0, "The completion list is null or empty!");
+            //Validate the first completion item in the list is preselected
+            Assert.True(completions[0].Preselect.HasValue && completions[0].Preselect.Value, "Preselect is not set properly in the first completion item by the completion extension!");
+            //Validate the Command object attached to the completion item by the extension
+            Assert.True(completions[0].Command != null && completions[0].Command.CommandStr == "vsintellicode.completionItemSelected", "Command is not set properly in the first completion item by the completion extension!");
+
+            //clean up the temp file
+            File.Delete(assemblyCopyPath);
+        }
+
+        /// <summary>
+        /// Make a copy of a file and update the last modified time
+        /// </summary>
+        /// <param name="filePath"></param>
+        /// <returns></returns>
+        private string CopyFileWithNewModifiedTime(string filePath)
+        {
+            var tempPath = Path.Combine(Path.GetTempPath(), Path.GetFileName(filePath));
+            if (File.Exists(tempPath))
+            {
+                File.Delete(tempPath);
+            }
+            File.Copy(filePath, tempPath);
+            File.SetLastWriteTimeUtc(tempPath, DateTime.UtcNow);
+            return tempPath;
+        }
+
         /// <summary>
         /// Verify that GetSignatureHelp returns not null when the provided TextDocumentPosition
         /// has an associated ScriptParseInfo and the provided query has a function that should
@@ -191,7 +303,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.LanguageServer
                     };
 
                 // First check that we don't have any items in the completion list as expected
-                var initialCompletionItems = langService.GetCompletionItems(
+                var initialCompletionItems = await langService.GetCompletionItems(
                     textDocumentPosition, connectionInfoResult.ScriptFile, connectionInfoResult.ConnectionInfo);
 
                 Assert.True(initialCompletionItems.Length == 0, $"Should not have any completion items initially. Actual : [{string.Join(',', initialCompletionItems.Select(ci => ci.Label))}]");
@@ -205,7 +317,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.LanguageServer
                     new TestEventContext());
 
                 // Now we should expect to see the item show up in the completion list
-                var afterTableCreationCompletionItems = langService.GetCompletionItems(
+                var afterTableCreationCompletionItems = await langService.GetCompletionItems(
                     textDocumentPosition, connectionInfoResult.ScriptFile, connectionInfoResult.ConnectionInfo);
 
                 Assert.True(afterTableCreationCompletionItems.Length == 1, $"Should only have a single completion item after rebuilding Intellisense cache. Actual : [{string.Join(',', initialCompletionItems.Select(ci => ci.Label))}]");

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/Microsoft.SqlTools.ServiceLayer.IntegrationTests.csproj

@@ -27,6 +27,7 @@
     <ProjectReference Include="../Microsoft.SqlTools.ServiceLayer.Test.Common/Microsoft.SqlTools.ServiceLayer.Test.Common.csproj" />
     <ProjectReference Include="../../src/Microsoft.SqlTools.ManagedBatchParser/Microsoft.SqlTools.ManagedBatchParser.csproj" />
     <ProjectReference Include="../Microsoft.SqlTools.ServiceLayer.UnitTests/Microsoft.SqlTools.ServiceLayer.UnitTests.csproj" />
+    <ProjectReference Include="..\CompletionExtSample\Microsoft.SqlTools.Test.CompletionExtension.csproj" />
   </ItemGroup>
   <ItemGroup>
     <PackageReference Include="System.Net.Http" Version="4.3.1" />

test/Microsoft.SqlTools.ServiceLayer.UnitTests/LanguageServer/AutocompleteTests.cs

@@ -117,7 +117,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.LanguageServer
         {
             InitializeTestObjects();
             textDocument.TextDocument.Uri = "somethinggoeshere";
-            Assert.True(langService.GetCompletionItems(textDocument, scriptFile.Object, null).Length > 0);
+            Assert.True(langService.GetCompletionItems(textDocument, scriptFile.Object, null).Result.Length > 0);
         }
 
         [Fact]

test/Microsoft.SqlTools.ServiceLayer.UnitTests/ServiceHost/ScriptFileTests.cs

@@ -22,7 +22,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ServiceHost
             "SELECT * FROM sys.objects as o2" + Environment.NewLine +
             "SELECT * FROM sys.objects as o3" + Environment.NewLine;
 
-        internal static ScriptFile GetTestScriptFile(string initialText = null)
+        public static ScriptFile GetTestScriptFile(string initialText = null)
         {
             if (initialText == null)
             {