From 4b9737042712cbd4c1500e9aa1d10ef8519eeef8 Mon Sep 17 00:00:00 2001 From: Nathan Lowe Date: Fri, 1 Apr 2022 14:19:20 -0700 Subject: [PATCH] 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 (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 | 3 ++ pkg/cloud/azure/azure_storage.go | 14 +++++- pkg/cloud/azure/azure_storage_test.go | 70 ++++++++++++++++++++++++--- pkg/roachpb/api.proto | 1 + 4 files changed, 81 insertions(+), 7 deletions(-) diff --git a/pkg/cloud/azure/BUILD.bazel b/pkg/cloud/azure/BUILD.bazel index a1339dca3678..bfbd2b84b232 100644 --- a/pkg/cloud/azure/BUILD.bazel +++ b/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", ], diff --git a/pkg/cloud/azure/azure_storage.go b/pkg/cloud/azure/azure_storage.go index 310124458493..c44d186c522a 100644 --- a/pkg/cloud/azure/azure_storage.go +++ b/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,6 +51,7 @@ 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) @@ -55,6 +59,10 @@ func parseAzureURL( 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") } diff --git a/pkg/cloud/azure/azure_storage_test.go b/pkg/cloud/azure/azure_storage_test.go index e120d98b6dc8..52447ed4e3d0 100644 --- a/pkg/cloud/azure/azure_storage_test.go +++ b/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()) + }) + } +} diff --git a/pkg/roachpb/api.proto b/pkg/roachpb/api.proto index a2fb4ece8002..89504c588d62 100644 --- a/pkg/roachpb/api.proto +++ b/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