feat: automatically add LIMIT clause (#792)

* feat: automatically add LIMIT clause Automatically add a LIMIT clause to queries that use an OFFSET clause with a query parameter and no LIMIT clause. That is currently not supported in Cloud Spanner PostgreSQL databases. This type of queries are generated by Prisma. * test: add more tests * test: add more tests and use fixed limit * perf: do sanity check before stepping through entire string * docs: add TODO and internal annotation * chore: run formatter

Author: olavloite

src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java

@@ -392,6 +392,19 @@ public boolean isForceAutocommit() {
  return getBoolSetting("spanner", "force_autocommit", false);
  }
 
+ /**
+ * Returns whether statements with an OFFSET clause that uses a parameter should be automatically
+ * appended with a LIMIT clause. The LIMIT clause will use the literal Long.MAX_VALUE for unbound
+ * statements, and Long.MAX_VALUE - offset for bound statements.
+ *
+ * <p>This method will be removed in the future.
+ */
+ // TODO: Remove when Cloud Spanner supports parametrized OFFSET clauses without a LIMIT clause.
+ @InternalApi
+ public boolean isAutoAddLimitClause() {
+ return getBoolSetting("spanner", "auto_add_limit_clause", false);
+ }
+
  /** Returns the current setting for replacing pg_catalog tables with common table expressions. */
  public boolean isReplacePgCatalogTables() {
  PGSetting setting = internalGet(toKey("spanner", "replace_pg_catalog_tables"), false);

src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java

@@ -295,7 +295,13 @@ void execute() {
  }
 
  Statement bindStatement(Statement statement) {
- return statementBinder.apply(statement);
+ Statement boundStatement = statementBinder.apply(statement);
+ // Add a LIMIT clause to the statement if it contains an OFFSET clause that uses a query
+ // parameter and there is no existing LIMIT clause in the query.
+ if (sessionState.isAutoAddLimitClause()) {
+ return SimpleParser.addLimitIfParameterizedOffset(boundStatement);
+ }
+ return boundStatement;
  }
 
  StatementResult analyzeOrExecute(Statement statement) {

src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleParser.java

@@ -24,6 +24,7 @@
 import com.google.common.collect.ImmutableList;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Objects;
@@ -225,6 +226,52 @@ static Statement copyStatement(Statement original, String sql) {
  return builder.build();
  }
 
+ /**
+ * Adds a LIMIT clause to the given statement if the statement contains a parameterized OFFSET
+ * clause and no LIMIT clause.
+ */
+ static Statement addLimitIfParameterizedOffset(Statement statement) {
+ String sqlLowerCase = statement.getSql().toLowerCase(Locale.ENGLISH);
+ // If there is no offset clause, then we know that we don't have to analyze any further.
+ if (!sqlLowerCase.contains("offset")) {
+ return statement;
+ }
+ SimpleParser parser = new SimpleParser(statement.getSql());
+ parser.parseExpressionUntilKeyword(ImmutableList.of("limit", "offset"), true, false, false);
+ if (parser.pos >= parser.getSql().length()) {
+ return statement;
+ }
+ if (parser.eatKeyword("limit")) {
+ // If the statement contains a LIMIT clause, then we're OK.
+ return statement;
+ }
+ String parameter;
+ if (parser.eatKeyword("offset") && (parameter = parser.readQueryParameter()) != null) {
+ // Check if we have a LIMIT clause after the OFFSET clause.
+ if (parser.peekKeyword("limit")) {
+ return statement;
+ }
+ // This is probably an invalid query. Do not modify it to prevent the user getting an error
+ // message for a query that they did not write.
+ if (parser.hasMoreTokens()) {
+ return statement;
+ }
+ // The statement contains an OFFSET clause using a query parameter and no LIMIT clause.
+ // Append a LIMIT clause equal to Long.MAX_VALUE / 2.
+ // We could also calculate it based on the actual OFFSET value, but that would make the
+ // query more prone to cache misses. Also, adding the LIMIT clause with a query parameter
+ // would structurally change the query, something that we also don't want.
+ long limit = Long.MAX_VALUE / 2;
+ return copyStatement(
+ statement,
+ parser.sql.substring(0, parser.pos)
+ + " limit "
+ + limit
+ + parser.sql.substring(parser.pos));
+ }
+ return statement;
+ }
+
  private String sql;
  private int pos;
 
@@ -661,6 +708,22 @@ boolean eatDotOperator() {
  return false;
  }
 
+ String readQueryParameter() {
+ if (eat(true, false, "$")) {
+ int startPos = pos - 1;
+ if (pos == sql.length() || !Character.isDigit(sql.charAt(pos))) {
+ return null;
+ }
+ while (pos < sql.length() && Character.isDigit(sql.charAt(pos))) {
+ pos++;
+ }
+ if (Character.isDigit(sql.charAt(pos - 1))) {
+ return sql.substring(startPos, pos);
+ }
+ }
+ return null;
+ }
+
  boolean peekJoinKeyword() {
  return peekKeyword("join")
  || peekKeyword("left")

src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java

@@ -195,6 +195,28 @@ public void testSelectHelloWorld() throws SQLException {
  }
  }
 
+ @Test
+ public void testSelectParameterizedOffsetWithoutLimit() throws SQLException {
+ try (Connection connection = DriverManager.getConnection(getConnectionUrl())) {
+ connection.createStatement().execute("set spanner.auto_add_limit_clause=true");
+ try (PreparedStatement statement =
+ connection.prepareStatement("select * from numbers offset ?")) {
+ for (long offset : new long[] {0L, 1L}) {
+ statement.setLong(1, offset);
+ try (ResultSet resultSet = statement.executeQuery()) {
+ if (offset == 0L) {
+ assertTrue(resultSet.next());
+ assertEquals(1L, resultSet.getLong(1));
+ assertFalse(resultSet.next());
+ } else {
+ assertFalse(resultSet.next());
+ }
+ }
+ }
+ }
+ }
+ }
+
  @Test
  public void testSelectCurrentSchema() throws SQLException {
  try (Connection connection = DriverManager.getConnection(getConnectionUrl())) {

src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java

@@ -1997,6 +1997,43 @@ public void testBatchedRewrittenPreparedStatementWithGenericBatchExecutionError(
  }
  }
 
+ @Test
+ public void testParameterizedOffsetWithoutLimit() throws SQLException {
+ // Add a result for the non-limited query that contains one row.
+ mockSpanner.putStatementResult(
+ StatementResult.query(
+ Statement.newBuilder("select * from foo offset $1").bind("p1").to(0L).build(),
+ SELECT1_RESULTSET));
+ // Add a result for the limited query that is empty.
+ mockSpanner.putStatementResult(
+ StatementResult.query(
+ Statement.newBuilder("select * from foo offset $1 limit 4611686018427387903")
+ .bind("p1")
+ .to(0L)
+ .build(),
+ EMPTY_RESULTSET));
+
+ try (Connection connection = DriverManager.getConnection(createUrl())) {
+ for (boolean addLimit : new boolean[] {true, false}) {
+ connection.createStatement().execute("set spanner.auto_add_limit_clause=" + addLimit);
+ try (PreparedStatement statement =
+ connection.prepareStatement("select * from foo offset ?")) {
+ statement.setLong(1, 0);
+ try (ResultSet resultSet = statement.executeQuery()) {
+ // We should get the empty result set when the auto-limit feature is enabled.
+ if (addLimit) {
+ assertFalse(resultSet.next());
+ } else {
+ assertTrue(resultSet.next());
+ assertEquals(1L, resultSet.getLong(1));
+ assertFalse(resultSet.next());
+ }
+ }
+ }
+ }
+ }
+ }
+
  @Test
  public void testTwoQueries() throws SQLException {
  try (Connection connection = DriverManager.getConnection(createUrl())) {