Add error handling for Azure Exceptions (#177)

* Add error handling for Azure Exceptions

* Add SRGen for string

* Add specific exception messages

* Move DefinitionResult class

* Add SqlLogin constant

* Add error scenarios

* revert timeout duration

* Modify tests

* Modify tests

* Add tests

* Revert live connection definition

* Modify DefinitionsHandlerWithNoConnectionTest

* fix test after merge

* Code review changes

* Code review changes

* Code review changes

src/Microsoft.SqlTools.ServiceLayer/Hosting/Protocol/Constants.cs

@@ -13,6 +13,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Hosting.Protocol
         public const string ContentLengthFormatString = "Content-Length: {0}\r\n\r\n";
         public static readonly JsonSerializerSettings JsonSerializerSettings;
 
+        public static readonly string SqlLoginAuthenticationType = "SqlLogin";
+
         static Constants()
         {
             JsonSerializerSettings = new JsonSerializerSettings();

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

@@ -57,7 +57,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// <summary>
         /// Queue a binding request item
         /// </summary>
-        public QueueItem QueueBindingOperation(
+        public virtual QueueItem QueueBindingOperation(
             string key,
             Func<IBindingContext, CancellationToken, object> bindOperation,
             Func<IBindingContext, object> timeoutOperation = null,

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

@@ -14,5 +14,16 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts
             RequestType<TextDocumentPosition, Location[]> Type =
             RequestType<TextDocumentPosition, Location[]>.Create("textDocument/definition");
     }
+
+    /// <summary>
+    /// Error object for Definition
+    /// </summary>
+    public class DefinitionError
+    {
+        /// <summary>
+        /// Error message 
+        /// </summary>
+        public string message { get; set; }
+    }
 }
 

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

@@ -236,6 +236,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
             // Store the SqlToolsContext for future use
             Context = context;
+
         }
 
         #endregion
@@ -317,11 +318,17 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 ConnectionInfo connInfo;
                 var scriptFile = LanguageService.WorkspaceServiceInstance.Workspace.GetFile(textDocumentPosition.TextDocument.Uri);
                 LanguageService.ConnectionServiceInstance.TryFindConnection(scriptFile.ClientFilePath, out connInfo);
-
-                Location[] locations = LanguageService.Instance.GetDefinition(textDocumentPosition, scriptFile, connInfo);
-                if (locations != null)
-                {
-                    await requestContext.SendResult(locations);
+                DefinitionResult definitionResult = LanguageService.Instance.GetDefinition(textDocumentPosition, scriptFile, connInfo);
+                if (definitionResult != null)
+                {   
+                    if (definitionResult.IsErrorResult)
+                    {
+                        await requestContext.SendError( new DefinitionError { message = definitionResult.Message });
+                    }
+                    else
+                    {
+                        await requestContext.SendResult(definitionResult.Locations);
+                    }                    
                 }
             }
             DocumentStatusHelper.SendStatusChange(requestContext, textDocumentPosition, DocumentStatusHelper.DefinitionRequestCompleted);
@@ -688,7 +695,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// <param name="scriptFile"></param>
         /// <param name="connInfo"></param>
         /// <returns> Location with the URI of the script file</returns>
-        internal Location[] GetDefinition(TextDocumentPosition textDocumentPosition, ScriptFile scriptFile, ConnectionInfo connInfo)
+        internal DefinitionResult GetDefinition(TextDocumentPosition textDocumentPosition, ScriptFile scriptFile, ConnectionInfo connInfo)
         {
             // Parse sql
             ScriptParseInfo scriptParseInfo = GetScriptParseInfo(textDocumentPosition.TextDocument.Uri);
@@ -726,27 +733,57 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                             int parserLine = textDocumentPosition.Position.Line + 1;
                             int parserColumn = textDocumentPosition.Position.Character + 1;
                             IEnumerable<Declaration> declarationItems = Resolver.FindCompletions(
-                                scriptParseInfo.ParseResult, 
-                                parserLine, parserColumn, 
+                                scriptParseInfo.ParseResult,
+                                parserLine, parserColumn,
                                 bindingContext.MetadataDisplayInfoProvider);
 
                             // Match token with the suggestions(declaration items) returned
-                            string schemaName = GetSchemaName(scriptParseInfo, textDocumentPosition.Position, scriptFile);
-                            PeekDefinition peekDefinition = new PeekDefinition(bindingContext.ServerConnection);
-                            return peekDefinition.GetScript(declarationItems, tokenText, schemaName);                        
+
+                            string schemaName = this.GetSchemaName(scriptParseInfo, textDocumentPosition.Position, scriptFile);
+                            PeekDefinition peekDefinition = new PeekDefinition(bindingContext.ServerConnection, connInfo);
+                            return peekDefinition.GetScript(declarationItems, tokenText, schemaName);
+                        },
+                        timeoutOperation: (bindingContext) =>
+                        {
+                            // return error result
+                            return new DefinitionResult
+                            {
+                                IsErrorResult = true,
+                                Message = SR.PeekDefinitionTimedoutError,
+                                Locations = null
+                            };
                         });
 
                     // wait for the queue item
                     queueItem.ItemProcessed.WaitOne();
-                    return queueItem.GetResultAsT<Location[]>();
+                    return queueItem.GetResultAsT<DefinitionResult>();
+                }
+                catch (Exception ex)
+                {
+                    // if any exceptions are raised return error result with message
+                    Logger.Write(LogLevel.Error, "Exception in GetDefinition " + ex.ToString());
+                    return new DefinitionResult
+                    {
+                        IsErrorResult = true,
+                        Message = SR.PeekDefinitionError(ex.Message),
+                        Locations = null
+                    };
                 }
                 finally
                 {
                     Monitor.Exit(scriptParseInfo.BuildingMetadataLock);
                 }
             }
-
-            return null;
+            else
+            {
+                // User is not connected.
+                return new DefinitionResult
+                {
+                    IsErrorResult = true,
+                    Message = SR.PeekDefinitionNotConnectedError,
+                    Locations = null
+                };
+            }
         }
 
         /// <summary>
@@ -755,7 +792,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// <param name="scriptParseInfo"></param>
         /// <param name="position"></param>
         /// <param name="scriptFile"></param>
-        /// <returns> schema nama</returns>
+        /// <returns> schema name</returns>
         private string GetSchemaName(ScriptParseInfo scriptParseInfo, Position position, ScriptFile scriptFile)
         {
             // Offset index by 1 for sql parser

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

@@ -7,12 +7,12 @@ using System.IO;
 using System.Collections.Generic;
 using System.Collections.Specialized;
 using System.Data.SqlClient;
-using System.Runtime.InteropServices;
 using Microsoft.SqlServer.Management.Smo;
 using Microsoft.SqlServer.Management.Common;
 using Microsoft.SqlServer.Management.SqlParser.Intellisense;
 using Microsoft.SqlTools.ServiceLayer.Utility;
 using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
@@ -23,10 +23,11 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
     /// </summary>
     internal class PeekDefinition
     {
+        private bool error;
+        private string errorMessage;
         private ServerConnection serverConnection;
-
+        private ConnectionInfo connectionInfo;
         private Database database;
-
         private string tempPath;
 
         internal delegate StringCollection ScriptGetter(string objectName, string schemaName);
@@ -42,9 +43,10 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// Initialize a Peek Definition helper object
         /// </summary>
         /// <param name="serverConnection">SMO Server connection</param>
-        internal PeekDefinition(ServerConnection serverConnection)
+        internal PeekDefinition(ServerConnection serverConnection, ConnectionInfo connInfo)
         {
             this.serverConnection = serverConnection;
+            this.connectionInfo = connInfo;
 
             DirectoryInfo tempScriptDirectory = Directory.CreateDirectory(Path.GetTempPath() + "mssql_definition");
             this.tempPath = tempScriptDirectory.FullName;
@@ -62,20 +64,29 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                         try
                         {
                             // Get server object from connection
-                            SqlConnection sqlConn = new SqlConnection(this.serverConnection.ConnectionString);                    
+                            SqlConnection sqlConn = new SqlConnection(this.serverConnection.ConnectionString);
                             sqlConn.Open();
                             ServerConnection peekConnection = new ServerConnection(sqlConn);
-                            Server server = new Server(peekConnection);                        
-                            this.database = new Database(server, peekConnection.DatabaseName);                        
+                            Server server = new Server(peekConnection);
+                            this.database = new Database(server, peekConnection.DatabaseName);
+                        }
+                        catch (ConnectionFailureException cfe)
+                        {
+                            Logger.Write(LogLevel.Error, "Exception at PeekDefinition Database.get() : " + cfe.Message);
+                            this.error = true;
+                            this.errorMessage = (connectionInfo != null && connectionInfo.IsAzure)? SR.PeekDefinitionAzureError(cfe.Message) : SR.PeekDefinitionError(cfe.Message);
+                            return null;
                         }
                         catch (Exception ex)
                         {
                             Logger.Write(LogLevel.Error, "Exception at PeekDefinition Database.get() : " + ex.Message);
-                        }   
+                            this.error = true;
+                            this.errorMessage = SR.PeekDefinitionError(ex.Message);
+                            return null;
+                        }
                     }
 
                 }
-                
                 return this.database;
             }
         }
@@ -114,7 +125,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         {
             if (Path.DirectorySeparatorChar.Equals('/'))
             {
-                tempFileName = "file:" + tempFileName;                 
+                tempFileName = "file:" + tempFileName;
             }
             else
             {
@@ -140,7 +151,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             string[] lines = script.Split(new string[] { Environment.NewLine }, StringSplitOptions.None);
             for (int lineNumber = 0; lineNumber < lines.Length; lineNumber++)
             {
-                if (lines[lineNumber].IndexOf( createString, StringComparison.OrdinalIgnoreCase) >= 0)
+                if (lines[lineNumber].IndexOf(createString, StringComparison.OrdinalIgnoreCase) >= 0)
                 {
                     return lineNumber;
                 }
@@ -155,7 +166,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// <param name="tokenText"></param>
         /// <param name="schemaName"></param>
         /// <returns>Location object of the script file</returns>
-        internal Location[] GetScript(IEnumerable<Declaration> declarationItems, string tokenText, string schemaName)
+        internal DefinitionResult GetScript(IEnumerable<Declaration> declarationItems, string tokenText, string schemaName)
         {
             foreach (Declaration declarationItem in declarationItems)
             {
@@ -167,29 +178,38 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 if (declarationItem.Title.Equals(tokenText))
                 {
                     // Script object using SMO based on type
-                    DeclarationType type  = declarationItem.Type;
+                    DeclarationType type = declarationItem.Type;
                     if (sqlScriptGetters.ContainsKey(type) && sqlObjectTypes.ContainsKey(type))
                     {
                         // On *nix and mac systems, the defaultSchema property throws an Exception when accessed.
                         // This workaround ensures that a schema name is present by attempting 
                         // to get the schema name from the declaration item 
                         // If all fails, the default schema name is assumed to be "dbo"
-                        if (!RuntimeInformation.IsOSPlatform(OSPlatform.Windows) && string.IsNullOrEmpty(schemaName))
+                        if ((connectionInfo != null && connectionInfo.ConnectionDetails.AuthenticationType.Equals(Constants.SqlLoginAuthenticationType)) && string.IsNullOrEmpty(schemaName))
                         {
                             string fullObjectName = declarationItem.DatabaseQualifiedName;
                             schemaName = this.GetSchemaFromDatabaseQualifiedName(fullObjectName, tokenText);
                         }
-                        return GetSqlObjectDefinition(
+                        Location[] locations = GetSqlObjectDefinition(
                                     sqlScriptGetters[type],
                                     tokenText,
                                     schemaName,
                                     sqlObjectTypes[type]
                                 );
+                        DefinitionResult result = new DefinitionResult
+                        {
+                            IsErrorResult = this.error,
+                            Message = this.errorMessage,
+                            Locations = locations
+                        };
+                        return result;
                     }
-                    return null;
+                    // sql object type is currently not supported
+                    return GetDefinitionErrorResult(SR.PeekDefinitionTypeNotSupportedError);
                 }
             }
-            return null;
+            // no definition found
+            return GetDefinitionErrorResult(SR.PeekDefinitionNoResultsError);
         }
 
         /// <summary>
@@ -203,9 +223,9 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             string[] tokens = fullObjectName.Split('.');
             for (int i = tokens.Length - 1; i > 0; i--)
             {
-                if(tokens[i].Equals(objectName))
+                if (tokens[i].Equals(objectName))
                 {
-                    return tokens[i-1];
+                    return tokens[i - 1];
                 }
             }
             return "dbo";
@@ -226,7 +246,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                     : new Table(this.Database, tableName, schemaName);
 
                 table.Refresh();
-                
+
                 return table.Script();
             }
             catch (Exception ex)
@@ -251,7 +271,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                     : new View(this.Database, viewName, schemaName);
 
                 view.Refresh();
-                
+
                 return view.Script();
             }
             catch (Exception ex)
@@ -276,7 +296,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                     : new StoredProcedure(this.Database, sprocName, schemaName);
 
                 sproc.Refresh();
-                
+
                 return sproc.Script();
             }
             catch (Exception ex)
@@ -300,30 +320,49 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 string schemaName,
                 string objectType)
         {
-                StringCollection scripts = sqlScriptGetter(objectName, schemaName);
-                string tempFileName = (schemaName != null) ?  Path.Combine(this.tempPath, string.Format("{0}.{1}.sql", schemaName, objectName))
-                                                    : Path.Combine(this.tempPath, string.Format("{0}.sql", objectName));
+            StringCollection scripts = sqlScriptGetter(objectName, schemaName);
+            string tempFileName = (schemaName != null) ? Path.Combine(this.tempPath, string.Format("{0}.{1}.sql", schemaName, objectName))
+                                                : Path.Combine(this.tempPath, string.Format("{0}.sql", objectName));
 
-                if (scripts != null)
+            if (scripts != null)
+            {
+                int lineNumber = 0;
+                using (StreamWriter scriptFile = new StreamWriter(File.Open(tempFileName, FileMode.Create, FileAccess.ReadWrite)))
                 {
-                    int lineNumber = 0;
-                    using (StreamWriter scriptFile = new StreamWriter(File.Open(tempFileName, FileMode.Create, FileAccess.ReadWrite)))
-                    {
 
-                        foreach (string script in scripts)
+                    foreach (string script in scripts)
+                    {
+                        string createSyntax = string.Format("CREATE {0}", objectType);
+                        if (script.IndexOf(createSyntax, StringComparison.OrdinalIgnoreCase) >= 0)
                         {
-                            string createSyntax = string.Format("CREATE {0}", objectType);
-                            if (script.IndexOf(createSyntax, StringComparison.OrdinalIgnoreCase) >= 0)
-                            {
-                                scriptFile.WriteLine(script);
-                                lineNumber = GetStartOfCreate(script, createSyntax);
-                            }
+                            scriptFile.WriteLine(script);
+                            lineNumber = GetStartOfCreate(script, createSyntax);
                         }
                     }
-                    return GetLocationFromFile(tempFileName, lineNumber);
                 }
+                return GetLocationFromFile(tempFileName, lineNumber);
+            }
+            else
+            {
+                this.error = true;
+                this.errorMessage = SR.PeekDefinitionNoResultsError;
+                return null;
+            }
+        }
 
-            return null;
+        /// <summary>
+        /// Helper method to create definition error result object
+        /// </summary>
+        /// <param name="errorMessage">Error message</param>
+        /// <returns> DefinitionResult</returns>
+        internal DefinitionResult GetDefinitionErrorResult(string errorMessage)
+        {
+            return new DefinitionResult
+            {
+                IsErrorResult = true,
+                Message = errorMessage,
+                Locations = null
+            };
         }
     }
 }

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

@@ -0,0 +1,28 @@
+//
+// 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.Workspace.Contracts;
+namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
+{
+    /// <summary>
+    /// /// Result object for PeekDefinition
+    /// </summary>
+    public class DefinitionResult
+    {
+        /// <summary>
+        /// True, if definition error occured
+        /// </summary>
+        public bool IsErrorResult;
+
+        /// <summary>
+        /// Error message, if any
+        /// </summary>
+        public string Message { get; set; }
+
+        /// <summary>
+        /// Location object representing the definition script file
+        /// </summary>
+        public Location[] Locations;
+    }
+}

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

@@ -39,7 +39,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// <summary>
         /// Gets or sets an event to signal when this queue item has been processed
         /// </summary>
-        public ManualResetEvent ItemProcessed { get; set; } 
+        public virtual ManualResetEvent ItemProcessed { get; set; } 
 
         /// <summary>
         /// Gets or sets the result of the queued task

src/Microsoft.SqlTools.ServiceLayer/sr.cs

@@ -341,6 +341,38 @@ namespace Microsoft.SqlTools.ServiceLayer
             }
         }
 
+        public static string PeekDefinitionNoResultsError
+        {
+            get
+            {
+                return Keys.GetString(Keys.PeekDefinitionNoResultsError);
+            }
+        }
+
+        public static string PeekDefinitionNotConnectedError
+        {
+            get
+            {
+                return Keys.GetString(Keys.PeekDefinitionNotConnectedError);
+            }
+        }
+
+        public static string PeekDefinitionTimedoutError
+        {
+            get
+            {
+                return Keys.GetString(Keys.PeekDefinitionTimedoutError);
+            }
+        }
+
+        public static string PeekDefinitionTypeNotSupportedError
+        {
+            get
+            {
+                return Keys.GetString(Keys.PeekDefinitionTypeNotSupportedError);
+            }
+        }
+
         public static string WorkspaceServicePositionLineOutOfRange
         {
             get
@@ -384,6 +416,16 @@ namespace Microsoft.SqlTools.ServiceLayer
             return Keys.GetString(Keys.QueryServiceQueryFailed, message);
         }
 
+        public static string PeekDefinitionAzureError(string errorMessage)
+        {
+            return Keys.GetString(Keys.PeekDefinitionAzureError, errorMessage);
+        }
+
+        public static string PeekDefinitionError(string errorMessage)
+        {
+            return Keys.GetString(Keys.PeekDefinitionError, errorMessage);
+        }
+
         public static string WorkspaceServicePositionColumnOutOfRange(int line)
         {
             return Keys.GetString(Keys.WorkspaceServicePositionColumnOutOfRange, line);
@@ -397,7 +439,7 @@ namespace Microsoft.SqlTools.ServiceLayer
         [System.Runtime.CompilerServices.CompilerGeneratedAttribute()]
         public class Keys
         {
-            static ResourceManager resourceManager = new ResourceManager(typeof(SR));
+            static ResourceManager resourceManager = new ResourceManager("Microsoft.SqlTools.ServiceLayer.SR", typeof(SR).GetTypeInfo().Assembly);
 
             static CultureInfo _culture = null;
 
@@ -540,6 +582,24 @@ namespace Microsoft.SqlTools.ServiceLayer
             public const string QueryServiceResultSetNoColumnSchema = "QueryServiceResultSetNoColumnSchema";
 
 
+            public const string PeekDefinitionAzureError = "PeekDefinitionAzureError";
+
+
+            public const string PeekDefinitionError = "PeekDefinitionError";
+
+
+            public const string PeekDefinitionNoResultsError = "PeekDefinitionNoResultsError";
+
+
+            public const string PeekDefinitionNotConnectedError = "PeekDefinitionNotConnectedError";
+
+
+            public const string PeekDefinitionTimedoutError = "PeekDefinitionTimedoutError";
+
+
+            public const string PeekDefinitionTypeNotSupportedError = "PeekDefinitionTypeNotSupportedError";
+
+
             public const string WorkspaceServicePositionLineOutOfRange = "WorkspaceServicePositionLineOutOfRange";
 
 

src/Microsoft.SqlTools.ServiceLayer/sr.resx

@@ -308,6 +308,32 @@
     <value>Could not retrieve column schema for result set</value>
     <comment></comment>
   </data>
+  <data name="PeekDefinitionAzureError" xml:space="preserve">
+    <value>This feature is currently not supported on Azure SQL DB and Data Warehouse: {0}</value>
+    <comment>.
+ Parameters: 0 - errorMessage (string) </comment>
+  </data>
+  <data name="PeekDefinitionError" xml:space="preserve">
+    <value>An unexpected error occurred during Peek Definition execution: {0}</value>
+    <comment>.
+ Parameters: 0 - errorMessage (string) </comment>
+  </data>
+  <data name="PeekDefinitionNoResultsError" xml:space="preserve">
+    <value>No results were found.</value>
+    <comment></comment>
+  </data>
+  <data name="PeekDefinitionNotConnectedError" xml:space="preserve">
+    <value>Please connect to a server.</value>
+    <comment></comment>
+  </data>
+  <data name="PeekDefinitionTimedoutError" xml:space="preserve">
+    <value>Operation timed out.</value>
+    <comment></comment>
+  </data>
+  <data name="PeekDefinitionTypeNotSupportedError" xml:space="preserve">
+    <value>This object type is currently not supported by this feature.</value>
+    <comment></comment>
+  </data>
   <data name="WorkspaceServicePositionLineOutOfRange" xml:space="preserve">
     <value>Position is outside of file line range</value>
     <comment></comment>

src/Microsoft.SqlTools.ServiceLayer/sr.strings

@@ -139,6 +139,21 @@ QueryServiceResultSetRowCountOutOfRange = Row count must be a positive integer
 
 QueryServiceResultSetNoColumnSchema = Could not retrieve column schema for result set
 
+############################################################################
+# Language Service
+
+PeekDefinitionAzureError(string errorMessage) = This feature is currently not supported on Azure SQL DB and Data Warehouse: {0}
+
+PeekDefinitionError(string errorMessage) = An unexpected error occurred during Peek Definition execution: {0}
+
+PeekDefinitionNoResultsError = No results were found.
+
+PeekDefinitionNotConnectedError = Please connect to a server.
+
+PeekDefinitionTimedoutError = Operation timed out.
+
+PeekDefinitionTypeNotSupportedError = This object type is currently not supported by this feature.
+
 ############################################################################
 # Workspace Service
 

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

@@ -3,11 +3,14 @@
 // 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.Runtime.InteropServices;
+using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
+using System.Runtime.InteropServices;
+using Microsoft.SqlServer.Management.Common;
 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.SqlTools.ServiceLayer.Connection;
@@ -91,8 +94,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
             requestContext = new Mock<RequestContext<Location[]>>();
             requestContext.Setup(rc => rc.SendResult(It.IsAny<Location[]>()))
                 .Returns(Task.FromResult(0));
-            requestContext.Setup(r => r.SendEvent(It.IsAny<EventType<TelemetryParams>>(), It.IsAny<TelemetryParams>()));
-            requestContext.Setup(r => r.SendEvent(It.IsAny<EventType<StatusChangeParams>>(), It.IsAny<StatusChangeParams>()));
+            requestContext.Setup(rc => rc.SendError(It.IsAny<DefinitionError>())).Returns(Task.FromResult(0));;
+            requestContext.Setup(r => r.SendEvent(It.IsAny<EventType<TelemetryParams>>(), It.IsAny<TelemetryParams>())).Returns(Task.FromResult(0));;
+            requestContext.Setup(r => r.SendEvent(It.IsAny<EventType<StatusChangeParams>>(), It.IsAny<StatusChangeParams>())).Returns(Task.FromResult(0));;
 
             // setup the IBinder mock
             binder = new Mock<IBinder>();
@@ -103,7 +107,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
 
             var testScriptParseInfo = new ScriptParseInfo();
             LanguageService.Instance.AddOrUpdateScriptParseInfo(this.testScriptUri, testScriptParseInfo);
-            testScriptParseInfo.IsConnected = true;
+            testScriptParseInfo.IsConnected = false;
             testScriptParseInfo.ConnectionKey = LanguageService.Instance.BindingQueue.AddConnectionContext(connectionInfo);
 
             // setup the binding context object
@@ -115,18 +119,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
 
 
         /// <summary>
-        /// Tests the definition event handler. When called with no active connection, no definition is sent
+        /// Tests the definition event handler. When called with no active connection, an error is sent
         /// </summary>
         [Fact]
         public async Task DefinitionsHandlerWithNoConnectionTest()
         {
             TestObjects.InitializeTestServices();
             InitializeTestObjects();
-            // request the completion list
-            await Task.WhenAny(LanguageService.HandleDefinitionRequest(textDocument, requestContext.Object), Task.Delay(TaskTimeout));
-            
-            // verify that send result was not called
+            // request definition
+            var definitionTask = await Task.WhenAny(LanguageService.HandleDefinitionRequest(textDocument, requestContext.Object), Task.Delay(TaskTimeout));
+            await definitionTask;
+            // verify that send result was not called and send error was called
             requestContext.Verify(m => m.SendResult(It.IsAny<Location[]>()), Times.Never());
+            requestContext.Verify(m => m.SendError(It.IsAny<DefinitionError>()), Times.Once());
         }
 
         /// <summary>
@@ -136,7 +141,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         public void GetLocationFromFileForValidFilePathTest()
         {
             String filePath = RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? "C:\\test\\script.sql" : "/test/script.sql";
-            PeekDefinition peekDefinition = new PeekDefinition(null);
+            PeekDefinition peekDefinition = new PeekDefinition(null, null);
             Location[] locations = peekDefinition.GetLocationFromFile(filePath, 0);
 
             String expectedFilePath = RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? "file:///C:/test/script.sql" : "file:/test/script.sql";
@@ -149,11 +154,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetSchemaFromDatabaseQualifiedNameWithValidNameTest()
         {
-            PeekDefinition peekDefinition = new PeekDefinition(null);
+            PeekDefinition peekDefinition = new PeekDefinition(null, null);
             string validDatabaseQualifiedName = "master.test.test_table";
             string objectName = "test_table";
             string expectedSchemaName = "test";
-        
+
             string actualSchemaName = peekDefinition.GetSchemaFromDatabaseQualifiedName(validDatabaseQualifiedName, objectName);
             Assert.Equal(actualSchemaName, expectedSchemaName);
         }
@@ -165,11 +170,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetSchemaFromDatabaseQualifiedNameWithNoSchemaTest()
         {
-            PeekDefinition peekDefinition = new PeekDefinition(null);
+            PeekDefinition peekDefinition = new PeekDefinition(null, null);
             string validDatabaseQualifiedName = "test_table";
             string objectName = "test_table";
             string expectedSchemaName = "dbo";
-        
+
             string actualSchemaName = peekDefinition.GetSchemaFromDatabaseQualifiedName(validDatabaseQualifiedName, objectName);
             Assert.Equal(actualSchemaName, expectedSchemaName);
         }
@@ -180,15 +185,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetSchemaFromDatabaseQualifiedNameWithInvalidNameTest()
         {
-            PeekDefinition peekDefinition = new PeekDefinition(null);
+            PeekDefinition peekDefinition = new PeekDefinition(null, null);
             string validDatabaseQualifiedName = "x.y.z";
             string objectName = "test_table";
             string expectedSchemaName = "dbo";
-        
+
             string actualSchemaName = peekDefinition.GetSchemaFromDatabaseQualifiedName(validDatabaseQualifiedName, objectName);
             Assert.Equal(actualSchemaName, expectedSchemaName);
         }
-         
+
 #if LIVE_CONNECTION_TESTS
         /// <summary>
         /// Test get definition for a table object with active connection
@@ -196,9 +201,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetValidTableDefinitionTest()
         {
-            // Get live connectionInfo
-            PeekDefinition peekDefinition = new PeekDefinition(TestObjects.InitLiveConnectionInfoForDefinition());
+            // Get live connectionInfo and serverConnection
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "spt_monitor";
+
             string schemaName = null;
             string objectType = "TABLE";
 
@@ -214,8 +223,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetTableDefinitionInvalidObjectTest()
         {
-            // Get live connectionInfo
-            PeekDefinition peekDefinition = new PeekDefinition(TestObjects.InitLiveConnectionInfoForDefinition());
+            // Get live connectionInfo and serverConnection
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "test_invalid";
             string schemaName = null;
             string objectType = "TABLE";
@@ -231,24 +243,102 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetTableDefinitionWithSchemaTest()
         {
-            // Get live connectionInfo
-            PeekDefinition peekDefinition = new PeekDefinition(TestObjects.InitLiveConnectionInfoForDefinition());
+            // Get live connectionInfo and serverConnection
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "spt_monitor";
+
             string schemaName = "dbo";
             string objectType = "TABLE";
 
             // Get locations for valid table object with schema name
             Location[] locations = peekDefinition.GetSqlObjectDefinition(peekDefinition.GetTableScripts, objectName, schemaName, objectType);
             Assert.NotNull(locations);
-            Cleanup(locations);            
+            Cleanup(locations);
         }
 
         /// <summary>
-        /// Test GetDefinition with an unsupported type(function)
+        /// Test GetDefinition with an unsupported type(schema - dbo). Expect a error result.
         /// </summary>
         [Fact]
-        public void GetUnsupportedDefinitionForFullScript()
+        public void GetUnsupportedDefinitionErrorTest()
         {
+            ScriptFile scriptFile;
+            TextDocumentPosition textDocument = new TextDocumentPosition
+            {
+                TextDocument = new TextDocumentIdentifier { Uri = OwnerUri },
+                Position = new Position
+                {
+                    Line = 0,
+                    // test for 'dbo'
+                    Character = 16
+                }
+            };
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfo(out scriptFile);
+            scriptFile.Contents = "select * from dbo.func ()";
+            var languageService = new LanguageService();
+            ScriptParseInfo scriptInfo = new ScriptParseInfo { IsConnected = true };
+            languageService.ScriptParseInfoMap.Add(OwnerUri, scriptInfo);
+
+            // When I call the language service
+            var result = languageService.GetDefinition(textDocument, scriptFile, connInfo);
+
+            // Then I expect null locations and an error to be reported
+            Assert.NotNull(result);
+            Assert.True(result.IsErrorResult);
+        }
+
+        /// <summary>
+        /// Get Definition for a object with no definition. Expect a error result
+        /// </summary>
+        [Fact]
+        public void GetDefinitionWithNoResultsFoundError()
+        {
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
+            string objectName = "from";
+
+            List<Declaration> declarations = new List<Declaration>();
+            DefinitionResult result = peekDefinition.GetScript(declarations, objectName, null);
+
+            Assert.NotNull(result);
+            Assert.True(result.IsErrorResult);
+            Assert.Equal(SR.PeekDefinitionNoResultsError, result.Message);
+        }
+    
+        /// <summary>
+        /// Test GetDefinition with a forced timeout. Expect a error result.
+        /// </summary>
+        [Fact]
+        public void GetDefinitionTimeoutTest()
+        {
+            // Given a binding queue that will automatically time out
+            var languageService = new LanguageService();
+            Mock<ConnectedBindingQueue> queueMock = new Mock<ConnectedBindingQueue>();
+            languageService.BindingQueue = queueMock.Object;
+            ManualResetEvent mre = new ManualResetEvent(true); // Do not block
+            Mock<QueueItem> itemMock = new Mock<QueueItem>();
+            itemMock.Setup(i => i.ItemProcessed).Returns(mre);
+            
+            DefinitionResult timeoutResult = null;
+            
+            queueMock.Setup(q => q.QueueBindingOperation(
+                It.IsAny<string>(),
+                It.IsAny<Func<IBindingContext, CancellationToken, object>>(),
+                It.IsAny<Func<IBindingContext, object>>(),
+                It.IsAny<int?>(),
+                It.IsAny<int?>()))
+            .Callback<string, Func<IBindingContext, CancellationToken, object>, Func<IBindingContext, object>, int?, int?>(
+                (key, bindOperation, timeoutOperation, blah, blah2) => 
+            {
+                timeoutResult = (DefinitionResult) timeoutOperation((IBindingContext)null);
+                itemMock.Object.Result = timeoutResult;
+            })
+            .Returns(() => itemMock.Object);
 
             ScriptFile scriptFile;
             TextDocumentPosition textDocument = new TextDocumentPosition
@@ -263,12 +353,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
             ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfo(out scriptFile);
             scriptFile.Contents = "select * from dbo.func ()";
 
-            var languageService = LanguageService.Instance;
             ScriptParseInfo scriptInfo = new ScriptParseInfo { IsConnected = true };
             languageService.ScriptParseInfoMap.Add(OwnerUri, scriptInfo);
 
-            var locations = languageService.GetDefinition(textDocument, scriptFile, connInfo);
-            Assert.Null(locations);
+            // When I call the language service
+            var result = languageService.GetDefinition(textDocument, scriptFile, connInfo);
+
+            // Then I expect null locations and an error to be reported
+            Assert.NotNull(result);
+            Assert.True(result.IsErrorResult);
+            // Check timeout message
+            Assert.Equal(SR.PeekDefinitionTimedoutError, result.Message);
         }
 
         /// <summary>
@@ -276,8 +371,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         /// </summary>
         [Fact]
         public void GetValidViewDefinitionTest()
-        {            
-            PeekDefinition peekDefinition = new PeekDefinition(TestObjects.InitLiveConnectionInfoForDefinition());
+        {
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "objects";
             string schemaName = "sys";
             string objectType = "VIEW";
@@ -293,7 +391,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetViewDefinitionInvalidObjectTest()
         {
-            PeekDefinition peekDefinition = new PeekDefinition(TestObjects.InitLiveConnectionInfoForDefinition());
+            // Get live connectionInfo and serverConnection
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "objects";
             string schemaName = null;
             string objectType = "VIEW";
@@ -308,8 +410,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetStoredProcedureDefinitionTest()
         {
-            PeekDefinition peekDefinition = new PeekDefinition(TestObjects.InitLiveConnectionInfoForDefinition());
+            // Get live connectionInfo and serverConnection
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "sp_MSrepl_startup";
+
             string schemaName = "dbo";
             string objectType = "PROCEDURE";
 
@@ -324,7 +431,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetStoredProcedureDefinitionFailureTest()
         {
-            PeekDefinition peekDefinition = new PeekDefinition(TestObjects.InitLiveConnectionInfoForDefinition());
+            // Get live connectionInfo and serverConnection
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "SP2";
             string schemaName = "dbo";
             string objectType = "PROCEDURE";
@@ -339,7 +450,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetStoredProcedureDefinitionWithoutSchemaTest()
         {
-            PeekDefinition peekDefinition = new PeekDefinition(TestObjects.InitLiveConnectionInfoForDefinition());
+            // Get live connectionInfo and serverConnection
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "sp_MSrepl_startup";
             string schemaName = null;
             string objectType = "PROCEDURE";

test/Microsoft.SqlTools.ServiceLayer.Test/Utility/TestObjects.cs

@@ -194,7 +194,7 @@ namespace Microsoft.SqlTools.Test.Utility
             return connInfo;
         }
 
-        public static ServerConnection InitLiveConnectionInfoForDefinition()
+        public static ConnectionInfo InitLiveConnectionInfoForDefinition()
         {
             TestObjects.InitializeTestServices();
 
@@ -212,7 +212,11 @@ namespace Microsoft.SqlTools.Test.Utility
 
             ConnectionInfo connInfo = null;
             connectionService.TryFindConnection(ownerUri, out connInfo);
-            
+            return connInfo;
+        }
+
+        public static ServerConnection InitLiveServerConnectionForDefinition(ConnectionInfo connInfo)
+        {
             SqlConnection sqlConn = new SqlConnection(ConnectionService.BuildConnectionString(connInfo.ConnectionDetails));                                
             return new ServerConnection(sqlConn);
         }