Allow custom MongoHandlerObservationConvention

Closes spring-projects#4321

add author

Author: francoiskha

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoObservationCommandListener.java

@@ -39,6 +39,7 @@
  * @author OpenZipkin Brave Authors
  * @author Marcin Grzejszczak
  * @author Greg Turnquist
+ * @author François Kha
  * @since 4.0
  */
 public class MongoObservationCommandListener implements CommandListener {
@@ -48,7 +49,7 @@ public class MongoObservationCommandListener implements CommandListener {
 	private final ObservationRegistry observationRegistry;
 	private final @Nullable ConnectionString connectionString;
 
-	private final MongoHandlerObservationConvention observationConvention = new DefaultMongoHandlerObservationConvention();
+	private final MongoHandlerObservationConvention observationConvention;
 
 	/**
 	 * Create a new {@link MongoObservationCommandListener} to record {@link Observation}s.
@@ -61,6 +62,7 @@ public MongoObservationCommandListener(ObservationRegistry observationRegistry)
 
 		this.observationRegistry = observationRegistry;
 		this.connectionString = null;
+		this.observationConvention = new DefaultMongoHandlerObservationConvention();
 	}
 
 	/**
@@ -77,6 +79,26 @@ public MongoObservationCommandListener(ObservationRegistry observationRegistry,
 
 		this.observationRegistry = observationRegistry;
 		this.connectionString = connectionString;
+		this.observationConvention = new DefaultMongoHandlerObservationConvention();
+	}
+
+	/**
+	 * Create a new {@link MongoObservationCommandListener} to record {@link Observation}s. This constructor attaches the
+	 * {@link ConnectionString} to every {@link Observation} and uses the given {@link MongoHandlerObservationConvention}
+	 *
+	 * @param observationRegistry must not be {@literal null}
+	 * @param connectionString must not be {@literal null}
+	 * @param observationConvention must not be {@literal null}
+	 */
+	public MongoObservationCommandListener(ObservationRegistry observationRegistry, ConnectionString connectionString, MongoHandlerObservationConvention observationConvention) {
+
+		Assert.notNull(observationRegistry, "ObservationRegistry must not be null");
+		Assert.notNull(connectionString, "ConnectionString must not be null");
+		Assert.notNull(observationConvention, "MongoHandlerObservationConvention must not be null");
+
+		this.observationRegistry = observationRegistry;
+		this.connectionString = connectionString;
+		this.observationConvention = observationConvention;
 	}
 
 	@Override

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/observability/MongoObservationCommandListenerTests.java

@@ -18,6 +18,7 @@
 import static io.micrometer.core.tck.MeterRegistryAssert.*;
 import static org.mockito.Mockito.*;
 
+import com.mongodb.ConnectionString;
 import io.micrometer.common.KeyValues;
 import io.micrometer.core.instrument.MeterRegistry;
 import io.micrometer.core.instrument.observation.DefaultMeterObservationHandler;
@@ -48,6 +49,7 @@
  * @author Marcin Grzejszczak
  * @author Greg Turnquist
  * @author Mark Paluch
+ * @author François Kha
  */
 class MongoObservationCommandListenerTests {
 
@@ -108,8 +110,7 @@ void successfullyCompletedCommandShouldCreateTimerWhenParentSampleInRequestConte
 				new ConnectionDescription( //
 						new ServerId( //
 								new ClusterId("description"), //
-								new ServerAddress("localhost", 1234))),
-				"database", "insert", //
+								new ServerAddress("localhost", 1234))), "database", "insert", //
 				new BsonDocument("collection", new BsonString("user"))));
 		listener.commandSucceeded(new CommandSucceededEvent(traceRequestContext, 0, null, "insert", null, 0));
 
@@ -179,7 +180,8 @@ void commandWithErrorShouldCreateTimerWhenParentSampleInRequestContext() {
 		assertThatTimerRegisteredWithTags();
 	}
 
-	@Test // GH-4481
+	@Test
+		// GH-4481
 	void completionShouldIgnoreIncompatibleObservationContext() {
 
 		// given
@@ -195,7 +197,8 @@ void completionShouldIgnoreIncompatibleObservationContext() {
 		verifyNoMoreInteractions(observation);
 	}
 
-	@Test // GH-4481
+	@Test
+		// GH-4481
 	void failureShouldIgnoreIncompatibleObservationContext() {
 
 		// given
@@ -211,6 +214,31 @@ void failureShouldIgnoreIncompatibleObservationContext() {
 		verifyNoMoreInteractions(observation);
 	}
 
+	@Test
+		// GH-4321
+	void shouldUseObservationConvention() {
+		//given
+		MongoHandlerObservationConvention customObservationConvention = new MongoHandlerObservationConvention() {
+			@Override
+			public boolean supportsContext(Observation.Context context) {
+				return MongoHandlerObservationConvention.super.supportsContext(context);
+			}
+
+			@Override
+			public String getName() {
+				return "custom.name";
+			}
+		};
+		this.listener = new MongoObservationCommandListener(observationRegistry, mock(ConnectionString.class),
+				customObservationConvention);
+
+		// when
+		listener.commandStarted(new CommandStartedEvent(new MapRequestContext(), 0, null, "some name", "", null));
+
+		// then
+		assertThat(meterRegistry).hasMeterWithName("custom.name.active");
+	}
+
 	private RequestContext getContext() {
 		return ((SynchronousContextProvider) ContextProviderFactory.create(observationRegistry)).getContext();
 	}