pkg/cloud/azure: Support specifying Azure environments in storage URLs

The Azure Storage cloud provider learned a new parameter, AZURE_ENVIRONMENT,
which specifies which azure environment the storage account in question
belongs to. This allows cockroach to backup and restore data to Azure
Storage Accounts outside the main Azure Public Cloud. For backwards
compatibility, this defaults to "AzurePublicCloud" if AZURE_ENVIRONMENT
is not specified.

Fixes cockroachdb#47163

pkg/cloud/azure/BUILD.bazel

@@ -15,6 +15,7 @@ go_library(
         "//pkg/util/ioctx",
         "//pkg/util/tracing",
         "@com_github_azure_azure_storage_blob_go//azblob",
+        "@com_github_azure_go_autorest_autorest//azure",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_gogo_protobuf//types",
     ],

pkg/cloud/azure/azure_storage.go

@@ -19,6 +19,7 @@ import (
 	"strings"
 
 	"github.com/Azure/azure-storage-blob-go/azblob"
+	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/cloud"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -36,6 +37,8 @@ const (
 	AzureAccountNameParam = "AZURE_ACCOUNT_NAME"
 	// AzureAccountKeyParam is the query parameter for account_key in an azure URI.
 	AzureAccountKeyParam = "AZURE_ACCOUNT_KEY"
+	// AzureEnvironmentKeyParam is the query parameter for the environment name in an azure URI.
+	AzureEnvironmentKeyParam = "AZURE_ENVIRONMENT"
 )
 
 func parseAzureURL(
@@ -48,13 +51,18 @@ func parseAzureURL(
 		Prefix:      uri.Path,
 		AccountName: uri.Query().Get(AzureAccountNameParam),
 		AccountKey:  uri.Query().Get(AzureAccountKeyParam),
+		Environment: uri.Query().Get(AzureEnvironmentKeyParam),
 	}
 	if conf.AzureConfig.AccountName == "" {
 		return conf, errors.Errorf("azure uri missing %q parameter", AzureAccountNameParam)
 	}
 	if conf.AzureConfig.AccountKey == "" {
 		return conf, errors.Errorf("azure uri missing %q parameter", AzureAccountKeyParam)
 	}
+	if conf.AzureConfig.Environment == "" {
+		// Default to AzurePublicCloud if not specified for backwards compatibility
+		conf.AzureConfig.Environment = azure.PublicCloud.Name
+	}
 	conf.AzureConfig.Prefix = strings.TrimLeft(conf.AzureConfig.Prefix, "/")
 	return conf, nil
 }
@@ -81,8 +89,12 @@ func makeAzureStorage(
 	if err != nil {
 		return nil, errors.Wrap(err, "azure credential")
 	}
+	env, err := azure.EnvironmentFromName(conf.Environment)
+	if err != nil {
+		return nil, errors.Wrap(err, "azure environment")
+	}
 	p := azblob.NewPipeline(credential, azblob.PipelineOptions{})
-	u, err := url.Parse(fmt.Sprintf("https://%s.blob.core.windows.net", conf.AccountName))
+	u, err := url.Parse(fmt.Sprintf("https://%s.blob.%s", conf.AccountName, env.StorageEndpointSuffix))
 	if err != nil {
 		return nil, errors.Wrap(err, "azure: account name is not valid")
 	}

pkg/cloud/azure/azure_storage_test.go

@@ -11,13 +11,17 @@
 package azure
 
 import (
+	"context"
+	"encoding/base64"
 	"fmt"
 	"net/url"
 	"os"
 	"testing"
 
+	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/cockroachdb/cockroach/pkg/cloud"
 	"github.com/cockroachdb/cockroach/pkg/cloud/cloudtestutils"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
@@ -27,25 +31,30 @@ import (
 )
 
 type azureConfig struct {
-	account, key, bucket string
+	account, key, bucket, environment string
 }
 
 func (a azureConfig) filePath(f string) string {
-	return fmt.Sprintf("azure://%s/%s?%s=%s&%s=%s",
+	return fmt.Sprintf("azure://%s/%s?%s=%s&%s=%s&%s=%s",
 		a.bucket, f,
 		AzureAccountKeyParam, url.QueryEscape(a.key),
-		AzureAccountNameParam, url.QueryEscape(a.account))
+		AzureAccountNameParam, url.QueryEscape(a.account),
+		AzureEnvironmentKeyParam, url.QueryEscape(a.environment))
 }
 
 func getAzureConfig() (azureConfig, error) {
 	cfg := azureConfig{
 		account: os.Getenv("AZURE_ACCOUNT_NAME"),
 		key:     os.Getenv("AZURE_ACCOUNT_KEY"),
 		bucket:  os.Getenv("AZURE_CONTAINER"),
+		environment: azure.PublicCloud.Name,
 	}
 	if cfg.account == "" || cfg.key == "" || cfg.bucket == "" {
 		return azureConfig{}, errors.New("AZURE_ACCOUNT_NAME, AZURE_ACCOUNT_KEY, AZURE_CONTAINER must all be set")
 	}
+	if v, ok := os.LookupEnv(AzureEnvironmentKeyParam); ok {
+		cfg.environment = v
+	}
 	return cfg, nil
 }
 func TestAzure(t *testing.T) {
@@ -80,3 +89,52 @@ func TestAntagonisticAzureRead(t *testing.T) {
 
 	cloudtestutils.CheckAntagonisticRead(t, conf, testSettings)
 }
+
+func TestParseAzureURL(t *testing.T) {
+	t.Run("Defaults to Public Cloud when AZURE_ENVIRONEMNT unset", func(t *testing.T) {
+		u, err := url.Parse("azure://container/path?AZURE_ACCOUNT_NAME=account&AZURE_ACCOUNT_KEY=key")
+		require.NoError(t, err)
+
+		sut, err := parseAzureURL(cloud.ExternalStorageURIContext{}, u)
+		require.NoError(t, err)
+
+		require.Equal(t, azure.PublicCloud.Name, sut.AzureConfig.Environment)
+	})
+
+	t.Run("Can Override AZURE_ENVIRONMENT", func(t *testing.T) {
+		u, err := url.Parse("azure://container/path?AZURE_ACCOUNT_NAME=account&AZURE_ACCOUNT_KEY=key&AZURE_ENVIRONMENT=AzureUSGovernmentCloud")
+		require.NoError(t, err)
+
+		sut, err := parseAzureURL(cloud.ExternalStorageURIContext{}, u)
+		require.NoError(t, err)
+
+		require.Equal(t, azure.USGovernmentCloud.Name, sut.AzureConfig.Environment)
+	})
+}
+
+func TestMakeAzureStorageURLFromEnvironment(t *testing.T) {
+	for _, tt := range []struct{
+		environment string
+		expected string
+	}{
+		{environment: azure.PublicCloud.Name, expected: "https://account.blob.core.windows.net/container"},
+		{environment: azure.USGovernmentCloud.Name, expected: "https://account.blob.core.usgovcloudapi.net/container"},
+	} {
+		t.Run(tt.environment, func(t *testing.T) {
+			sut, err := makeAzureStorage(context.Background(), cloud.ExternalStorageContext{}, roachpb.ExternalStorage{
+				AzureConfig: &roachpb.ExternalStorage_Azure{
+					Container:   "container",
+					Prefix:      "path",
+					AccountName: "account",
+					AccountKey:  base64.StdEncoding.EncodeToString([]byte("key")),
+					Environment: tt.environment,
+				},
+			})
+
+			require.NoError(t, err)
+
+			u := sut.(*azureStorage).container.URL()
+			require.Equal(t, tt.expected, u.String())
+		})
+	}
+}

pkg/roachpb/api.proto

@@ -1404,6 +1404,7 @@ message ExternalStorage {
 
     string account_name = 3;
     string account_key = 4;
+    string environment = 5;
   }
   message FileTable {
     // User interacting with the external storage. This is used to check access