Fix timeout of Invoke-AzOperationalInsightsQuery (#24882)

Author: isra-fel

src/OperationalInsights/OperationalInsights/ChangeLog.md

@@ -19,6 +19,7 @@
 -->
 
 ## Upcoming Release
+* Fixed an issue that `Invoke-AzOperationalInsightsQuery` timed out after 100 seconds. The timeout is now bound to the `-Wait` parameter. (#16553)
 * Removed the outdated deps.json file.
 
 ## Version 3.2.0

src/OperationalInsights/OperationalInsights/Query/InvokeOperationalInsightsQuery.cs

@@ -20,6 +20,8 @@
 using System.Management.Automation;
 using System.Net;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
+using System.Threading;
+using Microsoft.WindowsAzure.Commands.Utilities.Common;
 
 namespace Microsoft.Azure.Commands.OperationalInsights.Query
 {
@@ -28,6 +30,7 @@ public class InvokeOperationalInsightsQuery : ResourceManager.Common.AzureRmLong
     {
         private const string ParamSetNameByWorkspaceId = "ByWorkspaceId";
         private const string ParamSetNameByWorkspaceObject = "ByWorkspaceObject";
+        private readonly double _timeoutBufferInSeconds = 30;
 
         [Parameter(Mandatory = true, ParameterSetName = ParamSetNameByWorkspaceId, HelpMessage = "The workspace ID.")]
         [ValidateNotNullOrEmpty]
@@ -73,11 +76,13 @@ internal OperationalInsightsDataClient OperationalInsightsDataClient
 
                     this._operationalInsightsDataClient =
                         AzureSession.Instance.ClientFactory.CreateCustomArmClient<OperationalInsightsDataClient>(clientCredentials);
+                    ConfigureTimeoutForClient();
+
                     this._operationalInsightsDataClient.Preferences.IncludeRender = IncludeRender.IsPresent;
                     this._operationalInsightsDataClient.Preferences.IncludeStatistics = IncludeStatistics.IsPresent;
                     this._operationalInsightsDataClient.NameHeader = "LogAnalyticsPSClient";
 
-                    Uri targetUri= null;
+                    Uri targetUri = null;
                     DefaultContext.Environment.TryGetEndpointUrl(
                         AzureEnvironment.ExtendedEndpoint.OperationalInsightsEndpoint, out targetUri);
                     if (targetUri == null)
@@ -101,6 +106,14 @@ internal OperationalInsightsDataClient OperationalInsightsDataClient
             }
         }
 
+        private void ConfigureTimeoutForClient()
+        {
+            if (this.IsParameterBound(c => c.Wait) && Wait != null)
+            {
+                _operationalInsightsDataClient.HttpClient.Timeout = TimeSpan.FromSeconds(Convert.ToDouble(Wait)).Add(TimeSpan.FromSeconds(_timeoutBufferInSeconds));
+            }
+        }
+
         public override void ExecuteCmdlet()
         {
             if (ParameterSetName == ParamSetNameByWorkspaceId)

src/OperationalInsights/OperationalInsights/help/Invoke-AzOperationalInsightsQuery.md

@@ -33,6 +33,8 @@ You can retrieve the results of the search from the Value property of the return
 Please check detail of general query limits here: 
 https://learn.microsoft.com/azure/azure-monitor/service-limits#log-queries-and-language.
 
+Note: try setting `-Wait` to a larger value if you experience the error 'The request was canceled due to the configured HttpClient.Timeout of 100 seconds elapsing'.
+
 ## EXAMPLES
 
 ### Example 1: Get search results using a query