From 21490b839761c3a22fbe519d1b11e286204da26a Mon Sep 17 00:00:00 2001
From: Jens Schauder <jschauder@pivotal.io>
Date: Mon, 27 Feb 2017 08:22:19 +0100
Subject: [PATCH] DATAJDBC-96 - Database based ID generation.

No longer using batch inserts, which won't (easily) hold up for long anyway, since we need to decide between update and insert anyway. If one wants to support batch insert one also has to check if any autoid columns are present, because those seem not to work with batch inserts with many or at least some drivers.

Related ticket: SPR-1836.
Original pull request: #5.
---
 .../mapping/model/JdbcPersistentEntity.java   |  4 ++
 .../jdbc/repository/SimpleJdbcRepository.java | 22 +++++-----
 .../JdbcRepositoryIntegrationTests.java       | 40 +++++++++++++++++--
 .../createTable.sql                           |  2 +-
 4 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/src/main/java/org/springframework/data/jdbc/mapping/model/JdbcPersistentEntity.java b/src/main/java/org/springframework/data/jdbc/mapping/model/JdbcPersistentEntity.java
index 217d8d70..bbe2d69c 100644
--- a/src/main/java/org/springframework/data/jdbc/mapping/model/JdbcPersistentEntity.java
+++ b/src/main/java/org/springframework/data/jdbc/mapping/model/JdbcPersistentEntity.java
@@ -50,4 +50,8 @@ public class JdbcPersistentEntity<T> extends BasicPersistentEntity<T, JdbcPersis
 	public Object getIdValue(T instance) {
 		return getPropertyAccessor(instance).getProperty(getIdProperty());
 	}
+
+	public void setId(T instance, Object value) {
+		getPropertyAccessor(instance).setProperty(getIdProperty(),value);
+	}
 }
diff --git a/src/main/java/org/springframework/data/jdbc/repository/SimpleJdbcRepository.java b/src/main/java/org/springframework/data/jdbc/repository/SimpleJdbcRepository.java
index 556e2a1e..e66502ac 100644
--- a/src/main/java/org/springframework/data/jdbc/repository/SimpleJdbcRepository.java
+++ b/src/main/java/org/springframework/data/jdbc/repository/SimpleJdbcRepository.java
@@ -28,6 +28,8 @@ import org.springframework.data.repository.CrudRepository;
 import org.springframework.jdbc.core.namedparam.MapSqlParameterSource;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
+import org.springframework.jdbc.support.GeneratedKeyHolder;
+import org.springframework.jdbc.support.KeyHolder;
 
 /**
  * @author Jens Schauder
@@ -50,22 +52,24 @@ public class SimpleJdbcRepository<T, ID extends Serializable> implements CrudRep
 	}
 
 	@Override
-	public <S extends T> S save(S entity) {
+	public <S extends T> S save(S instance) {
 
-		template.update(sql.getInsert(), getPropertyMap(entity));
+		KeyHolder holder = new GeneratedKeyHolder();
 
-		return entity;
+		template.update(
+				sql.getInsert(),
+				new MapSqlParameterSource(getPropertyMap(instance)),
+				holder);
+
+		entity.setId(instance, holder.getKey());
+
+		return instance;
 	}
 
 	@Override
 	public <S extends T> Iterable<S> save(Iterable<S> entities) {
 
-		Map<String, ?>[] batchValues = StreamSupport
-				.stream(entities.spliterator(), false)
-				.map(i -> getPropertyMap(i))
-				.toArray(size -> new Map[size]);
-
-		template.batchUpdate(sql.getInsert(), batchValues);
+		entities.forEach(this::save);
 
 		return entities;
 	}
diff --git a/src/test/java/org/springframework/data/jdbc/repository/JdbcRepositoryIntegrationTests.java b/src/test/java/org/springframework/data/jdbc/repository/JdbcRepositoryIntegrationTests.java
index 8af0faf1..0251b976 100644
--- a/src/test/java/org/springframework/data/jdbc/repository/JdbcRepositoryIntegrationTests.java
+++ b/src/test/java/org/springframework/data/jdbc/repository/JdbcRepositoryIntegrationTests.java
@@ -19,9 +19,7 @@ import static java.util.Arrays.*;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.*;
 
-import java.sql.SQLException;
 import org.junit.After;
-import org.junit.Assert;
 import org.junit.Test;
 import org.springframework.data.annotation.Id;
 import org.springframework.data.jdbc.repository.support.JdbcRepositoryFactory;
@@ -91,6 +89,26 @@ public class JdbcRepositoryIntegrationTests {
 				reloadedEntity.getName());
 	}
 
+	@Test
+	public void canSaveAndLoadAnEntityWithDatabaseBasedIdGeneration() {
+
+		entity = createDummyEntity(null);
+
+		entity = repository.save(entity);
+
+		assertThat(entity).isNotNull();
+
+		DummyEntity reloadedEntity = repository.findOne(entity.getId());
+
+		assertEquals(
+				entity.getId(),
+				reloadedEntity.getId());
+		assertEquals(
+				entity.getName(),
+				reloadedEntity.getName());
+	}
+
+
 	@Test
 	public void saveMany() {
 
@@ -101,6 +119,21 @@ public class JdbcRepositoryIntegrationTests {
 		assertThat(repository.findAll()).extracting(DummyEntity::getId).containsExactlyInAnyOrder(23L, 24L);
 	}
 
+	@Test
+	public void saveManyWithIdGeneration() {
+
+		DummyEntity one = createDummyEntity(null);
+		DummyEntity two = createDummyEntity(null);
+
+		Iterable<DummyEntity> entities = repository.save(asList(one, two));
+
+		assertThat(entities).allMatch(e -> e.getId() != null);
+
+		assertThat(repository.findAll())
+				.extracting(DummyEntity::getId)
+				.containsExactlyInAnyOrder(new Long[]{one.getId(), two.getId()});
+	}
+
 	@Test
 	public void existsReturnsTrueIffEntityExists() {
 
@@ -195,13 +228,12 @@ public class JdbcRepositoryIntegrationTests {
 	}
 
 
-
 	private static DummyEntityRepository createRepository(EmbeddedDatabase db) {
 		return new JdbcRepositoryFactory(db).getRepository(DummyEntityRepository.class);
 	}
 
 
-	private static DummyEntity createDummyEntity(long id) {
+	private static DummyEntity createDummyEntity(Long id) {
 
 		DummyEntity entity = new DummyEntity();
 		entity.setId(id);
diff --git a/src/test/resources/org.springframework.data.jdbc.repository/createTable.sql b/src/test/resources/org.springframework.data.jdbc.repository/createTable.sql
index 6bfdd471..c818bad8 100644
--- a/src/test/resources/org.springframework.data.jdbc.repository/createTable.sql
+++ b/src/test/resources/org.springframework.data.jdbc.repository/createTable.sql
@@ -1 +1 @@
-CREATE TABLE dummyentity (ID BIGINT, NAME VARCHAR(100), PRIMARY KEY (ID))
\ No newline at end of file
+CREATE TABLE dummyentity (ID BIGINT GENERATED BY DEFAULT AS IDENTITY(START WITH 4711) PRIMARY KEY, NAME VARCHAR(100))
\ No newline at end of file