From 4b3fd8da15cf70d329dc76a636f0974bac8445ee Mon Sep 17 00:00:00 2001
From: Kazuki Shimizu <kazuki43zoo@gmail.com>
Date: Fri, 18 May 2018 03:58:55 +0900
Subject: [PATCH] DATAJDBC-216 : Support primitive type @Id property on
 auditing feature

---
 .../support/JdbcAuditingEventListener.java    | 25 ++++++++++++++++---
 .../config/JdbcAuditingRegistrar.java         |  2 ++
 ...nableJdbcAuditingHsqlIntegrationTests.java |  2 +-
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/springframework/data/jdbc/domain/support/JdbcAuditingEventListener.java b/src/main/java/org/springframework/data/jdbc/domain/support/JdbcAuditingEventListener.java
index bdb31ee21e..8de09cd655 100644
--- a/src/main/java/org/springframework/data/jdbc/domain/support/JdbcAuditingEventListener.java
+++ b/src/main/java/org/springframework/data/jdbc/domain/support/JdbcAuditingEventListener.java
@@ -19,6 +19,8 @@
 import org.springframework.context.ApplicationListener;
 import org.springframework.data.auditing.AuditingHandler;
 import org.springframework.data.jdbc.mapping.event.BeforeSaveEvent;
+import org.springframework.data.jdbc.mapping.model.JdbcMappingContext;
+import org.springframework.data.jdbc.mapping.model.JdbcPersistentEntityInformation;
 import org.springframework.data.jdbc.repository.config.EnableJdbcAuditing;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
@@ -35,6 +37,7 @@
 public class JdbcAuditingEventListener implements ApplicationListener<BeforeSaveEvent> {
 
 	@Nullable private AuditingHandler handler;
+	private JdbcMappingContext context;
 
 	/**
 	 * Configures the {@link AuditingHandler} to be used to set the current auditor on the domain types touched.
@@ -48,6 +51,17 @@ public void setAuditingHandler(ObjectFactory<AuditingHandler> auditingHandler) {
 		this.handler = auditingHandler.getObject();
 	}
 
+	/**
+	 * Configures a {@link JdbcMappingContext} that use for judging whether new object or not.
+	 * @param context must not be {@literal null}
+	 */
+	public void setJdbcMappingContext(JdbcMappingContext context) {
+
+		Assert.notNull(context, "JdbcMappingContext must not be null!");
+
+		this.context = context;
+	}
+
 	/**
 	 * {@inheritDoc}
 	 * 
@@ -59,11 +73,14 @@ public void onApplicationEvent(BeforeSaveEvent event) {
 		if (handler != null) {
 
 			event.getOptionalEntity().ifPresent(entity -> {
-
-				if (event.getId().getOptionalValue().isPresent()) {
-					handler.markModified(entity);
-				} else {
+				@SuppressWarnings("unchecked")
+				Class<Object> entityType = event.getChange().getEntityType();
+				JdbcPersistentEntityInformation<Object, ?> entityInformation =
+						context.getRequiredPersistentEntityInformation(entityType);
+				if (entityInformation.isNew(entity)) {
 					handler.markCreated(entity);
+				} else {
+					handler.markModified(entity);
 				}
 			});
 		}
diff --git a/src/main/java/org/springframework/data/jdbc/repository/config/JdbcAuditingRegistrar.java b/src/main/java/org/springframework/data/jdbc/repository/config/JdbcAuditingRegistrar.java
index 41e6a87a82..d40566a70f 100644
--- a/src/main/java/org/springframework/data/jdbc/repository/config/JdbcAuditingRegistrar.java
+++ b/src/main/java/org/springframework/data/jdbc/repository/config/JdbcAuditingRegistrar.java
@@ -17,6 +17,7 @@
 
 import java.lang.annotation.Annotation;
 
+import org.springframework.beans.factory.config.AutowireCapableBeanFactory;
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.support.BeanDefinitionBuilder;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
@@ -83,6 +84,7 @@ protected void registerAuditListenerBeanDefinition(BeanDefinition auditingHandle
 		BeanDefinitionBuilder builder = BeanDefinitionBuilder.rootBeanDefinition(listenerClass);
 		builder.addPropertyValue("auditingHandler",
 				ParsingUtils.getObjectFactoryBeanDefinition(getAuditingHandlerBeanName(), null));
+		builder.setAutowireMode(AutowireCapableBeanFactory.AUTOWIRE_BY_TYPE);
 		registerInfrastructureBeanWithId(builder.getRawBeanDefinition(), listenerClass.getName(), registry);
 	}
 
diff --git a/src/test/java/org/springframework/data/jdbc/repository/config/EnableJdbcAuditingHsqlIntegrationTests.java b/src/test/java/org/springframework/data/jdbc/repository/config/EnableJdbcAuditingHsqlIntegrationTests.java
index 24b8b0b494..8eb9df7b1e 100644
--- a/src/test/java/org/springframework/data/jdbc/repository/config/EnableJdbcAuditingHsqlIntegrationTests.java
+++ b/src/test/java/org/springframework/data/jdbc/repository/config/EnableJdbcAuditingHsqlIntegrationTests.java
@@ -199,7 +199,7 @@ interface AuditingAnnotatedDummyEntityRepository extends CrudRepository<Auditing
 	@Data
 	static class AuditingAnnotatedDummyEntity {
 
-		@Id Long id;
+		@Id long id;
 		@CreatedBy String createdBy;
 		@CreatedDate LocalDateTime createdDate;
 		@LastModifiedBy String lastModifiedBy;