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 {