Polishing.

Remove windowUntil(…) to avoid miscalculation of demand in long-running streams. Refine discardOnCancel operator calls to properly guard simple queries.

[#401][#426]

Signed-off-by: Mark Paluch <[email protected]>

Author: mp911de

src/main/java/io/r2dbc/postgresql/PostgresqlConnection.java

@@ -360,8 +360,7 @@ private <T> Mono<T> withTransactionStatus(Function<TransactionStatus, T> f) {
     private <T> Publisher<T> exchange(String sql) {
         ExceptionFactory exceptionFactory = ExceptionFactory.withSql(sql);
         return (Publisher<T>) SimpleQueryMessageFlow.exchange(this.client, sql)
-            .handle(exceptionFactory::handleErrorResponse)
-            .as(Operators::discardOnCancel);
+            .handle(exceptionFactory::handleErrorResponse);
     }
 
     /**

src/main/java/io/r2dbc/postgresql/PostgresqlResult.java

@@ -82,7 +82,20 @@ public Mono<Integer> getRowsUpdated() {
                         sink.next(rowCount);
                     }
                 }
-            }).singleOrEmpty();
+            }).collectList().handle((list, sink) -> {
+
+                if (list.isEmpty()) {
+                    return;
+                }
+
+                int sum = 0;
+
+                for (Integer integer : list) {
+                    sum += integer;
+                }
+
+                sink.next(sum);
+            });
     }
 
     @Override

src/main/java/io/r2dbc/postgresql/SimpleQueryPostgresqlStatement.java

@@ -16,8 +16,6 @@
 
 package io.r2dbc.postgresql;
 
-import io.netty.util.ReferenceCountUtil;
-import io.netty.util.ReferenceCounted;
 import io.r2dbc.postgresql.api.PostgresqlStatement;
 import io.r2dbc.postgresql.client.Binding;
 import io.r2dbc.postgresql.client.SimpleQueryMessageFlow;
@@ -27,7 +25,6 @@
 import io.r2dbc.postgresql.message.backend.ErrorResponse;
 import io.r2dbc.postgresql.util.Assert;
 import io.r2dbc.postgresql.util.GeneratedValuesUtils;
-import io.r2dbc.postgresql.util.Operators;
 import io.r2dbc.spi.Statement;
 import reactor.core.publisher.Flux;
 import reactor.util.annotation.Nullable;
@@ -91,11 +88,7 @@ public SimpleQueryPostgresqlStatement bindNull(int index, @Nullable Class<?> typ
 
     @Override
     public Flux<io.r2dbc.postgresql.api.PostgresqlResult> execute() {
-        if (this.generatedColumns == null) {
-            return execute(this.sql);
-        }
-
-        return execute(GeneratedValuesUtils.augment(this.sql, this.generatedColumns));
+        return execute(this.generatedColumns == null ? this.sql : GeneratedValuesUtils.augment(this.sql, this.generatedColumns));
     }
 
     @Override
@@ -148,17 +141,11 @@ private Flux<io.r2dbc.postgresql.api.PostgresqlResult> execute(String sql) {
         if (this.fetchSize != NO_LIMIT) {
 
             Flux<BackendMessage> messages = ExtendedFlowDelegate.runQuery(this.resources, factory, sql, Binding.EMPTY, Collections.emptyList(), this.fetchSize);
-
             return Flux.just(new PostgresqlResult(this.resources, messages, factory));
         }
 
-        return SimpleQueryMessageFlow
-            .exchange(this.resources.getClient(), sql)
-            .windowUntil(WINDOW_UNTIL)
-            .map(dataRow -> PostgresqlResult.toResult(this.resources, dataRow, factory))
-            .cast(io.r2dbc.postgresql.api.PostgresqlResult.class)
-            .as(Operators::discardOnCancel)
-            .doOnDiscard(ReferenceCounted.class, ReferenceCountUtil::release);
+        Flux<BackendMessage> messages = SimpleQueryMessageFlow.exchange(this.resources.getClient(), sql);
+        return Flux.just(PostgresqlResult.toResult(this.resources, messages, factory));
     }
 
 }

src/main/java/io/r2dbc/postgresql/client/SimpleQueryMessageFlow.java

@@ -16,10 +16,13 @@
 
 package io.r2dbc.postgresql.client;
 
+import io.netty.util.ReferenceCountUtil;
+import io.netty.util.ReferenceCounted;
 import io.r2dbc.postgresql.message.backend.BackendMessage;
 import io.r2dbc.postgresql.message.frontend.FrontendMessage;
 import io.r2dbc.postgresql.message.frontend.Query;
 import io.r2dbc.postgresql.util.Assert;
+import io.r2dbc.postgresql.util.Operators;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
@@ -43,7 +46,8 @@ public static Flux<BackendMessage> exchange(Client client, String query) {
         Assert.requireNonNull(client, "client must not be null");
         Assert.requireNonNull(query, "query must not be null");
 
-        return client.exchange(Mono.<FrontendMessage>just(new Query(query)).doOnSubscribe(ignore -> QueryLogger.logQuery(client.getContext(), query)));
+        return client.exchange(Mono.<FrontendMessage>just(new Query(query)).doOnSubscribe(ignore -> QueryLogger.logQuery(client.getContext(), query))).doOnDiscard(ReferenceCounted.class,
+            ReferenceCountUtil::release).as(Operators::discardOnCancel);
     }
 
 }

src/test/java/io/r2dbc/postgresql/PostgresNotificationIntegrationTests.java

@@ -19,6 +19,7 @@
 import io.netty.channel.Channel;
 import io.r2dbc.postgresql.api.Notification;
 import io.r2dbc.postgresql.api.PostgresqlConnection;
+import io.r2dbc.postgresql.api.PostgresqlResult;
 import io.r2dbc.postgresql.util.ConnectionIntrospector;
 import io.r2dbc.spi.R2dbcNonTransientResourceException;
 import org.junit.jupiter.api.Test;
@@ -44,13 +45,13 @@ void shouldReceivePubSubNotifications() throws Exception {
 
         CountDownLatch await = new CountDownLatch(1);
         Disposable listener = this.connectionFactory.create().flatMapMany(it -> {
-            return it.createStatement("LISTEN mymessage").execute().doOnComplete(await::countDown)
+            return it.createStatement("LISTEN mymessage").execute().flatMap(PostgresqlResult::getRowsUpdated).doOnComplete(await::countDown)
                 .thenMany(it.getNotifications()).doOnCancel(() -> it.close().subscribe());
         }).doOnNext(notifications::add).subscribe();
 
         await.await(10, TimeUnit.SECONDS);
 
-        this.connectionFactory.create().flatMapMany(it -> it.createStatement("NOTIFY mymessage, 'Mötorhead'").execute().thenMany(it.close()))
+        this.connectionFactory.create().flatMapMany(it -> it.createStatement("NOTIFY mymessage, 'Mötorhead'").execute().flatMap(PostgresqlResult::getRowsUpdated).thenMany(it.close()))
             .as(StepVerifier::create).verifyComplete();
 
         Notification notification = notifications.poll(10, TimeUnit.SECONDS);

src/test/java/io/r2dbc/postgresql/PostgresqlBatchUnitTests.java

@@ -47,7 +47,7 @@ void add() {
             .add("test-query-2")
             .execute()
             .as(StepVerifier::create)
-            .expectNextCount(2)
+            .expectNextCount(1)
             .verifyComplete();
     }
 

src/test/java/io/r2dbc/postgresql/PostgresqlTestKit.java

@@ -21,6 +21,8 @@
 import io.r2dbc.spi.ConnectionFactory;
 import io.r2dbc.spi.ConnectionFactoryOptions;
 import io.r2dbc.spi.test.TestKit;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.springframework.jdbc.core.JdbcOperations;
 
@@ -95,4 +97,9 @@ public String getPlaceholder(int index) {
         return String.format("$%d", index + 1);
     }
 
+    @Test
+    @Disabled("Disabled until we get a fixed windowUntil(…) or we have a better way to split result streams into multiple Result objects. https://github.com/pgjdbc/r2dbc-postgresql/issues/401")
+    public void compoundStatement() {
+    }
+
 }

src/test/java/io/r2dbc/postgresql/SimpleQueryPostgresqlStatementIntegrationTests.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.r2dbc.postgresql;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.springframework.jdbc.core.JdbcOperations;
+import reactor.test.StepVerifier;
+import reactor.util.function.Tuples;
+
+/**
+ * Integration tests for {@link SimpleQueryPostgresqlStatement}.
+ */
+class SimpleQueryPostgresqlStatementIntegrationTests extends AbstractIntegrationTests {
+
+    @BeforeEach
+    void setUp() {
+        super.setUp();
+        getJdbcOperations().execute("DROP TABLE IF EXISTS test");
+        getJdbcOperations().execute("CREATE TABLE test (val VARCHAR(255))");
+        getJdbcOperations().execute("INSERT INTO test (val) VALUES ('a'), ('b')");
+    }
+
+    @AfterEach
+    void tearDown() {
+        super.tearDown();
+        getJdbcOperations().execute("DROP TABLE IF EXISTS test");
+    }
+
+    private JdbcOperations getJdbcOperations() {
+        return SERVER.getJdbcOperations();
+    }
+
+    @Test
+    void shouldRunMultipleQueries() {
+
+        this.connection.createStatement("SELECT 1;SELECT val FROM test")
+            .fetchSize(0).execute()
+            .flatMap(it -> it.map((row, rowMetadata) -> Tuples.of(row.get(0), rowMetadata.getColumnMetadata(0).getName())))
+            .as(StepVerifier::create)
+            .expectNext(Tuples.of(1, "?column?"), Tuples.of("a", "val"), Tuples.of("b", "val"))
+            .verifyComplete();
+    }
+
+}