Skip to content

Instantly share code, notes, and snippets.

@SQLDBAWithABeard
Forked from JustinGrote/Trace-AICommand.ps1
Created June 27, 2022 09:15
Show Gist options
  • Save SQLDBAWithABeard/92818c366da94a0da4191cf47453eef0 to your computer and use it in GitHub Desktop.
Save SQLDBAWithABeard/92818c366da94a0da4191cf47453eef0 to your computer and use it in GitHub Desktop.
Report the results and performance of any scriptblock to Azure Application Insights
#requires -version 7
#You can load this script with $(iwr https://tinyurl.com/TraceAICommand | iex)
using namespace Microsoft.ApplicationInsights
using namespace Microsoft.ApplicationInsights.Extensibility
using namespace Microsoft.ApplicationInsights.DataContracts
using namespace System.Management.Automation
using namespace System.Collections.Generic
using namespace System.Net
#Reference: https://docs.microsoft.com/en-us/azure/azure-monitor/app/console
#Reference: https://docs.microsoft.com/en-us/azure/azure-monitor/app/custom-operations-tracking
function Trace-AICommand {
[CmdletBinding()]
param(
#The scriptblock to execute. It will run in the current scope.
[Parameter(Mandatory, ValueFromPipeline)][ScriptBlock]$ScriptBlock,
#A label for this operation to more easily identify it in the telemetry
[ValidateNotNullOrEmpty()][String]$Name = 'PowerShell ScriptBlock',
#Specify where this request is originating from. Defaults to the current hostname but you may want to define something custom like 'AzureAutomation'
[String]$InstanceName = [Environment]::MachineName,
#Specify the app name which will represent an individual instance in the app map. By default we use the $PSScriptRoot if this is a script and 'Powershell ScriptBlock' otherwise.
[ValidateNotNullOrEmpty()][String]$RoleName = $MyInvocation.ScriptName ? $(Split-Path $MyInvocation.ScriptName -Leaf) : 'PowerShell ScriptBlock',
#By default we log nonterminating errors as exceptions (but requests can still show OK). Set this if you want them tracked as traces instead.
[Switch]$ErrorAsTrace,
#By default we flush after each invocation but don't wait for it to complete. You will want to specify this if you want to ensure your telemetry uploads after the scriptblock before moving on.
[Switch]$Wait,
#Specify the connection string for your app insights instance. You can also specify this in the environment variable APPLICATIONINSIGHTS_CONNECTION_STRING
[String]$ConnectionString = $(Get-Content ENV:APPLICATIONINSIGHTS_CONNECTION_STRING -ErrorAction SilentlyContinue),
#By Default we don't log the script output however you can enable this. Be very careful with this setting, you can incur a lot of cost logging large outputs.
[Switch]$LogOutput,
#Specify if you would like the output to be serialized to JSON first before being sent to telemetry
[Switch]$OutputAsJson,
#Specify how many properties deep you would like to check.
[int]$OutputAsJsonDepth = 3,
#Specify an operation ID for a parent operation. This will be automatically recursively referenced if Trace-AICommand is referenced within itself
[string]$ParentId = $__TRACEAICOMMAND_CURRENTID,
#If you are supplying a parentId for an external process correlation (e.g. microservices call), specify this flag as well. You don't normally need to do this.
[string]$NoDependency,
#Send telemetry as fast as possible. Only use for testing, this is not good for performance.
[Switch]$DeveloperMode,
[SeverityLevel]$WarningAs = [SeverityLevel]::Warning,
[SeverityLevel]$ErrorAs = [SeverityLevel]::Error,
[SeverityLevel]$VerboseAs = [SeverityLevel]::Verbose,
[SeverityLevel]$DebugAs = [SeverityLevel]::Verbose,
[SeverityLevel]$InformationAs = [SeverityLevel]::Information
)
#BUG: ValidateNotNullOrEmpty doesnt apply if a default value might be null so we have to do this extra check
if ([string]::IsNullOrEmpty($connectionString)) { throw 'You must specify a connection string or set the environment variable APPLICATIONINSIGHTS_CONNECTION_STRING' }
#Instantiate a common telemetry client. This is the rare case GLOBAL makes sense as we want it to be process-level.
$GLOBAL:__AIClient ??= [Dictionary[guid, TelemetryClient]]@{}
if (-not $connectionString) {
Write-Warning 'You must specify a connection string for your app insights instance either via the ConnectionString parameter or by setting the environment variable APPLICATIONINSIGHTS_CONNECTION_STRING. The scriptblock will now be run without telemetry'
& $ScriptBlock
return
}
[guid]$InstrumentationKey = $connectionString -replace 'InstrumentationKey=([\w-]+)?;.+', '$1'
if (-not $GLOBAL:__AIClient.ContainsKey($InstrumentationKey)) {
$config = [TelemetryConfiguration]::CreateDefault()
$config.ConnectionString = $connectionString
$GLOBAL:__AIClient[$InstrumentationKey] = [TelemetryClient]::New($config)
}
[TelemetryClient]$client = $GLOBAL:__AIClient[$InstrumentationKey]
if ($DeveloperMode) { $client.TelemetryConfiguration.TelemetryChannel.DeveloperMode = $true }
$client.Context.Cloud.RoleInstance = $InstanceName
$client.Context.Cloud.RoleName = $RoleName
#This should expire when this goes out of scope and allows for recursive operation dependencies
#If we have a parentId and detect this is nested, a dependency item needs to be linked between our request and the parent to draw the
#Lines in application map. This will be disposed once the current activity is completed.
if ($ParentId -and $__TRACEAICOMMAND_CURRENTROLENAME -and $__TRACEAICOMMAND_CURRENTROLEINSTANCE) {
[IOperationHolder[DependencyTelemetry]]$LOCAL:dependency = Start-AIDependency -Client $client -OperationName 'ChildScriptBlock' -ParentId $ParentId
# $dependency.Telemetry.Type = 'script'
$context = $LOCAL:dependency.Telemetry.Context
$context.Cloud.RoleInstance = $__TRACEAICOMMAND_CURRENTROLEINSTANCE
$context.Cloud.RoleName = $__TRACEAICOMMAND_CURRENTROLENAME
#By replacing this, our Start-AIOperation will chain to the dependency link, not the original request
$ParentId = $LOCAL:dependency.Telemetry.Id
}
[IOperationHolder[RequestTelemetry]]$operation = Start-AIOperation -Client $client -OperationName $Name -ParentId $ParentId
# 'Tracking AppInsights Operation {0} OpId:{1} ParentId: {2}' -f $Name, $operation.Telemetry.Id, $operation.Telemetry.Context.Operation.ParentId | Write-Debug
[string]$LOCAL:__TRACEAICOMMAND_CURRENTID = $operation.Telemetry.Id
[string]$LOCAL:__TRACEAICOMMAND_CURRENTROLENAME = $operation.Telemetry.Context.Cloud.RoleName
[string]$LOCAL:__TRACEAICOMMAND_CURRENTROLEINSTANCE = $operation.Telemetry.Context.Cloud.RoleInstance
[bool]$isFailed = $false
try {
& $ScriptBlock *>&1 | ForEach-Object {
$object = $PSItem
#Should only happen once per object at its deepest scope
Write-AITelemetry $object
#Child Operations should pass thru on output stream
if ($ParentId) {
return $object
}
#The parent operation re-emits the object on the appropriate stream
if ($object.__TRACEAICOMMAND_PROCESSED) {
[void]$object.PSObject.Properties.Remove('__TRACEAICOMMAND_PROCESSED')
}
switch ($object.GetType().Name) {
'DebugRecord' { $PSCmdlet.WriteDebug($object) }
'VerboseRecord' { $PSCmdlet.WriteVerbose($object) }
'InformationRecord' { $PSCmdlet.WriteInformation($object) }
'WarningRecord' { $PSCmdlet.WriteWarning($object) }
'ErrorRecord' { $PSCmdlet.WriteError($object) }
default { $PSCmdlet.WriteObject($object) }
}
}
} catch {
if (-not $PSItem.Exception.Data['TraceIACommandProcessed']) {
$client.TrackException(
$PSItem.Exception,
(Get-AIErrorRecordCustomProperties $PSItem)
)
#This should prevent further throws up the chain from being logged
$psitem.exception.data['TraceIACommandProcessed'] = $true
}
$isFailed = $true
throw $PSItem
} finally {
Stop-AIOperation -Client $client -Operation $operation -Failed:$isFailed
if ($LOCAL:dependency) {
$LOCAL:dependency.Dispose()
}
if ($Wait) {
$client.Flush()
} else {
[void]$client.FlushAsync([System.Threading.CancellationTokenSource]::new().Token)
}
}
}
function Start-AIOperation ([TelemetryClient]$client, [String]$ParentId, [String]$OperationName, [String]$OperationId) {
#TODO: In 7.3 we can call this generic method more succinctly
$startOperationMethod = [TelemetryClientExtensions].GetMethod(
'StartOperation',
[type[]]@([telemetryclient], [string], [string], [string])).MakeGenericMethod([RequestTelemetry]
)
$operation = [IOperationHolder[RequestTelemetry]]$startOperationMethod.Invoke([TelemetryClientExtensions], @($client, $OperationName, $OperationId, $ParentId))
return $operation
}
function Start-AIDependency ([TelemetryClient]$client, [String]$ParentId, [String]$OperationName, [String]$OperationId) {
[OutputType([IOperationHolder[DependencyTelemetry]])]
$startDependency = [TelemetryClientExtensions].GetMethod(
'StartOperation',
[type[]]@(
[telemetryclient], [string], [string], [string]
)).MakeGenericMethod([DependencyTelemetry])
[IOperationHolder[DependencyTelemetry]]$startDependency.Invoke(
[TelemetryClientExtensions],
@(
$client,
$OperationName,
$OperationId,
$ParentId
)
)
}
function Stop-AIOperation ([TelemetryClient]$client, [IOperationHolder[RequestTelemetry]]$Operation, [Switch]$Failed) {
$Operation.Telemetry.Success = -not $Failed
$Operation.Telemetry.ResponseCode = $Failed ? 'TerminatingError' : 'OK'
$Operation.Dispose()
}
function Get-AIErrorRecordCustomProperties ([ErrorRecord]$record) {
[OutputType([Dictionary[String, String]])]
[Dictionary[String, String]]$customProperties = @{}
$customProperties.TargetObject = $record.TargetObject
$customProperties.FullyQualifiedErrorId = $record.FullyQualifiedErrorId
$customProperties.PositionMessage = $record.InvocationInfo.PositionMessage
$customProperties.Category = $record.CategoryInfo.Category
$customProperties.Reason = $record.CategoryInfo.Reason
$customProperties.ScriptStackTrace = $record.ScriptStackTrace
return $customProperties
}
function Write-AITelemetry ($object) {
<#
.SYNOPSIS
Does the heavy lifting of sending the telemetry to appinsights, and will flag the object as processed so it only happens once per object.
#>
if ($object.__TRACEAICOMMAND_PROCESSED) {
return
}
switch ($object.GetType().Name) {
'VerboseRecord' {
if (-not $__TRACEAICOMMAND_NOTRACE) {
$client.TrackTrace(
$object.Message,
$VerboseAs
)
}
}
'DebugRecord' {
if (-not $__TRACEAICOMMAND_NOTRACE) {
[Dictionary[String, String]]$customProperties = @{}
$customProperties.Severity = 'Debug'
$client.TrackTrace(
$object,
$DebugAs,
$customProperties
)
}
}
'WarningRecord' {
if (-not $__TRACEAICOMMAND_NOTRACE) {
$client.TrackTrace(
$object,
$WarningAs
)
}
}
'InformationRecord' {
if (-not $__TRACEAICOMMAND_NOTRACE) {
[InformationRecord]$record = $object
[Dictionary[String, String]]$customProperties = @{}
foreach ($tag in $record.Tags) {
$customProperties["TAG:$tag"] = 'true'
}
$client.TrackTrace(
$record.MessageData,
$InformationAs,
$customProperties
)
}
}
'ErrorRecord' {
if (-not $__TRACEAICOMMAND_NOTRACE) {
$err = [ErrorRecord]$object
if ($ErrorAsTrace) {
$client.TrackTrace(
('{0}: {1}' -f $err.InvocationInfo.MyCommand, $err.ToString()),
$ErrorAs,
(Get-AIErrorRecordCustomProperties $err)
)
} else {
$client.TrackException(
$err.Exception,
(Get-AIErrorRecordCustomProperties $err)
)
}
}
$PSCmdlet.WriteError($err)
}
default {
if (-not $__TRACEAICOMMAND_NOTRACE) {
if ($LogOutput) {
if ($PSStyle.OutputRendering) {
$currentStyle = $PSStyle.OutputRendering
}
[string]$objectAsString = try {
if ($OutputAsJson) {
$object | ConvertTo-Json -Depth $OutputAsJsonDepth -WarningAction silentlycontinue
} else {
$PSStyle.OutputRendering = 'PlainText'
$object | Out-String
}
} finally {
if (-not $OutputAsJson) {
$PSStyle.OutputRendering = $currentStyle
}
}
[Dictionary[String, String]]$customProperties = @{}
$customProperties.Severity = 'Output'
$client.TrackTrace(
('OUTPUT: ' + $objectAsString),
[SeverityLevel]::Information,
$customProperties
)
}
}
}
}
#Use ETS to flag the object as processed so parents do not duplicate the telemetry processing
$object | Add-Member -NotePropertyName __TRACEAICOMMAND_PROCESSED -NotePropertyValue $true
}

Powershell Application Insights Telemetry

This is an easy way to instantly add telemetry collection and reporting to your Powershell Script.

QuickStart

  1. Create an Application Insights instance

  2. SET A DAILY CAP OF 1GB/DAY WHEN TESTING. Avoid surprise bills

  3. Get your connection string

  4. Specify the connection string for your app insights instance either via the ConnectionString parameter or by setting the environment variable APPLICATIONINSIGHTS_CONNECTION_STRING

  5. iwr tinyurl.com/InvokeAICommand | iex (This links to a gist, you should review the code but can be generally trusted )

  6. Invoke-AICommand -Name AITest {Write-Warning "Hello World Warning!"} -ConnectionString 'ConnectionString'
    #ConnectionString optional if you set it as env var
    #Name is optional, will default to 'PowerShell Scriptblock'
  7. You should see telemetry in the transaction search tab Telemetry Demo

Notes

Only what the script outputs is logged, so if you don't have verbosepreference set to 'continue' and you do write-verbose, that won't get logged. If you don't see it show up on screen, it won't show up in appinsights as a rule.

Regular script output isn't logged by default, you must specify the -LogOutput switch and optionally the -OutputAsJson switch

You can nest Invoke-AICommands within each other, and the parent/child dependencies will be automatically tracked.

More Detailed Demo

function Get-UserSentiment {
  [CmdletBinding()]
  param()
  'Starting up the script!'
  Write-Debug "Things looking good for process $PID"
  Write-Verbose 'Alright lets ask the user something'
  Write-Host 'Hi User! Want to ask you something...'
  $mood = Read-Host -Prompt 'What is your Mood?'
  Write-Information -Tag 'Mood' -MessageData "The user is $Mood"
  if ($mood -eq 'Angry') {
    Write-Warning 'oh no! you are angry!'
    # Will throw an error saying we cant find happiness
    Resolve-Path 'happiness'
    Write-Verbose "Gonna sleep for a bit and see if the user's mood improves"
    Start-Sleep 5
  }

  $mood = Read-Host -Prompt 'How about now?'
  Write-Information -Tag 'Mood' -MessageData "The user is $Mood"
  if ($mood -eq 'Angry') {
    Resolve-Path 'I dont know what to do' -ErrorAction stop
  } else {
    Write-Host -Fore Green 'Well I am glad you arent angry anymore at least'
  }
}

Trace-AICommand -Name 'Get-UserSentiment' { Get-UserSentiment -Verbose -Debug -InformationAction continue }

OutputResult

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment