From 93741cc54238c1cb1fd4d1ac5fdd470b03a9f54d Mon Sep 17 00:00:00 2001 From: Tanmay Vartak Date: Tue, 2 Jun 2026 00:22:37 +0000 Subject: [PATCH 1/2] feat: add support for USE keyspace attachments --- .../adapter/DriverConnectionHandler.java | 18 ++++++- .../adapter/DriverConnectionHandlerTest.java | 47 +++++++++++++++++++ 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java index c812fb85..e1991ea4 100644 --- a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java +++ b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java @@ -71,6 +71,8 @@ final class DriverConnectionHandler implements Runnable { private static final char WRITE_ACTION_QUERY_ID_PREFIX = 'W'; private static final String ROUTE_TO_LEADER_HEADER_KEY = "x-goog-spanner-route-to-leader"; private static final String MAX_COMMIT_DELAY_ATTACHMENT_KEY = "max_commit_delay"; + private static final String KEYSPACE_ATTACHMENT_KEY = "keyspace"; + private static final ByteBufAllocator byteBufAllocator = ByteBufAllocator.DEFAULT; private static final FrameCodec serverFrameCodec = customServerCodec(byteBufAllocator); private static final FrameCodec clientFrameCodec = @@ -353,8 +355,11 @@ private PreparePayloadResult preparePayload(MessageContext ctx) { return prepareBatchMessage((Batch) decodeFrame(ctx.payload).message, ctx.streamId); case ProtocolConstants.Opcode.QUERY: return prepareQueryMessage((Query) decodeFrame(ctx.payload).message); + case ProtocolConstants.Opcode.PREPARE: + return preparePrepareMessage(); default: return new PreparePayloadResult(DEFAULT_CONTEXT); + } } @@ -396,19 +401,28 @@ private PreparePayloadResult prepareBatchMessage(Batch message, int streamId) { private PreparePayloadResult prepareQueryMessage(Query message) { ApiCallContext context; - Map attachments = Collections.emptyMap(); + Map attachments = new HashMap<>(); if (startsWith(message.query, "SELECT")) { context = DEFAULT_CONTEXT; } else { context = DEFAULT_CONTEXT_WITH_LAR; if (maxCommitDelayMillis.isPresent()) { - attachments = new HashMap<>(); attachments.put(MAX_COMMIT_DELAY_ATTACHMENT_KEY, maxCommitDelayMillis.get()); } } + Optional keyspace = adapterClientWrapper.getAttachmentsCache().get(KEYSPACE_ATTACHMENT_KEY); + keyspace.ifPresent(v -> attachments.put(KEYSPACE_ATTACHMENT_KEY, v)); return new PreparePayloadResult(context, attachments); } + private PreparePayloadResult preparePrepareMessage() { + Map attachments = new HashMap<>(); + Optional keyspace = adapterClientWrapper.getAttachmentsCache().get(KEYSPACE_ATTACHMENT_KEY); + keyspace.ifPresent(v -> attachments.put(KEYSPACE_ATTACHMENT_KEY, v)); + return new PreparePayloadResult(DEFAULT_CONTEXT, attachments); + } + + private Optional prepareAttachmentForQueryId( int streamId, Map attachments, byte[] queryId) { String key = constructKey(queryId); diff --git a/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/DriverConnectionHandlerTest.java b/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/DriverConnectionHandlerTest.java index 2a0d0d37..ddc9ea67 100644 --- a/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/DriverConnectionHandlerTest.java +++ b/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/DriverConnectionHandlerTest.java @@ -82,6 +82,7 @@ public void setUp() throws IOException { mockSocket = mock(Socket.class); outputStream = new ByteArrayOutputStream(); when(mockSocket.getOutputStream()).thenReturn(outputStream); + when(mockAdapterClient.getAttachmentsCache()).thenReturn(new AttachmentsCache(10)); } @Test @@ -319,6 +320,52 @@ public void shortBody_writesErrorMessageToSocket() throws IOException { verify(mockSocket).close(); } + @Test + public void successfulQueryMessageWithKeyspace() throws IOException { + byte[] validPayload = createQueryMessage(); + ByteString grpcResponse = ByteString.copyFromUtf8("gRPC response"); + when(mockSocket.getInputStream()).thenReturn(new ByteArrayInputStream(validPayload)); + when(mockAdapterClient.sendGrpcRequest(any(byte[].class), any(), any(), any(int.class))) + .thenReturn(grpcResponse); + + AttachmentsCache attachmentsCache = new AttachmentsCache(1); + attachmentsCache.put("keyspace", "test_keyspace"); + when(mockAdapterClient.getAttachmentsCache()).thenReturn(attachmentsCache); + + DriverConnectionHandler handler = + new DriverConnectionHandler(mockSocket, mockAdapterClient, mockMetricsRecorder); + handler.run(); + + assertThat(outputStream.toString(StandardCharsets.UTF_8.name())).isEqualTo("gRPC response"); + verify(mockSocket).close(); + verify(mockAdapterClient) + .sendGrpcRequest(any(), attachmentsCaptor.capture(), contextCaptor.capture(), any(int.class)); + assertThat(attachmentsCaptor.getValue()).containsExactly("keyspace", "test_keyspace"); + } + + @Test + public void successfulPrepareMessageWithKeyspace() throws IOException { + byte[] validPayload = createPrepareMessage(); + ByteString grpcResponse = ByteString.copyFromUtf8("gRPC response"); + when(mockSocket.getInputStream()).thenReturn(new ByteArrayInputStream(validPayload)); + when(mockAdapterClient.sendGrpcRequest(any(byte[].class), any(), any(), any(int.class))) + .thenReturn(grpcResponse); + + AttachmentsCache attachmentsCache = new AttachmentsCache(1); + attachmentsCache.put("keyspace", "test_keyspace"); + when(mockAdapterClient.getAttachmentsCache()).thenReturn(attachmentsCache); + + DriverConnectionHandler handler = + new DriverConnectionHandler(mockSocket, mockAdapterClient, mockMetricsRecorder); + handler.run(); + + assertThat(outputStream.toString(StandardCharsets.UTF_8.name())).isEqualTo("gRPC response"); + verify(mockSocket).close(); + verify(mockAdapterClient) + .sendGrpcRequest(any(), attachmentsCaptor.capture(), contextCaptor.capture(), any(int.class)); + assertThat(attachmentsCaptor.getValue()).containsExactly("keyspace", "test_keyspace"); + } + private static byte[] createQueryMessage() { return encodeMessage(new Query("SELECT * FROM ks.T")); } From e6a31213c8d53fea2a4284fc1cf72361aae5e384 Mon Sep 17 00:00:00 2001 From: Tanmay Vartak Date: Tue, 2 Jun 2026 00:31:43 +0000 Subject: [PATCH 2/2] style: format modified lines in DriverConnectionHandler and Test --- .../cloud/spanner/adapter/DriverConnectionHandler.java | 8 ++++---- .../spanner/adapter/DriverConnectionHandlerTest.java | 10 ++++++---- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java index e1991ea4..9bc1e50c 100644 --- a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java +++ b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java @@ -359,7 +359,6 @@ private PreparePayloadResult preparePayload(MessageContext ctx) { return preparePrepareMessage(); default: return new PreparePayloadResult(DEFAULT_CONTEXT); - } } @@ -410,19 +409,20 @@ private PreparePayloadResult prepareQueryMessage(Query message) { attachments.put(MAX_COMMIT_DELAY_ATTACHMENT_KEY, maxCommitDelayMillis.get()); } } - Optional keyspace = adapterClientWrapper.getAttachmentsCache().get(KEYSPACE_ATTACHMENT_KEY); + Optional keyspace = + adapterClientWrapper.getAttachmentsCache().get(KEYSPACE_ATTACHMENT_KEY); keyspace.ifPresent(v -> attachments.put(KEYSPACE_ATTACHMENT_KEY, v)); return new PreparePayloadResult(context, attachments); } private PreparePayloadResult preparePrepareMessage() { Map attachments = new HashMap<>(); - Optional keyspace = adapterClientWrapper.getAttachmentsCache().get(KEYSPACE_ATTACHMENT_KEY); + Optional keyspace = + adapterClientWrapper.getAttachmentsCache().get(KEYSPACE_ATTACHMENT_KEY); keyspace.ifPresent(v -> attachments.put(KEYSPACE_ATTACHMENT_KEY, v)); return new PreparePayloadResult(DEFAULT_CONTEXT, attachments); } - private Optional prepareAttachmentForQueryId( int streamId, Map attachments, byte[] queryId) { String key = constructKey(queryId); diff --git a/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/DriverConnectionHandlerTest.java b/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/DriverConnectionHandlerTest.java index ddc9ea67..12d5cb2c 100644 --- a/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/DriverConnectionHandlerTest.java +++ b/google-cloud-spanner-cassandra/src/test/java/com/google/cloud/spanner/adapter/DriverConnectionHandlerTest.java @@ -327,7 +327,7 @@ public void successfulQueryMessageWithKeyspace() throws IOException { when(mockSocket.getInputStream()).thenReturn(new ByteArrayInputStream(validPayload)); when(mockAdapterClient.sendGrpcRequest(any(byte[].class), any(), any(), any(int.class))) .thenReturn(grpcResponse); - + AttachmentsCache attachmentsCache = new AttachmentsCache(1); attachmentsCache.put("keyspace", "test_keyspace"); when(mockAdapterClient.getAttachmentsCache()).thenReturn(attachmentsCache); @@ -339,7 +339,8 @@ public void successfulQueryMessageWithKeyspace() throws IOException { assertThat(outputStream.toString(StandardCharsets.UTF_8.name())).isEqualTo("gRPC response"); verify(mockSocket).close(); verify(mockAdapterClient) - .sendGrpcRequest(any(), attachmentsCaptor.capture(), contextCaptor.capture(), any(int.class)); + .sendGrpcRequest( + any(), attachmentsCaptor.capture(), contextCaptor.capture(), any(int.class)); assertThat(attachmentsCaptor.getValue()).containsExactly("keyspace", "test_keyspace"); } @@ -350,7 +351,7 @@ public void successfulPrepareMessageWithKeyspace() throws IOException { when(mockSocket.getInputStream()).thenReturn(new ByteArrayInputStream(validPayload)); when(mockAdapterClient.sendGrpcRequest(any(byte[].class), any(), any(), any(int.class))) .thenReturn(grpcResponse); - + AttachmentsCache attachmentsCache = new AttachmentsCache(1); attachmentsCache.put("keyspace", "test_keyspace"); when(mockAdapterClient.getAttachmentsCache()).thenReturn(attachmentsCache); @@ -362,7 +363,8 @@ public void successfulPrepareMessageWithKeyspace() throws IOException { assertThat(outputStream.toString(StandardCharsets.UTF_8.name())).isEqualTo("gRPC response"); verify(mockSocket).close(); verify(mockAdapterClient) - .sendGrpcRequest(any(), attachmentsCaptor.capture(), contextCaptor.capture(), any(int.class)); + .sendGrpcRequest( + any(), attachmentsCaptor.capture(), contextCaptor.capture(), any(int.class)); assertThat(attachmentsCaptor.getValue()).containsExactly("keyspace", "test_keyspace"); }