Add a new observation for generated keys

This commit introduces a new observation for generated keys.
For Spring Boot, a new type `TraceType.KEYS` is added to determine
whether to create the new observation by `jdbc.includes` property.

Close #36

datasource-micrometer-spring-boot/src/main/java/net/ttddyy/observation/boot/autoconfigure/DataSourceObservationAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022-2023 the original author or authors.
+ * Copyright 2022-2024 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.
@@ -126,11 +126,13 @@ public static DataSourceObservationBeanPostProcessor dataSourceObservationBeanPo
 			ObjectProvider<ParameterTransformer> parameterTransformer,
 			ObjectProvider<QueryTransformer> queryTransformer,
 			ObjectProvider<ResultSetProxyLogicFactory> resultSetProxyLogicFactory,
+			ObjectProvider<ResultSetProxyLogicFactory> generatedKeysProxyLogicFactory,
 			ObjectProvider<DataSourceProxyConnectionIdManagerProvider> dataSourceProxyConnectionIdManagerProvider,
 			ObjectProvider<ProxyDataSourceBuilderCustomizer> proxyDataSourceBuilderCustomizers) {
 		return new DataSourceObservationBeanPostProcessor(jdbcProperties, dataSourceNameResolvers, listeners,
 				methodExecutionListeners, parameterTransformer, queryTransformer, resultSetProxyLogicFactory,
-				dataSourceProxyConnectionIdManagerProvider, proxyDataSourceBuilderCustomizers);
+				generatedKeysProxyLogicFactory, dataSourceProxyConnectionIdManagerProvider,
+				proxyDataSourceBuilderCustomizers);
 	}
 
 	@Configuration(proxyBeanMethods = false)

datasource-micrometer-spring-boot/src/main/java/net/ttddyy/observation/boot/autoconfigure/DataSourceObservationBeanPostProcessor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 the original author or authors.
+ * Copyright 2022-2024 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.
@@ -51,6 +51,8 @@ public class DataSourceObservationBeanPostProcessor implements BeanPostProcessor
 
 	private final ObjectProvider<ResultSetProxyLogicFactory> resultSetProxyLogicFactoryProvider;
 
+	private final ObjectProvider<ResultSetProxyLogicFactory> generatedKeysProxyLogicFactoryProvider;
+
 	private final ObjectProvider<DataSourceProxyConnectionIdManagerProvider> dataSourceProxyConnectionIdManagerProviderProvider;
 
 	private DataSourceProxyBuilderConfigurer dataSourceProxyBuilderConfigurer;
@@ -64,6 +66,7 @@ public DataSourceObservationBeanPostProcessor(ObjectProvider<JdbcProperties> jdb
 			ObjectProvider<ParameterTransformer> parameterTransformerProvider,
 			ObjectProvider<QueryTransformer> queryTransformerProvider,
 			ObjectProvider<ResultSetProxyLogicFactory> resultSetProxyLogicFactoryProvider,
+			ObjectProvider<ResultSetProxyLogicFactory> generatedKeysProxyLogicFactoryProvider,
 			ObjectProvider<DataSourceProxyConnectionIdManagerProvider> dataSourceProxyConnectionIdManagerProviderProvider,
 			ObjectProvider<ProxyDataSourceBuilderCustomizer> proxyDataSourceBuilderCustomizers) {
 		this.jdbcPropertiesProvider = jdbcPropertiesProvider;
@@ -73,6 +76,7 @@ public DataSourceObservationBeanPostProcessor(ObjectProvider<JdbcProperties> jdb
 		this.parameterTransformerProvider = parameterTransformerProvider;
 		this.queryTransformerProvider = queryTransformerProvider;
 		this.resultSetProxyLogicFactoryProvider = resultSetProxyLogicFactoryProvider;
+		this.generatedKeysProxyLogicFactoryProvider = generatedKeysProxyLogicFactoryProvider;
 		this.dataSourceProxyConnectionIdManagerProviderProvider = dataSourceProxyConnectionIdManagerProviderProvider;
 		this.proxyDataSourceBuilderCustomizers = proxyDataSourceBuilderCustomizers;
 	}
@@ -100,6 +104,7 @@ private DataSourceProxyBuilderConfigurer getConfigurer() {
 					this.methodExecutionListenersProvider.orderedStream().toList(),
 					this.parameterTransformerProvider.getIfAvailable(), this.queryTransformerProvider.getIfAvailable(),
 					this.resultSetProxyLogicFactoryProvider.getIfAvailable(),
+					this.generatedKeysProxyLogicFactoryProvider.getIfAvailable(),
 					this.dataSourceProxyConnectionIdManagerProviderProvider.getIfAvailable());
 		}
 		return this.dataSourceProxyBuilderConfigurer;

datasource-micrometer-spring-boot/src/main/java/net/ttddyy/observation/boot/autoconfigure/DataSourceProxyBuilderConfigurer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2022 the original author or authors.
+ * Copyright 2013-2024 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.
@@ -53,12 +53,19 @@ public class DataSourceProxyBuilderConfigurer {
 
 	private final List<MethodExecutionListener> methodExecutionListeners;
 
+	@Nullable
 	private final ParameterTransformer parameterTransformer;
 
+	@Nullable
 	private final QueryTransformer queryTransformer;
 
+	@Nullable
 	private final ResultSetProxyLogicFactory resultSetProxyLogicFactory;
 
+	@Nullable
+	private final ResultSetProxyLogicFactory generatedKeysProxyLogicFactory;
+
+	@Nullable
 	private final DataSourceProxyConnectionIdManagerProvider dataSourceProxyConnectionIdManagerProvider;
 
 	private final JdbcProperties jdbcProperties;
@@ -67,13 +74,15 @@ public DataSourceProxyBuilderConfigurer(JdbcProperties jdbcProperties, List<Quer
 			List<MethodExecutionListener> methodExecutionListeners, @Nullable ParameterTransformer parameterTransformer,
 			@Nullable QueryTransformer queryTransformer,
 			@Nullable ResultSetProxyLogicFactory resultSetProxyLogicFactory,
+			@Nullable ResultSetProxyLogicFactory generatedKeysProxyLogicFactory,
 			@Nullable DataSourceProxyConnectionIdManagerProvider dataSourceProxyConnectionIdManagerProvider) {
 		this.jdbcProperties = jdbcProperties;
 		this.listeners = listeners;
 		this.methodExecutionListeners = methodExecutionListeners;
 		this.parameterTransformer = parameterTransformer;
 		this.queryTransformer = queryTransformer;
 		this.resultSetProxyLogicFactory = resultSetProxyLogicFactory;
+		this.generatedKeysProxyLogicFactory = generatedKeysProxyLogicFactory;
 		this.dataSourceProxyConnectionIdManagerProvider = dataSourceProxyConnectionIdManagerProvider;
 	}
 
@@ -151,6 +160,11 @@ TimeUnit.SECONDS, toCommonsLogLevel(datasourceProxy.getSlowQuery().getLogLevel()
 					? ResultSetProxyLogicFactory.DEFAULT : this.resultSetProxyLogicFactory;
 			proxyDataSourceBuilder.proxyResultSet(factory);
 		}
+		if (this.jdbcProperties.getIncludes().contains(TraceType.KEYS)) {
+			ResultSetProxyLogicFactory factory = this.generatedKeysProxyLogicFactory == null
+					? ResultSetProxyLogicFactory.DEFAULT : this.generatedKeysProxyLogicFactory;
+			proxyDataSourceBuilder.proxyGeneratedKeys(factory);
+		}
 		ifAvailable(this.listeners, l -> l.forEach(proxyDataSourceBuilder::listener));
 		ifAvailable(this.methodExecutionListeners, m -> m.forEach(proxyDataSourceBuilder::methodListener));
 		ifAvailable(this.parameterTransformer, proxyDataSourceBuilder::parameterTransformer);

datasource-micrometer-spring-boot/src/main/java/net/ttddyy/observation/boot/autoconfigure/JdbcProperties.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022-2023 the original author or authors.
+ * Copyright 2022-2024 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.
@@ -36,7 +36,7 @@ public class JdbcProperties {
 	/**
 	 * Which types of tracing we would like to include.
 	 */
-	private Set<TraceType> includes = Set.of(TraceType.CONNECTION, TraceType.QUERY, TraceType.FETCH);
+	private Set<TraceType> includes = Set.of(TraceType.CONNECTION, TraceType.QUERY, TraceType.FETCH, TraceType.KEYS);
 
 	/**
 	 * List of DataSource bean names that will not be decorated.
@@ -308,7 +308,12 @@ public enum TraceType {
 		/**
 		 * Related to ResultSets.
 		 */
-		FETCH(JdbcObservationDocumentation.RESULT_SET);
+		FETCH(JdbcObservationDocumentation.RESULT_SET),
+
+		/**
+		 * Related to generated keys.
+		 */
+		KEYS(JdbcObservationDocumentation.GENERATED_KEYS);
 
 		final JdbcObservationDocumentation supportedDocumentation;
 

datasource-micrometer-spring-boot/src/test/java/net/ttddyy/observation/boot/autoconfigure/DataSourceObservationBeanPostProcessorTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 the original author or authors.
+ * Copyright 2022-2024 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.
@@ -60,6 +60,8 @@ class DataSourceObservationBeanPostProcessorTests {
 
 	private ObjectProvider<ResultSetProxyLogicFactory> resultSetProxyLogicFactoryProvider;
 
+	private ObjectProvider<ResultSetProxyLogicFactory> generatedKeysProxyLogicFactoryProvider;
+
 	private ObjectProvider<DataSourceProxyConnectionIdManagerProvider> dataSourceProxyConnectionIdManagerProviderProvider;
 
 	private ObjectProvider<ProxyDataSourceBuilderCustomizer> proxyDataSourceBuilderCustomizers;
@@ -76,14 +78,15 @@ void beforeEach() {
 		this.parameterTransformerProvider = mock(ObjectProvider.class);
 		this.queryTransformerProvider = mock(ObjectProvider.class);
 		this.resultSetProxyLogicFactoryProvider = mock(ObjectProvider.class);
+		this.generatedKeysProxyLogicFactoryProvider = mock(ObjectProvider.class);
 		this.dataSourceProxyConnectionIdManagerProviderProvider = mock(ObjectProvider.class);
 		this.proxyDataSourceBuilderCustomizers = mock(ObjectProvider.class);
 
 		this.processor = new DataSourceObservationBeanPostProcessor(this.jdbcPropertiesProvider,
 				this.dataSourceNameResolverProvider, this.listenersProvider, this.methodExecutionListenersProvider,
 				this.parameterTransformerProvider, this.queryTransformerProvider,
-				this.resultSetProxyLogicFactoryProvider, this.dataSourceProxyConnectionIdManagerProviderProvider,
-				this.proxyDataSourceBuilderCustomizers);
+				this.resultSetProxyLogicFactoryProvider, this.generatedKeysProxyLogicFactoryProvider,
+				this.dataSourceProxyConnectionIdManagerProviderProvider, this.proxyDataSourceBuilderCustomizers);
 	}
 
 	@Test

datasource-micrometer/src/main/java/net/ttddyy/observation/tracing/ConnectionAttributesManager.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 the original author or authors.
+ * Copyright 2022-2024 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.
@@ -69,9 +69,19 @@ class ConnectionAttributes {
 
 	class ResultSetAttributesManager {
 
-		Map<ResultSet, ResultSetAttributes> byResultSet = new ConcurrentHashMap<>();
+		private final Map<ResultSet, ResultSetAttributes> byResultSet = new ConcurrentHashMap<>();
 
-		Map<ResultSet, Statement> statements = new ConcurrentHashMap<>();
+		private final Map<ResultSet, Statement> statements = new ConcurrentHashMap<>();
+
+		private final Set<ResultSet> generatedKeys = new HashSet<>();
+
+		boolean isGeneratedKeys(ResultSet resultSet) {
+			return this.generatedKeys.contains(resultSet);
+		}
+
+		void addGeneratedKeys(ResultSet resultSet) {
+			this.generatedKeys.add(resultSet);
+		}
 
 		ResultSetAttributes add(ResultSet resultSet, @Nullable Statement statement, ResultSetAttributes attributes) {
 			this.byResultSet.put(resultSet, attributes);
@@ -88,6 +98,7 @@ ResultSetAttributes getByResultSet(ResultSet resultSet) {
 
 		@Nullable
 		ResultSetAttributes removeByResultSet(ResultSet resultSet) {
+			this.generatedKeys.remove(resultSet);
 			this.statements.remove(resultSet);
 			return this.byResultSet.remove(resultSet);
 		}
@@ -102,13 +113,14 @@ Set<ResultSetAttributes> removeByStatement(Statement statement) {
 					iter.remove();
 				}
 			}
-
+			this.generatedKeys.removeAll(resultSets);
 			return resultSets.stream().map(this.byResultSet::remove).filter(Objects::nonNull)
 					.collect(Collectors.toSet());
 		}
 
 		Set<ResultSetAttributes> removeAll() {
 			Set<ResultSetAttributes> attributes = new HashSet<>(this.byResultSet.values());
+			this.generatedKeys.clear();
 			this.byResultSet.clear();
 			this.statements.clear();
 			return attributes;

datasource-micrometer/src/main/java/net/ttddyy/observation/tracing/DataSourceObservationListener.java

@@ -42,6 +42,7 @@
 import net.ttddyy.dsproxy.listener.MethodExecutionContext;
 import net.ttddyy.dsproxy.listener.MethodExecutionListener;
 import net.ttddyy.dsproxy.listener.QueryExecutionListener;
+import net.ttddyy.dsproxy.proxy.ProxyJdbcObject;
 import net.ttddyy.observation.tracing.ConnectionAttributesManager.ConnectionAttributes;
 import net.ttddyy.observation.tracing.ConnectionAttributesManager.ResultSetAttributes;
 import net.ttddyy.observation.tracing.JdbcObservationDocumentation.JdbcEvents;
@@ -55,6 +56,14 @@ public class DataSourceObservationListener implements QueryExecutionListener, Me
 
 	private static final InternalLogger logger = InternalLoggerFactory.getInstance(DataSourceObservationListener.class);
 
+	private static final Set<String> METHODS_TO_IGNORE = new HashSet<>(Arrays.asList(
+			// java.lang.Object methods
+			"toString", "equals", "hashCode",
+			// java.sql.Wrapper methods
+			"unwrap", "isWrapperFor",
+			// datasource-proxy methods
+			"getTarget", "getProxyConfig", "getDataSourceName"));
+
 	private final Supplier<ObservationRegistry> observationRegistrySupplier;
 
 	private ConnectionAttributesManager connectionAttributesManager = new DefaultConnectionAttributesManager();
@@ -68,6 +77,9 @@ public class DataSourceObservationListener implements QueryExecutionListener, Me
 	private ResultSetObservationConvention resultSetObservationConvention = new ResultSetObservationConvention() {
 	};
 
+	private GeneratedKeysObservationConvention generatedKeysObservationConvention = new GeneratedKeysObservationConvention() {
+	};
+
 	private QueryParametersSpanTagProvider queryParametersSpanTagProvider = new DefaultQueryParametersSpanTagProvider();
 
 	/**
@@ -193,6 +205,9 @@ else if ("executeBatch".equals(methodName)) {
 	@Override
 	public void beforeMethod(MethodExecutionContext executionContext) {
 		String methodName = executionContext.getMethod().getName();
+		if (METHODS_TO_IGNORE.contains(methodName)) {
+			return;
+		}
 		Object target = executionContext.getTarget();
 		if (target instanceof DataSource && "getConnection".equals(methodName)) {
 			handleGetConnectionBefore(executionContext);
@@ -202,6 +217,9 @@ public void beforeMethod(MethodExecutionContext executionContext) {
 	@Override
 	public void afterMethod(MethodExecutionContext executionContext) {
 		String methodName = executionContext.getMethod().getName();
+		if (METHODS_TO_IGNORE.contains(methodName)) {
+			return;
+		}
 		Object target = executionContext.getTarget();
 		if (target instanceof DataSource && "getConnection".equals(methodName)) {
 			handleGetConnectionAfter(executionContext);
@@ -221,6 +239,18 @@ else if (target instanceof Statement) {
 			if ("close".equals(methodName)) {
 				handleStatementClose(executionContext);
 			}
+			else if ("getGeneratedKeys".equals(methodName) && executionContext.getResult() instanceof ResultSet) {
+				String connectionId = executionContext.getConnectionInfo().getConnectionId();
+				ConnectionAttributes connectionAttributes = this.connectionAttributesManager.get(connectionId);
+				if (connectionAttributes != null) {
+					ResultSet resultSet = (ResultSet) executionContext.getResult();
+					// result could be a proxy, unwrap it.
+					if (resultSet instanceof ProxyJdbcObject) {
+						resultSet = (ResultSet) ((ProxyJdbcObject) resultSet).getTarget();
+					}
+					connectionAttributes.resultSetAttributesManager.addGeneratedKeys(resultSet);
+				}
+			}
 		}
 		else if (target instanceof ResultSet) {
 			handleResultSet(executionContext);
@@ -345,7 +375,8 @@ private void handleResultSet(MethodExecutionContext executionContext) {
 		ResultSetAttributes resultSetAttributes = connectionAttributes.resultSetAttributesManager
 				.getByResultSet(resultSet);
 		if (resultSetAttributes == null) {
-			resultSetAttributes = createResultSetAttributesAndStartObservation(executionContext);
+			boolean isGeneratedKey = connectionAttributes.resultSetAttributesManager.isGeneratedKeys(resultSet);
+			resultSetAttributes = createResultSetAttributesAndStartObservation(executionContext, isGeneratedKey);
 
 			Statement statement = null;
 			try {
@@ -361,7 +392,8 @@ private void handleResultSet(MethodExecutionContext executionContext) {
 			connectionAttributes.resultSetAttributesManager.add(resultSet, statement, resultSetAttributes);
 		}
 
-		ResultSetOperation operation = new ResultSetOperation(executionContext.getMethod(), executionContext.getResult());
+		ResultSetOperation operation = new ResultSetOperation(executionContext.getMethod(),
+				executionContext.getResult());
 		resultSetAttributes.context.addOperation(operation);
 
 		String methodName = executionContext.getMethod().getName();
@@ -377,12 +409,21 @@ else if ("next".equals(methodName)) {
 		}
 	}
 
-	private ResultSetAttributes createResultSetAttributesAndStartObservation(MethodExecutionContext executionContext) {
+	private ResultSetAttributes createResultSetAttributesAndStartObservation(MethodExecutionContext executionContext,
+			boolean isGeneratedKeys) {
 		// new ResultSet observation
 		ResultSetContext resultSetContext = new ResultSetContext();
 		populateFromConnectionAttributes(resultSetContext, executionContext.getConnectionInfo().getConnectionId());
-		Observation observation = createAndStartObservation(JdbcObservationDocumentation.RESULT_SET, resultSetContext,
-				this.resultSetObservationConvention);
+
+		Observation observation;
+		if (isGeneratedKeys) {
+			observation = createAndStartObservation(JdbcObservationDocumentation.GENERATED_KEYS, resultSetContext,
+					this.generatedKeysObservationConvention);
+		}
+		else {
+			observation = createAndStartObservation(JdbcObservationDocumentation.RESULT_SET, resultSetContext,
+					this.resultSetObservationConvention);
+		}
 
 		if (logger.isDebugEnabled()) {
 			logger.debug("Created a new result-set observation [" + observation + "]");
@@ -458,6 +499,11 @@ public void setResultSetObservationConvention(ResultSetObservationConvention res
 		this.resultSetObservationConvention = resultSetObservationConvention;
 	}
 
+	public void setGeneratedKeysObservationConvention(
+			GeneratedKeysObservationConvention generatedKeysObservationConvention) {
+		this.generatedKeysObservationConvention = generatedKeysObservationConvention;
+	}
+
 	public void setQueryParametersSpanTagProvider(QueryParametersSpanTagProvider queryParametersSpanTagProvider) {
 		this.queryParametersSpanTagProvider = queryParametersSpanTagProvider;
 	}