Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Powershell script to upload a file incrementally to an Azure Data Lake Storage using OAuth2 authentication
[CmdletBinding()]
Param(
# OAuth parameters
[Parameter(Mandatory = $false)] [string] $TenantID = 'a guid',
[Parameter(Mandatory = $false)] [string] $ClientID = 'a client id',
[Parameter(Mandatory = $false)] [string] $Scope = 'https://storage.azure.com/.default',
[Parameter(Mandatory = $false)] [string] $GrantType = 'client_credentials',
[Parameter(Mandatory = $false)] [string] $ClientSecret = '',
[Parameter(Mandatory = $false)] [string] $ClientSecretKeyVaultSecret = 'the secret name',
[Parameter(Mandatory = $false)] [string] $KeyVaultName = 'the key vault name',
# File parameters
[Parameter(Mandatory = $false)] [string] $StorageAccountName = 'storage account name',
[Parameter(Mandatory = $false)] [string] $DestinationPath = 'temp/config/sample.jpg',
[Parameter(Mandatory = $false)] [string] $SourcePath = 'C:\Temp\sample.jpg',
# Other
[Parameter(Mandatory = $false)] [string] $AzContextPath = 'C:\Temp\AzContext.json'
)
# begin functions
function Get-OAuth2Token {
<#
.SYNOPSIS
Get OAuth 2.0 bearer token from Azure.
#>
[CmdletBinding()]
Param(
[Parameter(Mandatory = $true, Position = 1)] [string] $ClientID ,
[Parameter(Mandatory = $True, Position = 2)] [string] $ClientSecret,
[Parameter(Mandatory = $True, Position = 3)] [string] $Scope,
[Parameter(Mandatory = $True, Position = 4)] [string] $GrantType,
[Parameter(Mandatory = $True, Position = 5)] [string] $TenantID
)
$formData = @{
client_id = $ClientId;
client_secret = $ClientSecret;
scope = $Scope;
grant_type = $GrantType;
}
$uri = "https://login.microsoftonline.com/$TenantID/oauth2/v2.0/token"
try {
$response = Invoke-RestMethod -Uri $uri -Method Post -Body $formData -ContentType "application/x-www-form-urlencoded"
$response.access_token
}
catch {
$ErrorMessage = $_.Exception.Message
$StatusDescription = $_.Exception.Response.StatusDescription
$false
Throw $ErrorMessage + " " + $StatusDescription + " ClientID: $ClientID, Scope: $Scope, GrantType: $GrantType, TenantID: $TenantID"
}
}
function Create-FileInADLS {
<#
.SYNOPSIS
Create file in Data Lake Generation 2 using REST API and bearer token.
#>
[CmdletBinding()]
Param(
[Parameter(Mandatory = $True, Position = 1)] [string] $BearerToken,
[Parameter(Mandatory = $true, Position = 2)] [string] $StorageAccountName,
[Parameter(Mandatory = $True, Position = 3)] [string] $DestinationPath
)
# Rest documentation:
# https://docs.microsoft.com/en-us/rest/api/storageservices/datalakestoragegen2/path/create
$PathToCreate = $DestinationPath.trim("/") # remove all "//path" or "path/"
$method = "PUT"
$headers = @{ }
$headers.Add("x-ms-version", "2018-11-09")
$headers.Add("Authorization", "Bearer $BearerToken")
#$headers.Add("If-None-Match", "*") # To fail if the destination already exists, use a conditional request with If-None-Match: "*"
$URI = "https://$StorageAccountName.dfs.core.windows.net/" + $PathToCreate + "?resource=file"
try {
Invoke-RestMethod -method $method -Uri $URI -Headers $headers # returns empty response
}
catch {
$ErrorMessage = $_.Exception.Message
$StatusDescription = $_.Exception.Response.StatusDescription
Throw $ErrorMessage + " " + $StatusDescription + " Path: $PathToCreate"
}
}
function Upload-FileToADLS {
<#
.SYNOPSIS
Uplaoad file in Data Lake Generation 2 using REST API and bearer token.
#>
[CmdletBinding()]
Param(
[Parameter(Mandatory = $True, Position = 1)] [string] $BearerToken,
[Parameter(Mandatory = $true, Position = 2)] [string] $StorageAccountName,
[Parameter(Mandatory = $True, Position = 3)] [string] $DestinationPath,
[Parameter(Mandatory = $True, Position = 4)] [string] $SourcePath
)
# Rest documentation:
# https://docs.microsoft.com/en-us/rest/api/storageservices/datalakestoragegen2/path/update
$result = @{ }
$filesize = (Get-Item $SourcePath).length
if ($filesize -gt 0kb)
{
$PathToUpload = $DestinationPath.trim("/") # remove all "//path" or "path/"
$destinationURI = "https://$StorageAccountName.dfs.core.windows.net/" + $PathToUpload
$method = "PATCH"
$headers = @{ }
$headers.Add("x-ms-version", "2018-11-09")
$headers.Add("Authorization", "Bearer $BearerToken")
$offset = 0
$chunkSize = 256kb
$chunk = New-Object byte[] $chunkSize
try
{
$fileStream = [System.IO.File]::OpenRead($SourcePath)
while ($bytesRead = $fileStream.Read($chunk, 0, $chunkSize))
{
$percentComplete = [math]::Round($offset*100/$filesize)
Write-Progress -Activity "Uploading $SourcePath to $destinationURI" -Status "$percentComplete% Complete" -PercentComplete $percentComplete
$URI = $destinationURI + "?action=append&position=$offset"
if ($bytesRead -eq $chunkSize)
{
$body = $chunk
}
else
{
$body = New-Object byte[] $bytesRead
for($i=0; $i -lt $bytesRead; $i++)
{
$body[$i] = $chunk[$i]
}
}
try {
$response = Invoke-WebRequest -method $method -Uri $URI -Headers $headers -Body $body
}
catch {
$ErrorMessage = $_.Exception.Message
$StatusDescription = $_.Exception.Response.StatusDescription
Throw $ErrorMessage + " " + $StatusDescription + " Path: $PathToUpload, SourcePath: $SourcePath"
}
$offset += $bytesRead
}
}
finally
{
$fileStream.Close()
}
}
$result.Add("filesize", $filesize)
$result
}
function Flush-FileCacheInADLS {
<#
.SYNOPSIS
Flush file in Data Lake Generation 2 using REST API and bearer token.
#>
[CmdletBinding()]
Param(
[Parameter(Mandatory = $True, Position = 1)] [string] $BearerToken,
[Parameter(Mandatory = $true, Position = 2)] [string] $StorageAccountName,
[Parameter(Mandatory = $True, Position = 3)] [string] $DestinationPath,
[Parameter(Mandatory = $True, Position = 4)] [string] $FlushPosition
)
# Rest documentation:
# https://docs.microsoft.com/en-us/rest/api/storageservices/datalakestoragegen2/path/update
$PathToFlush = $DestinationPath.trim("/") # remove all "//path" or "path/"
$method = "PATCH"
$headers = @{ }
$headers.Add("x-ms-version", "2018-11-09")
$headers.Add("Authorization", "Bearer $BearerToken")
$URI = "https://$StorageAccountName.dfs.core.windows.net/" + $PathToFlush + "?action=flush&position=$FlushPosition"
# PATCH http://{accountName}.{dnsSuffix}/{filesystem}/{path}?action={action}&position={position}
try {
$response = Invoke-WebRequest -method $method -Uri $URI -Headers $headers # returns empty response
$response
}
catch {
$ErrorMessage = $_.Exception.Message
$StatusDescription = $_.Exception.Response.StatusDescription
Throw $ErrorMessage + " " + $StatusDescription + " Path: $DestinationPath, FlushPosition: $FlushPosition"
}
}
# end functions
# begin
if ($AzContextPath -ne '')
{
$context = Import-AzContext -Path $AzContextPath
}
else
{
$response = Connect-AzAccount -Tenant $TenantID
}
if ($ClientSecret -eq '')
{
$ClientSecret = Get-AzKeyVaultSecret -vaultName $KeyVaultName -name $ClientSecretKeyVaultSecret -AsPlainText
}
$BearerToken = Get-OAuth2Token -ClientID $ClientID -ClientSecret $ClientSecret -Scope $Scope -GrantType $GrantType -TenantID $TenantID
Create-FileInADLS -BearerToken $BearerToken -StorageAccountName $StorageAccountName -DestinationPath $DestinationPath
$response = Upload-FileToADLS -BearerToken $BearerToken -StorageAccountName $StorageAccountName -DestinationPath $DestinationPath -SourcePath $SourcePath
$response = Flush-FileCacheInADLS -BearerToken $BearerToken -StorageAccountName $StorageAccountName -DestinationPath $DestinationPath -FlushPosition $response.filesize
# end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment