Merge pull request #54 from Microsoft/feature/autocompFixes

Merge autocomplete improvements to dev

src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs

@@ -73,7 +73,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// <summary>
         // Callback for ondisconnect handler
         /// </summary>
-        public delegate Task OnDisconnectHandler(ConnectionSummary summary);
+        public delegate Task OnDisconnectHandler(ConnectionSummary summary, string ownerUri);
 
         /// <summary>
         /// List of onconnection handlers
@@ -241,7 +241,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             // Invoke callback notifications
             foreach (var activity in this.onDisconnectActivities)
             {
-                activity(info.ConnectionDetails);
+                activity(info.ConnectionDetails, disconnectParams.OwnerUri);
             }
 
             // Success

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

@@ -4,8 +4,13 @@
 //
 
 using System.Collections.Generic;
+using Microsoft.SqlServer.Management.SqlParser.Binder;
 using Microsoft.SqlServer.Management.SqlParser.Intellisense;
+using Microsoft.SqlServer.Management.SqlParser.Parser;
+using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts;
+using Microsoft.SqlTools.ServiceLayer.SqlContext;
+using Microsoft.SqlTools.ServiceLayer.Workspace;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
@@ -487,7 +492,6 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                     Label = autoCompleteItem.Title,
                     Kind = CompletionItemKind.Variable,
                     Detail = autoCompleteItem.Title,
-                    Documentation = autoCompleteItem.Description,
                     TextEdit = new TextEdit
                     {
                         NewText = autoCompleteItem.Title,
@@ -510,5 +514,76 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
             return completions.ToArray();
         }
+
+        /// <summary>
+        /// Preinitialize the parser and binder with common metadata.
+        /// This should front load the long binding wait to the time the
+        /// connection is established.  Once this is completed other binding
+        /// requests should be faster.
+        /// </summary>
+        /// <param name="info"></param>
+        /// <param name="scriptInfo"></param>
+        internal static void PrepopulateCommonMetadata(ConnectionInfo info, ScriptParseInfo scriptInfo)
+        {
+            if (scriptInfo.IsConnected)
+            {
+                var scriptFile = WorkspaceService<SqlToolsSettings>.Instance.Workspace.GetFile(info.OwnerUri);                                
+                LanguageService.Instance.ParseAndBind(scriptFile, info);
+
+                if (scriptInfo.BuildingMetadataEvent.WaitOne(LanguageService.OnConnectionWaitTimeout))
+                {
+                    try
+                    {
+                        scriptInfo.BuildingMetadataEvent.Reset();
+
+                        // parse a simple statement that returns common metadata
+                        ParseResult parseResult = Parser.Parse(
+                            "select ", 
+                            scriptInfo.ParseOptions);
+
+                        List<ParseResult> parseResults = new List<ParseResult>();
+                        parseResults.Add(parseResult);
+                        scriptInfo.Binder.Bind(
+                            parseResults, 
+                            info.ConnectionDetails.DatabaseName, 
+                            BindMode.Batch);
+
+                        // get the completion list from SQL Parser
+                        var suggestions = Resolver.FindCompletions(
+                            parseResult, 1, 8, 
+                            scriptInfo.MetadataDisplayInfoProvider); 
+
+                        // this forces lazy evaluation of the suggestion metadata
+                        AutoCompleteHelper.ConvertDeclarationsToCompletionItems(suggestions, 1, 8, 8);
+
+                        parseResult = Parser.Parse(
+                            "exec ", 
+                            scriptInfo.ParseOptions);
+
+                        parseResults = new List<ParseResult>();
+                        parseResults.Add(parseResult);
+                        scriptInfo.Binder.Bind(
+                            parseResults, 
+                            info.ConnectionDetails.DatabaseName, 
+                            BindMode.Batch);
+
+                        // get the completion list from SQL Parser
+                        suggestions = Resolver.FindCompletions(
+                            parseResult, 1, 6, 
+                            scriptInfo.MetadataDisplayInfoProvider); 
+
+                        // this forces lazy evaluation of the suggestion metadata
+                        AutoCompleteHelper.ConvertDeclarationsToCompletionItems(suggestions, 1, 6, 6);
+                    }
+                    catch
+                    {
+                    }
+                    finally
+                    {
+                        scriptInfo.BuildingMetadataEvent.Set();
+                    }
+                }
+            }
+        }
     }
 }

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

@@ -8,6 +8,13 @@ using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.SqlServer.Management.Common;
+using Microsoft.SqlServer.Management.SqlParser;
+using Microsoft.SqlServer.Management.SqlParser.Binder;
+using Microsoft.SqlServer.Management.SqlParser.Intellisense;
+using Microsoft.SqlServer.Management.SqlParser.MetadataProvider;
+using Microsoft.SqlServer.Management.SqlParser.Parser;
+using Microsoft.SqlServer.Management.SmoMetadataProvider;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
@@ -15,16 +22,9 @@ using Microsoft.SqlTools.ServiceLayer.Hosting;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
+using Microsoft.SqlTools.ServiceLayer.Utility;
 using Microsoft.SqlTools.ServiceLayer.Workspace;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
-using Microsoft.SqlServer.Management.Common;
-using Microsoft.SqlServer.Management.SqlParser;
-using Microsoft.SqlTools.ServiceLayer.Utility;
-using Microsoft.SqlServer.Management.SqlParser.Binder;
-using Microsoft.SqlServer.Management.SqlParser.Intellisense;
-using Microsoft.SqlServer.Management.SqlParser.MetadataProvider;
-using Microsoft.SqlServer.Management.SqlParser.Parser;
-using Microsoft.SqlServer.Management.SmoMetadataProvider;
 using Location = Microsoft.SqlTools.ServiceLayer.Workspace.Contracts.Location;
 
 namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
@@ -35,17 +35,21 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
     /// </summary>
     public sealed class LanguageService
     {
-        public const string DefaultBatchSeperator = "GO";
+        internal const string DefaultBatchSeperator = "GO";
 
-        private const int DiagnosticParseDelay = 750;
+        internal const int DiagnosticParseDelay = 750;
 
-        private const int FindCompletionsTimeout = 3000;
+        internal const int FindCompletionsTimeout = 3000;
 
-        private const int FindCompletionStartTimeout = 50;
+        internal const int FindCompletionStartTimeout = 50;
 
-        private const int OnConnectionWaitTimeout = 30000;
+        internal const int OnConnectionWaitTimeout = 300000;
 
-        private bool ShouldEnableAutocomplete()
+        private object parseMapLock = new object();
+
+        private ScriptParseInfo currentCompletionParseInfo;
+
+        internal bool ShouldEnableAutocomplete()
         {
             return true;
         }
@@ -105,12 +109,12 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
         private static CancellationTokenSource ExistingRequestCancellation { get; set; }
 
-        private SqlToolsSettings CurrentSettings
+        internal SqlToolsSettings CurrentSettings
         {
             get { return WorkspaceService<SqlToolsSettings>.Instance.CurrentSettings; }
         }
 
-        private Workspace.Workspace CurrentWorkspace
+        internal Workspace.Workspace CurrentWorkspace
         {
             get { return WorkspaceService<SqlToolsSettings>.Instance.Workspace; }
         }
@@ -119,7 +123,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// Gets or sets the current SQL Tools context
         /// </summary>
         /// <returns></returns>
-        private SqlToolsContext Context { get; set; }
+        internal SqlToolsContext Context { get; set; }
 
         #endregion
 
@@ -196,11 +200,25 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             await requestContext.SendResult(completionItems); 
         }
 
+        /// <summary>
+        /// Handle the resolve completion request event to provide additional
+        /// autocomplete metadata to the currently select completion item
+        /// </summary>
+        /// <param name="completionItem"></param>
+        /// <param name="requestContext"></param>
+        /// <returns></returns>
+        private static async Task HandleCompletionResolveRequest(
+            CompletionItem completionItem,
+            RequestContext<CompletionItem> requestContext)
+        {
+            completionItem = LanguageService.Instance.ResolveCompletionItem(completionItem);
+            await requestContext.SendResult(completionItem); 
+        }
+
         private static async Task HandleDefinitionRequest(
             TextDocumentPosition textDocumentPosition,
             RequestContext<Location[]> requestContext)
         {
-            Logger.Write(LogLevel.Verbose, "HandleDefinitionRequest");
             await Task.FromResult(true);
         }
 
@@ -208,15 +226,6 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             ReferencesParams referencesParams,
             RequestContext<Location[]> requestContext)
         {
-            Logger.Write(LogLevel.Verbose, "HandleReferencesRequest");
-            await Task.FromResult(true);
-        }
-
-        private static async Task HandleCompletionResolveRequest(
-            CompletionItem completionItem,
-            RequestContext<CompletionItem> requestContext)
-        {
-            Logger.Write(LogLevel.Verbose, "HandleCompletionResolveRequest");
             await Task.FromResult(true);
         }
 
@@ -224,7 +233,6 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             TextDocumentPosition textDocumentPosition,
             RequestContext<SignatureHelp> requestContext)
         {
-            Logger.Write(LogLevel.Verbose, "HandleSignatureHelpRequest");
             await Task.FromResult(true);
         }
 
@@ -232,7 +240,6 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             TextDocumentPosition textDocumentPosition,
             RequestContext<DocumentHighlight[]> requestContext)
         {
-            Logger.Write(LogLevel.Verbose, "HandleDocumentHighlightRequest");
             await Task.FromResult(true);
         }
 
@@ -240,7 +247,6 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             TextDocumentPosition textDocumentPosition,
             RequestContext<Hover> requestContext)
         {
-            Logger.Write(LogLevel.Verbose, "HandleHoverRequest");
             await Task.FromResult(true);
         }
 
@@ -258,9 +264,12 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             ScriptFile scriptFile, 
             EventContext eventContext)
         {
-            await this.RunScriptDiagnostics( 
+            if (!IsPreviewWindow(scriptFile))
+            {
+                await RunScriptDiagnostics( 
                     new ScriptFile[] { scriptFile },
                     eventContext); 
+            }
 
             await Task.FromResult(true);             
         }
@@ -327,8 +336,10 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// it is the last URI connected to a particular connection,
         /// then remove the cache.
         /// </summary>
-        public async Task RemoveAutoCompleteCacheUriReference(ConnectionSummary summary)
+        public async Task RemoveAutoCompleteCacheUriReference(ConnectionSummary summary, string ownerUri)
         {
+            RemoveScriptParseInfo(ownerUri);
+
             // currently this method is disabled, but we need to reimplement now that the 
             // implementation of the 'cache' has changed.
             await Task.FromResult(0);
@@ -342,16 +353,8 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// <returns></returns>
         public ParseResult ParseAndBind(ScriptFile scriptFile, ConnectionInfo connInfo)
         {
-            ScriptParseInfo parseInfo = null;
+            // get or create the current parse info object
-            if (this.ScriptParseInfoMap.ContainsKey(scriptFile.ClientFilePath))
+            ScriptParseInfo parseInfo = GetScriptParseInfo(scriptFile.ClientFilePath, createIfNotExists: true);
-            {
-                parseInfo = this.ScriptParseInfoMap[scriptFile.ClientFilePath];   
-            }
-            else
-            {
-                parseInfo = new ScriptParseInfo();
-                this.ScriptParseInfoMap.Add(scriptFile.ClientFilePath, parseInfo);
-            }
 
             if (parseInfo.BuildingMetadataEvent.WaitOne(LanguageService.FindCompletionsTimeout))
             {
@@ -398,7 +401,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         }
 
         /// <summary>
-        /// Update the cached autocomplete candidate list when the user connects to a database
+        /// Update the autocomplete metadata provider when the user connects to a database
         /// </summary>
         /// <param name="info"></param>
         public async Task UpdateLanguageServiceOnConnection(ConnectionInfo info)
@@ -407,16 +410,12 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             {
                 if (ShouldEnableAutocomplete())
                 {
-                    ScriptParseInfo scriptInfo = 
+                    ScriptParseInfo scriptInfo = GetScriptParseInfo(info.OwnerUri, createIfNotExists: true);
-                            this.ScriptParseInfoMap.ContainsKey(info.OwnerUri)
+                    if (scriptInfo.BuildingMetadataEvent.WaitOne(LanguageService.OnConnectionWaitTimeout))
-                                ? this.ScriptParseInfoMap[info.OwnerUri]
+                    {
-                                : new ScriptParseInfo();
-
                         try
                         {
-                        scriptInfo.BuildingMetadataEvent.WaitOne(LanguageService.OnConnectionWaitTimeout);
                             scriptInfo.BuildingMetadataEvent.Reset();
-
                             var sqlConn = info.SqlConnection as ReliableSqlConnection;
                             if (sqlConn != null)
                             {
@@ -425,7 +424,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                                 scriptInfo.MetadataProvider = SmoMetadataProvider.CreateConnectedProvider(serverConn);
                                 scriptInfo.Binder = BinderProvider.CreateBinder(scriptInfo.MetadataProvider);                           
                                 scriptInfo.ServerConnection = new ServerConnection(sqlConn.GetUnderlyingConnection());
-                            this.ScriptParseInfoMap[info.OwnerUri] = scriptInfo;
+                                scriptInfo.IsConnected = true;
                             }
                         }
                         catch (Exception)
@@ -438,12 +437,10 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                             // (Tell Language Service that I am ready with Metadata Provider Object)
                             scriptInfo.BuildingMetadataEvent.Set();
                         }
-
-                    if (scriptInfo.IsConnected)
-                    {
-                        var scriptFile = WorkspaceService<SqlToolsSettings>.Instance.Workspace.GetFile(info.OwnerUri);                                
-                        ParseAndBind(scriptFile, info);
                     }
+
+                    // populate SMO metadata provider with most common info
+                    AutoCompleteHelper.PrepopulateCommonMetadata(info, scriptInfo);
                 }
             });
         }
@@ -466,6 +463,28 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 || !string.Equals(prevSqlText, currentSqlText);
         }
 
+        /// <summary>
+        /// Resolves the details and documentation for a completion item
+        /// </summary>
+        /// <param name="completionItem"></param>
+        internal CompletionItem ResolveCompletionItem(CompletionItem completionItem)
+        {
+            var scriptParseInfo = LanguageService.Instance.currentCompletionParseInfo;
+            if (scriptParseInfo != null && scriptParseInfo.CurrentSuggestions != null)
+            {
+                foreach (var suggestion in scriptParseInfo.CurrentSuggestions)
+                {
+                    if (string.Equals(suggestion.Title, completionItem.Label))
+                    {
+                        completionItem.Detail = suggestion.DatabaseQualifiedName;
+                        completionItem.Documentation = suggestion.Description;
+                        break;
+                    }
+                }
+            }
+            return completionItem;
+        }
+
         /// <summary>
         /// Return the completion item list for the current text position.
         /// This method does not await cache builds since it expects to return quickly
@@ -484,15 +503,17 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                                 textDocumentPosition.Position.Character);
             int endColumn = textDocumentPosition.Position.Character;
 
+            this.currentCompletionParseInfo = null;
+
             // Take a reference to the list at a point in time in case we update and replace the list
-            if (connInfo == null 
+
-                || !LanguageService.Instance.ScriptParseInfoMap.ContainsKey(textDocumentPosition.TextDocument.Uri))
+            ScriptParseInfo scriptParseInfo = GetScriptParseInfo(textDocumentPosition.TextDocument.Uri);
+            if (connInfo == null || scriptParseInfo == null)
             {
                 return AutoCompleteHelper.GetDefaultCompletionItems(startLine, startColumn, endColumn);
             }
 
             // reparse and bind the SQL statement if needed
-            var scriptParseInfo = ScriptParseInfoMap[textDocumentPosition.TextDocument.Uri];
             if (RequiresReparse(scriptParseInfo, scriptFile))
             {       
                 ParseAndBind(scriptFile, connInfo);
@@ -511,15 +532,18 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                     try
                     {
                         // get the completion list from SQL Parser
-                        var suggestions = Resolver.FindCompletions(
+                        scriptParseInfo.CurrentSuggestions = Resolver.FindCompletions(
                             scriptParseInfo.ParseResult, 
                             textDocumentPosition.Position.Line + 1, 
                             textDocumentPosition.Position.Character + 1, 
                             scriptParseInfo.MetadataDisplayInfoProvider); 
 
+                        // cache the current script parse info object to resolve completions later
+                        this.currentCompletionParseInfo = scriptParseInfo;
+
                         // convert the suggestion list to the VS Code format
                         return AutoCompleteHelper.ConvertDeclarationsToCompletionItems(
-                            suggestions, 
+                            scriptParseInfo.CurrentSuggestions, 
                             startLine, 
                             startColumn, 
                             endColumn);
@@ -685,5 +709,76 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         }
 
         #endregion
+
+        private void AddOrUpdateScriptParseInfo(string uri, ScriptParseInfo scriptInfo)
+        {
+            lock (this.parseMapLock)
+            {
+                if (this.ScriptParseInfoMap.ContainsKey(uri))
+                {
+                    this.ScriptParseInfoMap[uri] = scriptInfo;
+                }
+                else
+                {
+                    this.ScriptParseInfoMap.Add(uri, scriptInfo);
+                }
+
+            }
+        }
+
+        private ScriptParseInfo GetScriptParseInfo(string uri, bool createIfNotExists = false)
+        {
+            lock (this.parseMapLock)
+            {
+                if (this.ScriptParseInfoMap.ContainsKey(uri))
+                {
+                    return this.ScriptParseInfoMap[uri];
+                }
+                else if (createIfNotExists)
+                {
+                    ScriptParseInfo scriptInfo = new ScriptParseInfo();
+                    this.ScriptParseInfoMap.Add(uri, scriptInfo);
+                    return scriptInfo;
+                }
+                else
+                {
+                    return null;
+                }
+            }
+        }
+
+        private bool RemoveScriptParseInfo(string uri)
+        {
+            lock (this.parseMapLock)
+            {
+                if (this.ScriptParseInfoMap.ContainsKey(uri))
+                {
+                    var scriptInfo = this.ScriptParseInfoMap[uri];
+                    scriptInfo.ServerConnection.Disconnect();
+                    scriptInfo.ServerConnection = null;
+                    return this.ScriptParseInfoMap.Remove(uri);
+                }
+                else
+                {
+                    return false;
+                }
+            }
+        }
+
+        /// <summary>
+        /// Returns a flag indicating if the ScriptFile refers to the output window.
+        /// </summary>
+        /// <param name="scriptFile"></param>
+        private bool IsPreviewWindow(ScriptFile scriptFile)
+        {
+            if (scriptFile != null && !string.IsNullOrWhiteSpace(scriptFile.ClientFilePath))
+            {
+                return scriptFile.ClientFilePath.StartsWith("tsqloutput:");
+            }
+            else
+            {
+                return false;
+            }
+        }
     }
 }

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

@@ -3,14 +3,16 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System;
+using System.Collections.Generic;
+using System.Threading;
 using Microsoft.SqlServer.Management.Common;
 using Microsoft.SqlServer.Management.SmoMetadataProvider;
 using Microsoft.SqlServer.Management.SqlParser.Binder;
 using Microsoft.SqlServer.Management.SqlParser.Common;
+using Microsoft.SqlServer.Management.SqlParser.Intellisense;
 using Microsoft.SqlServer.Management.SqlParser.MetadataProvider;
 using Microsoft.SqlServer.Management.SqlParser.Parser;
-using System;
-using System.Threading;
 
 namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 {
@@ -33,7 +35,6 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             get { return this.buildingMetadataEvent; }
         }
 
-
         /// <summary>
         /// Gets or sets a flag determining is the LanguageService is connected
         /// </summary>
@@ -56,7 +57,6 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                     isQuotedIdentifierSet: true, 
                     compatibilityLevel: DatabaseCompatibilityLevel, 
                     transactSqlVersion: TransactSqlVersion);
-                this.IsConnected = true;
             }
         }
 
@@ -143,6 +143,11 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// </summary>
         public MetadataDisplayInfoProvider MetadataDisplayInfoProvider { get; set; }
         
+        /// <summary>
+        /// Gets or sets the current autocomplete suggestion list
+        /// </summary>
+        public IEnumerable<Declaration> CurrentSuggestions { get; set; }
+
         /// <summary>
         /// Gets the database compatibility level from a server version
         /// </summary>

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

@@ -131,6 +131,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Utility
 
     internal class LogWriter : IDisposable
     {
+        private object logLock = new object();
         private TextWriter textWriter;
         private LogLevel minimumLogLevel = LogLevel.Verbose;
 
@@ -169,6 +170,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Utility
         {
             if (this.textWriter != null &&
                 logLevel >= this.minimumLogLevel)
+            {
+                // System.IO is not thread safe
+                lock (this.logLock)
                 {
                     // Print the timestamp and log level
                     this.textWriter.WriteLine(
@@ -190,6 +194,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Utility
                     this.textWriter.Flush();
                 }
             }
+        }
 
         public void Dispose()
         {

test/Microsoft.SqlTools.ServiceLayer.Test/Connection/ConnectionServiceTests.cs

@@ -138,8 +138,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Connection
 
             // register disconnect callback
             connectionService.RegisterOnDisconnectTask(
-                (result) => { 
+                (result, uri) => { 
                     callbackInvoked = true;
+                    Assert.True(uri.Equals(ownerUri));
                     return Task.FromResult(true);
                 }
             );
@@ -433,8 +434,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Connection
 
             // register disconnect callback
             connectionService.RegisterOnDisconnectTask(
-                (result) => { 
+                (result, uri) => { 
                     callbackInvoked = true;
+                    Assert.True(uri.Equals(ownerUri));
                     return Task.FromResult(true);
                 }
             );

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

@@ -21,9 +21,13 @@ using Microsoft.SqlServer.Management.SqlParser.MetadataProvider;
 using Microsoft.SqlServer.Management.SqlParser.Parser;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
+using Microsoft.SqlTools.ServiceLayer.Credentials;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution;
+using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Test.QueryExecution;
 using Microsoft.SqlTools.ServiceLayer.Test.Utility;
+using Microsoft.SqlTools.ServiceLayer.Workspace;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 using Microsoft.SqlTools.Test.Utility;
 using Moq;
@@ -148,6 +152,150 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
 
         #endregion
 
+        #region "General Language Service tests"
+
+        /// <summary>
+        /// Check that autocomplete is enabled by default
+        /// </summary>
+        [Fact]
+        public void CheckAutocompleteEnabledByDefault()
+        {
+            // get test service
+            LanguageService service = TestObjects.GetTestLanguageService();
+            Assert.True(service.ShouldEnableAutocomplete());
+        }
+
+        /// <summary>
+        /// Test the service initialization code path and verify nothing throws
+        /// </summary>
+        [Fact]
+        public void ServiceInitiailzation()
+        {
+            InitializeTestServices();
+
+            Assert.True(LanguageService.Instance.Context != null);
+            Assert.True(LanguageService.Instance.ConnectionServiceInstance != null);
+            Assert.True(LanguageService.Instance.CurrentSettings != null);
+            Assert.True(LanguageService.Instance.CurrentWorkspace != null);
+
+            LanguageService.Instance.ConnectionServiceInstance = null;
+            Assert.True(LanguageService.Instance.ConnectionServiceInstance == null);
+        }        
+        
+        /// <summary>
+        /// Test the service initialization code path and verify nothing throws
+        /// </summary>
+        [Fact]
+        public void UpdateLanguageServiceOnConnection()
+        {
+            string ownerUri = "file://my/sample/file.sql";
+            var connectionService = TestObjects.GetTestConnectionService();
+            var connectionResult =
+                connectionService
+                .Connect(new ConnectParams()
+                {
+                    OwnerUri = ownerUri,
+                    Connection = TestObjects.GetTestConnectionDetails()
+                });
+            
+            ConnectionInfo connInfo = null;
+            connectionService.TryFindConnection(ownerUri, out connInfo);
+            
+            var task = LanguageService.Instance.UpdateLanguageServiceOnConnection(connInfo);
+            task.Wait();
+        }
+
+        /// <summary>
+        /// Test the service initialization code path and verify nothing throws
+        /// </summary>
+        [Fact]
+        public void PrepopulateCommonMetadata()
+        {
+            InitializeTestServices();
+
+            string sqlFilePath = GetTestSqlFile();            
+            ScriptFile scriptFile = WorkspaceService<SqlToolsSettings>.Instance.Workspace.GetFile(sqlFilePath);
+
+            string ownerUri = scriptFile.ClientFilePath;
+            var connectionService = TestObjects.GetTestConnectionService();
+            var connectionResult =
+                connectionService
+                .Connect(new ConnectParams()
+                {
+                    OwnerUri = ownerUri,
+                    Connection = TestObjects.GetTestConnectionDetails()
+                });
+            
+            ConnectionInfo connInfo = null;
+            connectionService.TryFindConnection(ownerUri, out connInfo);
+            
+            ScriptParseInfo scriptInfo = new ScriptParseInfo();
+            scriptInfo.IsConnected = true;
+
+            AutoCompleteHelper.PrepopulateCommonMetadata(connInfo, scriptInfo);
+        }
+
+        private string GetTestSqlFile()
+        {
+            string filePath = Path.Combine(
+                Path.GetDirectoryName(Assembly.GetEntryAssembly().Location),
+                "sqltest.sql");
+            
+            if (File.Exists(filePath))
+            {
+                File.Delete(filePath);
+            }
+
+            File.WriteAllText(filePath, "SELECT * FROM sys.objects\n");
+
+            return filePath;
+        }
+
+        private void InitializeTestServices()
+        {
+            const string hostName = "SQL Tools Service Host";
+            const string hostProfileId = "SQLToolsService";
+            Version hostVersion = new Version(1,0); 
+
+            // set up the host details and profile paths 
+            var hostDetails = new HostDetails(hostName, hostProfileId, hostVersion);     
+            var profilePaths = new ProfilePaths(hostProfileId, "baseAllUsersPath", "baseCurrentUserPath");
+            SqlToolsContext sqlToolsContext = new SqlToolsContext(hostDetails, profilePaths);
+
+            // Grab the instance of the service host
+            Hosting.ServiceHost serviceHost = Hosting.ServiceHost.Instance;
+
+            // Start the service
+            serviceHost.Start().Wait();
+
+            // Initialize the services that will be hosted here
+            WorkspaceService<SqlToolsSettings>.Instance.InitializeService(serviceHost);
+            LanguageService.Instance.InitializeService(serviceHost, sqlToolsContext);
+            ConnectionService.Instance.InitializeService(serviceHost);
+            CredentialService.Instance.InitializeService(serviceHost);
+            QueryExecutionService.Instance.InitializeService(serviceHost);
+
+            serviceHost.Initialize();
+        }
+
+        private Hosting.ServiceHost GetTestServiceHost()
+        {
+            // set up the host details and profile paths 
+            var hostDetails = new HostDetails("Test Service Host", "SQLToolsService", new Version(1,0));     
+            var profilePaths = new ProfilePaths("SQLToolsService", "baseAllUsersPath", "baseCurrentUserPath");
+            SqlToolsContext context = new SqlToolsContext(hostDetails, profilePaths);
+
+            // Grab the instance of the service host
+            Hosting.ServiceHost host = Hosting.ServiceHost.Instance;
+
+            // Start the service
+            host.Start().Wait();
+
+            return host;
+        }
+
+        #endregion
+
         #region "Autocomplete Tests"
 
         // This test currently requires a live database connection to initialize