Managed Identity to Connect AKS with Azure SQL

#This script creates and assigns workload identities in the AKS cluster. Workload identities need to be enabled, and also the managed identity needs
#to be created. The managed identity created by Teraform will then be "federated" for workload identity assignment to pods.
param(
    [Parameter(Mandatory=$true)]
    [string]
    $clusterName,
    [Parameter(Mandatory=$true)]
    [string]
    $rg,
    [Parameter(Mandatory=$true)]
    [string]
    $serviceName,
    [Parameter(Mandatory=$true)]
    [string]
    $namespace,
    [string]
    $prefix="test-",
    [string]
    $identityName,
    [string]
    $tenantId
)
class CustomException : Exception {
    [string] $additionalData

    CustomException($Message, $additionalData) : base($Message) {
        $this.additionalData = $additionalData
    }
}

$clientID = "$(az identity show --name $identityName --resource-group $rg --query 'clientId' -otsv)"

$serviceAccountName = "$prefix$serviceName-sa"
$yaml = `apiVersion: v1
kind: ServiceAccount
metadata:
  annotations:
azure.workload.identity/client-id: __clientID__ 
azure.workload.identity/tenant-id: __tenantId__
  labels:
azure.workload.identity/use: "true"
  name: __serviceAccountName__ 
  namespace: __namespace__ `

#overwrite tokens in the template and write out to a new file
$yaml = $yaml.Replace("__clientID__", $clientID);
$yaml = $yaml.Replace("__serviceAccountName__", $serviceAccountName);
$yaml = $yaml.Replace("__namespace__", $namespace);
$yaml = $yaml.Replace("__tenantId__", $tenantId);
Set-Content -Path .service-account.yaml -Value $yaml

#create the service account that needs to be bound to the federated identity
kubectl apply -f .service-account.yaml -n $namespace

#Clear content
Set-Content -Path .service-account.yaml -Value ""

try {
    $oidcIssuer="$(az aks show -n $clusterName -g $rg --query "oidcIssuerProfile.issuerUrl" -otsv)"
    if($null -eq $oidcIssuer -or "" -eq $oidcIssuer)
    {
        #oidc does not exist in the cluster. It needs to be enabled. Check if cluster addon is enabled. First check az version
        $azVersions = az version | ConvertFrom-Json
        $cliVersion = $azVersions.'azure-cli';
        #cliVersion needs to be >= 2.40.0

        if(([Version] $cliVersion) -lt ([Version] "2.40.0")){
            Write-Host "CLI needs upgraded based on the dependencies here... https://learn.microsoft.com/en-us/azure/aks/workload-identity-overview";
            #az upgrade --yes --all --quiet
            throw [CustomException]::new("Azure CLI $cliVersion needs to be upgraded to >=2.40.0 to leverage Workload Identities and Federated credentials with the azure cli.","") > $PSScriptRooterror.log
        }
        
        #see if the WorkloadIdentityPreview feature is enabled and registered
        $featureState = az feature show --namespace "Microsoft.ContainerService" --name "EnableWorkloadIdentityPreview" | ConvertFrom-Json
        if($featureState.properties.state -ne "Registered"){
            #try to enable and register the 'EnableWorkloadIdentityPreview' feature flag
            Write-Host "Feature flag EnableWorkloadIdentityPreview is not registered... attempting to register..."
            az feature register --namespace "Microsoft.ContainerService" --name "EnableWorkloadIdentityPreview";
            #We have to wait a minute or two for the features to register before we can continue.
            Start-Sleep -Seconds 600;
            az provider register --namespace Microsoft.ContainerService
            #Verify feature is in a registered state
            $featureState = az feature show --namespace "Microsoft.ContainerService" --name "EnableWorkloadIdentityPreview" | ConvertFrom-Json
            if($featureState.properties.state -ne "Registered"){
                throw [CustomException]::new("Feature EnableWorkloadIdentityPreview did not finish registering after 120 seconds, please wait and try again. If this fails a second time the verify that Microsoft.ContainerService is enabled in the subscription under 'Features'.","") > $PSScriptRooterror.log
            }
            else{
                Write-Host "Feature was successfully registered... continuing..."
            }
        }
        #update the cluster to use OIDC and workload identities
        az aks update -g $rg -n $clusterName --enable-oidc-issuer --enable-workload-identity
        #now finally get the OIDC URL
        $oidcIssuer="$(az aks show -n $clusterName -g $rg --query "oidcIssuerProfile.issuerUrl" -otsv)"
        if($null -eq $oidcIssuer -or "" -eq $oidcIssuer){
            throw [CustomException]::new("Still cannot retrieve OIDC URLs. Please wait an hour and try again. Cluster might still be updating.",""); > $PSScriptRooterror.log
        }
    }

    az identity federated-credential create --name "$serviceName-$namespace" --identity-name $identityName --resource-group $rg --issuer $oidcIssuer --subject "system:serviceaccount:${namespace}:${serviceAccountName}"
}
catch {
    throw [CustomException]::new("Federation setup failed.","") > $PSScriptRooterror.log
}
#$oidcIssuer="$(az aks show -n $clusterName -g $rg --query "oidcIssuerProfile.issuerUrl" -otsv)"

try{
    #see if all helm charts exist for workload identity
    $charts = helm repo list -o json | ConvertFrom-Json
    $url = $charts | Where-Object {$_.Name -eq "azure-workload-identity"} | Select-Object -ExpandProperty url
    if($null -eq $url -or "" -eq $url){
        write-host "azure-workload-identity is missing from the helm charts... attempting to add it."
        helm repo add azure-workload-identity https://azure.github.io/azure-workload-identity/charts
        helm repo update
        helm install workload-identity-webhook azure-workload-identity/workload-identity-webhook --namespace azure-workload-identity-system --create-namespace --set azureTenantID="$tenantId"
        Write-Host "Helm charts 'workload-identity-webhook' and 'azure-workload-identity/workload-identity-webhook' successfully added!"
    }
}
catch {
    throw [CustomException]::new("Helm setup failed.","") > $PSScriptRooterror.log
}