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

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