WIP - buffer accumulates csv text; Broken resultHandler

Author: psolomin

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgConnectionImpl.java

@@ -19,6 +19,7 @@
 import io.vertx.core.Future;
 import io.vertx.core.Handler;
 import io.vertx.core.Promise;
+import io.netty.buffer.ByteBuf;
 import io.vertx.core.buffer.Buffer;
 import io.vertx.core.buffer.impl.BufferImpl;
 import io.vertx.core.impl.ContextInternal;
@@ -142,22 +143,8 @@ public Future<SqlResult<Buffer>> copyToBytes(String sql) {
     QueryResultBuilder<Buffer, SqlResultImpl<Buffer>, SqlResult<Buffer>> resultHandler =
       new QueryResultBuilder<>(factory, promise);
 
-    Collector<Row, Buffer, Buffer> collector = Collector.of(
-      BufferImpl::new,
-      (v, row) -> {
-        System.out.println(row);
-      },
-      (v1, v2) -> null,
-      Function.identity()
-    );
-
-    SimpleQueryCommand<Buffer> cmd = new SimpleQueryCommand<>(
-      sql, true, false, collector, resultHandler);
-    // this.schedule(promise.context(), cmd);
-
-    QueryExecutor executor = new QueryExecutor(factory, collector);
-    executor.executeSimpleQuery(this, sql, false, false, promise);
-    return promise.future();
+    CopyOutCommand cmd = new CopyOutCommand(sql, resultHandler);
+    return this.schedule(promise.context(), cmd);
   }
 
   @Override

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyOutCommand.java

@@ -1,17 +1,43 @@
 package io.vertx.pgclient.impl.codec;
 
+import io.netty.buffer.ByteBuf;
 import io.vertx.core.buffer.Buffer;
 import io.vertx.sqlclient.SqlResult;
 import io.vertx.sqlclient.impl.QueryResultBuilder;
 import io.vertx.sqlclient.impl.SqlResultImpl;
 import io.vertx.sqlclient.impl.command.CommandBase;
 
-public class CopyOutCommand extends CommandBase<Buffer> {
+import java.util.function.Function;
+import java.util.stream.Collector;
+
+public class CopyOutCommand extends CommandBase<SqlResult<Buffer>> {
   private final String sql;
+  private final Collector<ByteBuf, Buffer, Buffer> collector;
   private final QueryResultBuilder<Buffer, SqlResultImpl<Buffer>, SqlResult<Buffer>> resultHandler;
 
-  public CopyOutCommand(String sql, QueryResultBuilder<Buffer, SqlResultImpl<Buffer>, SqlResult<Buffer>> resultHandler) {
+  public CopyOutCommand(
+    String sql,
+    QueryResultBuilder<Buffer, SqlResultImpl<Buffer>, SqlResult<Buffer>> resultHandler
+  ) {
     this.sql = sql;
     this.resultHandler = resultHandler;
+    this.collector = Collector.of(
+      Buffer::buffer,
+      (v, chunk) -> v.appendBuffer(Buffer.buffer(chunk)),
+      (v1, v2) -> null,
+      Function.identity()
+    );
+  }
+
+  QueryResultBuilder<Buffer, SqlResultImpl<Buffer>, SqlResult<Buffer>> resultHandler() {
+    return resultHandler;
+  }
+
+  String sql() {
+    return sql;
+  }
+
+  Collector<ByteBuf, Buffer, Buffer> collector() {
+    return collector;
   }
 }

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyOutCommandCodec.java

@@ -1,4 +1,42 @@
 package io.vertx.pgclient.impl.codec;
 
-public class CopyOutCommandCodec {
+import io.vertx.core.buffer.Buffer;
+import io.vertx.core.buffer.impl.BufferImpl;
+import io.vertx.sqlclient.SqlResult;
+
+class CopyOutCommandCodec extends PgCommandCodec<SqlResult<Buffer>, CopyOutCommand> {
+  CopyOutDataDecoder decoder;
+
+  CopyOutCommandCodec(CopyOutCommand cmd) {
+    super(cmd);
+    decoder = new CopyOutDataDecoder(cmd.collector());
+  }
+
+  @Override
+  public void handleCommandComplete(int updated) {
+    this.result = null;
+    Buffer result;
+    Throwable failure;
+    int size;
+    if (decoder != null) {
+      failure = decoder.complete();
+      result = decoder.result();
+      size = decoder.size();
+      decoder.reset();
+    } else {
+      failure = null;
+      result = new BufferImpl();
+      size = 0;
+    }
+    cmd.resultHandler().handleResult(updated, size, null, result, failure);
+  }
+
+  @Override
+  public void handleErrorResponse(ErrorResponse errorResponse) {
+    failure = errorResponse.toException();
+  }
+
+  void encode(PgEncoder encoder) {
+    encoder.writeQuery(new Query(cmd.sql()));
+  }
 }

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyOutDataDecoder.java

@@ -0,0 +1,71 @@
+package io.vertx.pgclient.impl.codec;
+
+import io.netty.buffer.ByteBuf;
+import io.vertx.core.buffer.Buffer;
+import java.util.function.BiConsumer;
+import java.util.stream.Collector;
+
+public class CopyOutDataDecoder {
+
+  private final Collector<ByteBuf, Buffer, Buffer> collector;
+  private BiConsumer<Buffer, ByteBuf> accumulator;
+
+  private int size;
+  private Buffer container;
+  private Throwable failure;
+  private Buffer result;
+
+  protected CopyOutDataDecoder(Collector<ByteBuf, Buffer, Buffer> collector) {
+    this.collector = collector;
+    reset();
+  }
+
+  public int size() {
+    return size;
+  }
+
+  public void handleChunk(int len, ByteBuf in) {
+    if (failure != null) {
+      return;
+    }
+    if (accumulator == null) {
+      try {
+        accumulator = collector.accumulator();
+      } catch (Exception e) {
+        failure = e;
+        return;
+      }
+    }
+    try {
+      accumulator.accept(container, in);
+    } catch (Exception e) {
+      failure = e;
+      return;
+    }
+    size++;
+  }
+
+  public Buffer result() {
+    return result;
+  }
+
+  public Throwable complete() {
+    try {
+      result = collector.finisher().apply(container);
+    } catch (Exception e) {
+      failure = e;
+    }
+    return failure;
+  }
+
+  public void reset() {
+    size = 0;
+    failure = null;
+    result = null;
+    try {
+      this.container = collector.supplier().get();
+    } catch (Exception e) {
+      failure = e;
+    }
+  }
+}

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgDecoder.java

@@ -472,7 +472,12 @@ private void decodeNotificationResponse(ChannelHandlerContext ctx, ByteBuf in) {
 
   private void decodeCopyOutResponse(ChannelHandlerContext ctx, ByteBuf in) {}
 
-  private void decodeCopyData(ChannelHandlerContext ctx, ByteBuf in) {}
+  private void decodeCopyData(ChannelHandlerContext ctx, ByteBuf in) {
+    PgCommandCodec<?, ?> codec = inflight.peek();
+    CopyOutCommandCodec cmdCodec = (CopyOutCommandCodec) codec;
+    int len = in.readUnsignedShort();
+    cmdCodec.decoder.handleChunk(len, in);
+  }
 
   private void decodeCopyCompletion(ChannelHandlerContext ctx, ByteBuf in) {}
 }

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgEncoder.java

@@ -119,6 +119,8 @@ void write(CommandBase<?> cmd) {
       return new ClosePortalCommandCodec((CloseCursorCommand) cmd);
     } else if (cmd instanceof CloseStatementCommand) {
       return new CloseStatementCommandCodec((CloseStatementCommand) cmd);
+    } else if (cmd instanceof CopyOutCommand) {
+      return new CopyOutCommandCodec((CopyOutCommand) cmd);
     }
     throw new AssertionError();
   }

vertx-pg-client/src/test/java/io/vertx/pgclient/PgConnectionCopyTest.java

@@ -80,19 +80,4 @@ public void testSimpleQuery(TestContext ctx) {
         }));
     }));
   }
-
-  @Test
-  public void testMakeSureCopyOutProtocolIsDefined(TestContext ctx) {
-    Async async = ctx.async();
-    connector.accept(ctx.asyncAssertSuccess(conn -> {
-      conn
-        .query("copy fortune to stdout (format csv)")
-        .execute()
-        .onComplete(ctx.asyncAssertSuccess(result1 -> {
-          // nothing comes out
-          ctx.assertEquals(0, result1.size());
-          async.complete();
-        }));
-    }));
-  }
 }