Reinstate result splitting for composite simple statements.

Reactor 3.4 has fixed the underlying issue so it's safe now to apply windowUntil(…).

[resolves #440]

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

Author: mp911de

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

@@ -25,6 +25,7 @@
 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;
@@ -140,12 +141,16 @@ 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 ExtendedFlowDelegate.runQuery(this.resources, factory, sql, Binding.EMPTY, Collections.emptyList(), this.fetchSize)
+                .windowUntil(WINDOW_UNTIL)
+                .map(messages -> PostgresqlResult.toResult(this.resources, messages, factory))
+                .as(Operators::discardOnCancel);
         }
 
-        Flux<BackendMessage> messages = SimpleQueryMessageFlow.exchange(this.resources.getClient(), sql);
-        return Flux.just(PostgresqlResult.toResult(this.resources, messages, factory));
+        return SimpleQueryMessageFlow.exchange(this.resources.getClient(), sql)
+            .windowUntil(WINDOW_UNTIL)
+            .map(messages -> PostgresqlResult.toResult(this.resources, messages, factory))
+            .as(Operators::discardOnCancel);
     }
 
 }

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

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

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

@@ -21,8 +21,6 @@
 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;
 
@@ -97,9 +95,4 @@ 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() {
-    }
-
 }