-
Notifications
You must be signed in to change notification settings - Fork 57
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
SNOW-1497358 Support multiple stage for new table format (#812)
- Loading branch information
1 parent
66501f9
commit b9d31b6
Showing
12 changed files
with
601 additions
and
9 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
53 changes: 53 additions & 0 deletions
53
src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureRequest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,53 @@ | ||
| /* | ||
| * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. | ||
| */ | ||
|
|
||
| package net.snowflake.ingest.streaming.internal; | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonProperty; | ||
|
|
||
| /** Class used to serialize the channel configure request. */ | ||
| class ChannelConfigureRequest extends ClientConfigureRequest { | ||
| @JsonProperty("database") | ||
| private String database; | ||
|
|
||
| @JsonProperty("schema") | ||
| private String schema; | ||
|
|
||
| @JsonProperty("table") | ||
| private String table; | ||
|
|
||
| /** | ||
| * Constructor for channel configure request | ||
| * | ||
| * @param role Role to be used for the request. | ||
| * @param database Database name. | ||
| * @param schema Schema name. | ||
| * @param table Table name. | ||
| */ | ||
| ChannelConfigureRequest(String role, String database, String schema, String table) { | ||
| super(role); | ||
| this.database = database; | ||
| this.schema = schema; | ||
| this.table = table; | ||
| } | ||
|
|
||
| String getDatabase() { | ||
| return database; | ||
| } | ||
|
|
||
| String getSchema() { | ||
| return schema; | ||
| } | ||
|
|
||
| String getTable() { | ||
| return table; | ||
| } | ||
|
|
||
| @Override | ||
| public String getStringForLogging() { | ||
| return String.format( | ||
| "ChannelConfigureRequest(role=%s, db=%s, schema=%s, table=%s, file_name=%s)", | ||
| getRole(), database, schema, table, getFileName()); | ||
| } | ||
| } |
46 changes: 46 additions & 0 deletions
46
src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureResponse.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,46 @@ | ||
| /* | ||
| * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. | ||
| */ | ||
|
|
||
| package net.snowflake.ingest.streaming.internal; | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonIgnoreProperties; | ||
| import com.fasterxml.jackson.annotation.JsonProperty; | ||
|
|
||
| /** Class used to deserialize responses from channel configure endpoint */ | ||
| @JsonIgnoreProperties(ignoreUnknown = true) | ||
| class ChannelConfigureResponse extends StreamingIngestResponse { | ||
| @JsonProperty("status_code") | ||
| private Long statusCode; | ||
|
|
||
| @JsonProperty("message") | ||
| private String message; | ||
|
|
||
| @JsonProperty("stage_location") | ||
| private FileLocationInfo stageLocation; | ||
|
|
||
| @Override | ||
| Long getStatusCode() { | ||
| return statusCode; | ||
| } | ||
|
|
||
| void setStatusCode(Long statusCode) { | ||
| this.statusCode = statusCode; | ||
| } | ||
|
|
||
| String getMessage() { | ||
| return message; | ||
| } | ||
|
|
||
| void setMessage(String message) { | ||
| this.message = message; | ||
| } | ||
|
|
||
| FileLocationInfo getStageLocation() { | ||
| return stageLocation; | ||
| } | ||
|
|
||
| void setStageLocation(FileLocationInfo stageLocation) { | ||
| this.stageLocation = stageLocation; | ||
| } | ||
| } |
176 changes: 176 additions & 0 deletions
176
src/main/java/net/snowflake/ingest/streaming/internal/ExternalVolumeManager.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,176 @@ | ||
| /* | ||
| * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. | ||
| */ | ||
|
|
||
| package net.snowflake.ingest.streaming.internal; | ||
|
|
||
| import java.io.IOException; | ||
| import java.util.Calendar; | ||
| import java.util.Map; | ||
| import java.util.Optional; | ||
| import java.util.concurrent.ConcurrentHashMap; | ||
| import net.snowflake.client.jdbc.SnowflakeSQLException; | ||
| import net.snowflake.ingest.connection.IngestResponseException; | ||
| import net.snowflake.ingest.utils.ErrorCode; | ||
| import net.snowflake.ingest.utils.SFException; | ||
| import net.snowflake.ingest.utils.Utils; | ||
|
|
||
| class ExternalVolumeLocation { | ||
| public final String dbName; | ||
| public final String schemaName; | ||
| public final String tableName; | ||
|
|
||
| public ExternalVolumeLocation(String dbName, String schemaName, String tableName) { | ||
| this.dbName = dbName; | ||
| this.schemaName = schemaName; | ||
| this.tableName = tableName; | ||
| } | ||
| } | ||
|
|
||
| /** Class to manage multiple external volumes */ | ||
| class ExternalVolumeManager<T> implements IStorageManager<T, ExternalVolumeLocation> { | ||
| // Reference to the external volume per table | ||
| private final Map<String, StreamingIngestStorage<T, ExternalVolumeLocation>> externalVolumeMap; | ||
|
|
||
| // name of the owning client | ||
| private final String clientName; | ||
|
|
||
| // role of the owning client | ||
| private final String role; | ||
|
|
||
| // Reference to the Snowflake service client used for configure calls | ||
| private final SnowflakeServiceClient snowflakeServiceClient; | ||
|
|
||
| // Client prefix generated by the Snowflake server | ||
| private final String clientPrefix; | ||
|
|
||
| /** | ||
| * Constructor for ExternalVolumeManager | ||
| * | ||
| * @param isTestMode whether the manager in test mode | ||
| * @param role the role of the client | ||
| * @param clientName the name of the client | ||
| * @param snowflakeServiceClient the Snowflake service client used for configure calls | ||
| */ | ||
| ExternalVolumeManager( | ||
| boolean isTestMode, | ||
| String role, | ||
| String clientName, | ||
| SnowflakeServiceClient snowflakeServiceClient) { | ||
| this.role = role; | ||
| this.clientName = clientName; | ||
| this.snowflakeServiceClient = snowflakeServiceClient; | ||
| this.externalVolumeMap = new ConcurrentHashMap<>(); | ||
| try { | ||
| this.clientPrefix = | ||
| isTestMode | ||
| ? "testPrefix" | ||
| : this.snowflakeServiceClient | ||
| .clientConfigure(new ClientConfigureRequest(role)) | ||
| .getClientPrefix(); | ||
| } catch (IngestResponseException | IOException e) { | ||
| throw new SFException(e, ErrorCode.CLIENT_CONFIGURE_FAILURE, e.getMessage()); | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Given a fully qualified table name, return the target storage by looking up the table name | ||
| * | ||
| * @param fullyQualifiedTableName the target fully qualified table name | ||
| * @return target storage | ||
| */ | ||
| @Override | ||
| public StreamingIngestStorage<T, ExternalVolumeLocation> getStorage( | ||
| String fullyQualifiedTableName) { | ||
| // Only one chunk per blob in Iceberg mode. | ||
| StreamingIngestStorage<T, ExternalVolumeLocation> stage = | ||
| this.externalVolumeMap.get(fullyQualifiedTableName); | ||
|
|
||
| if (stage == null) { | ||
| throw new SFException( | ||
| ErrorCode.INTERNAL_ERROR, | ||
| String.format("No external volume found for table %s", fullyQualifiedTableName)); | ||
| } | ||
|
|
||
| return stage; | ||
| } | ||
|
|
||
| /** | ||
| * Add a storage to the manager by looking up the table name from the open channel response | ||
| * | ||
| * @param dbName the database name | ||
| * @param schemaName the schema name | ||
| * @param tableName the table name | ||
| * @param fileLocationInfo response from open channel | ||
| */ | ||
| @Override | ||
| public void addStorage( | ||
| String dbName, String schemaName, String tableName, FileLocationInfo fileLocationInfo) { | ||
| String fullyQualifiedTableName = | ||
| Utils.getFullyQualifiedTableName(dbName, schemaName, tableName); | ||
|
|
||
| try { | ||
| this.externalVolumeMap.put( | ||
| fullyQualifiedTableName, | ||
| new StreamingIngestStorage<T, ExternalVolumeLocation>( | ||
| this, | ||
| this.clientName, | ||
| fileLocationInfo, | ||
| new ExternalVolumeLocation(dbName, schemaName, tableName), | ||
| DEFAULT_MAX_UPLOAD_RETRIES)); | ||
| } catch (SnowflakeSQLException | IOException err) { | ||
| throw new SFException( | ||
| err, | ||
| ErrorCode.UNABLE_TO_CONNECT_TO_STAGE, | ||
| String.format("fullyQualifiedTableName=%s", fullyQualifiedTableName)); | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Gets the latest file location info (with a renewed short-lived access token) for the specified | ||
| * location | ||
| * | ||
| * @param location A reference to the target location | ||
| * @param fileName optional filename for single-file signed URL fetch from server | ||
| * @return the new location information | ||
| */ | ||
| @Override | ||
| public FileLocationInfo getRefreshedLocation( | ||
| ExternalVolumeLocation location, Optional<String> fileName) { | ||
| try { | ||
| ChannelConfigureRequest request = | ||
| new ChannelConfigureRequest( | ||
| this.role, location.dbName, location.schemaName, location.tableName); | ||
| fileName.ifPresent(request::setFileName); | ||
| ChannelConfigureResponse response = this.snowflakeServiceClient.channelConfigure(request); | ||
| return response.getStageLocation(); | ||
| } catch (IngestResponseException | IOException e) { | ||
| throw new SFException(e, ErrorCode.CHANNEL_CONFIGURE_FAILURE, e.getMessage()); | ||
| } | ||
| } | ||
|
|
||
| // TODO: SNOW-1502887 Blob path generation for external volume | ||
| @Override | ||
| public String generateBlobPath() { | ||
| return "snow_dummy_file_name.parquet"; | ||
| } | ||
|
|
||
| // TODO: SNOW-1502887 Blob path generation for iceberg table | ||
| @Override | ||
| public void decrementBlobSequencer() {} | ||
|
|
||
| // TODO: SNOW-1502887 Blob path generation for iceberg table | ||
| public String getBlobPath(Calendar calendar, String clientPrefix) { | ||
| return ""; | ||
| } | ||
|
|
||
| /** | ||
| * Get the client prefix from first external volume in the map | ||
| * | ||
| * @return the client prefix | ||
| */ | ||
| @Override | ||
| public String getClientPrefix() { | ||
| return this.clientPrefix; | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.