Private/Azure/Get-PolicySetDefinitionCached.ps1

function Get-PolicySetDefinitionCached {
    <#
    .SYNOPSIS
        Retrieve Azure Policy Set definition (initiative) with caching support.
     
    .DESCRIPTION
        Wrapper around Get-AzPolicySetDefinition with in-memory caching and retry logic.
        Supports lookup by Name, Id, or retrieving all policy sets.
        Tracks cache statistics for monitoring.
     
    .PARAMETER Name
        Name of the policy set definition to retrieve.
     
    .PARAMETER Id
        Resource ID of the policy set definition to retrieve.
     
    .PARAMETER ManagementGroupId
        Management Group scope for the query.
     
    .PARAMETER Cache
        Hashtable for caching policy set definitions (shared across calls).
     
    .EXAMPLE
        $cache = @{}
        $policySet = Get-PolicySetDefinitionCached -Name "MCSB" -Cache $cache
     
    .OUTPUTS
        Microsoft.Azure.Commands.ResourceManager.Cmdlets.Implementation.Policy.PsPolicySetDefinition
    #>

    [CmdletBinding()]
    param(
        [string]$Name,
        
        [string]$Id,
        
        [string]$ManagementGroupId,
        
        [hashtable]$Cache = @{}
    )
    
    # Initialize statistics if not already done
    if (-not $script:ApiCallStats) {
        Initialize-ApiStatistics
    }
    
    # Build cache key
    $cacheKey = if ($Id) { 
        $Id 
    } elseif ($Name) { 
        "name:$Name" 
    } else { 
        "all" 
    }
    
    # Check cache first
    if ($Cache.ContainsKey($cacheKey)) {
        Write-Debug "Cache HIT for policy set: $cacheKey"
        $script:ApiCallStats.CacheHits++
        return $Cache[$cacheKey]
    }
    
    # Cache miss
    Write-Debug "Cache MISS for policy set: $cacheKey"
    $script:ApiCallStats.CacheMisses++
    
    # Fetch with retry logic
    try {
        $definition = Invoke-AzCommandWithRetry -Command {
            if ($Id) {
                Write-Debug "Fetching policy set by ID: $Id"
                Get-AzPolicySetDefinition -Id $Id -ErrorAction Stop
            }
            elseif ($Name) {
                Write-Debug "Fetching policy set by Name: $Name"
                
                if ($ManagementGroupId) {
                    Get-AzPolicySetDefinition -Name $Name -ManagementGroupName $ManagementGroupId -ErrorAction Stop
                }
                else {
                    Get-AzPolicySetDefinition -Name $Name -ErrorAction Stop
                }
            }
            else {
                Write-Debug "Fetching all policy sets"
                
                if ($ManagementGroupId) {
                    Get-AzPolicySetDefinition -ManagementGroupName $ManagementGroupId -ErrorAction Stop
                }
                else {
                    Get-AzPolicySetDefinition -ErrorAction Stop -WarningAction SilentlyContinue
                }
            }
        } -OperationName "PolicySet"
        
        # Track API type
        $script:ApiCallStats.PolicySetDefinitionCalls++
        
        # Cache the result
        $Cache[$cacheKey] = $definition
        
        return $definition
    }
    catch {
        Write-Warning "Failed to get policy set definition: $($_.Exception.Message)"
        throw
    }
}