Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNOW-1789757 Support GCP region specific endpoint #1064

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 63 additions & 0 deletions Snowflake.Data.Tests/UnitTests/PutGetStageInfoTest.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
/*
* Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
*/

using System.Collections.Generic;
sfc-gh-knozderko marked this conversation as resolved.
Show resolved Hide resolved
using NUnit.Framework;
using Snowflake.Data.Core;
using Snowflake.Data.Core.FileTransfer;

namespace Snowflake.Data.Tests.UnitTests
{
[TestFixture]
public class PutGetStageInfoTest
{
[Test]
[TestCaseSource(nameof(TestCases))]
public void TestGcsRegionalUrl(string region, bool useRegionalUrl, string endPoint, string expectedGcsEndpoint)
{
// arrange
var stageInfo = CreateGcsStageInfo(region, useRegionalUrl, endPoint);

// act
var gcsCustomEndpoint = stageInfo.GcsCustomEndpoint();

// assert
Assert.AreEqual(expectedGcsEndpoint, gcsCustomEndpoint);
}

internal static IEnumerable<object[]> TestCases()
{
yield return new object[] { "US-CENTRAL1", false, null, null };
yield return new object[] { "US-CENTRAL1", false, "", null };
yield return new object[] { "US-CENTRAL1", false, "null", null };
yield return new object[] { "US-CENTRAL1", false, " ", null };
yield return new object[] { "US-CENTRAL1", false, "example.com", "example.com" };
yield return new object[] { "ME-CENTRAL2", false, null, "storage.me-central2.rep.googleapis.com" };
yield return new object[] { "ME-CENTRAL2", true, null, "storage.me-central2.rep.googleapis.com" };
yield return new object[] { "ME-CENTRAL2", true, "", "storage.me-central2.rep.googleapis.com" };
yield return new object[] { "ME-CENTRAL2", true, " ", "storage.me-central2.rep.googleapis.com" };
yield return new object[] { "ME-CENTRAL2", true, "example.com", "example.com" };
yield return new object[] { "US-CENTRAL1", true, null, "storage.us-central1.rep.googleapis.com" };
yield return new object[] { "US-CENTRAL1", true, "", "storage.us-central1.rep.googleapis.com" };
yield return new object[] { "US-CENTRAL1", true, " ", "storage.us-central1.rep.googleapis.com" };
yield return new object[] { "US-CENTRAL1", true, "null", "storage.us-central1.rep.googleapis.com" };
yield return new object[] { "US-CENTRAL1", true, "example.com", "example.com" };
}

private PutGetStageInfo CreateGcsStageInfo(string region, bool useRegionalUrl, string endPoint) =>
new PutGetStageInfo
{
locationType = SFRemoteStorageUtil.GCS_FS,
location = "some location",
path = "some path",
region = region,
storageAccount = "some storage account",
isClientSideEncrypted = true,
stageCredentials = new Dictionary<string, string>(),
presignedUrl = "some pre-signed url",
endPoint = endPoint,
useRegionalUrl = useRegionalUrl
};
}
}
35 changes: 33 additions & 2 deletions Snowflake.Data.Tests/UnitTests/SFGCSClientTest.cs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
* Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
*/

using System;
Expand All @@ -18,7 +18,7 @@ namespace Snowflake.Data.Tests.UnitTests
using Snowflake.Data.Tests.Mock;
using Moq;

[TestFixture]
[TestFixture, NonParallelizable]
class SFGCSClientTest : SFBaseTest
{
// Mock data for file metadata
Expand Down Expand Up @@ -340,6 +340,37 @@ public async Task TestDownloadFileAsync(HttpStatusCode? httpStatusCode, ResultSt
AssertForDownloadFileTests(expectedResultStatus);
}

[Test]
[TestCase("us-central1", null, null, "https://storage.googleapis.com/mock-customer-stage/mock-id/tables/mock-key/")]
[TestCase("us-central1", "example.com", null, "https://example.com/mock-customer-stage/mock-id/tables/mock-key/")]
[TestCase("us-central1", "https://example.com", null, "https://example.com/mock-customer-stage/mock-id/tables/mock-key/")]
[TestCase("us-central1", null, true, "https://storage.us-central1.rep.googleapis.com/mock-customer-stage/mock-id/tables/mock-key/")]
[TestCase("me-central2", null, null, "https://storage.me-central2.rep.googleapis.com/mock-customer-stage/mock-id/tables/mock-key/")]
public void TestUseUriWithRegionsWhenNeeded(string region, string endPoint, bool useRegionalUrl, string expectedRequestUri)
{
var fileMetadata = new SFFileMetadata()
{
stageInfo = new PutGetStageInfo()
{
endPoint = endPoint,
location = Location,
locationType = SFRemoteStorageUtil.GCS_FS,
path = LocationPath,
presignedUrl = null,
region = region,
stageCredentials = _stageCredentials,
storageAccount = null,
useRegionalUrl = useRegionalUrl
}
};

// act
var uri = _client.FormBaseRequest(fileMetadata, "PUT").RequestUri.ToString();

// assert
Assert.AreEqual(expectedRequestUri, uri);
}

private void AssertForDownloadFileTests(ResultStatus expectedResultStatus)
{
if (expectedResultStatus == ResultStatus.DOWNLOADED)
Expand Down
50 changes: 39 additions & 11 deletions Snowflake.Data/Core/FileTransfer/StorageClient/SFGCSClient.cs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
using Newtonsoft.Json;
using Snowflake.Data.Log;
using System.Net;
using Google.Apis.Storage.v1;
using Google.Cloud.Storage.V1;

namespace Snowflake.Data.Core.FileTransfer.StorageClient
{
Expand Down Expand Up @@ -52,6 +54,8 @@
/// </summary>
private WebRequest _customWebRequest = null;

private static readonly string[] s_scopes = new[] { StorageService.Scope.DevstorageFullControl };

/// <summary>
/// GCS client with access token.
/// </summary>
Expand All @@ -65,15 +69,32 @@
Logger.Debug("Constructing client using access token");
AccessToken = accessToken;
GoogleCredential creds = GoogleCredential.FromAccessToken(accessToken, null);
StorageClient = Google.Cloud.Storage.V1.StorageClient.Create(creds);
var storageClientBuilder = new StorageClientBuilder
{
Credential = creds?.CreateScoped(s_scopes),
EncryptionKey = null
};
StorageClient = BuildStorageClient(storageClientBuilder, stageInfo);
}
else
{
Logger.Info("No access token received from GS, constructing anonymous client with no encryption support");
StorageClient = Google.Cloud.Storage.V1.StorageClient.CreateUnauthenticated();
var storageClientBuilder = new StorageClientBuilder
{
UnauthenticatedAccess = true
};
StorageClient = BuildStorageClient(storageClientBuilder, stageInfo);

Check warning on line 86 in Snowflake.Data/Core/FileTransfer/StorageClient/SFGCSClient.cs

View check run for this annotation

Codecov / codecov/patch

Snowflake.Data/Core/FileTransfer/StorageClient/SFGCSClient.cs#L82-L86

Added lines #L82 - L86 were not covered by tests
}
}

private Google.Cloud.Storage.V1.StorageClient BuildStorageClient(StorageClientBuilder builder, PutGetStageInfo stageInfo)
{
var gcsCustomEndpoint = stageInfo.GcsCustomEndpoint();
if (!string.IsNullOrEmpty(gcsCustomEndpoint))
builder.BaseUri = gcsCustomEndpoint;
return builder.Build();
}

internal void SetCustomWebRequest(WebRequest mockWebRequest)
{
_customWebRequest = mockWebRequest;
Expand Down Expand Up @@ -112,7 +133,7 @@
internal WebRequest FormBaseRequest(SFFileMetadata fileMetadata, string method)
{
string url = string.IsNullOrEmpty(fileMetadata.presignedUrl) ?
generateFileURL(fileMetadata.stageInfo.location, fileMetadata.RemoteFileName()) :
generateFileURL(fileMetadata.stageInfo, fileMetadata.RemoteFileName()) :
fileMetadata.presignedUrl;

WebRequest request = WebRequest.Create(url);
Expand Down Expand Up @@ -219,19 +240,26 @@
return null;
}

/// <summary>
/// Generate the file URL.
/// </summary>
/// <param name="stageLocation">The GCS file metadata.</param>
/// <param name="fileName">The GCS file metadata.</param>
internal string generateFileURL(string stageLocation, string fileName)
internal string generateFileURL(PutGetStageInfo stageInfo, string fileName)
{
var gcsLocation = ExtractBucketNameAndPath(stageLocation);
var storageHostPath = ExtractStorageHostPath(stageInfo);
var gcsLocation = ExtractBucketNameAndPath(stageInfo.location);
var fullFilePath = gcsLocation.key + fileName;
var link = "https://storage.googleapis.com/" + gcsLocation.bucket + "/" + fullFilePath;
var link = storageHostPath + gcsLocation.bucket + "/" + fullFilePath;
return link;
}

private string ExtractStorageHostPath(PutGetStageInfo stageInfo)
{
var gcsEndpoint = stageInfo.GcsCustomEndpoint();
var storageHostPath = string.IsNullOrEmpty(gcsEndpoint) ? "https://storage.googleapis.com/" : gcsEndpoint;
if (!storageHostPath.StartsWith("https://", StringComparison.OrdinalIgnoreCase))
storageHostPath = "https://" + storageHostPath;
if (!storageHostPath.EndsWith("/"))
storageHostPath = storageHostPath + "/";
return storageHostPath;
}

/// <summary>
/// Upload the file to the GCS location.
/// </summary>
Expand Down
17 changes: 17 additions & 0 deletions Snowflake.Data/Core/RestResponse.cs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
using Newtonsoft.Json.Converters;
using Newtonsoft.Json.Linq;
using Snowflake.Data.Client;
using Snowflake.Data.Core.FileTransfer;

namespace Snowflake.Data.Core
{
Expand Down Expand Up @@ -439,6 +440,22 @@

[JsonProperty(PropertyName = "endPoint", NullValueHandling = NullValueHandling.Ignore)]
internal string endPoint { get; set; }

[JsonProperty(PropertyName = "useRegionalUrl", NullValueHandling = NullValueHandling.Ignore)]
internal bool useRegionalUrl { get; set; }

private const string GcsRegionMeCentral2 = "me-central2";

internal string GcsCustomEndpoint()
{
if (!(locationType ?? string.Empty).Equals(SFRemoteStorageUtil.GCS_FS, StringComparison.OrdinalIgnoreCase))
return null;

Check warning on line 452 in Snowflake.Data/Core/RestResponse.cs

View check run for this annotation

Codecov / codecov/patch

Snowflake.Data/Core/RestResponse.cs#L452

Added line #L452 was not covered by tests
if (!string.IsNullOrWhiteSpace(endPoint) && endPoint != "null")
return endPoint;
if (GcsRegionMeCentral2.Equals(region, StringComparison.OrdinalIgnoreCase) || useRegionalUrl)
return $"storage.{region.ToLower()}.rep.googleapis.com";
return null;
}
}

internal class PutGetEncryptionMaterial
Expand Down
Loading