Return context after it has been generated (#2194)

* Emit generate scripts complete event to client

* Rename Message to ErrorMessage

* Sets owner URI for complete params obj

* Setting complete flag explicitly

* Making errorMessage prop nullable

* Localizes error messages

* Return context scripts and remove script tabs

* Send event when script gen isn't needed

* Change notification to request endpoint

* test get context when context doesn't exist

* Stop reading old context files

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.cs

@@ -437,6 +437,14 @@ namespace Microsoft.SqlTools.ServiceLayer
             }
         }
 
+        public static string FailedToGenerateServerContextualizationScripts
+        {
+            get
+            {
+                return Keys.GetString(Keys.FailedToGenerateServerContextualizationScripts);
+            }
+        }
+
         public static string PeekDefinitionNoResultsError
         {
             get
@@ -10862,6 +10870,11 @@ namespace Microsoft.SqlTools.ServiceLayer
             return Keys.GetString(Keys.ScriptTooLarge, maxChars, currentChars);
         }
 
+        public static string WritingServerContextualizationToCacheError(string message)
+        {
+            return Keys.GetString(Keys.WritingServerContextualizationToCacheError, message);
+        }
+
         public static string SerializationServiceUnsupportedFormat(string formatName)
         {
             return Keys.GetString(Keys.SerializationServiceUnsupportedFormat, formatName);
@@ -11431,6 +11444,12 @@ namespace Microsoft.SqlTools.ServiceLayer
             public const string ScriptTooLarge = "ScriptTooLarge";
 
 
+            public const string WritingServerContextualizationToCacheError = "WritingServerContextualizationToCacheError";
+
+
+            public const string FailedToGenerateServerContextualizationScripts = "FailedToGenerateServerContextualizationScripts";
+
+
             public const string SerializationServiceUnsupportedFormat = "SerializationServiceUnsupportedFormat";
 
 

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.resx

@@ -426,6 +426,15 @@
     <comment>.
  Parameters: 0 - maxChars (int), 1 - currentChars (int) </comment>
   </data>
+  <data name="WritingServerContextualizationToCacheError" xml:space="preserve">
+    <value>An error was encountered while writing server contextualization scripts to the cache. Error: {0}</value>
+    <comment>.
+ Parameters: 0 - message (string) </comment>
+  </data>
+  <data name="FailedToGenerateServerContextualizationScripts" xml:space="preserve">
+    <value>Failed to generate server contextualization scripts</value>
+    <comment></comment>
+  </data>
   <data name="SerializationServiceUnsupportedFormat" xml:space="preserve">
     <value>Unsupported Save Format: {0}</value>
     <comment>.

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.strings

@@ -188,6 +188,13 @@ ParsingErrorHeader (int lineNumber, int columnNumber) = Line {0}, column {1}
 
 ScriptTooLarge (int maxChars, int currentChars) = The current script is too large for Parameterization for Always Encrypted, please disable Parameterization for Always Encrypted in Query Options (Query > Query Options > Execution > Advanced). Maximum allowable length: {0} characters, Current script length: {1} characters
 
+############################################################################
+# Server Contextualization Service
+
+WritingServerContextualizationToCacheError (string message) = An error was encountered while writing server contextualization scripts to the cache. Error: {0}
+
+FailedToGenerateServerContextualizationScripts = Failed to generate server contextualization scripts
+
 ############################################################################
 # Serialization Service
 

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.xlf

@@ -1,4 +1,4 @@
-<?xml version="1.0" encoding="utf-8"?>
+<?xml version="1.0" encoding="utf-8"?>
 <xliff version="1.2" xmlns="urn:oasis:names:tc:xliff:document:1.2" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="urn:oasis:names:tc:xliff:document:1.2 xliff-core-1.2-transitional.xsd">
   <file datatype="xml" original="sr.resx" source-language="en">
     <body>
@@ -7248,6 +7248,17 @@ The Query Processor estimates that implementing the following index could improv
         <target state="new">Error parsing ScriptingListObjectsCompleteParams.ConnectionString property.</target>
         <note></note>
       </trans-unit>
+      <trans-unit id="WritingServerContextualizationToCacheError">
+        <source>An error was encountered while writing server contextualization scripts to the cache. Error: {0}</source>
+        <target state="new">An error was encountered while writing server contextualization scripts to the cache. Error: {0}</target>
+        <note>.
+ Parameters: 0 - message (string) </note>
+      </trans-unit>
+      <trans-unit id="FailedToGenerateServerContextualizationScripts">
+        <source>Failed to generate server contextualization scripts</source>
+        <target state="new">Failed to generate server contextualization scripts</target>
+        <note></note>
+      </trans-unit>
     </body>
   </file>
 </xliff>

src/Microsoft.SqlTools.ServiceLayer/Metadata/Contracts/GenerateServerContextualizationNotification.cs

@@ -1,26 +0,0 @@
-//
-// 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.Hosting.Protocol.Contracts;
-
-namespace Microsoft.SqlTools.ServiceLayer.Metadata.Contracts
-{
-    public class GenerateServerContextualizationParams
-    {
-        /// <summary>
-        /// The URI of the connection to generate context for.
-        /// </summary>
-        public string OwnerUri { get; set; }
-    }
-
-    /// <summary>
-    /// Event set after a connection to a server is completed.
-    /// </summary>
-    public class GenerateServerContextualizationNotification
-    {
-        public static readonly EventType<GenerateServerContextualizationParams> Type =
-            EventType<GenerateServerContextualizationParams>.Create("metadata/generateServerContext");
-    }
-}

src/Microsoft.SqlTools.ServiceLayer/Metadata/Contracts/GenerateServerContextualizationRequest.cs

@@ -0,0 +1,34 @@
+//
+// 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.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.Metadata.Contracts
+{
+    public class GenerateServerContextualizationParams
+    {
+        /// <summary>
+        /// The URI of the connection to generate context for.
+        /// </summary>
+        public string OwnerUri { get; set; }
+    }
+
+    public class GenerateServerContextualizationResult
+    {
+        /// <summary>
+        /// The generated server context.
+        /// </summary>
+        public string? Context { get; set; }
+    }
+
+    /// <summary>
+    /// Generate server context request endpoint.
+    /// </summary>
+    public class GenerateServerContextualizationRequest
+    {
+        public static readonly RequestType<GenerateServerContextualizationParams, GenerateServerContextualizationResult> Type =
+            RequestType<GenerateServerContextualizationParams, GenerateServerContextualizationResult>.Create("metadata/generateServerContext");
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/Metadata/Contracts/GetServerContextualizationRequest.cs

@@ -18,9 +18,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Metadata.Contracts
     public class GetServerContextualizationResult
     {
         /// <summary>
-        /// An array containing the generated server context.
+        /// The generated server context.
         /// </summary>
-        public string[] Context { get; set; }
+        public string? Context { get; set; }
     }
 
     public class GetServerContextualizationRequest

src/Microsoft.SqlTools.ServiceLayer/Metadata/MetadataService.cs

@@ -59,7 +59,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Metadata
             serviceHost.SetRequestHandler(MetadataListRequest.Type, HandleMetadataListRequest, true);
             serviceHost.SetRequestHandler(TableMetadataRequest.Type, HandleGetTableRequest, true);
             serviceHost.SetRequestHandler(ViewMetadataRequest.Type, HandleGetViewRequest, true);
-            serviceHost.SetEventHandler(GenerateServerContextualizationNotification.Type, HandleGenerateServerContextualizationNotification, true);
+            serviceHost.SetRequestHandler(GenerateServerContextualizationRequest.Type, HandleGenerateServerContextualizationNotification, true);
             serviceHost.SetRequestHandler(GetServerContextualizationRequest.Type, HandleGetServerContextualizationRequest, true);
         }
 
@@ -125,11 +125,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Metadata
         /// Handles the event for generating server contextualization scripts.
         /// </summary>
         internal static Task HandleGenerateServerContextualizationNotification(GenerateServerContextualizationParams contextualizationParams,
-            EventContext eventContext)
+            RequestContext<GenerateServerContextualizationResult> requestContext)
         {
-            _ = Task.Factory.StartNew(() =>
+            _ = Task.Factory.StartNew(async () =>
             {
-                GenerateServerContextualization(contextualizationParams);
+                await GenerateServerContextualization(contextualizationParams, requestContext);
             },
             CancellationToken.None,
             TaskCreationOptions.None,
@@ -143,7 +143,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Metadata
         /// database objects like tables and views.
         /// </summary>
         /// <param name="contextualizationParams">The contextualization parameters.</param>
-        internal static void GenerateServerContextualization(GenerateServerContextualizationParams contextualizationParams)
+        internal static async Task GenerateServerContextualization(GenerateServerContextualizationParams contextualizationParams, RequestContext<GenerateServerContextualizationResult> requestContext)
         {
             MetadataService.ConnectionServiceInstance.TryFindConnection(contextualizationParams.OwnerUri, out ConnectionInfo connectionInfo);
 
@@ -153,26 +153,39 @@ namespace Microsoft.SqlTools.ServiceLayer.Metadata
                 {
                     // If scripts have been generated within the last 30 days then there isn't a need to go through the process
                     // of generating scripts again.
-                    if (!MetadataScriptTempFileStream.IsScriptTempFileUpdateNeeded(connectionInfo.ConnectionDetails.ServerName))
+                    if (MetadataScriptTempFileStream.IsScriptTempFileUpdateNeeded(connectionInfo.ConnectionDetails.ServerName))
                     {
-                        return;
+                        var scripts = SmoScripterHelpers.GenerateAllServerTableScripts(sqlConn)?.ToArray();
-                    }
+                        if (scripts != null)
+                        {
+                            try
+                            {
+                                await requestContext.SendResult(new GenerateServerContextualizationResult()
+                                {
+                                    Context = string.Join('\n', scripts)
+                                });
 
-                    var scripts = SmoScripterHelpers.GenerateAllServerTableScripts(sqlConn);
+                                MetadataScriptTempFileStream.Write(connectionInfo.ConnectionDetails.ServerName, scripts);
-                    if (scripts != null)
+                            }
-                    {
+                            catch (Exception ex)
-                        try
+                            {
-                        {
+                                Logger.Error($"An error was encountered while writing server contextualization scripts to the cache. Error: {ex.Message}");
-                            MetadataScriptTempFileStream.Write(connectionInfo.ConnectionDetails.ServerName, scripts);
+                                await requestContext.SendError(ex);
+                            }
                         }
-                        catch (Exception ex)
+                        else
                         {
-                            Logger.Error($"An error was encountered while writing to the cache. Error: {ex.Message}");
+                            Logger.Error("Failed to generate server contextualization scripts");
+                            await requestContext.SendError(SR.FailedToGenerateServerContextualizationScripts);
                         }
                     }
                     else
                     {
-                        Logger.Error("Failed to generate server scripts");
+                        var generateContextResult = new GenerateServerContextualizationResult()
+                        {
+                            Context = null
+                        };
+                        await requestContext.SendResult(generateContextResult);
                     }
                 }
             }
@@ -204,27 +217,37 @@ namespace Microsoft.SqlTools.ServiceLayer.Metadata
         internal static async Task GetServerContextualization(GetServerContextualizationParams contextualizationParams, RequestContext<GetServerContextualizationResult> requestContext)
         {
             MetadataService.ConnectionServiceInstance.TryFindConnection(contextualizationParams.OwnerUri, out ConnectionInfo connectionInfo);
-
+            // When the filed context is too old don't read it
-            if (connectionInfo != null)
+            if (MetadataScriptTempFileStream.IsScriptTempFileUpdateNeeded(connectionInfo.ConnectionDetails.ServerName))
             {
-                try
+                await requestContext.SendResult(new GetServerContextualizationResult
                 {
-                    var scripts = MetadataScriptTempFileStream.Read(connectionInfo.ConnectionDetails.ServerName);
+                    Context = null
-                    await requestContext.SendResult(new GetServerContextualizationResult
+                });
-                    {
-                        Context = scripts.ToArray()
-                    });
-                }
-                catch (Exception ex)
-                {
-                    Logger.Error("Failed to read scripts from the script cache");
-                    await requestContext.SendError(ex);
-                }
             }
             else
             {
-                Logger.Error("Failed to find connection info about the server.");
+                if (connectionInfo != null)
-                await requestContext.SendError("Failed to find connection info about the server.");
+                {
+                    try
+                    {
+                        var scripts = MetadataScriptTempFileStream.Read(connectionInfo.ConnectionDetails.ServerName).ToArray();
+                        await requestContext.SendResult(new GetServerContextualizationResult
+                        {
+                            Context = string.Join('\n', scripts)
+                        });
+                    }
+                    catch (Exception ex)
+                    {
+                        Logger.Error("Failed to read scripts from the script cache");
+                        await requestContext.SendError(ex);
+                    }
+                }
+                else
+                {
+                    Logger.Error("Failed to find connection info about the server.");
+                    await requestContext.SendError("Failed to find connection info about the server.");
+                }
             }
         }
 

src/Microsoft.SqlTools.ServiceLayer/Metadata/SmoScripterHelpers.cs

@@ -161,11 +161,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Metadata
             var scripts = new List<string>();
             foreach (var s in generatedScripts)
             {
-                // Needed to remove '\r' and '\n' characters from script, so that an entire create script
+                // Needed to remove '\r', '\n', and '\t' characters from script, so that an entire create script
                 // can be written and read as a single line to and from a temp file. Since scripts aren't
-                // going to be read by people, and mainly sent to Copilot to generate accurate suggestions,
+                // going to be read by people, and sent to Copilot to generate accurate suggestions,
                 // a lack of formatting is fine.
-                var script = s.Replace("\r", string.Empty).Replace("\n", string.Empty);
+                var script = s.Replace("\r", string.Empty).Replace("\n", string.Empty).Replace("\t", string.Empty);
                 scripts.Add(script);
             }
 

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/Metadata/MetadataServiceTests.cs

@@ -150,11 +150,22 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Metadata
                 OwnerUri = connectionResult.ConnectionInfo.OwnerUri
             };
 
-            MetadataService.GenerateServerContextualization(generateServerContextualizationParams);
+            var actualGenerateServerContextResponse = new GenerateServerContextualizationResult();
+            var mockGenerateRequestContext = new Mock<RequestContext<GenerateServerContextualizationResult>>();
+            mockGenerateRequestContext.Setup(x => x.SendResult(It.IsAny<GenerateServerContextualizationResult>()))
+                .Callback<GenerateServerContextualizationResult>(result => actualGenerateServerContextResponse = result)
+                .Returns(Task.CompletedTask);
+            await MetadataService.GenerateServerContextualization(generateServerContextualizationParams, mockGenerateRequestContext.Object);
 
             DeleteTestTable(sqlConn, this.testTableSchema, this.testTableName);
             DeleteTestTable(sqlConn, this.testTableSchema, this.testTableName2);
 
+            var firstCreateTableScript = $"CREATE TABLE [{this.testTableSchema}].[{this.testTableName}]([id] [int] NULL)";
+            var secondCreateTableScript = $"CREATE TABLE [{this.testTableSchema}].[{this.testTableName2}]([id] [int] NULL)";
+
+            Assert.IsTrue(actualGenerateServerContextResponse.Context.Contains(firstCreateTableScript));
+            Assert.IsTrue(actualGenerateServerContextResponse.Context.Contains(secondCreateTableScript));
+
             DeleteServerContextualizationTempFile(sqlConn.DataSource);
         }
 
@@ -170,18 +181,8 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Metadata
             CreateTestTable(sqlConn, this.testTableSchema, this.testTableName);
             CreateTestTable(sqlConn, this.testTableSchema, this.testTableName2);
 
-            var generateServerContextualizationParams = new GenerateServerContextualizationParams
+            var firstCreateTableScript = $"CREATE TABLE [{this.testTableSchema}].[{this.testTableName}]([id] [int] NULL)";
-            {
+            var secondCreateTableScript = $"CREATE TABLE [{this.testTableSchema}].[{this.testTableName2}]([id] [int] NULL)";
-                OwnerUri = connectionResult.ConnectionInfo.OwnerUri
-            };
-
-            MetadataService.GenerateServerContextualization(generateServerContextualizationParams);
-
-            DeleteTestTable(sqlConn, this.testTableSchema, this.testTableName);
-            DeleteTestTable(sqlConn, this.testTableSchema, this.testTableName2);
-
-            var firstCreateTableScript = $"CREATE TABLE [{this.testTableSchema}].[{this.testTableName}](\t[id] [int] NULL)";
-            var secondCreateTableScript = $"CREATE TABLE [{this.testTableSchema}].[{this.testTableName2}](\t[id] [int] NULL)";
 
             var mockGetServerContextualizationRequestContext = new Mock<RequestContext<GetServerContextualizationResult>>();
             var actualGetServerContextualizationResponse = new GetServerContextualizationResult();
@@ -193,6 +194,25 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Metadata
             {
                 OwnerUri = connectionResult.ConnectionInfo.OwnerUri
             };
+            await MetadataService.GetServerContextualization(getServerContextualizationParams, mockGetServerContextualizationRequestContext.Object);
+
+            Assert.IsNull(actualGetServerContextualizationResponse.Context);
+            Assert.IsNull(actualGetServerContextualizationResponse.Context);
+
+            var generateServerContextualizationParams = new GenerateServerContextualizationParams
+            {
+                OwnerUri = connectionResult.ConnectionInfo.OwnerUri
+            };
+
+            var actualGenerateServerContextResponse = new GenerateServerContextualizationResult();
+            var mockGenerateRequestContext = new Mock<RequestContext<GenerateServerContextualizationResult>>();
+            mockGenerateRequestContext.Setup(x => x.SendResult(It.IsAny<GenerateServerContextualizationResult>()))
+                .Callback<GenerateServerContextualizationResult>(actual => actualGenerateServerContextResponse = actual)
+                .Returns(Task.CompletedTask);
+            await MetadataService.GenerateServerContextualization(generateServerContextualizationParams, mockGenerateRequestContext.Object);
+
+            DeleteTestTable(sqlConn, this.testTableSchema, this.testTableName);
+            DeleteTestTable(sqlConn, this.testTableSchema, this.testTableName2);
 
             await MetadataService.GetServerContextualization(getServerContextualizationParams, mockGetServerContextualizationRequestContext.Object);