Add support for unquoted NEW and OLD keywords in RETURNING clauses

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/expressions/UpdateExpression.java

@@ -58,8 +58,8 @@
  */
 public class UpdateExpression extends AbstractRelationalExpressionWithChildren implements PlannerGraphRewritable {
 
-    private static final String OLD_FIELD_NAME = "old";
-    private static final String NEW_FIELD_NAME = "new";
+    private static final String OLD_FIELD_NAME = "OLD";
+    private static final String NEW_FIELD_NAME = "NEW";
 
     @Nonnull
     private final Quantifier.ForEach inner;

fdb-relational-core/src/main/antlr/RelationalLexer.g4

@@ -583,6 +583,7 @@ NAME:                                'NAME';
 NAMES:                               'NAMES';
 NCHAR:                               'NCHAR';
 NEVER:                               'NEVER';
+NEW:                                 'NEW';
 NEXT:                                'NEXT';
 NO:                                  'NO';
 NOCOPY:                              'NOCOPY';
@@ -596,6 +597,7 @@ OFFLINE:                             'OFFLINE';
 OFFSET:                              'OFFSET';
 OF:                                  'OF';
 OJ:                                  'OJ';
+OLD:                                 'OLD';
 OLD_PASSWORD:                        'OLD_PASSWORD';
 ONE:                                 'ONE';
 ONLINE:                              'ONLINE';

fdb-relational-core/src/main/antlr/RelationalParser.g4

@@ -1252,8 +1252,8 @@ keywordsCanBeId
     | MAX_USER_CONNECTIONS | MEDIUM | MEMBER | MEMORY | MERGE | MESSAGE | MESSAGE_TEXT
     | MID | MIGRATE
     | MIN | MIN_ROWS | MODE | MODIFY | MUTEX | MYSQL | MYSQL_ERRNO | NAME | NAMES
-    | NCHAR | NDB_STORED_USER | NEVER | NEXT | NO | NOCOPY | NODEGROUP | NOCACHE | NONE | NOWAIT | NUMBER | ODBC | OFFLINE | OFFSET
-    | OF | OJ | OLD_PASSWORD | ONE | ONLINE | ONLY | OPEN | OPTIMIZER_COSTS
+    | NCHAR | NDB_STORED_USER | NEVER | NEW | NEXT | NO | NOCOPY | NODEGROUP | NOCACHE | NONE | NOWAIT | NUMBER | ODBC | OFFLINE | OFFSET
+    | OF | OJ | OLD | OLD_PASSWORD | ONE | ONLINE | ONLY | OPEN | OPTIMIZER_COSTS
     | OPTIONAL | OPTIONS | ORDER | OWNER | PACK_KEYS | PAGE | PARSER | PARTIAL
     | PARTITIONING | PARTITIONS | PASSWORD | PERSIST_RO_VARIABLES_ADMIN | PHASE | PLUGINS
     | PLUGIN_DIR | PLUGIN | PORT | PRECEDES | PREPARE | PRESERVE | PREV

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/LogicalPlanFragment.java

@@ -23,6 +23,7 @@
 import com.apple.foundationdb.annotation.API;
 
 import com.apple.foundationdb.record.query.plan.cascades.CorrelationIdentifier;
+import com.apple.foundationdb.record.query.plan.cascades.Quantifier;
 import com.apple.foundationdb.record.query.plan.cascades.typing.Type;
 import com.apple.foundationdb.relational.util.Assert;
 
@@ -151,10 +152,15 @@
         @Nonnull
         private final Optional<StringTrieNode> targetTypeReorderings;
 
+        @Nonnull
+        private final Optional<Quantifier.ForEach> updateQuantifier;
+
         private State(@Nonnull Optional<Type> targetType,
-                      @Nonnull Optional<StringTrieNode> targetTypeReorderings) {
+                      @Nonnull Optional<StringTrieNode> targetTypeReorderings,
+                      @Nonnull Optional<Quantifier.ForEach> updateQuantifier) {
             this.targetType = targetType;
             this.targetTypeReorderings = targetTypeReorderings;
+            this.updateQuantifier = updateQuantifier;
         }
 
         @Nonnull
@@ -172,6 +178,11 @@
             return targetTypeReorderings;
         }
 
+        @Nonnull
+        public Optional<Quantifier.ForEach> getUpdateQuantifier() {
+            return updateQuantifier;
+        }
+
         public static final class Builder {
 
             @Nullable
@@ -180,6 +191,9 @@
             @Nullable
             StringTrieNode targetTypeReorderings;
 
+            @Nullable
+            private Quantifier.ForEach updateQuantifier;
+
             private Builder() {
             }
 
@@ -195,9 +209,15 @@
                 return this;
             }
 
+            @Nonnull
+            public Builder withUpdateQuantifier(@Nonnull Quantifier.ForEach updateQuantifier) {
+                this.updateQuantifier = updateQuantifier;
+                return this;
+            }
+
             @Nonnull
             public State build() {
-                return new State(Optional.ofNullable(targetType), Optional.ofNullable(targetTypeReorderings));
+                return new State(Optional.ofNullable(targetType), Optional.ofNullable(targetTypeReorderings), Optional.ofNullable(updateQuantifier));
             }
         }
     }

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/SemanticAnalyzer.java

@@ -292,6 +292,23 @@
     @Nonnull
     public Star expandStar(@Nonnull Optional<Identifier> optionalQualifier,
                            @Nonnull LogicalOperators operators) {
+        return expandStar(optionalQualifier, operators, Optional.empty());
+    }
+
+    @Nonnull
+    public Star expandStar(@Nonnull Optional<Identifier> optionalQualifier,
+                           @Nonnull LogicalOperators operators,
+                           @Nonnull Optional<LogicalPlanFragment> planFragmentMaybe) {
+        // In UPDATE RETURNING context, unqualified * should expand to NEW.*
+        if (planFragmentMaybe.isPresent() && optionalQualifier.isEmpty()) {
+            final var stateMaybe = planFragmentMaybe.get().getStateMaybe();
+            final var updateQuantifierMaybe = stateMaybe.flatMap(LogicalPlanFragment.State::getUpdateQuantifier);
+            if (updateQuantifierMaybe.isPresent()) {
+                // Treat unqualified * as NEW.* in UPDATE RETURNING
+                return expandStar(Optional.of(Identifier.of("NEW")), operators, planFragmentMaybe);
+            }
+        }
+
         final var forEachOperators = operators.forEachOnly();
         // Case 1: no qualifier, e.g. SELECT * FROM T, R;
         if (optionalQualifier.isEmpty()) {
@@ -331,6 +348,32 @@
     @Nonnull
     public Expression resolveIdentifier(@Nonnull Identifier identifier,
                                         @Nonnull LogicalPlanFragment planFragment) {
+        // Check if we're in UPDATE RETURNING context
+        final var stateMaybe = planFragment.getStateMaybe();
+        final var updateQuantifierMaybe = stateMaybe.flatMap(LogicalPlanFragment.State::getUpdateQuantifier);
+
+        if (updateQuantifierMaybe.isPresent() && !identifier.isQualified()) {
+            // Unqualified identifier in UPDATE RETURNING - automatically access NEW.<field>
+            final var updateQuantifier = updateQuantifierMaybe.get();
+            final var updateValue = updateQuantifier.getFlowedObjectValue();
+            final var fieldValue = com.apple.foundationdb.record.query.plan.cascades.values.FieldValue.ofFieldNames(
+                    updateValue, ImmutableList.of("NEW", identifier.getName()));
+            return Expression.fromUnderlying(fieldValue).withName(identifier);
+        }
+
+        if (updateQuantifierMaybe.isPresent() && identifier.isQualified() && identifier.getQualifier().size() == 1) {
+            // Check if qualified with NEW or OLD
+            final var qualifier = identifier.getQualifier().get(0);
+            if ("NEW".equals(qualifier) || "OLD".equals(qualifier)) {
+                final var updateQuantifier = updateQuantifierMaybe.get();
+                final var updateValue = updateQuantifier.getFlowedObjectValue();
+                final var fieldValue = com.apple.foundationdb.record.query.plan.cascades.values.FieldValue.ofFieldNames(
+                        updateValue, ImmutableList.of(qualifier, identifier.getName()));
+                return Expression.fromUnderlying(fieldValue).withName(Identifier.of(identifier.getName()));
+            }
+        }
+
+        // Normal identifier resolution
         // search throw all visible plan fragments:
         // - in each plan fragment, search operators left to right.
         // - if identifier is not resolve, go to parent plan fragment.

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/ExpressionVisitor.java

@@ -134,7 +134,8 @@ public Expression visitContinuationAtom(@Nonnull RelationalParser.ContinuationAt
     @Nonnull
     @Override
     public Expression visitSelectStarElement(@Nonnull RelationalParser.SelectStarElementContext ignored) {
-        return getDelegate().getSemanticAnalyzer().expandStar(Optional.empty(), getDelegate().getLogicalOperators());
+        return getDelegate().getSemanticAnalyzer().expandStar(Optional.empty(), getDelegate().getLogicalOperators(),
+                getDelegate().getCurrentPlanFragmentMaybe());
     }
 
     @Nonnull

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/QueryVisitor.java

@@ -479,26 +479,30 @@
                 table.getName(),
                 tableType,
                 transformMapBuilder.build());
         final var updateQuantifier = Quantifier.forEach(Reference.initialOf(updateExpression));
         final var resultingUpdate = LogicalOperator.newUnnamedOperator(Expressions.fromQuantifier(updateQuantifier), updateQuantifier);
 
         getDelegate().getCurrentPlanFragment().setOperator(resultingUpdate);
 
         //        if (ctx.CONTINUATION() != null) {
         //            getDelegate().getPlanGenerationContext().setContinuation((byte[]) visit(ctx.continuationAtom()));
         //        }
 
         if (ctx.RETURNING() != null) {
+            // Store the updateQuantifier in State so SemanticAnalyzer can access it
+            final var stateBuilder = LogicalPlanFragment.State.newBuilder().withUpdateQuantifier(updateQuantifier);
+            getDelegate().getCurrentPlanFragment().setState(stateBuilder.build());
+
             final var selectExpressions = visitSelectElements(ctx.selectElements());
             final var result = LogicalOperator.generateSelect(selectExpressions, getDelegate().getLogicalOperators(),
                     Optional.empty(), List.of(), Optional.empty(),
                     getDelegate().getCurrentPlanFragment().getOuterCorrelations(), getDelegate().isTopLevel(), false);
             getDelegate().getCurrentPlanFragment().setOperator(result);
             return result;
         }
         final var result = LogicalOperator.generateSort(resultingUpdate, List.of(), Set.of(), Optional.empty());
         getDelegate().popPlanFragment();
         return result;
     }
 
     @Nonnull

fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/AutoCommitTests.java

@@ -210,7 +210,7 @@ public void simpleUpdateWithReturningExhaustedResultSet(TransactionType transact
         }
         setAutoCommit(conn, transactionType);
         try (final var statement = conn.createStatement()) {
-            final var hasResultSet = statement.execute("UPDATE RESTAURANT SET NAME = 'aa' WHERE REST_NO = 1 RETURNING \"new\".* ");
+            final var hasResultSet = statement.execute("UPDATE RESTAURANT SET NAME = 'aa' WHERE REST_NO = 1 RETURNING NEW.* ");
             Assertions.assertTrue(hasResultSet);
             try (final var resultSet = statement.getResultSet()) {
                 ResultSetAssert.assertThat(resultSet)
@@ -244,7 +244,7 @@ public void simpleUpdateWithReturningNonExhaustedResultSet(TransactionType trans
         }
         setAutoCommit(conn, transactionType);
         try (final var statement = conn.createStatement()) {
-            final var hasResultSet = statement.execute("UPDATE RESTAURANT SET NAME = 'aa' WHERE REST_NO < 3 RETURNING \"new\".* ");
+            final var hasResultSet = statement.execute("UPDATE RESTAURANT SET NAME = 'aa' WHERE REST_NO < 3 RETURNING NEW.* ");
             Assertions.assertTrue(hasResultSet);
             // resultSet not retrieved, hence it is still owned by the statement and will be closed.
         }
@@ -530,7 +530,7 @@ public void switchOffAutoCommitBetweenOngoingTransaction() throws SQLException {
         Assertions.assertFalse(conn.inActiveTransaction());
         Assertions.assertTrue(conn.getAutoCommit());
         try (final var statement = conn.createStatement()) {
-            try (final var rs = statement.executeQuery("UPDATE RESTAURANT SET name = 'aa' WHERE REST_NO < 3 RETURNING \"new\".* ")) {
+            try (final var rs = statement.executeQuery("UPDATE RESTAURANT SET name = 'aa' WHERE REST_NO < 3 RETURNING NEW.* ")) {
                 Assertions.assertTrue(conn.inActiveTransaction());
                 ResultSetAssert.assertThat(rs)
                         .hasNextRow()
@@ -571,7 +571,7 @@ public void switchOnAutoCommitBetweenOngoingTransaction() throws SQLException {
         conn.setAutoCommit(false);
         Assertions.assertFalse(conn.getAutoCommit());
         try (final var statement = conn.createStatement()) {
-            try (final var rs = statement.executeQuery("UPDATE RESTAURANT SET name = 'aa' WHERE REST_NO < 3 RETURNING \"new\".* ")) {
+            try (final var rs = statement.executeQuery("UPDATE RESTAURANT SET name = 'aa' WHERE REST_NO < 3 RETURNING NEW.* ")) {
                 Assertions.assertTrue(conn.inActiveTransaction());
                 ResultSetAssert.assertThat(rs)
                         .hasNextRow()
@@ -613,7 +613,7 @@ public void transactionClosesWithStatement() throws SQLException {
 
         RelationalResultSet rs;
         try (final var statement = conn.createStatement()) {
-            rs = statement.executeQuery("UPDATE RESTAURANT SET name = 'aa' WHERE REST_NO < 3 RETURNING \"new\".* ");
+            rs = statement.executeQuery("UPDATE RESTAURANT SET name = 'aa' WHERE REST_NO < 3 RETURNING NEW.* ");
             Assertions.assertTrue(conn.inActiveTransaction());
             ResultSetAssert.assertThat(rs).hasNextRow();
 

fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/PreparedStatementTests.java

@@ -416,8 +416,7 @@ void setArrayTypeOfByte() throws Exception {
                 final var count = statement.executeUpdate("INSERT INTO RestaurantReviewer(id) VALUES (1)");
                 Assertions.assertThat(count).isEqualTo(1);
             }
-            // "new" should not be quoted. TODO ([Post] Fix identifiers case-sensitivity matching in plan generator)
-            try (var ps = ddl.setSchemaAndGetConnection().prepareStatement("UPDATE RestaurantReviewer SET secrets = ?param WHERE id = 1 RETURNING \"new\".*")) {
+            try (var ps = ddl.setSchemaAndGetConnection().prepareStatement("UPDATE RestaurantReviewer SET secrets = ?param WHERE id = 1 RETURNING NEW.*")) {
                 final var array = List.of(new byte[]{1, 2, 3, 4}, new byte[]{5, 6, 7, 8});
                 final var arrayObject = ddl.getConnection().createArrayOf("BINARY", array.toArray());
                 ps.setArray("param", arrayObject);
@@ -442,8 +441,7 @@ void setByteType() throws Exception {
                 final var count = statement.executeUpdate("INSERT INTO RestaurantComplexRecord(rest_no) VALUES (1)");
                 Assertions.assertThat(count).isEqualTo(1);
             }
-            // "new" should not be quoted. TODO ([Post] Fix identifiers case-sensitivity matching in plan generator)
-            try (var ps = ddl.setSchemaAndGetConnection().prepareStatement("UPDATE RestaurantComplexRecord SET key = ?param WHERE rest_no = 1 RETURNING \"new\".*")) {
+            try (var ps = ddl.setSchemaAndGetConnection().prepareStatement("UPDATE RestaurantComplexRecord SET key = ?param WHERE rest_no = 1 RETURNING NEW.*")) {
                 ps.setBytes("param", new byte[]{1, 2, 3, 4});
                 try (final var resultSet = ps.executeQuery()) {
                     ResultSetAssert.assertThat(resultSet)
@@ -668,8 +666,7 @@ void prepareUpdateWithStruct() throws Exception {
             try (var statement = ddl.setSchemaAndGetConnection().createStatement()) {
                 statement.execute("INSERT INTO RestaurantReviewer(id, stats) VALUES (1, (2, 'a', 'b')), (2, (3, 'b', 'c')), (3, (4, 'c', 'd')), (4, (5, 'd', 'e')), (5, (6, 'e', 'f'))");
             }
-            // "new" should not be quoted. TODO ([Post] Fix identifiers case-sensitivity matching in plan generator)
-            final var query = "UPDATE RestaurantReviewer SET stats = ?param WHERE id = 1 RETURNING \"new\".stats";
+            final var query = "UPDATE RestaurantReviewer SET stats = ?param WHERE id = 1 RETURNING NEW.stats";
             try (var ps = ddl.setSchemaAndGetConnection().prepareStatement(query)) {
                 ps.setObject("param", ddl.getConnection().createStruct("blah", statsAttributes));
                 final var expectedStats = EmbeddedRelationalStruct.newBuilder()
@@ -705,8 +702,7 @@ void prepareUpdateWithNestedStruct(Object[] attributes, boolean succeed) throws
             try (var statement = ddl.setSchemaAndGetConnection().createStatement()) {
                 statement.execute("INSERT INTO RestaurantComplexRecord(rest_no, name) VALUES (1, 'mango & miso'), (2, 'basil & brawn'), (3, 'peach & pepper'), (4, 'smoky skillet'), (5, 'the tin pot')");
             }
-            // "new" should not be quoted. TODO ([Post] Fix identifiers case-sensitivity matching in plan generator)
-            final var query = "UPDATE RestaurantComplexRecord SET location = ?param WHERE rest_no = 1 RETURNING \"new\".location";
+            final var query = "UPDATE RestaurantComplexRecord SET location = ?param WHERE rest_no = 1 RETURNING NEW.location";
             try (var ps = ddl.setSchemaAndGetConnection().prepareStatement(query)) {
                 final var latLong = ddl.getConnection().createStruct("LATLONG", attributes);
                 final var location = ddl.getConnection().createStruct("LOCATION", new Object[]{"next door", 217, latLong});
@@ -740,8 +736,7 @@ void prepareUpdateWithArrayOfPrimitives() throws Exception {
             try (var statement = ddl.setSchemaAndGetConnection().createStatement()) {
                 statement.execute("INSERT INTO RestaurantComplexRecord(rest_no, name) VALUES (1, 'mango & miso'), (2, 'basil & brawn'), (3, 'peach & pepper'), (4, 'smoky skillet'), (5, 'the tin pot')");
             }
-            // "new" should not be quoted. TODO ([Post] Fix identifiers case-sensitivity matching in plan generator)
-            final var query = "UPDATE RestaurantComplexRecord SET customer = ?param WHERE rest_no = 1 RETURNING \"new\".customer";
+            final var query = "UPDATE RestaurantComplexRecord SET customer = ?param WHERE rest_no = 1 RETURNING NEW.customer";
             try (var ps = ddl.setSchemaAndGetConnection().prepareStatement(query)) {
 
                 final var customer = ddl.getConnection().createArrayOf("STRING", customerAttributes);
@@ -767,8 +762,7 @@ void prepareUpdateWithArrayOfStructs() throws Exception {
             try (var statement = ddl.setSchemaAndGetConnection().createStatement()) {
                 statement.execute("INSERT INTO RestaurantComplexRecord(rest_no, name) VALUES (1, 'mango & miso'), (2, 'basil & brawn'), (3, 'peach & pepper'), (4, 'smoky skillet'), (5, 'the tin pot')");
             }
-            // "new" should not be quoted. TODO ([Post] Fix identifiers case-sensitivity matching in plan generator)
-            final var query = "UPDATE RestaurantComplexRecord SET tags = ?param WHERE rest_no = 1 RETURNING \"new\".tags OPTIONS (LOG QUERY)";
+            final var query = "UPDATE RestaurantComplexRecord SET tags = ?param WHERE rest_no = 1 RETURNING NEW.tags OPTIONS (LOG QUERY)";
             final var restaurantTagAttributes = new Object[][]{{"chinese", 343}, {"top-rated", 2356}, {"exotic", 10}};
             try (var ps = ddl.setSchemaAndGetConnection().prepareStatement(query)) {
                 ps.setArray("param", createTagArray(ddl.getConnection(), restaurantTagAttributes));

fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/UpdateTest.java

@@ -180,7 +180,7 @@ public void insertRecords(int numRecords) throws RelationalException, SQLExcepti
 
     private static RelationalPreparedStatement prepareUpdate(RelationalConnection conn, String updateField, Object param, Continuation continuation) throws SQLException {
         if (continuation.atBeginning()) {
-            final var statement = conn.prepareStatement("UPDATE RestaurantReviewer SET " + updateField + " = ?param WHERE id >= 0 RETURNING \"new\"." + updateField + ", \"new\".id");
+            final var statement = conn.prepareStatement("UPDATE RestaurantReviewer SET " + updateField + " = ?param WHERE id >= 0 RETURNING NEW." + updateField + ", NEW.id");
             statement.setObject("param", param);
             return statement;
         } else {

yaml-tests/src/test/resources/arrays.yamsql

@@ -66,12 +66,12 @@ test_block:
       - supported_version: 4.5.13.0
       - count: 1
     -
-      # "new" should not be quoted. TODO ([Post] Fix identifiers case-sensitivity matching in plan generator)
-      - query: UPDATE A SET X = [11, 12, 13] WHERE PK = 1 RETURNING "new".X
+      - query: UPDATE A SET X = [11, 12, 13] WHERE PK = 1 RETURNING NEW.X
+      - supported_version: !current_version
       - result: [{[11, 12, 13]}]
     -
-      # "new" should not be quoted. TODO ([Post] Fix identifiers case-sensitivity matching in plan generator)
-      - query: UPDATE D SET X = [(11), (12), (13)] WHERE PK = 1 RETURNING "new".X
+      - query: UPDATE D SET X = [(11), (12), (13)] WHERE PK = 1 RETURNING NEW.X
+      - supported_version: !current_version
       - result: [{[{F: 11}, {F: 12}, {F: 13}]}]
     -
       - query: select X[3] from A where PK = 1