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 #47163

Release note (sql change): 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.
This parameter should not be used when the cluster is in a mixed version
or upgrading state, as nodes that have not been upgraded will continue to
send requests to the AzurePublicCloud even in the presence of this parameter.

pkg/cloud/azure/BUILD.bazel

@@ -14,6 +14,7 @@ go_library(
         "//pkg/util/contextutil",
         "//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",
     ],
@@ -26,10 +27,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"
@@ -35,6 +36,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(
@@ -47,13 +50,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
 }
@@ -80,8 +88,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_ENVIRONMENT 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())
+		})
+	}
+}