feat: Refactor error handling and update parameter retrieval methods

infra/main.bicepparam

@@ -8,7 +8,7 @@ param environmentName = readEnvironmentVariable('AZURE_ENV_NAME', '')
 param location = readEnvironmentVariable('AZURE_LOCATION', '')
 param cosmosLocation = readEnvironmentVariable('AZURE_COSMOS_LOCATION', '')
 // Entra object ID of the identity to grant RBAC (user, group, service principal, or UAI). Set this if Graph lookup is blocked.
-param principalId = ''
+param principalId = readEnvironmentVariable('AZURE_PRINCIPAL_ID', '')
 param principalType = 'User'
 
 // ========================================

scripts/automationScripts/FabricWorkspace/CreateWorkspace/create_fabric_workspace.ps1

@@ -341,12 +341,12 @@ if ($CapacityName) {
 $capName = Get-CapacityLookupName -ResolvedCapacityId $CapacityId -ResolvedCapacityName $CapacityName
 if ($capName) {
   Log "Deriving Fabric capacity GUID for name: $capName"
-  
+
   try { 
     $caps = Invoke-SecureRestMethod -Uri "$apiRoot/capacities" -Headers $apiHeaders -Method Get
     if ($caps.value) { 
       Log "Searching through $($caps.value.Count) capacities for: '$capName'"
-      
+
       # Use a simple foreach loop instead of Where-Object to debug comparison issues
       foreach ($cap in $caps.value) {
         $capDisplayName = if ($cap.PSObject.Properties['displayName']) { $cap.displayName } else { '' }
@@ -361,15 +361,15 @@ if ($capName) {
           Log "EXACT MATCH FOUND: Using capacity '$capDisplayName' with GUID: $capacityGuid"
           break
         }
-        
+
         # Case-insensitive fallback
         if (([string]$capDisplayName).ToLowerInvariant() -eq $capName.ToLowerInvariant() -or ([string]$capName2).ToLowerInvariant() -eq $capName.ToLowerInvariant() -or ([string]$capId).ToLowerInvariant() -eq $capName.ToLowerInvariant()) {
           $capacityGuid = $capId
           Log "CASE-INSENSITIVE MATCH FOUND: Using capacity '$capDisplayName' with GUID: $capacityGuid"
           break
         }
       }
-      
+
       if (-not $capacityGuid) {
         Log "NO MATCH FOUND. Available capacities:"
         foreach ($cap in $caps.value) {
@@ -386,7 +386,7 @@ if ($capName) {
   } catch { 
     Fail "Failed to query capacities: $_"
   }
-  
+
   if ($capacityGuid) {
     Log "Resolved capacity GUID: $capacityGuid"
   } else {
@@ -398,6 +398,15 @@ if ($capName) {
 $workspaceId = $null
 $workspaceId = Resolve-WorkspaceIdByName -Name $WorkspaceName
 
+# Fallback: also check via /workspaces endpoint (uses displayName instead of name)
+if (-not $workspaceId) {
+  try {
+    $wsListResp = Invoke-SecureRestMethod -Uri "$apiRoot/workspaces?%24top=5000" -Headers $apiHeaders -Method Get -ErrorAction Stop
+    $wsMatch = $wsListResp.value | Where-Object { $_.displayName -eq $WorkspaceName }
+    if ($wsMatch) { $workspaceId = $wsMatch.id }
+  } catch { }
+}
+
 if ($workspaceId) {
   Log "Workspace '$WorkspaceName' already exists (id=$workspaceId). Ensuring capacity assignment & admins."
   if ($capacityGuid) {
@@ -524,7 +533,7 @@ if ($capacityGuid) {
     Log "Assigning workspace to capacity GUID: $capacityGuid"
     $assignResp = Invoke-SecureWebRequest -Uri "$apiRoot/workspaces/$workspaceId/assignToCapacity" -Method Post -Headers ($apiHeaders) -Body (@{ capacityId = $capacityGuid } | ConvertTo-Json) -ErrorAction Stop
     Log "Capacity assignment response: $($assignResp.StatusCode)"
-    
+
     # Verify assignment worked
     Start-Sleep -Seconds 3
     $workspace = Invoke-SecureRestMethod -Uri "$apiRoot/workspaces/$workspaceId" -Headers $apiHeaders -Method Get -ErrorAction Stop
@@ -582,4 +591,4 @@ Log "Workspace ID: $workspaceId"
 
 # Clean up sensitive variables
 Clear-SensitiveVariables -VariableNames @('accessToken')
-exit 0
+exit 0

scripts/automationScripts/OneLakeIndex/02_create_onelake_skillsets.ps1

@@ -53,7 +53,225 @@ if (-not $aiSearchName -or -not $resourceGroup -or -not $subscription) {
     exit 1
 }
 
-. "$PSScriptRoot/SearchHelpers.ps1"
+function Get-SearchPublicNetworkAccess {
+    try {
+        return az search service show --name $aiSearchName --resource-group $resourceGroup --subscription $subscription --query "publicNetworkAccess" -o tsv
+    } catch {
+        return $null
+    }
+}
+
+function Get-SearchResourceId {
+    try {
+        return az search service show --name $aiSearchName --resource-group $resourceGroup --subscription $subscription --query "id" -o tsv
+    } catch {
+        return $null
+    }
+}
+
+function Get-ArmAccessToken {
+    try {
+        return az account get-access-token --resource https://management.azure.com/ --subscription $subscription --query accessToken -o tsv
+    } catch {
+        return $null
+    }
+}
+
+function Invoke-AzCliWithTimeout {
+    param(
+        [string[]]$Args,
+        [int]$TimeoutSeconds = 120
+    )
+
+    $escapedArgs = $Args | ForEach-Object {
+        if ($_ -match '\s|"') { '"' + ($_ -replace '"', '\"') + '"' } else { $_ }
+    }
+
+    $azPath = (Get-Command az -ErrorAction SilentlyContinue).Source
+    if (-not $azPath) {
+        throw "Azure CLI (az) not found on PATH."
+    }
+
+    $psi = New-Object System.Diagnostics.ProcessStartInfo
+    $psi.FileName = $azPath
+    $psi.Arguments = ($escapedArgs -join ' ')
+    $psi.RedirectStandardOutput = $true
+    $psi.RedirectStandardError = $true
+    $psi.UseShellExecute = $false
+    $psi.CreateNoWindow = $true
+
+    $process = [System.Diagnostics.Process]::Start($psi)
+    if (-not $process.WaitForExit($TimeoutSeconds * 1000)) {
+        try { $process.Kill() } catch { }
+        throw "Azure CLI command timed out after $TimeoutSeconds seconds: az $($psi.Arguments)"
+    }
+
+    $stdout = $process.StandardOutput.ReadToEnd()
+    $stderr = $process.StandardError.ReadToEnd()
+
+    if ($process.ExitCode -ne 0) {
+        throw "Azure CLI failed with exit code $($process.ExitCode): $stderr"
+    }
+
+    return $stdout
+}
+
+function Set-SearchPublicNetworkAccess {
+    param([string]$Mode)
+
+    $timeoutSeconds = 120
+    if ($env:AI_SEARCH_PUBLIC_ACCESS_TIMEOUT_SECONDS) {
+        [int]::TryParse($env:AI_SEARCH_PUBLIC_ACCESS_TIMEOUT_SECONDS, [ref]$timeoutSeconds) | Out-Null
+    }
+
+    $maxRetries = 3
+    if ($env:AI_SEARCH_PUBLIC_ACCESS_MAX_RETRIES) {
+        [int]::TryParse($env:AI_SEARCH_PUBLIC_ACCESS_MAX_RETRIES, [ref]$maxRetries) | Out-Null
+    }
+
+    $waitSeconds = 120
+    if ($env:AI_SEARCH_PUBLIC_ACCESS_POLL_SECONDS) {
+        [int]::TryParse($env:AI_SEARCH_PUBLIC_ACCESS_POLL_SECONDS, [ref]$waitSeconds) | Out-Null
+    }
+
+    for ($attempt = 1; $attempt -le $maxRetries; $attempt++) {
+        try {
+            $resourceId = Get-SearchResourceId
+            if (-not $resourceId) {
+                throw "Unable to resolve AI Search resource ID."
+            }
+            $armToken = Get-ArmAccessToken
+            if (-not $armToken) {
+                throw "Unable to acquire ARM access token."
+            }
+            $body = @{ properties = @{ publicNetworkAccess = $Mode } } | ConvertTo-Json -Compress
+            $url = "https://management.azure.com${resourceId}?api-version=2023-11-01"
+            $headers = @{ Authorization = "Bearer $armToken"; 'Content-Type' = 'application/json' }
+            Invoke-RestMethod -Method Patch -Uri $url -Headers $headers -Body $body | Out-Null
+            $deadline = (Get-Date).AddSeconds($waitSeconds)
+            do {
+                $current = Get-SearchPublicNetworkAccess
+                if ($current -eq $Mode) { return }
+                Start-Sleep -Seconds 5
+            } while ((Get-Date) -lt $deadline)
+
+            throw "Timed out waiting for AI Search public network access to become '$Mode'."
+        } catch {
+            if ($attempt -lt $maxRetries) {
+                Write-Warning "Failed to set AI Search public network access (attempt $attempt of $maxRetries): $($_.Exception.Message)"
+                Start-Sleep -Seconds 10
+                continue
+            }
+            throw
+        }
+    }
+}
+
+function Ensure-SearchPublicAccess {
+    if ($env:AI_SEARCH_SKIP_PUBLIC_ACCESS_TOGGLE -and $env:AI_SEARCH_SKIP_PUBLIC_ACCESS_TOGGLE.ToLowerInvariant() -eq 'true') {
+        Write-Host "Skipping temporary public network access toggle for AI Search."
+        return $null
+    }
+
+    $current = Get-SearchPublicNetworkAccess
+    if (-not $current) { return $null }
+
+    if ($current -eq 'Disabled') {
+        Write-Warning "AI Search public network access is Disabled. Enabling temporarily for OneLake setup."
+        Set-SearchPublicNetworkAccess -Mode 'Enabled'
+    }
+
+    return $current
+}
+
+function Restore-SearchPublicAccess {
+    param([string]$OriginalAccess)
+
+    if (-not $OriginalAccess) { return }
+    if ($env:AI_SEARCH_SKIP_PUBLIC_ACCESS_TOGGLE -and $env:AI_SEARCH_SKIP_PUBLIC_ACCESS_TOGGLE.ToLowerInvariant() -eq 'true') { return }
+
+    $current = Get-SearchPublicNetworkAccess
+    if ($current -and $current -ne $OriginalAccess) {
+        Write-Host "Restoring AI Search public network access to '$OriginalAccess'."
+        try {
+            Set-SearchPublicNetworkAccess -Mode $OriginalAccess
+        } catch {
+            Write-Warning "Failed to restore AI Search public network access: $($_.Exception.Message)"
+        }
+    }
+}
+
+function Get-SearchAccessToken {
+    try {
+        return az account get-access-token --resource https://search.azure.com --subscription $subscription --query accessToken -o tsv
+    } catch {
+        return $null
+    }
+}
+
+function New-SearchHeaders {
+    param(
+        [string]$AccessToken
+    )
+
+    if ($AccessToken) {
+        return @{
+            'Authorization' = "Bearer $AccessToken"
+            'Content-Type' = 'application/json'
+        }
+    }
+
+    return $null
+}
+
+function Invoke-SearchRequest {
+    param(
+        [string]$Method,
+        [string]$Uri,
+        [string]$Body
+    )
+
+    $maxAttempts = 6
+    $delaySeconds = 30
+
+    for ($attempt = 1; $attempt -le $maxAttempts; $attempt++) {
+        $accessToken = Get-SearchAccessToken
+        $headers = New-SearchHeaders -AccessToken $accessToken
+
+        if (-not $headers) {
+            Write-Error "Failed to acquire Azure AI Search access token via Microsoft Entra ID"
+            exit 1
+        }
+
+        try {
+            if ($Body) {
+                return Invoke-RestMethod -Uri $Uri -Headers $headers -Method $Method -Body $Body
+            }
+            return Invoke-RestMethod -Uri $Uri -Headers $headers -Method $Method
+        } catch {
+            $statusCode = $null
+            try { $statusCode = $_.Exception.Response.StatusCode.value__ } catch { }
+
+            if (($statusCode -eq 401 -or $statusCode -eq 403) -and $attempt -lt $maxAttempts) {
+                Write-Warning "Search request denied (HTTP $statusCode). Waiting ${delaySeconds}s for RBAC propagation (attempt $attempt of $maxAttempts)."
+                Start-Sleep -Seconds $delaySeconds
+                continue
+            }
+
+            # Retry on connection timeouts (search service data plane may not be reachable yet)
+            $isConnTimeout = ($_.Exception -is [System.Net.WebException] -and $_.Exception.Status -eq [System.Net.WebExceptionStatus]::ConnectFailure) -or
+                             ($_.Exception.Message -match 'A connection attempt failed|No connection could be made|actively refused') -or
+                             ($_.Exception.InnerException -and $_.Exception.InnerException.Message -match 'A connection attempt failed|No connection could be made|actively refused')
+            if ($isConnTimeout -and $attempt -lt $maxAttempts) {
+                Write-Warning "Connection timeout to search service. Waiting ${delaySeconds}s for data plane availability (attempt $attempt of $maxAttempts)."
+                Start-Sleep -Seconds $delaySeconds
+                continue
+            }
+
+            throw
+        }
+    }
+}
 
 $originalPublicAccess = Ensure-SearchPublicAccess
 try {