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.

Sign up for GitHub

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-1373151: Proactively refresh token to avoid token expiration exception from JDBC #780

Merged
merged 9 commits into from
Jul 16, 2024
Merged

SNOW-1373151: Proactively refresh token to avoid token expiration exception from JDBC #780

Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
75a0727
fix cve
sfc-gh-tzhang Jun 15, 2024
85b5f74
Merge branch 'master' into tzhang-si-random
sfc-gh-tzhang Jun 20, 2024
f2c4e0b
Revert "SNOW-1457523: Fix CVE for snowflake-ingest-java io.airlift:ai…
sfc-gh-tzhang Jun 20, 2024
e0334c6
update pom
sfc-gh-tzhang Jun 20, 2024
8bddd2c
refresh token
sfc-gh-tzhang Jun 21, 2024
e6c4413
Merge branch 'master' into tzhang-si-refresh
sfc-gh-xhuang Jul 10, 2024
f2025e8
merge master
sfc-gh-tzhang Jul 13, 2024
521d696
merge master
sfc-gh-tzhang Jul 13, 2024
7eadb4a
fix pom
sfc-gh-tzhang Jul 13, 2024
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
18 changes: 15 additions & 3 deletions src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Paths;
import java.time.Duration;
import java.time.Instant;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
Expand Down Expand Up @@ -47,6 +49,11 @@ class StreamingIngestStage {
private static final long REFRESH_THRESHOLD_IN_MS =
TimeUnit.MILLISECONDS.convert(1, TimeUnit.MINUTES);

// Stage credential refresh interval, currently the token will expire in 1hr for GCS and 2hr for
sfc-gh-psaha marked this conversation as resolved.
Show resolved Hide resolved
// AWS/Azure, so set it a bit smaller than 1hr
private static final Duration refreshDuration = Duration.ofMinutes(58);
private static Instant prevRefresh = Instant.EPOCH;

private static final Logging logger = new Logging(StreamingIngestStage.class);

/**
Expand Down Expand Up @@ -180,6 +187,12 @@ private void putRemote(String fullFilePath, byte[] data, int retryCount)
InputStream inStream = new ByteArrayInputStream(data);

try {
// Proactively refresh the credential if it's going to expire, to avoid the token expiration
// error from JDBC which confuses customer
if (Instant.now().isAfter(prevRefresh.plus(refreshDuration))) {
refreshSnowflakeMetadata();
}

SnowflakeFileTransferAgent.uploadWithoutConnection(
SnowflakeFileTransferConfig.Builder.newInstance()
.setSnowflakeFileTransferMetadata(fileTransferMetadataCopy)
Expand All @@ -194,9 +207,6 @@ private void putRemote(String fullFilePath, byte[] data, int retryCount)
} catch (Exception e) {
if (retryCount == 0) {
// for the first exception, we always perform a metadata refresh.
logger.logInfo(
"Stage metadata need to be refreshed due to upload error: {} on first retry attempt",
e.getMessage());
this.refreshSnowflakeMetadata();
}
if (retryCount >= maxUploadRetries) {
Expand Down Expand Up @@ -281,6 +291,8 @@ synchronized SnowflakeFileTransferMetadataWithAge refreshSnowflakeMetadata(boole
SnowflakeFileTransferAgent.getFileTransferMetadatas(responseNode).get(0),
Optional.of(System.currentTimeMillis()));
}

prevRefresh = Instant.now();
return this.fileTransferMetadataWithAge;
}

Expand Down
Loading