wip

Author: ArjanSchouten

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

@@ -31,11 +31,8 @@
 import io.r2dbc.postgresql.message.frontend.Query;
 import io.r2dbc.postgresql.util.Assert;
 import org.reactivestreams.Publisher;
-import reactor.core.Disposable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
-import reactor.core.publisher.Sinks;
-import reactor.core.scheduler.Schedulers;
 
 import java.nio.ByteBuffer;
 
@@ -48,8 +45,6 @@ final class PostgresqlCopyIn {
 
     private final ConnectionResources context;
 
-    private volatile boolean cancelled = false;
-
     PostgresqlCopyIn(ConnectionResources context) {
         this.context = Assert.requireNonNull(context, "context must not be null");
     }
@@ -66,14 +61,15 @@ private Mono<Long> copyIn(String sql, Flux<FrontendMessage> frontendMessages) {
 
         Flux<BackendMessage> backendMessages = frontendMessages
             .doOnNext(client::send)
-            .doOnError(e -> !(e instanceof IllegalArgumentException), (e) -> {
-                copyFail(e.getMessage()).subscribe();
-            })
+            .doOnError(e -> !(e instanceof IllegalArgumentException), (e) -> sendCopyFail(e.getMessage()))
             .doOnDiscard(ReferenceCounted.class, ReferenceCountUtil::release)
             .thenMany(client.exchange(Mono.just(CopyDone.INSTANCE)));
 
         return startCopy(sql)
             .concatWith(backendMessages)
+            .doOnCancel(() -> {
+                sendCopyFail("Cancelled");
+            })
             .as(messages -> toResult(context, messages, ExceptionFactory.INSTANCE).getRowsUpdated());
     }
 
@@ -90,21 +86,15 @@ private Flux<BackendMessage> startCopy(String sql) {
             });
     }
 
-    private Flux<BackendMessage> copyFail(String message) {
-        if (!cancelled) {
-            cancelled = true;
-            return context.getClient()
-                .exchange(backendMessage -> backendMessage instanceof CommandComplete, Mono.just(new CopyFail("Copy operation failed: " + message)));
-        }
-
-        return Flux.empty();
+    private void sendCopyFail(String message) {
+        context.getClient().exchange(m -> m instanceof CommandComplete, Mono.just(new CopyFail("Copy operation failed: " + message)))
+            .subscribe();
     }
 
     @Override
     public String toString() {
         return "PostgresqlCopyIn{" +
             "context=" + this.context +
-            ", cancelled=" + this.cancelled +
             '}';
     }
 

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

@@ -17,10 +17,8 @@
 package io.r2dbc.postgresql;
 
 import io.r2dbc.postgresql.ExceptionFactory.PostgresqlBadGrammarException;
-import org.junit.Ignore;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.springframework.jdbc.core.JdbcOperations;
 import reactor.core.publisher.Flux;