Public/Invoke-AzTableAPIBatch.ps1

function Invoke-AzTableAPIBatch {
    <#
    .SYNOPSIS
        Executes an atomic batch (entity group transaction) against Azure Table Storage.

    .DESCRIPTION
        Submits up to 100 entity operations as a single atomic unit using the Azure Table
        Storage $batch endpoint (OData entity group transactions). All operations must
        target the same table and the same PartitionKey.

        Supported actions per operation:
          - Insert : POST (add entity; fails if it already exists)
          - Replace : PUT (insert or replace; removes unlisted properties)
          - Merge : PATCH (insert or merge; preserves unlisted properties)
          - Delete : DELETE (remove entity)

        Each operation is a hashtable or PSCustomObject with:
          - Action : 'Insert', 'Replace', 'Merge', or 'Delete' (required)
          - Entity : hashtable or PSCustomObject with at least PartitionKey and RowKey.
                         For Insert/Replace/Merge the entity data is included.
                         For Delete only PartitionKey and RowKey are used.
          - ETag : optional; defaults to '*' for Replace, Merge, and Delete.

    .PARAMETER Context
        Connection context created by New-AzTableAPIContext.

    .PARAMETER TableName
        Name of the target table.

    .PARAMETER Operations
        Array of 1-100 operation objects. Each must have an Action and an Entity property.
        All entities must share the same PartitionKey.

    .OUTPUTS
        [PSCustomObject[]] - one result object per operation, in the same order, with:
          - Index : zero-based position in the Operations array
          - StatusCode : HTTP status code for this operation
          - Content : parsed JSON response body, or $null for empty responses
          - ETag : ETag header value returned for this operation, or $null

    .EXAMPLE
        $ops = @(
            @{ Action = 'Insert'; Entity = @{ PartitionKey = 'Sales'; RowKey = '001'; Amount = 99.99 } }
            @{ Action = 'Merge'; Entity = @{ PartitionKey = 'Sales'; RowKey = '002'; Status = 'Done' } }
            @{ Action = 'Delete'; Entity = @{ PartitionKey = 'Sales'; RowKey = '003' } }
        )
        Invoke-AzTableAPIBatch -Context $ctx -TableName 'Orders' -Operations $ops
    #>

    [CmdletBinding()]
    [OutputType([PSCustomObject[]])]
    param (
        [Parameter(Mandatory)]
        [ValidateScript({
            ($_.PSObject.Properties.Name -contains 'Endpoint') -and
            ($_.PSObject.Properties.Name -contains 'AuthType')
        }, ErrorMessage = 'The -Context parameter requires a context object created by New-AzTableAPIContext.')]
        [PSCustomObject]$Context,

        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [string]$TableName,

        [Parameter(Mandatory)]
        [ValidateNotNull()]
        [ValidateCount(1, 100)]
        [object[]]$Operations
    )

    # -------------------------------------------------------------------------
    # Validate operations and extract entity data
    # -------------------------------------------------------------------------
    $validActions = @('Insert', 'Replace', 'Merge', 'Delete')
    $resolvedOps  = [System.Collections.Generic.List[hashtable]]::new()

    foreach ($op in $Operations) {
        $action = if ($op -is [hashtable]) { $op['Action'] } else { $op.Action }
        if ($action -notin $validActions) {
            throw "Invalid batch operation Action '$action'. Must be one of: $($validActions -join ', ')."
        }

        $entity = if ($op -is [hashtable]) { $op['Entity'] } else { $op.Entity }
        if ($null -eq $entity) {
            throw "Each batch operation must have a non-null Entity property."
        }

        $entityKeys = Get-AzTableEntityKey -Entity $entity -ErrorMessage 'Each batch operation Entity must have non-empty PartitionKey and RowKey.'
        $pk = $entityKeys.PartitionKey
        $rk = $entityKeys.RowKey

        $etag = if ($op -is [hashtable]) { $op['ETag'] } else { $op.ETag }
        if ([string]::IsNullOrEmpty($etag)) { $etag = '*' }

        $resolvedOps.Add(@{
            Action = $action
            Entity = $entity
            PK     = $pk
            RK     = $rk
            ETag   = $etag
        })
    }

    # All operations must share the same PartitionKey
    $uniquePartitions = $resolvedOps | ForEach-Object { $_.PK } | Select-Object -Unique
    if ($uniquePartitions.Count -gt 1) {
        throw "All batch operations must target the same PartitionKey. Found: $($uniquePartitions -join ', ')."
    }

    # -------------------------------------------------------------------------
    # Build the multipart/mixed batch body
    # -------------------------------------------------------------------------
    $batchId     = [System.Guid]::NewGuid().ToString()
    $changesetId = [System.Guid]::NewGuid().ToString()
    $CRLF        = "`r`n"

    $sb = [System.Text.StringBuilder]::new()

    $null = $sb.Append("--batch_$batchId$CRLF")
    $null = $sb.Append("Content-Type: multipart/mixed; boundary=changeset_$changesetId$CRLF")
    $null = $sb.Append($CRLF)

    foreach ($op in $resolvedOps) {
        $entityResource = Format-AzTableEntityResource -TableName $TableName -PartitionKey $op.PK -RowKey $op.RK

        switch ($op.Action) {
            'Insert' {
                $httpMethod = 'POST'
                $opUrl      = "$($Context.Endpoint)/$TableName"
            }
            'Replace' {
                $httpMethod = 'PUT'
                $opUrl      = "$($Context.Endpoint)/$entityResource"
            }
            'Merge' {
                $httpMethod = 'PATCH'
                $opUrl      = "$($Context.Endpoint)/$entityResource"
            }
            'Delete' {
                $httpMethod = 'DELETE'
                $opUrl      = "$($Context.Endpoint)/$entityResource"
            }
        }

        $null = $sb.Append("--changeset_$changesetId$CRLF")
        $null = $sb.Append("Content-Type: application/http$CRLF")
        $null = $sb.Append("Content-Transfer-Encoding: binary$CRLF")
        $null = $sb.Append($CRLF)
        $null = $sb.Append("$httpMethod $opUrl HTTP/1.1$CRLF")
        $null = $sb.Append("DataServiceVersion: 3.0;NetFx$CRLF")
        $null = $sb.Append("MaxDataServiceVersion: 3.0;NetFx$CRLF")
        $null = $sb.Append("Accept: application/json;odata=nometadata$CRLF")

        if ($op.Action -in 'Delete', 'Replace', 'Merge') {
            $null = $sb.Append("If-Match: $($op.ETag)$CRLF")
        }

        if ($op.Action -in 'Insert', 'Replace', 'Merge') {
            $bodyJson   = $op.Entity | ConvertTo-Json -Depth 10 -Compress
            $byteCount  = [System.Text.Encoding]::UTF8.GetByteCount($bodyJson)
            $null = $sb.Append("Content-Type: application/json$CRLF")
            $null = $sb.Append("Content-Length: $byteCount$CRLF")
            $null = $sb.Append($CRLF)
            $null = $sb.Append($bodyJson)
        } else {
            $null = $sb.Append($CRLF)
        }

        $null = $sb.Append($CRLF)
    }

    $null = $sb.Append("--changeset_$changesetId--$CRLF")
    $null = $sb.Append("--batch_$batchId--$CRLF")

    $batchBody        = $sb.ToString()
    $batchContentType = "multipart/mixed; boundary=batch_$batchId"

    # -------------------------------------------------------------------------
    # Build request headers (auth + OData)
    # -------------------------------------------------------------------------
    $authHeaders = Get-AzTableAuthorizationHeader `
        -Context     $Context `
        -Method      'POST' `
        -Resource    '$batch' `
        -ContentType $batchContentType

    $requestHeaders = @{
        'DataServiceVersion'    = '3.0;NetFx'
        'MaxDataServiceVersion' = '3.0;NetFx'
        'Accept'                = 'application/json;odata=nometadata'
    }

    foreach ($key in $authHeaders.Keys) {
        $requestHeaders[$key] = $authHeaders[$key]
    }

    # -------------------------------------------------------------------------
    # Build the request URL (append SAS token if applicable)
    # -------------------------------------------------------------------------
    $url = "$($Context.Endpoint)/`$batch"
    if ($Context.AuthType -eq 'SasToken') {
        $url += $Context.SasToken
    }

    # -------------------------------------------------------------------------
    # Execute the batch request
    # -------------------------------------------------------------------------
    $bodyBytes = [System.Text.Encoding]::UTF8.GetBytes($batchBody)

    try {
        $response = Invoke-WebRequest `
            -Uri             $url `
            -Method          'POST' `
            -Headers         $requestHeaders `
            -Body            $bodyBytes `
            -ContentType     $batchContentType `
            -UseBasicParsing `
            -ErrorAction     Stop

        $rawContent = $response.Content
        if ($rawContent -is [byte[]]) {
            $rawContent = [System.Text.Encoding]::UTF8.GetString($rawContent)
        }

        $responseContentType = $response.Headers['Content-Type']
        if ($responseContentType -is [array]) { $responseContentType = $responseContentType[0] }

        return ConvertFrom-AzTableBatchResponse -RawContent $rawContent -ContentType $responseContentType
    } catch {
        $statusCode   = $null
        $errorMessage = $_.Exception.Message

        if ($_.Exception.Response) {
            $statusCode = [int]$_.Exception.Response.StatusCode
        }

        if ($_.ErrorDetails -and $_.ErrorDetails.Message) {
            $errorBody = $_.ErrorDetails.Message
            try {
                $errorObj = $errorBody | ConvertFrom-Json
                if ($errorObj.'odata.error'.message.value) {
                    $errorMessage = $errorObj.'odata.error'.message.value
                }
            } catch {
                $errorMessage = $errorBody
            }
        }

        throw [System.Exception]::new(
            "Azure Table Storage batch request failed (HTTP $statusCode): $errorMessage",
            $_.Exception
        )
    }
}

function ConvertFrom-AzTableBatchResponse {
    <#
    .SYNOPSIS
        Parses a multipart/mixed Azure Table Storage batch response into result objects.
    #>

    [CmdletBinding()]
    [OutputType([PSCustomObject[]])]
    param (
        [Parameter(Mandatory)]
        [string]$RawContent,

        [Parameter(Mandatory)]
        [string]$ContentType
    )

    # Extract outer batch boundary
    if ($ContentType -match 'boundary=(.+)') {
        $batchBoundary = $Matches[1].Trim()
    } else {
        throw "Cannot determine batch response boundary from Content-Type: $ContentType"
    }

    # Split by batch boundary to find the changeset part
    $batchParts    = $RawContent -split "--$([regex]::Escape($batchBoundary))"
    $changesetPart = $batchParts | Where-Object { $_ -match 'Content-Type:\s*multipart/mixed' } | Select-Object -First 1

    if (-not $changesetPart) {
        return @()
    }

    # Extract changeset boundary from the changeset part's Content-Type line
    if ($changesetPart -match 'boundary=(\S+?)[\r\n]') {
        $changesetBoundary = $Matches[1].Trim()
    } else {
        throw "Cannot determine changeset boundary from batch response."
    }

    # Split by changeset boundary to get individual operation responses
    $changesetParts = $changesetPart -split "--$([regex]::Escape($changesetBoundary))"

    $results = [System.Collections.Generic.List[PSCustomObject]]::new()
    $index   = 0

    foreach ($part in $changesetParts) {
        if ($part -notmatch 'HTTP/1\.1') { continue }

        # Extract HTTP status code and text
        $statusMatch = [regex]::Match($part, 'HTTP/1\.1 (\d{3}) (.+?)[\r\n]')
        if (-not $statusMatch.Success) { continue }

        $statusCode = [int]$statusMatch.Groups[1].Value
        $statusText = $statusMatch.Groups[2].Value.Trim()

        # Extract ETag response header (if present)
        $etagMatch = [regex]::Match($part, '(?i)ETag:\s*(.+?)[\r\n]')
        $etag      = if ($etagMatch.Success) { $etagMatch.Groups[1].Value.Trim() } else { $null }

        # Extract response body (content after the blank line separating headers from body)
        $bodyMatch    = [regex]::Match($part, '(?s)\r?\n\r?\n(.+)$')
        $parsedContent = $null

        if ($bodyMatch.Success) {
            $bodyText = $bodyMatch.Groups[1].Value.Trim()
            if (-not [string]::IsNullOrWhiteSpace($bodyText)) {
                try {
                    $parsedContent = $bodyText | ConvertFrom-Json
                } catch {
                    Write-Debug "Response body is not JSON: $_"
                }
            }
        }

        $results.Add([PSCustomObject]@{
            Index      = $index
            StatusCode = $statusCode
            StatusText = $statusText
            Content    = $parsedContent
            ETag       = $etag
        })
        $index++
    }

    return $results.ToArray()
}