diff --git a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/AdapterClientWrapper.java b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/AdapterClientWrapper.java index a6710342..9978e9c8 100644 --- a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/AdapterClientWrapper.java +++ b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/AdapterClientWrapper.java @@ -86,6 +86,10 @@ ByteString sendGrpcRequest( for (AdaptMessageResponse adaptMessageResponse : serverStream) { adaptMessageResponse.getStateUpdatesMap().forEach(attachmentsCache::put); collectedPayloads.add(adaptMessageResponse.getPayload()); + if (adaptMessageResponse.getLast()) { + // This is the last message in the stream. Stop processing immediately. + break; + } } } catch (RuntimeException e) { LOG.error("Error executing AdaptMessage request: ", e); diff --git a/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/AdapterClientWrapperTest.java b/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/AdapterClientWrapperTest.java index 78a3c496..dadb78ad 100644 --- a/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/AdapterClientWrapperTest.java +++ b/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/AdapterClientWrapperTest.java @@ -74,11 +74,7 @@ public void sendGrpcRequest_SuccessfulResponse() { Map stateUpdates = new HashMap<>(); stateUpdates.put("k1", "v1"); stateUpdates.put("k2", "v2"); - AdaptMessageResponse mockResponse = - AdaptMessageResponse.newBuilder() - .setPayload(ByteString.copyFromUtf8("test response")) - .putAllStateUpdates(stateUpdates) - .build(); + AdaptMessageResponse mockResponse = createAdaptMessageResponse("test response", stateUpdates); Iterator mockResponseIterator = Collections.singletonList(mockResponse).iterator(); AdaptMessageRequest expectedRequest = @@ -106,21 +102,13 @@ public void sendGrpcRequest_MultipleResponses() { stateUpdates1.put("k1", "v1"); stateUpdates1.put("k2", "v2"); AdaptMessageResponse mockResponse1 = - AdaptMessageResponse.newBuilder() - .setPayload(ByteString.copyFromUtf8(" test response 1")) - .putAllStateUpdates(stateUpdates1) - .build(); + createAdaptMessageResponse(" test response 1", stateUpdates1); Map stateUpdates2 = new HashMap<>(); stateUpdates2.put("k3", "v3"); AdaptMessageResponse mockResponse2 = - AdaptMessageResponse.newBuilder() - .setPayload(ByteString.copyFromUtf8(" test response 2")) - .putAllStateUpdates(stateUpdates2) - .build(); + createAdaptMessageResponse(" test response 2", stateUpdates2); AdaptMessageResponse mockResponse3 = - AdaptMessageResponse.newBuilder() - .setPayload(ByteString.copyFromUtf8("test header")) - .build(); + createAdaptMessageResponse("test header", Collections.emptyMap()); Iterator mockResponseIterator = Arrays.asList(mockResponse1, mockResponse2, mockResponse3).iterator(); when(mockServerStream.iterator()).thenReturn(mockResponseIterator); @@ -142,6 +130,42 @@ public void sendGrpcRequest_MultipleResponses() { assertThat(attachmentsCache.get("k3")).hasValue("v3"); } + @Test + public void sendGrpcRequest_MultipleResponsesWithLastBitSet() { + int streamId = 1; + byte[] payload = "test payload".getBytes(); + Map stateUpdates1 = new HashMap<>(); + stateUpdates1.put("k1", "v1"); + stateUpdates1.put("k2", "v2"); + AdaptMessageResponse mockResponse1 = + createAdaptMessageResponse(" test response 1", stateUpdates1); + Map stateUpdates2 = new HashMap<>(); + stateUpdates2.put("k3", "v3"); + AdaptMessageResponse mockResponse2 = + createAdaptMessageResponse(" test response 2", stateUpdates2); + AdaptMessageResponse mockResponse3 = + createAdaptMessageResponse("test header (last)", Collections.emptyMap(), true); + Iterator mockResponseIterator = + Arrays.asList(mockResponse1, mockResponse2, mockResponse3).iterator(); + when(mockServerStream.iterator()).thenReturn(mockResponseIterator); + AdaptMessageRequest expectedRequest = + AdaptMessageRequest.newBuilder() + .setName("test-session") + .setProtocol("cassandra") + .setPayload(ByteString.copyFrom(payload)) + .build(); + + ByteString response = + adapterClientWrapper.sendGrpcRequest(payload, new HashMap<>(), context, streamId); + + verify(mockCallable).call(expectedRequest, context); + assertThat(response) + .isEqualTo(ByteString.copyFromUtf8("test header (last) test response 1 test response 2")); + assertThat(attachmentsCache.get("k1")).hasValue("v1"); + assertThat(attachmentsCache.get("k2")).hasValue("v2"); + assertThat(attachmentsCache.get("k3")).hasValue("v3"); + } + @Test public void sendGrpcRequest_NoResponse() { int streamId = 1; @@ -156,8 +180,7 @@ public void sendGrpcRequest_NoResponse() { when(mockServerStream.iterator()).thenReturn(mockResponseIterator); when(mockSession.getName()).thenReturn("test-session"); - ByteString response = - adapterClientWrapper.sendGrpcRequest(payload, new HashMap<>(), context, streamId); + adapterClientWrapper.sendGrpcRequest(payload, new HashMap<>(), context, streamId); verify(mockCallable).call(expectedRequest, context); } @@ -174,4 +197,18 @@ public void sendGrpcRequest_SessionCreationFailure() { adapterClientWrapper.sendGrpcRequest( payload, new HashMap<>(), GrpcCallContext.createDefault(), streamId)); } + + private static AdaptMessageResponse createAdaptMessageResponse( + String text, Map stateUpdates, boolean isLast) { + return AdaptMessageResponse.newBuilder() + .setPayload(ByteString.copyFromUtf8(text)) + .putAllStateUpdates(stateUpdates) + .setLast(isLast) + .build(); + } + + private static AdaptMessageResponse createAdaptMessageResponse( + String text, Map stateUpdates) { + return createAdaptMessageResponse(text, stateUpdates, false); + } } diff --git a/pom.xml b/pom.xml index 73858464..cd50545b 100644 --- a/pom.xml +++ b/pom.xml @@ -5,7 +5,7 @@ com.google.cloud sdk-platform-java-config - 3.47.0 + 3.54.2 4.0.0 @@ -79,7 +79,7 @@ com.google.cloud google-cloud-spanneradapter - 0.0.1 + 0.16.0 com.google.cloud @@ -110,12 +110,12 @@ io.opentelemetry opentelemetry-api - 1.47.0 + 1.52.0 io.opentelemetry opentelemetry-sdk - 1.47.0 + 1.52.0 com.google.cloud.opentelemetry