added support for connection and command timeout

GreyCorbel · Nov 7, 2024 · f60de63 · f60de63
1 parent 7324050
commit f60de63
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 5 deletions.
diff --git a/Commands/Public/Invoke-ExoCommand.ps1 b/Commands/Public/Invoke-ExoCommand.ps1
@@ -60,6 +60,13 @@ This command creates connection for IPPS REST API, retrieves list of sensitivity
             #Default is 100
         $PageSize = 100,
 
+        [Parameter()]
+        [System.Nullable[timespan]]
+            #Timeout for the command execution
+            #Default is timeout of the connection
+            #If specified, must be lower than connection timeout
+        $Timeout,
+
         [switch]
             #If we want to write any warnings returned by EXO REST API
         $ShowWarnings,
@@ -80,6 +87,14 @@ This command creates connection for IPPS REST API, retrieves list of sensitivity
 
     begin
     {
+        if($null -ne $Timeout)
+        {
+            $cts = new-object System.Threading.CancellationTokenSource($Timeout)
+        }
+        else {
+            $cts = new-object System.Threading.CancellationTokenSource($Connection.HttpClient.Timeout)
+        }
+
         $body = @{}
         if($PageSize -gt 1000)
         {
@@ -139,7 +154,7 @@ This command creates connection for IPPS REST API, retrieves list of sensitivity
                 #provide up to date token for each request of commands returning paged results that may take long to complete
                 $headers['Authorization'] = (Get-ExoToken -Connection $Connection).CreateAuthorizationHeader()
                 $requestMessage = GetRequestMessage -Uri $pageUri -Headers $headers -Body ($body | ConvertTo-Json -Depth 9)
-                $response = $Connection.HttpClient.SendAsync($requestMessage).GetAwaiter().GetResult()
+                $response = $Connection.HttpClient.SendAsync($requestMessage, $cts.Token).GetAwaiter().GetResult()
                 $requestMessage.Dispose()
 
                 if($null -ne $response.Content -and $response.StatusCode -ne [System.Net.HttpStatusCode]::NoContent)
@@ -243,6 +258,10 @@ This command creates connection for IPPS REST API, retrieves list of sensitivity
             }
             finally
             {
+                if($null -ne $cts)
+                {
+                    $cts.Dispose()
+                }
                 if($ShowRateLimits)
                 {
                     $val = $null

diff --git a/Commands/Public/New-ExoConnection.ps1 b/Commands/Public/New-ExoConnection.ps1
@@ -47,6 +47,11 @@ param
         [string]
         $AnchorMailbox,
 
+        [Parameter()]
+        [System.Nullable[timespan]]
+            #Default timeout for the EXO command execution
+        $DefaultTimeout,
+
         [switch]
         #Connection is specialized to call IPPS commands
         #If not present, connection is specialized to call Exchange Online commands
@@ -65,7 +70,9 @@ param
             IsIPPS = $IPPS.IsPresent
             HttpClient = new-object System.Net.Http.HttpClient
         }
-
+        $Connection.HttpClient.DefaultRequestHeaders.Add("User-Agent", "ExoHelper")
+        $Connection.HttpClient.Timeout = [timespan]::FromMinutes(60)
+
         #explicitly authenticate when establishing connection to catch any authentication problems early
         Get-ExoToken -Connection $Connection | Out-Null
         if([string]::IsNullOrEmpty($TenantId))

diff --git a/Module/ExoHelper/ExoHelper.psd1 b/Module/ExoHelper/ExoHelper.psd1
@@ -12,7 +12,7 @@
 RootModule = 'ExoHelper.psm1'
 
 # Version number of this module.
-ModuleVersion = '3.0.2'
+ModuleVersion = '3.0.3'
 
 # Supported PSEditions
 CompatiblePSEditions = @('Desktop', 'Core')

diff --git a/Module/ExoHelper/ExoHelper.psm1 b/Module/ExoHelper/ExoHelper.psm1
@@ -120,6 +120,13 @@ This command creates connection for IPPS REST API, retrieves list of sensitivity
             #Default is 100
         $PageSize = 100,
 
+        [Parameter()]
+        [System.Nullable[timespan]]
+            #Timeout for the command execution
+            #Default is timeout of the connection
+            #If specified, must be lower than connection timeout
+        $Timeout,
+
         [switch]
             #If we want to write any warnings returned by EXO REST API
         $ShowWarnings,
@@ -140,6 +147,14 @@ This command creates connection for IPPS REST API, retrieves list of sensitivity
 
     begin
     {
+        if($null -ne $Timeout)
+        {
+            $cts = new-object System.Threading.CancellationTokenSource($Timeout)
+        }
+        else {
+            $cts = new-object System.Threading.CancellationTokenSource($Connection.HttpClient.Timeout)
+        }
+
         $body = @{}
         if($PageSize -gt 1000)
         {
@@ -199,7 +214,7 @@ This command creates connection for IPPS REST API, retrieves list of sensitivity
                 #provide up to date token for each request of commands returning paged results that may take long to complete
                 $headers['Authorization'] = (Get-ExoToken -Connection $Connection).CreateAuthorizationHeader()
                 $requestMessage = GetRequestMessage -Uri $pageUri -Headers $headers -Body ($body | ConvertTo-Json -Depth 9)
-                $response = $Connection.HttpClient.SendAsync($requestMessage).GetAwaiter().GetResult()
+                $response = $Connection.HttpClient.SendAsync($requestMessage, $cts.Token).GetAwaiter().GetResult()
                 $requestMessage.Dispose()
 
                 if($null -ne $response.Content -and $response.StatusCode -ne [System.Net.HttpStatusCode]::NoContent)
@@ -303,6 +318,10 @@ This command creates connection for IPPS REST API, retrieves list of sensitivity
             }
             finally
             {
+                if($null -ne $cts)
+                {
+                    $cts.Dispose()
+                }
                 if($ShowRateLimits)
                 {
                     $val = $null
@@ -387,7 +406,9 @@ param
             IsIPPS = $IPPS.IsPresent
             HttpClient = new-object System.Net.Http.HttpClient
         }
-
+        $Connection.HttpClient.DefaultRequestHeaders.Add("User-Agent", "ExoHelper")
+        $Connection.HttpClient.Timeout = [timespan]::FromMinutes(60)
+
         #explicitly authenticate when establishing connection to catch any authentication problems early
         Get-ExoToken -Connection $Connection | Out-Null
         if([string]::IsNullOrEmpty($TenantId))