Created
August 21, 2024 13:39
-
-
Save JustinGrote/c0295b5683552d0d14c4e62c10fdaea7 to your computer and use it in GitHub Desktop.
PowerShell Azure Monitor Log Ingestion
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| using namespace System.Collections.Generic | |
| using namespace System.Management.Automation | |
| using namespace System.Text | |
| Function Send-AzMonitorLog { | |
| [CmdletBinding()] | |
| param( | |
| [Parameter(Mandatory)] | |
| [string]$Endpoint, | |
| [Parameter(Mandatory)] | |
| [string]$Id, | |
| [Parameter(Mandatory)] | |
| [AllowEmptyString()] | |
| [string]$StreamName, | |
| #The access token to use for authentication. If not specified, will use the AzAuth module to retrieve a token if present. | |
| [SecureString]$AccessToken, | |
| #The objects that represent logs to convert to JSON and send to the DCR. | |
| [Parameter(Mandatory, ValueFromPipeline)] | |
| [PSObject[]]$Data, | |
| #Show the content of the logs being sent in the debug stream. Only do this for troubleshooting purposes. | |
| [switch]$Trace = $ENV:__DCR_TRACE, | |
| #Dont wait for the results of log submission. This will return job objects you can follow up on later or discard. | |
| [switch]$AsJobs, | |
| #Run batches sequentially rather than in parallel. Useful for troubleshooting/debugging | |
| [switch]$Sequential | |
| ) | |
| begin { | |
| $ErrorActionPreference = 'Stop' | |
| $thisName = $MyInvocation.MyCommand.Name | |
| [uint]$MAX_BATCH_BYTES = 1MB | |
| [uint]$BATCH_PAYLOAD_OVERHEAD = 50KB | |
| [int]$TotalLogs = 0 | |
| [int]$TotalBatches = 0 | |
| [uint32]$BatchBytes = 0 | |
| [List[PSObject]]$BatchBuffer = @() | |
| [List[Management.Automation.Job2]]$jobs = @() | |
| if (-not $AccessToken) { | |
| try { | |
| $getAzToken = Get-Command -Name Get-AzToken -CommandType Cmdlet -Module AzAuth | |
| } catch { | |
| throw '-AccessToken was not specified. The AzAuth module can be installed to automatically retrieve a token for you.' | |
| } | |
| if (-not $getAzToken) { | |
| throw '-AccessToken was not specified. The AzAuth module can be installed to automatically retrieve a token for you.' | |
| } | |
| $resourceUrl = 'https://monitor.azure.com/' | |
| $tokenInfo = try { | |
| & $getAzToken -ManagedIdentity -ResourceUrl $resourceUrl | |
| } catch { | |
| & $getAzToken -ResourceUrl $resourceUrl | |
| } | |
| if (-not $tokenInfo) { throw "Failed to retrieve token automatically for resource $resourceUrl." } | |
| $AccessToken = $tokenInfo.token | ConvertTo-SecureString -AsPlainText | |
| $tokenInfo = $null | |
| } | |
| $iwrParams = @{ | |
| Uri = "$Endpoint/dataCollectionRules/$Id/streams/${StreamName}?api-version=2023-01-01" | |
| Method = 'POST' | |
| Authentication = 'Bearer' | |
| Token = $AccessToken | |
| ContentType = 'application/json' | |
| } | |
| } | |
| process { | |
| foreach ($logItem in $Data) { | |
| #TODO: Concatenate this to a stringbuilder with a comma and check the length of the stringbuilder | |
| $jsonObject = ConvertTo-Json -InputObject $logItem -Depth 3 -Compress -EnumsAsStrings | |
| $BatchBytes += $jsonObject.length + 1 #For comma added during concatenation | |
| if ($BatchBytes -gt ($MAX_BATCH_BYTES - $BATCH_PAYLOAD_OVERHEAD)) { | |
| Write-Debug "${thisName}: Batch limit reached" | |
| if ($BatchBuffer.Count -eq 0) { | |
| Write-Error "${thisName}: Individual Log entry is too large to send ($BatchBytes bytes). Max size is $MAX_BATCH_BYTES bytes." | |
| } | |
| $batch = $BatchBuffer.ToArray() | |
| $job = Submit-Logs $iwrParams $batch -Sequential:$Sequential -Trace:$Trace | |
| if ($job) { $jobs.Add($job) } | |
| $TotalLogs += $BatchBuffer.Count | |
| $TotalBatches++ | |
| $BatchBuffer.Clear() | |
| $BatchBytes = 0 | |
| } | |
| $BatchBuffer.Add($logItem) | |
| } | |
| } | |
| end { | |
| if ($BatchBuffer) { | |
| $job = Submit-Logs $iwrParams $batchBuffer -Sequential:$Sequential | |
| if ($job) { $jobs.Add($job) } | |
| $TotalLogs += $BatchBuffer.Count | |
| $TotalBatches++ | |
| } | |
| if (-not $TotalLogs) { | |
| Write-Verbose "${thisName}: No logs were provided and so none were sent." | |
| return | |
| } | |
| Write-Verbose "${thisName}: Submitted $TotalLogs logs in $TotalBatches batches." | |
| if (-not $Sequential) { | |
| if ($AsJobs) { | |
| return $jobs | |
| } | |
| $jobs | Receive-Job -AutoRemoveJob -Wait -ErrorAction 'Continue' #Continue is so we surface multiple batch errors if encountered. | |
| } | |
| } | |
| # clean { | |
| # if (-not $Sequential -and -not $AsJobs) { | |
| # $jobs | Remove-Job -Force | |
| # } | |
| # } | |
| } | |
| function Submit-Logs ([Parameter(Mandatory)][hashtable]$iwrParams, [Parameter(Mandatory)]$Batch, [switch]$Sequential, [switch]$Trace) { | |
| $thisName = 'Send-AzMonitorLog' #Workaround | |
| $requestGUID = (New-Guid).Guid | |
| Write-Debug "${thisName} [$requestGUID]: Submitting $($Batch.Count) records." | |
| $submitTaskParams = [ordered]@{ | |
| iwrParams = $iwrParams | |
| Batch = $Batch | |
| requestGUID = $requestGUID | |
| Trace = $Trace.IsPresent | |
| } | |
| if ($Sequential) { | |
| SubmitLogsTask @submitTaskParams | |
| } else { | |
| Start-ThreadJob -Name "Send-AzMonitorLog-$requestGUID" -ScriptBlock (Get-Item Function:\SubmitLogsTask).ScriptBlock -ArgumentList $submitTaskParams.Values -Verbose -Debug | |
| } | |
| } | |
| function SubmitLogsTask { | |
| param( | |
| [Parameter(Mandatory)] | |
| [hashtable]$iwrParams, | |
| [Parameter(Mandatory)] | |
| [PSObject[]]$Batch, | |
| [Parameter(Mandatory)] | |
| [string]$requestGUID, | |
| [Parameter(Mandatory)] | |
| [bool]$Trace | |
| ) | |
| $VerbosePreference = 'Continue' | |
| $DebugPreference = 'Continue' | |
| $thisName = 'Send-AzMonitorLog' #Workaround | |
| Write-Debug "${thisName} [$requestGUID]: $($iwrParams.Method) $($iwrParams.Uri)" | |
| $jsonBody = ConvertTo-Json $Batch -Depth 3 -Compress -EnumsAsStrings | |
| if ($Trace) { | |
| Write-Debug "${thisName}: ===REQUEST BODY===`n $jsonBody`n ===END REQUEST BODY===" | |
| } | |
| $iwrParams.Headers = @{ 'x-ms-client-request-id' = $requestGUID } | |
| try { | |
| $response = Invoke-WebRequest @iwrParams -Body $jsonBody -Verbose:$false | |
| } catch { | |
| $err = $PSItem | |
| try { | |
| $message = $err.ErrorDetails.Message | |
| if (-not $message.Trim()?.StartsWith('{')) { | |
| throw $err | |
| } | |
| $jsonMessage = ConvertFrom-Json $message | |
| if (-not $jsonMessage.error.code -or -not $jsonMessage.error.message) { | |
| throw $err | |
| } | |
| $PSItem.ErrorDetails = "[$requestGUID]: $($jsonMessage.error.code) - $($jsonMessage.error.message)" | |
| $PSItem.FullyQualifiedErrorId = $jsonMessage.error.code | |
| $PSCmdlet.ThrowTerminatingError($PSItem) | |
| } catch { | |
| throw $err | |
| } | |
| } | |
| if ($response.StatusCode -eq 204) { | |
| if ($response.Headers['x-ms-client-request-id'] -ne $requestGUID) { | |
| Write-Error "Request ID mismatch. Sent: $requestGUID, Received: $($response.Headers['x-ms-client-request-id'])" | |
| } | |
| Write-Debug "${thisName} [$requestGUID]: $($batch.count) Logs ingested [ID: $($response.Headers['x-ms-request-id'])]" | |
| } else { | |
| Write-Error "${thisName} [$requestGUID] [$($response.Headers['x-ms-request-id'])]: Failed to send logs. Status code: $($response.StatusCode). [ID: $($response.Headers['x-ms-request-id'])]" | |
| } | |
| } | |
| Export-ModuleMember Send-AzMonitorLog |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment