From e9e07bd4a3db10091834c42465a009d91739d849 Mon Sep 17 00:00:00 2001
From: Jens Schauder <jschauder@vmware.com>
Date: Tue, 2 Nov 2021 14:18:38 +0100
Subject: [PATCH] Render `LOCK` clause through `SelectRenderContext`.

Closes #1078
---
 .../core/sql/render/SelectRenderContext.java  | 32 +++++++++-
 .../sql/render/SelectRendererUnitTests.java   | 61 ++++++++++++++++++-
 2 files changed, 89 insertions(+), 4 deletions(-)

diff --git a/spring-data-relational/src/main/java/org/springframework/data/relational/core/sql/render/SelectRenderContext.java b/spring-data-relational/src/main/java/org/springframework/data/relational/core/sql/render/SelectRenderContext.java
index b420fefe..c78f3791 100644
--- a/spring-data-relational/src/main/java/org/springframework/data/relational/core/sql/render/SelectRenderContext.java
+++ b/spring-data-relational/src/main/java/org/springframework/data/relational/core/sql/render/SelectRenderContext.java
@@ -15,8 +15,10 @@
  */
 package org.springframework.data.relational.core.sql.render;
 
+import java.util.OptionalLong;
 import java.util.function.Function;
 
+import org.springframework.data.relational.core.sql.LockMode;
 import org.springframework.data.relational.core.sql.Select;
 
 /**
@@ -26,6 +28,7 @@ import org.springframework.data.relational.core.sql.Select;
  *
  * @author Mark Paluch
  * @author Myeonghyeon Lee
+ * @author Jens Schauder
  * @since 1.1
  */
 public interface SelectRenderContext {
@@ -51,12 +54,35 @@ public interface SelectRenderContext {
 
 	/**
 	 * Customization hook: Rendition of a part after {@code ORDER BY}. The rendering function is called always, regardless
-	 * whether {@code ORDER BY} exists or not. Renders an empty string by default.
-	 *
+	 * whether {@code ORDER BY} exists or not.
+	 * <p>
+	 * Renders lock, limit and offset clause as appropriate.
+	 * </p>
+	 * 
 	 * @param hasOrderBy the actual value whether the {@link Select} statement has a {@code ORDER BY} clause.
 	 * @return render {@link Function} invoked after rendering {@code ORDER BY}.
 	 */
 	default Function<Select, ? extends CharSequence> afterOrderBy(boolean hasOrderBy) {
-		return select -> "";
+
+		return select -> {
+
+			OptionalLong limit = select.getLimit();
+			OptionalLong offset = select.getOffset();
+			LockMode lockMode = select.getLockMode();
+
+			String lockPrefix = (lockMode == null) ? "" : " FOR UPDATE";
+
+			if (limit.isPresent() && offset.isPresent()) {
+				return lockPrefix
+						+ String.format(" OFFSET %d ROWS FETCH FIRST %d ROWS ONLY", offset.getAsLong(), limit.getAsLong());
+			}
+			if (limit.isPresent()) {
+				return lockPrefix + String.format(" FETCH FIRST %d ROWS ONLY", limit.getAsLong());
+			}
+			if (offset.isPresent()) {
+				return lockPrefix + String.format(" OFFSET %d ROWS", offset.getAsLong());
+			}
+			return lockPrefix;
+		};
 	}
 }
diff --git a/spring-data-relational/src/test/java/org/springframework/data/relational/core/sql/render/SelectRendererUnitTests.java b/spring-data-relational/src/test/java/org/springframework/data/relational/core/sql/render/SelectRendererUnitTests.java
index 7893bd5f..7e0f22d4 100644
--- a/spring-data-relational/src/test/java/org/springframework/data/relational/core/sql/render/SelectRendererUnitTests.java
+++ b/spring-data-relational/src/test/java/org/springframework/data/relational/core/sql/render/SelectRendererUnitTests.java
@@ -37,7 +37,7 @@ class SelectRendererUnitTests {
 		Table bar = SQL.table("bar");
 		Column foo = bar.column("foo");
 
-		Select select = Select.builder().select(foo).from(bar).limitOffset(1, 2).build();
+		Select select = Select.builder().select(foo).from(bar).build();
 
 		assertThat(SqlRenderer.toString(select)).isEqualTo("SELECT bar.foo FROM bar");
 	}
@@ -480,6 +480,65 @@ class SelectRendererUnitTests {
 		assertThat(rendered).isEqualTo("SELECT CAST(User.name AS VARCHAR2) FROM User");
 	}
 
+	@Test // GH-1076
+	void rendersLimitAndOffset() {
+
+		Table table_user = SQL.table("User");
+		Select select = StatementBuilder.select(table_user.column("name")).from(table_user).limitOffset(10, 5).build();
+
+		final String rendered = SqlRenderer.toString(select);
+		assertThat(rendered).isEqualTo("SELECT User.name FROM User OFFSET 5 ROWS FETCH FIRST 10 ROWS ONLY");
+	}
+
+	@Test // GH-1076
+	void rendersLimit() {
+
+		Table table_user = SQL.table("User");
+		Select select = StatementBuilder.select(table_user.column("name")).from(table_user) //
+				.limit(3) //
+				.build();
+
+		final String rendered = SqlRenderer.toString(select);
+		assertThat(rendered).isEqualTo("SELECT User.name FROM User FETCH FIRST 3 ROWS ONLY");
+	}
+
+	@Test // GH-1076
+	void rendersLock() {
+
+		Table table_user = SQL.table("User");
+		Select select = StatementBuilder.select(table_user.column("name")).from(table_user) //
+				.lock(LockMode.PESSIMISTIC_READ) //
+				.build();
+
+		final String rendered = SqlRenderer.toString(select);
+		assertThat(rendered).isEqualTo("SELECT User.name FROM User FOR UPDATE");
+	}
+
+	@Test // GH-1076
+	void rendersLockAndOffset() {
+
+		Table table_user = SQL.table("User");
+		Select select = StatementBuilder.select(table_user.column("name")).from(table_user).offset(3) //
+				.lock(LockMode.PESSIMISTIC_WRITE) //
+				.build();
+
+		final String rendered = SqlRenderer.toString(select);
+		assertThat(rendered).isEqualTo("SELECT User.name FROM User FOR UPDATE OFFSET 3 ROWS");
+	}
+
+	@Test // GH-1076
+	void rendersLockAndOffsetUsingDialect() {
+
+		Table table_user = SQL.table("User");
+		Select select = StatementBuilder.select(table_user.column("name")).from(table_user).limitOffset(3, 6) //
+				.lock(LockMode.PESSIMISTIC_WRITE) //
+				.build();
+
+		String rendered = SqlRenderer.create(new RenderContextFactory(PostgresDialect.INSTANCE).createRenderContext())
+				.render(select);
+		assertThat(rendered).isEqualTo("SELECT User.name FROM User LIMIT 3 OFFSET 6 FOR UPDATE OF User");
+	}
+
 	@Test // GH-1007
 	void shouldRenderConditionAsExpression() {