Autoscale/ScaleUpRunbook.ps1 (88 lines of code) (raw):

$connectionName = "AzureRunAsConnection" $subId = Get-AutomationVariable -Name 'subId' $resourceGroupName = Get-AutomationVariable -Name 'resourceGroupName' $jobName = Get-AutomationVariable -Name 'jobName' $targetSU = Get-AutomationVariable -Name 'increasedSU' "Executing with subscription id: $subId; resource group name: $resourceGroupName; job name: $jobName; target SU: $targetSU" $ErrorActionPreference = 'Stop' function Get-AzureRmCachedAccessToken() { $ErrorActionPreference = 'Stop' if(-not (Get-Module AzureRm.Profile)) { Import-Module AzureRm.Profile } $azureRmProfileModuleVersion = (Get-Module AzureRm.Profile).Version # refactoring performed in AzureRm.Profile v3.0 or later if($azureRmProfileModuleVersion.Major -ge 3) { $azureRmProfile = [Microsoft.Azure.Commands.Common.Authentication.Abstractions.AzureRmProfileProvider]::Instance.Profile if(-not $azureRmProfile.Accounts.Count) { Write-Error "Ensure you have logged in before calling this function." } } else { # AzureRm.Profile < v3.0 $azureRmProfile = [Microsoft.WindowsAzure.Commands.Common.AzureRmProfileProvider]::Instance.Profile if(-not $azureRmProfile.Context.Account.Count) { Write-Error "Ensure you have logged in before calling this function." } } $currentAzureContext = Get-AzureRmContext $profileClient = New-Object Microsoft.Azure.Commands.ResourceManager.Common.RMProfileClient($azureRmProfile) Write-Debug ("Getting access token for tenant" + $currentAzureContext.Tenant.TenantId) $token = $profileClient.AcquireAccessToken($currentAzureContext.Tenant.TenantId) $token.AccessToken } try { # Get the connection "AzureRunAsConnection " $servicePrincipalConnection = Get-AutomationConnection -Name $connectionName "Logging in to Azure..." Add-AzureRmAccount ` -ServicePrincipal ` -TenantId $servicePrincipalConnection.TenantId ` -ApplicationId $servicePrincipalConnection.ApplicationId ` -CertificateThumbprint $servicePrincipalConnection.CertificateThumbprint } catch { if (!$servicePrincipalConnection) { $ErrorMessage = "Connection $connectionName not found." throw $ErrorMessage } else{ Write-Error -Message $_.Exception throw $_.Exception } } $accessToken = Get-AzureRmCachedAccessToken $headers = @{ 'Content-Type' = 'application/json' 'Authorization' = 'Bearer ' + $accessToken } $getJobUri = "https://management.azure.com/subscriptions/$subId/resourceGroups/$resourceGroupName/providers/Microsoft.StreamAnalytics/streamingjobs/$jobName" + '?$expand=transformation&api-version=2017-04-01-preview' $jobState = Invoke-RestMethod -Uri $getJobUri -Method Get -Headers $headers $curState = $jobState.properties.jobState $curSU = $jobState.properties.transformation.properties.streamingUnits "Current Job state: $curState; Current SU: $curSU" $scaleUri = "https://management.azure.com/subscriptions/$subId/resourceGroups/$resourceGroupName/providers/Microsoft.StreamAnalytics/streamingjobs/$jobName/scale?api-version=2017-04-01-preview" $scaleForm = @{ "streamingUnits" = $targetSU } $jsonBody = $scaleForm | ConvertTo-Json "Sending scale request..." Invoke-RestMethod -Uri $scaleUri -Method Post -Headers $headers -Body $jsonBody # Check scale result Start-Sleep -Seconds 3 $pollingCount = 0 while ($pollingCount -ne 30) { $pollingCount++ $jobState = Invoke-RestMethod -Uri $getJobUri -Method Get -Headers $headers if ($jobState.properties.jobState.equals("Running") -and $jobState.properties.transformation.properties.streamingUnits -eq $targetSU.ToString()) { break } "Waiting for the job to scale." Start-Sleep -Seconds 5 } $jobState = Invoke-RestMethod -Uri $getJobUri -Method Get -Headers $headers $finalState = $jobState.properties.jobState $finalSU = $jobState.properties.transformation.properties.streamingUnits "Final Job state: $finalState; Final SU: $finalSU" if ($finalState.equals("Running") -and $finalSU -eq $targetSU.ToString()) { "Scaling has completed." } else { throw 'Failed to scale the job' }