From d47b5c1053ff4ae46abb95660606de0bc8528c9b Mon Sep 17 00:00:00 2001
From: Jay Patel <jay.patel@snowflake.com>
Date: Fri, 7 Jun 2024 14:05:26 -0700
Subject: [PATCH] Fix tests and log requestId

---
 .../ingest/connection/RequestBuilder.java     | 19 ++++++++++++-------
 .../SnowflakeStreamingIngestClientTest.java   | 12 ++++++++----
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java b/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java
index 2e8b13553..fa2ac3ba7 100644
--- a/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java
+++ b/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java
@@ -678,18 +678,23 @@ public HttpGet generateHistoryRangeRequest(
    */
   public HttpPost generateStreamingIngestPostRequest(
       String payload, String endPoint, String message) {
-    LOGGER.debug("Generate Snowpipe streaming request: endpoint={}, payload={}", endPoint, payload);
+    final String requestId = UUID.randomUUID().toString();
+    LOGGER.debug(
+        "Generate Snowpipe streaming request: endpoint={}, payload={}, requestId={}",
+        endPoint,
+        payload,
+        requestId);
     // Make the corresponding URI
     URI uri = null;
     try {
       uri =
           new URIBuilder()
-                  .setScheme(scheme)
-                  .setHost(host)
-                  .setPort(port)
-                  .setPath(endPoint)
-                  .setParameter(REQUEST_ID, UUID.randomUUID().toString())
-                  .build();
+              .setScheme(scheme)
+              .setHost(host)
+              .setPort(port)
+              .setPath(endPoint)
+              .setParameter(REQUEST_ID, requestId)
+              .build();
     } catch (URISyntaxException e) {
       throw new SFException(e, ErrorCode.BUILD_REQUEST_FAILURE, message);
     }
diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java
index 1693e1520..5b24bcc7f 100644
--- a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java
+++ b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java
@@ -368,7 +368,6 @@ public void testGetChannelsStatusWithRequest() throws Exception {
     ChannelsStatusRequest.ChannelStatusRequestDTO dto =
         new ChannelsStatusRequest.ChannelStatusRequestDTO(channel);
     ChannelsStatusRequest request = new ChannelsStatusRequest();
-    request.setRequestId("null_0");
     request.setChannels(Collections.singletonList(dto));
     ChannelsStatusResponse result = client.getChannelsStatus(Collections.singletonList(channel));
     Assert.assertEquals(response.getMessage(), result.getMessage());
@@ -547,9 +546,15 @@ public void testRegisterBlobRequestCreationSuccess() throws Exception {
     HttpPost request =
         requestBuilder.generateStreamingIngestPostRequest(
             payload, REGISTER_BLOB_ENDPOINT, "register blob");
+    String expectedUrlPattern =
+        String.format("%s%s", urlStr, REGISTER_BLOB_ENDPOINT) + "(\\?requestId=[a-f0-9\\-]{36})?";
+
+    Assert.assertTrue(
+        String.format(
+            "Expected URL to match pattern: %s but was: %s",
+            expectedUrlPattern, request.getRequestLine().getUri()),
+        request.getRequestLine().getUri().matches(expectedUrlPattern));
 
-    Assert.assertEquals(
-        String.format("%s%s", urlStr, REGISTER_BLOB_ENDPOINT), request.getRequestLine().getUri());
     Assert.assertNotNull(request.getFirstHeader(HttpHeaders.USER_AGENT));
     Assert.assertNotNull(request.getFirstHeader(HttpHeaders.AUTHORIZATION));
     Assert.assertEquals("POST", request.getMethod());
@@ -1432,7 +1437,6 @@ public void testGetLatestCommittedOffsetTokens() throws Exception {
     ChannelsStatusRequest.ChannelStatusRequestDTO dto =
         new ChannelsStatusRequest.ChannelStatusRequestDTO(channel);
     ChannelsStatusRequest request = new ChannelsStatusRequest();
-    request.setRequestId("null_0");
     request.setChannels(Collections.singletonList(dto));
     Map<String, String> result =
         client.getLatestCommittedOffsetTokens(Collections.singletonList(channel));