Return proxied Statement from ObservationStatement mutator calls.

Add test to verify proxy preserves Observation

Signed-off-by: Ben Efrati <BenEfrati@users.noreply.github.com>

Closes #1651
Original pull request: #1652

Author: BenEfrati

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/observability/ObservationStatement.java

@@ -34,6 +34,7 @@
  *
  * @author Mark Paluch
  * @author Marcin Grzejszczak
+ * @author Ben Efrati
  * @since 4.0
  */
 final class ObservationStatement implements MethodInterceptor {
@@ -103,9 +104,10 @@ public static boolean isObservationStatement(Statement<?> statement) {
 			}
 		}
 
-		Object result = invocation.proceed();
-		if (result instanceof Statement<?>) {
-			this.delegate = (Statement<?>) result;
+		Object result = invocation.proceed(); 
+		if (result instanceof Statement<?> statement) { 
+			this.delegate = statement;
+			return createProxy(this.observation, statement);
 		}
 
 		return result;

spring-data-cassandra/src/test/java/org/springframework/data/cassandra/observability/ObservationStatementUnitTests.java

@@ -15,11 +15,17 @@
  */
 package org.springframework.data.cassandra.observability;
 
-import static org.assertj.core.api.Assertions.*;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
 
 import io.micrometer.observation.Observation;
 import io.micrometer.observation.tck.TestObservationRegistry;
 
+import org.assertj.core.api.Assertions;
+import org.assertj.core.api.InstanceOfAssertFactory;
 import org.junit.jupiter.api.Test;
 
 import com.datastax.oss.driver.api.core.cql.SimpleStatement;
@@ -28,6 +34,7 @@
  * Unit test for {@link ObservationStatement}
  *
  * @author Mark Paluch
+ * @author Ben Efrati
  */
 class ObservationStatementUnitTests {
 
@@ -50,4 +57,25 @@ void equalsAndHashCodeShouldBeEqual() {
 		assertThat(statement1).isNotEqualTo(statement3);
 		assertThat(statement1.hashCode()).isNotEqualTo(statement3.hashCode());
 	}
+
+	@Test // GH-1651
+	void mutatingCallsReturnProxyPreservingObservation() {
+
+		TestObservationRegistry registry = TestObservationRegistry.create();
+
+		Observation observation = Observation.start("foo", registry);
+
+		SimpleStatement statement = ObservationStatement.createProxy(observation,
+				SimpleStatement.newInstance("SELECT * FROM foo"));
+		Map<String, ByteBuffer> customPayload = Map.of("key",
+				ByteBuffer.wrap("value".getBytes(StandardCharsets.UTF_8)));
+		SimpleStatement customPayloadStatement = statement.setCustomPayload(customPayload);
+		assertThat(customPayloadStatement)
+				.isNotEqualTo(statement)
+				.doesNotHaveSameHashCodeAs(statement)
+				.asInstanceOf(new InstanceOfAssertFactory<>(CassandraObservationSupplier.class, Assertions::assertThat))
+				.extracting(CassandraObservationSupplier::getObservation)
+				.isEqualTo(observation);
+
+	}
 }