Polishing.

Introduce factory beans for easier wrapping of CqlSession.

See #1321
Original pull request: #1322

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

@@ -0,0 +1,98 @@
+/*
+ * Copyright 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 org.springframework.beans.factory.config.AbstractFactoryBean;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+import org.springframework.util.ObjectUtils;
+
+import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.CqlSessionBuilder;
+
+import io.micrometer.observation.ObservationRegistry;
+
+/**
+ * Factory bean to construct a {@link CqlSession} integrated with given {@link ObservationRegistry}. This factory bean
+ * registers also {@link ObservationRequestTracker#INSTANCE ObservationRequestTracker.INSTANCE} with the builder to
+ * ensure full integration with the required infrastructure.
+ *
+ * @author Mark Paluch
+ * @since 4.0
+ * @see ObservationRequestTracker
+ * @see ObservableCqlSessionFactory
+ */
+public class ObservableCqlSessionFactoryBean extends AbstractFactoryBean<CqlSession> {
+
+	private final CqlSessionBuilder cqlSessionBuilder;
+
+	private final ObservationRegistry observationRegistry;
+
+	private @Nullable String remoteServiceName;
+
+	/**
+	 * Construct a new {@link ObservableCqlSessionFactoryBean}.
+	 *
+	 * @param cqlSessionBuilder must not be {@literal null}.
+	 * @param observationRegistry must not be {@literal null}.
+	 */
+	public ObservableCqlSessionFactoryBean(CqlSessionBuilder cqlSessionBuilder, ObservationRegistry observationRegistry) {
+
+		Assert.notNull(cqlSessionBuilder, "CqlSessionBuilder must not be null");
+		Assert.notNull(observationRegistry, "ObservationRegistry must not be null");
+
+		this.cqlSessionBuilder = cqlSessionBuilder;
+		this.observationRegistry = observationRegistry;
+	}
+
+	@Override
+	protected CqlSession createInstance() {
+
+		cqlSessionBuilder.addRequestTracker(ObservationRequestTracker.INSTANCE);
+
+		if (ObjectUtils.isEmpty(getRemoteServiceName())) {
+			return ObservableCqlSessionFactory.wrap(cqlSessionBuilder.build(), observationRegistry);
+		}
+
+		return ObservableCqlSessionFactory.wrap(cqlSessionBuilder.build(), getRemoteServiceName(), observationRegistry);
+	}
+
+	@Override
+	protected void destroyInstance(@Nullable CqlSession instance) {
+		if (instance != null) {
+			instance.close();
+		}
+	}
+
+	@Override
+	public Class<?> getObjectType() {
+		return CqlSession.class;
+	}
+
+	@Nullable
+	public String getRemoteServiceName() {
+		return remoteServiceName;
+	}
+
+	/**
+	 * Set the remote service name.
+	 *
+	 * @param remoteServiceName
+	 */
+	public void setRemoteServiceName(@Nullable String remoteServiceName) {
+		this.remoteServiceName = remoteServiceName;
+	}
+}

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

@@ -0,0 +1,92 @@
+/*
+ * Copyright 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 org.springframework.beans.factory.config.AbstractFactoryBean;
+import org.springframework.data.cassandra.ReactiveSession;
+import org.springframework.data.cassandra.core.cql.session.DefaultBridgedReactiveSession;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+import org.springframework.util.ObjectUtils;
+
+import com.datastax.oss.driver.api.core.CqlSession;
+
+import io.micrometer.observation.ObservationRegistry;
+
+/**
+ * Factory bean to construct a {@link ReactiveSession} integrated with given {@link ObservationRegistry}. The required
+ * {@link CqlSession} must be associated with {@link ObservationRequestTracker#INSTANCE
+ * ObservationRequestTracker.INSTANCE} to ensure proper integration with all observability components. You can use
+ * {@link ObservableCqlSessionFactoryBean} to obtain a properly configured {@link CqlSession}.
+ *
+ * @author Mark Paluch
+ * @since 4.0
+ * @see ObservationRequestTracker
+ * @see ObservableReactiveSessionFactory
+ */
+public class ObservableReactiveSessionFactoryBean extends AbstractFactoryBean<ReactiveSession> {
+
+	private final CqlSession cqlSession;
+
+	private final ObservationRegistry observationRegistry;
+
+	private @Nullable String remoteServiceName;
+
+	/**
+	 * Construct a new {@link ObservableReactiveSessionFactoryBean}.
+	 *
+	 * @param cqlSession must not be {@literal null}.
+	 * @param observationRegistry must not be {@literal null}.
+	 */
+	public ObservableReactiveSessionFactoryBean(CqlSession cqlSession, ObservationRegistry observationRegistry) {
+
+		Assert.notNull(cqlSession, "CqlSession must not be null");
+		Assert.notNull(observationRegistry, "ObservationRegistry must not be null");
+
+		this.cqlSession = cqlSession;
+		this.observationRegistry = observationRegistry;
+	}
+
+	@Override
+	protected ReactiveSession createInstance() {
+
+		if (ObjectUtils.isEmpty(getRemoteServiceName())) {
+			return ObservableReactiveSessionFactory.wrap(new DefaultBridgedReactiveSession(cqlSession), observationRegistry);
+		}
+
+		return ObservableReactiveSessionFactory.wrap(new DefaultBridgedReactiveSession(cqlSession), getRemoteServiceName(),
+				observationRegistry);
+	}
+
+	@Override
+	public Class<?> getObjectType() {
+		return ReactiveSession.class;
+	}
+
+	@Nullable
+	public String getRemoteServiceName() {
+		return remoteServiceName;
+	}
+
+	/**
+	 * Set the remote service name.
+	 *
+	 * @param remoteServiceName
+	 */
+	public void setRemoteServiceName(@Nullable String remoteServiceName) {
+		this.remoteServiceName = remoteServiceName;
+	}
+}

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

@@ -45,7 +45,7 @@ import reactor.util.context.Context;
 @ContextConfiguration(classes = TestConfig.class)
 public class ReactiveIntegrationTests extends SampleTestRunner {
 
-	@Autowired ReactiveSession session;
+	@Autowired ReactiveSession observableSession;
 
 	ReactiveIntegrationTests() {
 		super(SampleRunnerConfig.builder().build());
@@ -68,7 +68,6 @@ public class ReactiveIntegrationTests extends SampleTestRunner {
 
 			Observation intermediate = Observation.start("intermediate", createObservationRegistry());
 
-			ReactiveSession observableSession = ObservableReactiveSessionFactory.wrap(session, createObservationRegistry());
 
 			Mono<ReactiveResultSet> drop = observableSession.execute("DROP KEYSPACE IF EXISTS ObservationTest");
 			Mono<ReactiveResultSet> create = observableSession.execute("CREATE KEYSPACE ObservationTest " + "WITH "

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

@@ -18,9 +18,7 @@ package org.springframework.data.cassandra.observability;
 import org.springframework.context.annotation.Bean;
 import org.springframework.core.io.ClassPathResource;
 import org.springframework.core.io.Resource;
-import org.springframework.data.cassandra.ReactiveSession;
 import org.springframework.data.cassandra.config.SessionBuilderConfigurer;
-import org.springframework.data.cassandra.core.cql.session.DefaultBridgedReactiveSession;
 import org.springframework.data.cassandra.support.AbstractTestJavaConfig;
 import org.springframework.lang.Nullable;
 
@@ -71,7 +69,7 @@ class TestConfig extends AbstractTestJavaConfig {
 	}
 
 	@Bean
-	ReactiveSession reactiveSession(CqlSession session) {
-		return new DefaultBridgedReactiveSession(session);
+	ObservableReactiveSessionFactoryBean reactiveSession(CqlSession session, ObservationRegistry observationRegistry) {
+		return new ObservableReactiveSessionFactoryBean(session, observationRegistry);
 	}
 }

src/main/asciidoc/reference/observability.adoc

@@ -14,41 +14,22 @@ To enable the instrumentation, apply the following configuration to your applica
 class ObservabilityConfiguration {
 
   @Bean
-  public ObservationBeanPostProcessor observationBeanPostProcessor(ObservationRegistry observationRegistry) {
-    return new ObservationBeanPostProcessor(observationRegistry);                                           <1>
+  public ObservableCqlSessionFactoryBean observableCqlSession(CqlSessionBuilder builder,
+                                                              ObservationRegistry registry) {
+    return new ObservableCqlSessionFactoryBean(builder, registry);                        <1>
   }
 
   @Bean
-  public SessionBuilderConfigurer getSessionBuilderConfigurer() {
-    return sessionBuilder -> sessionBuilder.addRequestTracker(ObservationRequestTracker.INSTANCE);          <2>
-  }
-
-  class ObservationBeanPostProcessor implements BeanPostProcessor {
-
-    public final ObservationRegistry observationRegistry;
-
-    public ObservationBeanPostProcessor(ObservationRegistry observationRegistry) {
-      this.observationRegistry = observationRegistry;
-    }
-
-    @Override
-    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
-
-      if (bean instanceof CqlSession) {
-        return ObservableCqlSessionFactory.wrap((CqlSession) bean, observationRegistry);
-      }
-
-      if (bean instanceof ReactiveSession) {
-        return ObservableReactiveSessionFactory.wrap((ReactiveSession) bean, observationRegistry);
-      }
-
-      return BeanPostProcessor.super.postProcessAfterInitialization(bean, beanName);
-    }
+  public ObservableReactiveSessionFactoryBean observableReactiveSession(CqlSession session,
+                                                              ObservationRegistry registry) {
+    return new ObservableReactiveSessionFactoryBean(session, registry);                   <2>
   }
 }
 ----
-<1> Wraps all CQL session objects (imperative/reactive Session API) to observe Cassandra statement execution.
-<2> Integrate with the Cassandra driver to obtain success/error callbacks.
+
+<1> Wraps the CQL session object to observe Cassandra statement execution.
+Also, registers `ObservationRequestTracker.INSTANCE` with the `CqlSessionBuilder`.
+<1> Wraps a CQL session object to observe reactive Cassandra statement execution.
 ====
 
 include::../../../../target/_conventions.adoc[]