Fix tests and log requestId

snowflakedb · Jun 11, 2024 · 168f3d3 · 168f3d3
1 parent fe45ef5
commit 168f3d3
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 19 deletions.
diff --git 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/...est/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java b/...est/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java
@@ -1,12 +1,7 @@
 package net.snowflake.ingest.streaming.internal;
 
 import static java.time.ZoneOffset.UTC;
-import static net.snowflake.ingest.utils.Constants.ACCOUNT_URL;
-import static net.snowflake.ingest.utils.Constants.OPEN_CHANNEL_ENDPOINT;
-import static net.snowflake.ingest.utils.Constants.PRIVATE_KEY;
-import static net.snowflake.ingest.utils.Constants.RESPONSE_SUCCESS;
-import static net.snowflake.ingest.utils.Constants.ROLE;
-import static net.snowflake.ingest.utils.Constants.USER;
+import static net.snowflake.ingest.utils.Constants.*;
 import static org.mockito.ArgumentMatchers.argThat;
 
 import java.security.KeyPair;
@@ -313,8 +308,14 @@ public void testOpenChannelPostRequest() throws Exception {
         requestBuilder.generateStreamingIngestPostRequest(
             payload, OPEN_CHANNEL_ENDPOINT, "open channel");
 
-    Assert.assertEquals(
-        String.format("%s%s", urlStr, OPEN_CHANNEL_ENDPOINT), request.getRequestLine().getUri());
+    String expectedUrlPattern =
+            String.format("%s%s", urlStr, OPEN_CHANNEL_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.assertNotNull(request.getFirstHeader(HttpHeaders.USER_AGENT));
     Assert.assertNotNull(request.getFirstHeader(HttpHeaders.AUTHORIZATION));
     Assert.assertEquals("POST", request.getMethod());

diff --git a/...test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java b/...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));