diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/AbstractCloudStorage.java b/src/main/java/net/snowflake/ingest/streaming/internal/AbstractCloudStorage.java index 10d8f9177..200c0fb95 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/AbstractCloudStorage.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/AbstractCloudStorage.java @@ -253,8 +253,7 @@ SnowflakeFileTransferMetadataV1 fetchSignedURL(String fileName) payload.put("file_name", fileName); Map response = this.makeConfigureCall(payload); - JsonNode responseNode = - this.parseStageLocation(mapper.valueToTree(response).get("stage_location")); + JsonNode responseNode = this.parseConfigureResponse(response); SnowflakeFileTransferMetadataV1 metadata = (SnowflakeFileTransferMetadataV1) @@ -264,14 +263,15 @@ SnowflakeFileTransferMetadataV1 fetchSignedURL(String fileName) return metadata; } - protected JsonNode parseStageLocation(JsonNode stageLocation) { - ObjectNode responseNode = mapper.createObjectNode(); + protected JsonNode parseConfigureResponse(Map response) { + JsonNode responseNode = mapper.valueToTree(response); // Currently there are a few mismatches between the client/configure response and what // SnowflakeFileTransferAgent expects - responseNode.putObject("data"); - ObjectNode dataNode = (ObjectNode) responseNode.get("data"); - dataNode.set("stageInfo", stageLocation); + ObjectNode mutable = (ObjectNode) responseNode; + mutable.putObject("data"); + ObjectNode dataNode = (ObjectNode) mutable.get("data"); + dataNode.set("stageInfo", responseNode.get("stage_location")); // JDBC expects this field which maps to presignedFileUrlName. We will set this later dataNode.putArray("src_locations").add("placeholder"); @@ -300,8 +300,7 @@ synchronized SnowflakeFileTransferMetadataWithAge refreshCloudStorageMetadata(bo Map payload = getConfigurePayload(); Map response = this.makeConfigureCall(payload); - JsonNode responseNode = - this.parseStageLocation(mapper.valueToTree(response).get("stage_location")); + JsonNode responseNode = this.parseConfigureResponse(response); // Do not change the prefix everytime we have to refresh credentials if (Utils.isNullOrEmpty(this.clientPrefix)) { this.clientPrefix = createClientPrefix(responseNode); diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestExternalVolume.java b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestExternalVolume.java index 5abfc0559..8b6986c23 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestExternalVolume.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestExternalVolume.java @@ -123,4 +123,10 @@ void put(String filePath, byte[] blob) { public String getVolumeHash() { return volumeHash; } + + private JsonNode parseStageLocation(JsonNode stageLocation) { + Map response = new HashMap<>(); + response.put("stage_location", stageLocation); + return parseConfigureResponse(response); + } }