Provide smoother usage of Observabilty.

pom.xml

@@ -11,7 +11,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-cassandra-parent</artifactId>
-	<version>4.0.0-SNAPSHOT</version>
+	<version>4.0.0-gh-1321-SNAPSHOT</version>
 	<packaging>pom</packaging>
 
 	<name>Spring Data for Apache Cassandra</name>

spring-data-cassandra-distribution/pom.xml

@@ -8,7 +8,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-cassandra-parent</artifactId>
-		<version>4.0.0-SNAPSHOT</version>
+		<version>4.0.0-gh-1321-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-cassandra/pom.xml

@@ -8,7 +8,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-cassandra-parent</artifactId>
-		<version>4.0.0-SNAPSHOT</version>
+		<version>4.0.0-gh-1321-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

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

@@ -19,7 +19,7 @@
 import io.micrometer.observation.docs.ObservationDocumentation;
 
 /**
- * Cassandra-based implementation of {@link DocumentedObservation}.
+ * Cassandra-based implementation of {@link ObservationDocumentation}.
  *
  * @author Mark Paluch
  * @author Marcin Grzejszczak

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

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2013-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.cassandra.observability;
+
+import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.tracing.Tracer;
+
+import java.util.Optional;
+
+import org.springframework.context.annotation.Bean;
+
+/**
+ * Set of beans enable observability with Spring Data Cassandra.
+ *
+ * @author Greg Turnquist
+ * @since 4.0.0
+ */
+public class CassandraObservationConfiguration {
+
+	@Bean
+	CqlSessionObservationConvention observationConvention() {
+		return new DefaultCassandraObservationConvention();
+	}
+
+	@Bean
+	CqlSessionTracingObservationHandler cqlSessionTracingObservationHandler(
+			Optional<ObservationRegistry> observationRegistry, Tracer tracer) {
+
+		CqlSessionTracingObservationHandler observationHandler = new CqlSessionTracingObservationHandler(tracer);
+		observationRegistry.ifPresent(registry -> registry.observationConfig().observationHandler(observationHandler));
+		return observationHandler;
+	}
+
+	@Bean
+	CqlSessionTracingBeanPostProcessor traceCqlSessionBeanPostProcessor(ObservationRegistry observationRegistry,
+			CqlSessionObservationConvention observationConvention) {
+		return new CqlSessionTracingBeanPostProcessor(observationRegistry, observationConvention);
+	}
+}

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

@@ -102,7 +102,7 @@ private Object tracedCall(Statement<?> statement, String methodName,
 			Function<Statement<?>, Object> statementExecutor) {
 
 		if (this.observationRegistry.getCurrentObservation() == null) {
-			return null;
+			return statementExecutor.apply(statement);
 		}
 
 		Observation observation = childObservation(statement, methodName, this.delegateSession);
@@ -111,14 +111,7 @@ private Object tracedCall(Statement<?> statement, String methodName,
 			log.debug("Created a new child observation before query [" + observation + "]");
 		}
 
-		try (Observation.Scope scope = observation.openScope()) {
-			return statementExecutor.apply(statement);
-		} catch (Exception e) {
-			observation.error(e);
-			throw e;
-		} finally {
-			observation.stop();
-		}
+		return observation.observe(() -> statementExecutor.apply(statement));
 	}
 
 	/**
@@ -154,7 +147,7 @@ private Observation childObservation(Statement<?> statement, String methodName,
 				.observation(this.observationRegistry, () -> observationContext) //
 				.contextualName(CassandraObservation.CASSANDRA_QUERY_OBSERVATION.getContextualName()) //
 				.highCardinalityKeyValues(this.observationConvention.getHighCardinalityKeyValues(observationContext)) //
-				.lowCardinalityKeyValues(this.observationConvention.getLowCardinalityKeyValues(observationContext)) //
-				.start();
+				.lowCardinalityKeyValues(this.observationConvention.getLowCardinalityKeyValues(observationContext));
+		// .start();
 	}
 }

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/observability/DefaultCassandraObservationContention.java → spring-data-cassandra/src/main/java/org/springframework/data/cassandra/observability/DefaultCassandraObservationConvention.java

@@ -32,16 +32,16 @@
  * @author Greg Turnquist
  * @since 4.0.0
  */
-public class DefaultCassandraObservationContention implements CqlSessionObservationConvention {
+public class DefaultCassandraObservationConvention implements CqlSessionObservationConvention {
 
 	@Override
 	public KeyValues getLowCardinalityKeyValues(CqlSessionContext context) {
 
 		KeyValues keyValues = KeyValues.of(
 				LowCardinalityKeyNames.SESSION_NAME
 						.withValue(Optional.ofNullable(context.getDelegateSession().getName()).orElse("unknown")),
-				LowCardinalityKeyNames.KEYSPACE_NAME.withValue(
-						Optional.ofNullable(context.getStatement().getKeyspace()).map(CqlIdentifier::asInternal).orElse("unknown")),
+				LowCardinalityKeyNames.KEYSPACE_NAME
+						.withValue(context.getDelegateSession().getKeyspace().map(CqlIdentifier::asInternal).orElse("unknown")),
 				LowCardinalityKeyNames.METHOD_NAME.withValue(context.getMethodName()));
 
 		if (context.getStatement().getNode() != null) {

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

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2013-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.cassandra.observability;
+
+import java.lang.annotation.*;
+
+import org.springframework.context.annotation.Import;
+
+/**
+ * Annotation to enable Cassandra observability.
+ *
+ * @author Greg Turnquist
+ * @since 4.0.0
+ */
+@Inherited
+@Documented
+@Target(ElementType.TYPE)
+@Retention(RetentionPolicy.RUNTIME)
+@Import(CassandraObservationConfiguration.class)
+public @interface EnableCassandraObservability {
+}

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

@@ -18,7 +18,13 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
 
+import io.micrometer.observation.ObservationHandler;
 import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.tracing.Tracer;
+import io.micrometer.tracing.test.simple.SimpleTracer;
+
+import java.lang.reflect.Method;
+import java.util.Collection;
 
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -29,6 +35,7 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit.jupiter.SpringExtension;
+import org.springframework.util.ReflectionUtils;
 
 import com.datastax.oss.driver.api.core.CqlSession;
 
@@ -44,6 +51,8 @@
 public class CqlSessionTracingBeanPostProcessorTests {
 
 	@Autowired CqlSession session;
+	@Autowired ObservationRegistry registry;
+	@Autowired CqlSessionTracingObservationHandler handler;
 
 	@Test
 	void injectedCqlSessionShouldBeWrapped() throws Exception {
@@ -58,7 +67,22 @@ void injectedCqlSessionShouldBeWrapped() throws Exception {
 		assertThat(advised.getTargetSource().getTarget()).isEqualTo(TestConfig.originalSession);
 	}
 
+	@Test
+	void injectedObservationHandlerIsRegisteredWithRegistry() {
+
+		ObservationRegistry.ObservationConfig config = registry.observationConfig();
+
+		Method getObservationHandlers = ReflectionUtils.findMethod(ObservationRegistry.ObservationConfig.class,
+				"getObservationHandlers");
+		ReflectionUtils.makeAccessible(getObservationHandlers);
+		Collection<ObservationHandler<?>> handlers = (Collection<ObservationHandler<?>>) ReflectionUtils
+				.invokeMethod(getObservationHandlers, config);
+
+		assertThat(handlers).contains(handler);
+	}
+
 	@Configuration
+	@EnableCassandraObservability
 	static class TestConfig {
 
 		static CqlSession originalSession = mock(CqlSession.class);
@@ -74,14 +98,8 @@ ObservationRegistry meterRegistry() {
 		}
 
 		@Bean
-		CqlSessionObservationConvention observationConvention() {
-			return new DefaultCassandraObservationContention();
-		}
-
-		@Bean
-		CqlSessionTracingBeanPostProcessor traceCqlSessionBeanPostProcessor(ObservationRegistry observationRegistry,
-				CqlSessionObservationConvention tagsProvider) {
-			return new CqlSessionTracingBeanPostProcessor(observationRegistry, tagsProvider);
+		Tracer tracer() {
+			return new SimpleTracer();
 		}
 	}
 }

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

@@ -92,7 +92,7 @@ void tracingNoStatementsShouldProduceNoMetrics() {
 		ObservationRegistry observationRegistry = ObservationRegistry.create();
 		observationRegistry.observationConfig().observationHandler(new DefaultMeterObservationHandler(meterRegistry));
 
-		CqlSessionObservationConvention observationContention = new DefaultCassandraObservationContention();
+		CqlSessionObservationConvention observationContention = new DefaultCassandraObservationConvention();
 
 		SimpleTracer tracer = new SimpleTracer();
 		observationRegistry.observationConfig().observationHandler(new CqlSessionTracingObservationHandler(tracer));
@@ -109,7 +109,7 @@ void shouldCreateObservationForCqlSessionOperations() {
 		ObservationRegistry observationRegistry = ObservationRegistry.create();
 		observationRegistry.observationConfig().observationHandler(new DefaultMeterObservationHandler(meterRegistry));
 
-		CqlSessionObservationConvention observationContention = new DefaultCassandraObservationContention();
+		CqlSessionObservationConvention observationContention = new DefaultCassandraObservationConvention();
 		SimpleTracer tracer = new SimpleTracer();
 		observationRegistry.observationConfig().observationHandler(new CqlSessionTracingObservationHandler(tracer));
 
@@ -125,7 +125,7 @@ void shouldCreateObservationForCqlSessionOperations() {
 
 		MeterRegistryAssert.then(meterRegistry).hasTimerWithNameAndTags(CASSANDRA_QUERY_OBSERVATION.getName(), KeyValues.of( //
 				LowCardinalityKeyNames.SESSION_NAME.withValue("s5"), //
-				LowCardinalityKeyNames.KEYSPACE_NAME.withValue("unknown"), //
+				LowCardinalityKeyNames.KEYSPACE_NAME.withValue("system"), //
 				KeyValue.of("error", "none") //
 		));
 
@@ -140,7 +140,7 @@ void shouldCreateObservationForCqlSessionOperations() {
 				.hasRemoteServiceNameEqualTo("cassandra-s5") //
 				.hasNameEqualTo(CASSANDRA_QUERY_OBSERVATION.getContextualName()) //
 				.hasTag(LowCardinalityKeyNames.SESSION_NAME, "s5") //
-				.hasTag(LowCardinalityKeyNames.KEYSPACE_NAME, "unknown") //
+				.hasTag(LowCardinalityKeyNames.KEYSPACE_NAME, "system") //
 				.hasTag(HighCardinalityKeyNames.CQL_TAG, CREATE_KEYSPACE) //
 				.hasIpThatIsBlank() //
 				.hasPortEqualTo(0) //

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

@@ -27,7 +27,6 @@
 import java.util.Deque;
 import java.util.function.BiConsumer;
 
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Bean;
@@ -46,7 +45,6 @@
  * @author Greg Turnquist
  * @since 4.0.0
  */
-@Disabled("Run this manually to visually test spans in Zipkin")
 @ExtendWith({ SpringExtension.class, CassandraExtension.class })
 @TestKeyspaceName
 public class ZipkinIntegrationTests extends SampleTestRunner {
@@ -91,6 +89,9 @@ public SampleTestRunnerConsumer yourCode() {
 
 		return (tracer, meterRegistry) -> {
 
+			OBSERVATION_REGISTRY.observationConfig()
+					.observationHandler(new CqlSessionTracingObservationHandler(tracer.getTracer()));
+
 			session.execute("DROP KEYSPACE IF EXISTS ConfigTest");
 			session.execute("CREATE KEYSPACE ConfigTest " + "WITH "
 					+ "REPLICATION = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 };");
@@ -117,13 +118,13 @@ ObservationRegistry registry() {
 
 		@Bean
 		CqlSessionObservationConvention observationContention() {
-			return new DefaultCassandraObservationContention();
+			return new DefaultCassandraObservationConvention();
 		}
 
 		@Bean
 		CqlSessionTracingBeanPostProcessor traceCqlSessionBeanPostProcessor(ObservationRegistry observationRegistry,
-				CqlSessionObservationConvention tagsProvider) {
-			return new CqlSessionTracingBeanPostProcessor(observationRegistry, tagsProvider);
+				CqlSessionObservationConvention observationConvention) {
+			return new CqlSessionTracingBeanPostProcessor(observationRegistry, observationConvention);
 		}
 	}
 }