Add support for AzureUSGovernment - Issue #332 (#346)

PlagueHO · Mar 24, 2020 · 5f63d7c · 5f63d7c
1 parent 8cb3622
commit 5f63d7c
Show file tree

Hide file tree

Showing 8 changed files with 232 additions and 27 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Updated `README.MD` to documentation to reduce focus on collections without
  partition keys - fixes [Issue #342](https://github.com/PlagueHO/CosmosDB/issues/342).
 
+### Added
+
+- Added support for `Environment` parameter in `New-CosmosDbContext` to allow
+ using Azure US Government Cloud - fixes [Issue #322](https://github.com/PlagueHO/CosmosDB/issues/322).
+
 ## [3.6.1] - 2020-03-19
 
 ### Changed

diff --git a/README.md b/README.md
@@ -23,6 +23,7 @@
 - [Getting Started](#getting-started)
  - [Working with Contexts](#working-with-contexts)
  - [Create a Context specifying the Key Manually](#create-a-context-specifying-the-key-manually)
+ - [Create a Context for a Cosmos DB in Azure US Government Cloud](#create-a-context-for-a-cosmos-db-in-azure-us-government-cloud)
  - [Use CosmosDB Module to Retrieve Key from Azure Management Portal](#use-cosmosdb-module-to-retrieve-key-from-azure-management-portal)
  - [Create a Context for a Cosmos DB Emulator](#create-a-context-for-a-cosmos-db-emulator)
  - [Create a Context from Resource Authorization Tokens](#create-a-context-from-resource-authorization-tokens)
@@ -138,6 +139,16 @@ create a context variable:
 $cosmosDbContext = New-CosmosDbContext -Account 'MyAzureCosmosDB' -Database 'MyDatabase' -Key $primaryKey
 ```
 
+#### Create a Context for a Cosmos DB in Azure US Government Cloud
+
+Use the key secure string, Azure Cosmos DB account name and database to
+create a context variable and set the `Environment` parameter to
+`AzureUSGovernment`:
+
+```powershell
+$cosmosDbContext = New-CosmosDbContext -Account 'MyAzureCosmosDB' -Database 'MyDatabase' -Key $primaryKey -Environment AzureUSGovernment
+```
+
 #### Use CosmosDB Module to Retrieve Key from Azure Management Portal
 
 To create a context object so that the _CosmosDB PowerShell module_

diff --git a/docs/Get-CosmosDbUri.md b/docs/Get-CosmosDbUri.md
@@ -15,7 +15,8 @@ be sent to.
 ## SYNTAX
 
 ```powershell
-Get-CosmosDbUri [-Account] <String> [[-BaseUri] <String>] [<CommonParameters>]
+Get-CosmosDbUri [-Account] <String> [[-BaseUri] <String>] [[-Environment] <Environment>]
+ [<CommonParameters>]
 ```
 
 ## DESCRIPTION
@@ -41,6 +42,15 @@ PS C:\>$uri = Get-CosmosDbUri -Account 'MyAzureCosmosDB' -BaseUri 'localhost'
 Generates the URI for accessing a Cosmos DB emulator account
 on the localhost.
 
+### EXAMPLE 3
+
+```powershell
+PS C:\>$uri = Get-CosmosDbUri -Account 'MyAzureCosmosDB' -Environment 'AzureUSGovernment'
+```
+
+Generates the URI for accessing a Cosmos DB account in the
+US Government cloud.
+
 ## PARAMETERS
 
 ### -Account
@@ -67,7 +77,7 @@ If not specified it will default to 'documents.azure.com'.
 
 ```yaml
 Type: String
-Parameter Sets: (All)
+Parameter Sets: Uri
 Aliases:
 
 Required: False
@@ -77,6 +87,27 @@ Accept pipeline input: False
 Accept wildcard characters: False
 ```
 
+### -Environment
+
+This is the Azure environment hosting the Cosmos DB account.
+
+The supported values are:
+
+- AzureCloud
+- AzureUSGovernment
+
+```yaml
+Type: Environment
+Parameter Sets: Environment
+Aliases:
+
+Required: False
+Position: Named
+Default value: AzureCloud
+Accept pipeline input: False
+Accept wildcard characters: False
+```
+
 ### CommonParameters
 
 This cmdlet supports the common parameters: -Debug, -ErrorAction, -ErrorVariable, -InformationAction, -InformationVariable, -OutVariable, -OutBuffer, -PipelineVariable, -Verbose, -WarningAction, and -WarningVariable.

diff --git a/docs/New-CosmosDbContext.md b/docs/New-CosmosDbContext.md
@@ -17,29 +17,32 @@ to connect to a Cosmos DB.
 ### Account (Default)
 
 ```powershell
-New-CosmosDbContext -Account <String> [-Database <String>] -Key <SecureString> [-KeyType <String>]
- [-BackoffPolicy <BackoffPolicy>] [<CommonParameters>]
+New-CosmosDbContext -Account <String> [-Database <String>]
+ -Key <SecureString> [-KeyType <String>] [-BackoffPolicy <BackoffPolicy>]
+ [-Environment <Environment>] [<CommonParameters>]
 ```
 
 ### AzureAccount
 
 ```powershell
 New-CosmosDbContext -Account <String> [-Database <String>] -ResourceGroupName <String>
- [-MasterKeyType <String>] [-BackoffPolicy <BackoffPolicy>] [<CommonParameters>]
+ [-MasterKeyType <String>] [-BackoffPolicy <BackoffPolicy>] [-Environment <Environment>]
+ [<CommonParameters>]
 ```
 
 ### Token
 
 ```powershell
 New-CosmosDbContext -Account <String> [-Database <String>] -Token <ContextToken[]>
- [-BackoffPolicy <BackoffPolicy>] [<CommonParameters>]
+ [-BackoffPolicy <BackoffPolicy>] [-Environment <Environment>] [<CommonParameters>]
 ```
 
 ### Emulator
 
 ```powershell
-New-CosmosDbContext [-Database <String>] [-Key <SecureString>] [-Emulator] [-Port <Int16>] [-URI <String>]
- [-Token <ContextToken[]>] [-BackoffPolicy <BackoffPolicy>] [<CommonParameters>]
+New-CosmosDbContext [-Database <String>] [-Key <SecureString>] [-Emulator]
+ [-Port <Int16>] [-URI <String>] [-Token <ContextToken[]>]
+ [-BackoffPolicy <BackoffPolicy>] [<CommonParameters>]
 ```
 
 ## DESCRIPTION
@@ -106,6 +109,16 @@ PS C:\> $cosmosDbContext = New-CosmosDbContext -Emulator -URI 'mycosmosdb' -Key
 Creates a Cosmos DB context by using a Cosmos DB Emulator installed onto
 the machine 'mycosmosdb' with a custom key emulator key.
 
+### EXAMPLE 6
+
+```powershell
+PS C:\> $primaryKey = ConvertTo-SecureString -String 'your master key' -AsPlainText -Force
+PS C:\> $cosmosDbContext = New-CosmosDbContext -Account 'MyAzureGovCosmosDB' -Database 'MyDatabase' -Key $primaryKey -Environment 'AzureUSGovernment'
+```
+
+Creates a CosmosDB context specifying the master key manually connecting
+to the Azure US Government cloud.
+
 ## PARAMETERS
 
 ### -Account
@@ -126,7 +139,7 @@ Accept wildcard characters: False
 
 ### -BackoffPolicy
 
-This a Back-off Policy object that controls the retry behaviour for
+This a Back-off Policy object that controls the retry behavior for
 requests using this context.
 
 ```yaml
@@ -174,6 +187,27 @@ Accept pipeline input: False
 Accept wildcard characters: False
 ```
 
+### -Environment
+
+This is the Azure environment hosting the Cosmos DB account.
+
+The supported values are:
+
+- AzureCloud
+- AzureUSGovernment
+
+```yaml
+Type: Environment
+Parameter Sets: Account, AzureAccount, Token
+Aliases:
+
+Required: False
+Position: Named
+Default value: AzureCloud
+Accept pipeline input: False
+Accept wildcard characters: False
+```
+
 ### -Key
 
 The key to be used to access the Cosmos DB account or Cosmos DB emulator.

diff --git a/source/Private/utils/Get-CosmosDbUri.ps1 b/source/Private/utils/Get-CosmosDbUri.ps1
@@ -1,18 +1,35 @@
 function Get-CosmosDbUri
 {
 
- [CmdletBinding()]
- [OutputType([uri])]
+ [CmdletBinding(
+ DefaultParameterSetName = 'Environment'
+ )]
+ [OutputType([System.Uri])]
  param
  (
  [Parameter(Mandatory = $true)]
  [System.String]
  $Account,
 
- [Parameter()]
+ [Parameter(ParameterSetName = 'Uri')]
  [System.String]
- $BaseUri = 'documents.azure.com'
+ $BaseUri = 'documents.azure.com',
+
+ [Parameter(ParameterSetName = 'Environment')]
+ [CosmosDB.Environment]
+ $Environment = [CosmosDB.Environment]::AzureCloud
  )
 
- return [uri]::new(('https://{0}.{1}' -f $Account, $BaseUri))
+ if ($PSCmdlet.ParameterSetName -eq 'Environment')
+ {
+ switch ($Environment)
+ {
+ 'AzureUSGovernment'
+ {
+ $BaseUri = 'documents.azure.us'
+ }
+ }
+ }
+
+ return [System.Uri]::new(('https://{0}.{1}' -f $Account, $BaseUri))
 }
diff --git a/source/Public/utils/New-CosmosDbContext.ps1 b/source/Public/utils/New-CosmosDbContext.ps1
@@ -63,7 +63,13 @@ function New-CosmosDbContext
  [Parameter()]
  [ValidateNotNullOrEmpty()]
  [CosmosDB.BackoffPolicy]
- $BackoffPolicy
+ $BackoffPolicy,
+
+ [Parameter(ParameterSetName = 'Account')]
+ [Parameter(ParameterSetName = 'Token')]
+ [Parameter(ParameterSetName = 'AzureAccount')]
+ [CosmosDB.Environment]
+ $Environment = [CosmosDB.Environment]::AzureCloud
  )
 
  switch ($PSCmdlet.ParameterSetName)
@@ -92,25 +98,25 @@ function New-CosmosDbContext
  }
  catch
  {
- $null = Connect-AzAccount
+ $null = Connect-AzAccount -Environment $Environment
  }
 
  $Key = Get-CosmosDbAccountMasterKey `
  -ResourceGroupName $ResourceGroupName `
  -Name $Account `
  -MasterKeyType $MasterKeyType
 
- $BaseUri = (Get-CosmosDbUri -Account $Account)
+ $BaseUri = Get-CosmosDbUri -Account $Account -Environment $Environment
  }
 
  'Account'
  {
- $BaseUri = (Get-CosmosDbUri -Account $Account)
+ $BaseUri = Get-CosmosDbUri -Account $Account -Environment $Environment
  }
 
  'Token'
  {
- $BaseUri = (Get-CosmosDbUri -Account $Account)
+ $BaseUri = Get-CosmosDbUri -Account $Account -Environment $Environment
  }
  }
 

diff --git a/source/classes/CosmosDB/CosmosDB.cs b/source/classes/CosmosDB/CosmosDB.cs
@@ -1,6 +1,12 @@
 using System;
 
 namespace CosmosDB {
+ public enum Environment {
+ AzureCloud,
+ AzureUSGovernment
+ }
+
+
  public class ContextToken
  {
  public System.String Resource;
@@ -25,6 +31,7 @@ public class Context
  public System.String BaseUri;
  public CosmosDB.ContextToken[] Token;
  public CosmosDB.BackoffPolicy BackoffPolicy;
+ public CosmosDB.Environment Environment;
  }
 
  namespace IndexingPolicy {