Provide better support for activating observability with Spring Data Cassandra.

Resolves #1321.
Original pull request: #1322

Author: gregturn

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 renamed to 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);
 		}
 	}
 }

src/main/asciidoc/reference/observability.adoc

@@ -8,3 +8,77 @@ include::{root-target}_conventions.adoc[]
 include::{root-target}_metrics.adoc[]
 
 include::{root-target}_spans.adoc[]
+
+[[observability.registration]]
+== Observability Registration
+
+Spring Data Cassandra is not yet supported in Spring Boot to automatically enable observability.
+Don't worry, we've got you covered!
+Simply add the `@EnableCassandraObservability` annotation to your Spring Boot application, and you'll be all set.
+
+.Activating observability for your Spring Boot application
+====
+[source,java]
+----
+@SpringBootApplication
+@EnableCassandraObservability // <1>
+public class SpringDataCassandraObservabilityApplication {
+
+	public static void main(String[] args) {
+		SpringApplication.run(SpringDataCassandraObservabilityApplication.class, args);
+	}
+}
+----
+<1> This annotation will activate the bits needed start wrapping CQL calls and register them with your tracer of choice.
+====
+
+By the way, Spring Boot DOES have autoconfigured hooks into various parts of the system.
+For example, there is an observation filter that Spring Boot will apply to Spring MVC ensuring all your calls are wrapped properly.
+And if anywhere in the midst of that web call, you invoke Spring Data Cassandra (either through `CqlTemplate` or a custom repository), it will get captured properly.
+
+Something that is NOT covered are situations where your code runs independently.
+For example, if you have some block that is run during startup inside a `CommandLineRunner`, there is no way for Spring Boot to know that this should observed.
+Check out the code block below:
+
+.Loading data for a sample application
+====
+[source,java]
+----
+@Bean
+CommandLineRunner initData(EmployeeRepository repository) {
+    return args -> {
+        repository.save(new Employee("1", "Frodo", "ring bearer"));
+        repository.save(new Employee("2", "Bilbo", "burglar"));
+    };
+}
+----
+This tactic is used all the time in demos. These calls to a Spring Data Cassandra repository will NOT be observed.
+====
+
+If you DO want to observe such a code block, you must wrap it yourself, as shown below:
+
+.Observing the loading of data in a sample application
+====
+[source,java]
+----
+@Bean
+CommandLineRunner initData(EmployeeRepository repository, ObservationRegistry registry) { // <1>
+    return args -> {
+        Observation.createNotStarted("init-database", registry).observe(() -> { // <2>
+            repository.save(new Employee("1", "Frodo", "ring bearer"));
+            repository.save(new Employee("2", "Bilbo", "burglar"));
+        });
+    };
+}
+----
+<1> Your `CommandLineRunner` requires access to the app context's `ObservationRegistry`
+<2> You need to create your own `Observation` using the `createNotStarted()` method. Give it any contextual name you like, but be sure to include the `registry`.
+====
+
+The `observe()` method takes a Java 8 lambda function which is invoked inside a common Micrometer pattern of:
+
+* Starting the observation.
+* Invoking your callback.
+* Properly ending the observation by either reporting an error if it fails or stopping the observation if it succeeds.
+
+This will allow you to observe chunks of code that may fall outside of currently autoconfigured operations.