INT-4566: UPDATE for R2DBC In Channel Adapter

rohanmukesh · artembilan · commit b54085c99dff · 2020-08-10T14:12:28.000-04:00
JIRA: https://jira.spring.io/browse/INT-4566 * Rework UPDATE logic according deprecations * Use `ColumnMapRowMapper` for default `Map` payload type * Clean up tests
diff --git a/spring-integration-r2dbc/src/main/java/org/springframework/integration/r2dbc/inbound/R2dbcMessageSource.java b/spring-integration-r2dbc/src/main/java/org/springframework/integration/r2dbc/inbound/R2dbcMessageSource.java
@@ -17,7 +17,9 @@
 package org.springframework.integration.r2dbc.inbound;
 
 
+import java.util.HashMap;
 import java.util.Map;
+import java.util.function.BiFunction;
 
 import org.reactivestreams.Publisher;
 
@@ -30,9 +32,13 @@
 import org.springframework.expression.spel.support.StandardTypeLocator;
 import org.springframework.integration.endpoint.AbstractMessageSource;
 import org.springframework.integration.expression.ExpressionUtils;
+import org.springframework.r2dbc.core.ColumnMapRowMapper;
+import org.springframework.r2dbc.core.DatabaseClient;
 import org.springframework.r2dbc.core.RowsFetchSpec;
 import org.springframework.util.Assert;
 
+import io.r2dbc.spi.Row;
+import io.r2dbc.spi.RowMetadata;
 import reactor.core.publisher.Mono;
 
 /**
@@ -54,18 +60,24 @@
  */
 public class R2dbcMessageSource extends AbstractMessageSource<Publisher<?>> {
 
-	private final R2dbcEntityOperations r2dbcEntityOperations;
+	private final DatabaseClient databaseClient;
 
 	private final ReactiveDataAccessStrategy dataAccessStrategy;
 
 	private final Expression queryExpression;
 
 	private Class<?> payloadType = Map.class;
 
+	private BiFunction<Row, RowMetadata, ?> rowMapper = ColumnMapRowMapper.INSTANCE;
+
 	private boolean expectSingleResult = false;
 
 	private StandardEvaluationContext evaluationContext;
 
+	private String updateSql;
+
+	private BiFunction<DatabaseClient.GenericExecuteSpec, Object, DatabaseClient.GenericExecuteSpec> bindFunction;
+
 	private volatile boolean initialized = false;
 
 	/**
@@ -91,8 +103,8 @@ public R2dbcMessageSource(R2dbcEntityOperations r2dbcEntityOperations, String qu
 	public R2dbcMessageSource(R2dbcEntityOperations r2dbcEntityOperations, Expression queryExpression) {
 		Assert.notNull(r2dbcEntityOperations, "'r2dbcEntityOperations' must not be null");
 		Assert.notNull(queryExpression, "'queryExpression' must not be null");
-		this.r2dbcEntityOperations = r2dbcEntityOperations;
-		this.dataAccessStrategy = this.r2dbcEntityOperations.getDataAccessStrategy();
+		this.databaseClient = r2dbcEntityOperations.getDatabaseClient();
+		this.dataAccessStrategy = r2dbcEntityOperations.getDataAccessStrategy();
 		this.queryExpression = queryExpression;
 	}
 
@@ -108,10 +120,33 @@ public void setPayloadType(Class<?> payloadType) {
 	}
 
 	/**
-	 * Provide a way to return all the records matching criteria or only and only a one otherwise.
+	 * Provide a way to set update query that will be passed to the
+	 * {@link org.springframework.data.r2dbc.core.DatabaseClient#execute(String)}
+	 * method.
+	 * @param updateSql Update query string.
+	 */
+	public void setUpdateSql(String updateSql) {
+		this.updateSql = updateSql;
+	}
+
+	/**
+	 * Provide a way to set BindFunction which will be used to bind parameters
+	 * in the update query.
+	 * @param bindFunction The bindFunction.
+	 */
+	@SuppressWarnings("unchecked")
+	public void setBindFunction(
+			BiFunction<DatabaseClient.GenericExecuteSpec, ?, DatabaseClient.GenericExecuteSpec> bindFunction) {
+
+		this.bindFunction =
+				(BiFunction<DatabaseClient.GenericExecuteSpec, Object, DatabaseClient.GenericExecuteSpec>) bindFunction;
+	}
+
+	/**
+	 * Provide a way to manage which find* method to invoke on {@link R2dbcEntityOperations}.
 	 * Default is 'false', which means the {@link #receive()} method will use
-	 * the {@link org.springframework.data.r2dbc.core.DatabaseClient#execute(String)} method and will fetch all. If set
-	 * to 'true'{@link #receive()} will use {@link org.springframework.data.r2dbc.core.DatabaseClient#execute(String)}
+	 * the {@link DatabaseClient#sql(String)} method and will fetch all. If set
+	 * to 'true'{@link #receive()} will use {@link DatabaseClient#sql(String)}
 	 * and will fetch one and the payload of the returned {@link org.springframework.messaging.Message}
 	 * will be the returned target Object of type
 	 * identified by {@link #payloadType} instead of a List.
@@ -136,6 +171,9 @@ protected void onInit() {
 			 */
 			((StandardTypeLocator) typeLocator).registerImport("org.springframework.data.relational.core.query");
 		}
+		if (!Map.class.isAssignableFrom(this.payloadType)) {
+			this.rowMapper = this.dataAccessStrategy.getRowMapper(this.payloadType);
+		}
 		this.initialized = true;
 	}
 
@@ -155,17 +193,31 @@ protected Object doReceive() {
 				Mono.fromSupplier(() -> this.queryExpression.getValue(this.evaluationContext))
 						.map(this::prepareFetch);
 		if (this.expectSingleResult) {
-			return queryMono.flatMap(RowsFetchSpec::one);
+			return queryMono.flatMap(RowsFetchSpec::one)
+					.flatMap(this::executeUpdate);
+		}
+
+		return queryMono.flatMapMany(RowsFetchSpec::all)
+				.flatMap(this::executeUpdate);
+	}
+
+	private Mono<Object> executeUpdate(Object result) {
+		if (this.updateSql != null) {
+			DatabaseClient.GenericExecuteSpec genericExecuteSpec = this.databaseClient.sql(this.updateSql);
+			if (this.bindFunction != null) {
+				genericExecuteSpec = this.bindFunction.apply(genericExecuteSpec, result);
+			}
+			return genericExecuteSpec.then()
+					.thenReturn(result);
 		}
-		return queryMono.flatMapMany(RowsFetchSpec::all);
+		return Mono.just(result);
 	}
 
 	private RowsFetchSpec<?> prepareFetch(Object queryObject) {
 		String queryString = evaluateQueryObject(queryObject);
-		return this.r2dbcEntityOperations
-				.getDatabaseClient()
+		return this.databaseClient
 				.sql(queryString)
-				.map(this.dataAccessStrategy.getRowMapper(this.payloadType));
+				.map(this.rowMapper);
 	}
 
 	private String evaluateQueryObject(Object queryObject) {
diff --git a/spring-integration-r2dbc/src/test/java/org/springframework/integration/r2dbc/inbound/R2dbcMessageSourceTests.java b/spring-integration-r2dbc/src/test/java/org/springframework/integration/r2dbc/inbound/R2dbcMessageSourceTests.java
@@ -58,6 +58,9 @@ public class R2dbcMessageSourceTests {
 
 	R2dbcEntityTemplate entityTemplate;
 
+	@Autowired
+	R2dbcMessageSource defaultR2dbcMessageSource;
+
 	@Autowired
 	R2dbcMessageSource r2dbcMessageSourceSelectOne;
 
@@ -70,9 +73,13 @@ public class R2dbcMessageSourceTests {
 	@BeforeEach
 	public void setup() {
 		this.entityTemplate = new R2dbcEntityTemplate(this.client, H2Dialect.INSTANCE);
-		List<String> statements = Arrays.asList(
-				"DROP TABLE IF EXISTS person;",
-				"CREATE table person (id INT AUTO_INCREMENT NOT NULL, name VARCHAR2, age INT NOT NULL);");
+		r2dbcMessageSourceSelectMany.setExpectSingleResult(false);
+		defaultR2dbcMessageSource.setBindFunction(null);
+
+		List<String> statements =
+				Arrays.asList(
+						"DROP TABLE IF EXISTS person;",
+						"CREATE table person (id INT AUTO_INCREMENT NOT NULL, name VARCHAR2, age INT NOT NULL);");
 
 		statements.forEach(it -> this.client.sql(it)
 				.fetch()
@@ -129,6 +136,136 @@ public void validateSuccessfulQueryWithMultipleElementOfFluxDBObject() {
 
 	}
 
+	@Test
+	public void validateSuccessfulUpdateWithSingleElementOfMonoDBObject() {
+		this.entityTemplate.insert(new Person("Bob", 35))
+				.then()
+				.as(StepVerifier::create)
+				.verifyComplete();
+
+		r2dbcMessageSourceSelectMany.setUpdateSql("UPDATE Person SET name='Foo' where age = :age");
+		r2dbcMessageSourceSelectMany.setBindFunction(
+				(DatabaseClient.GenericExecuteSpec bindSpec, Person o) -> bindSpec.bind("age", o.getAge()));
+		r2dbcMessageSourceSelectMany.setExpectSingleResult(true);
+
+		StepVerifier.create(r2dbcMessageSourceSelectMany.receive().getPayload())
+				.assertNext(person -> assertThat(((Person) person).getName()).isEqualTo("Bob"))
+				.verifyComplete();
+
+		this.entityTemplate.select(Person.class)
+				.all()
+				.as(StepVerifier::create)
+				.assertNext(person -> assertThat(person.getName()).isEqualTo("Foo"))
+				.verifyComplete();
+
+	}
+
+	@Test
+	public void validateSuccessfulUpdateWithMultiplesElementsOfFluxDBObject() {
+		this.entityTemplate.insert(new Person("Bob", 35))
+				.then()
+				.as(StepVerifier::create)
+				.verifyComplete();
+
+		this.entityTemplate.insert(new Person("Tom", 40))
+				.then()
+				.as(StepVerifier::create)
+				.verifyComplete();
+
+		r2dbcMessageSourceSelectMany.setUpdateSql("UPDATE person SET name='Foo' where id = :id");
+		r2dbcMessageSourceSelectMany.setBindFunction(
+				(DatabaseClient.GenericExecuteSpec bindSpec, Person o) -> bindSpec.bind("id", o.getId()));
+		StepVerifier.create(r2dbcMessageSourceSelectMany.receive().getPayload())
+				.expectNextCount(2)
+				.verifyComplete();
+
+		this.entityTemplate.select(Person.class)
+				.all()
+				.as(StepVerifier::create)
+				.assertNext(person -> assertThat(person.getName()).isEqualTo("Foo"))
+				.assertNext(person -> assertThat(person.getName()).isEqualTo("Foo"))
+				.verifyComplete();
+
+	}
+
+	@Test
+	public void validateSuccessfulUpdateWithoutBindFunction() {
+		this.entityTemplate.insert(new Person("Bob", 35))
+				.then()
+				.as(StepVerifier::create)
+				.verifyComplete();
+
+		this.entityTemplate.insert(new Person("Tom", 40))
+				.then()
+				.as(StepVerifier::create)
+				.verifyComplete();
+
+		r2dbcMessageSourceSelectMany.setUpdateSql("UPDATE person SET name='Foo' where id = 1");
+
+		StepVerifier.create(r2dbcMessageSourceSelectMany.receive().getPayload())
+				.expectNextCount(2)
+				.verifyComplete();
+
+		this.entityTemplate.select(Person.class)
+				.all()
+				.as(StepVerifier::create)
+				.assertNext(person -> assertThat(person.getName()).isEqualTo("Foo"))
+				.assertNext(person -> assertThat(person.getName()).isEqualTo("Tom"))
+				.verifyComplete();
+
+	}
+
+	@Test
+	public void validateSuccessfulUpdateWithoutPayloadType() {
+		this.entityTemplate.insert(new Person("Bob", 35))
+				.then()
+				.as(StepVerifier::create)
+				.verifyComplete();
+
+		this.entityTemplate.insert(new Person("Tom", 40))
+				.then()
+				.as(StepVerifier::create)
+				.verifyComplete();
+
+		defaultR2dbcMessageSource.setUpdateSql("UPDATE person SET name='Foo' where id = 1");
+
+		StepVerifier.create(defaultR2dbcMessageSource.receive().getPayload())
+				.expectNextCount(2)
+				.verifyComplete();
+
+		this.client.sql("select * from person")
+				.fetch()
+				.all()
+				.as(StepVerifier::create)
+				.assertNext(person -> assertThat(person.get("name")).isEqualTo("Foo"))
+				.assertNext(person -> assertThat(person.get("name")).isEqualTo("Tom"))
+				.verifyComplete();
+
+	}
+
+	@Test
+	public void testWrongPayloadTypeInBindFunction() {
+		this.entityTemplate.insert(new Person("Bob", 35))
+				.then()
+				.as(StepVerifier::create)
+				.verifyComplete();
+
+		this.entityTemplate.insert(new Person("Tom", 40))
+				.then()
+				.as(StepVerifier::create)
+				.verifyComplete();
+
+		defaultR2dbcMessageSource.setUpdateSql("UPDATE person SET name='Foo' where id = 1");
+		defaultR2dbcMessageSource.setBindFunction(
+				(DatabaseClient.GenericExecuteSpec bindSpec, Person o) -> bindSpec.bind("id", o.getId()));
+
+		StepVerifier.create(defaultR2dbcMessageSource.receive().getPayload())
+				.expectErrorMatches(throwable -> throwable instanceof ClassCastException)
+				.verify();
+
+	}
+
+
 	@Test
 	public void testAnyOtherObjectQueryExpression() {
 
@@ -145,6 +282,12 @@ static class R2dbcMessageSourceConfiguration {
 		@Autowired
 		R2dbcEntityTemplate r2dbcEntityTemplate;
 
+		@Bean
+		public R2dbcMessageSource defaultR2dbcMessageSource() {
+			return new R2dbcMessageSource(r2dbcEntityTemplate, "select * from " +
+					"person");
+		}
+
 		@Bean
 		public R2dbcMessageSource r2dbcMessageSourceSelectOne() {
 			R2dbcMessageSource r2dbcMessageSource = new R2dbcMessageSource(this.r2dbcEntityTemplate,
