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

Release note (general): When using Azure Cloud Storage for data operations,
cockroach now calculates the Storage Account URL from the provided
AZURE_ENVIRONMENT query parameter. This defaults to AzurePublicCloud if not
specified to maintain backwards compatibility.

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",
     ],
@@ -27,10 +28,12 @@ go_test(
     deps = [
         "//pkg/cloud",
         "//pkg/cloud/cloudtestutils",
+        "//pkg/roachpb",
         "//pkg/security",
         "//pkg/settings/cluster",
         "//pkg/testutils/skip",
         "//pkg/util/leaktest",
+        "@com_github_azure_go_autorest_autorest//azure",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],

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"),
+		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