Abort with CopyFail and raise when copyOut is given a non-COPY-TO-STD…

async_postgres/pg_client/copy.nim

@@ -478,23 +478,55 @@ proc copyOutImpl*(conn: PgConnection, sql: string): Future[CopyResult] {.async.}
 
   var cr = CopyResult()
   var queryError: ref PgQueryError
+  var sawCopyOut = false
 
   block recvLoop:
     while true:
       while (let opt = conn.nextMessage(); opt.isSome):
         let msg = opt.get
         case msg.kind
         of bmkCopyOutResponse:
+          sawCopyOut = true
           cr.format = msg.copyFormat
           cr.columnFormats = msg.copyColumnFormats
+        of bmkCopyInResponse:
+          # Wrong direction: the backend waits for our CopyData, so send
+          # CopyFail before draining — a plain drain would deadlock.
+          if queryError == nil:
+            queryError = newException(
+              PgQueryError, "COPY OUT got a COPY ... FROM STDIN statement; use copyIn"
+            )
+          try:
+            await conn.sendMsg(
+              encodeCopyFail("COPY OUT got a COPY FROM STDIN statement")
+            )
+          except CatchableError as e:
+            # CopyFail undelivered: protocol out of sync, invalidate.
+            conn.state = csClosed
+            raise e
         of bmkCopyData:
           cr.data.add(msg.copyData)
         of bmkCopyDone:
           discard
+        of bmkRowDescription, bmkEmptyQueryResponse:
+          # Not a COPY statement; keep draining to ReadyForQuery, then raise there.
+          if queryError == nil:
+            queryError = newException(
+              PgQueryError, "COPY OUT requires a COPY ... TO STDOUT statement"
+            )
         of bmkCommandComplete:
-          cr.commandTag = msg.commandTag
+          if sawCopyOut:
+            cr.commandTag = msg.commandTag
+          elif queryError == nil:
+            # Completed without entering copy-out mode (e.g. INSERT/SET).
+            queryError = newException(
+              PgQueryError, "COPY OUT requires a COPY ... TO STDOUT statement"
+            )
         of bmkErrorResponse:
-          queryError = newPgQueryError(msg.errorFields)
+          # Keep an earlier misuse error; after CopyFail the server's
+          # ErrorResponse just echoes the abort.
+          if queryError == nil:
+            queryError = newPgQueryError(msg.errorFields)
         of bmkReadyForQuery:
           conn.txStatus = msg.txStatus
           if conn.state != csClosed:
@@ -514,6 +546,10 @@ proc copyOut*(
   ## Execute COPY ... TO STDOUT via simple query protocol.
   ## Collects all CopyData messages and returns them in a CopyResult.
   ## On timeout, the connection is marked csClosed (protocol out of sync).
+  ##
+  ## A statement that is not ``COPY ... TO STDOUT`` raises ``PgQueryError``;
+  ## a ``COPY ... FROM STDIN`` statement is aborted with CopyFail instead of
+  ## deadlocking.
   var cr: CopyResult
   withConnTracing(
     conn,
@@ -541,16 +577,37 @@ proc copyOutStreamImpl*(
 
   var info = CopyOutInfo()
   var queryError: ref PgQueryError
+  var sawCopyOut = false
 
   block recvLoop:
     while true:
       while (let opt = conn.nextMessage(); opt.isSome):
         let msg = opt.get
         case msg.kind
         of bmkCopyOutResponse:
+          sawCopyOut = true
           info.format = msg.copyFormat
           info.columnFormats = msg.copyColumnFormats
+        of bmkCopyInResponse:
+          # Wrong direction: the backend waits for our CopyData, so send
+          # CopyFail before draining — a plain drain would deadlock.
+          if queryError == nil:
+            queryError = newException(
+              PgQueryError,
+              "COPY OUT got a COPY ... FROM STDIN statement; use copyInStream",
+            )
+          try:
+            await conn.sendMsg(
+              encodeCopyFail("COPY OUT got a COPY FROM STDIN statement")
+            )
+          except CatchableError as e:
+            # CopyFail undelivered: protocol out of sync, invalidate.
+            conn.state = csClosed
+            raise e
         of bmkCopyData:
+          if queryError != nil:
+            # Misuse error pending; don't feed a doomed callback.
+            continue
           try:
             await callback(msg.copyData)
           except CancelledError as e:
@@ -571,10 +628,25 @@ proc copyOutStreamImpl*(
             raise e
         of bmkCopyDone:
           discard
+        of bmkRowDescription, bmkEmptyQueryResponse:
+          # Not a COPY statement; keep draining to ReadyForQuery, then raise there.
+          if queryError == nil:
+            queryError = newException(
+              PgQueryError, "COPY OUT requires a COPY ... TO STDOUT statement"
+            )
         of bmkCommandComplete:
-          info.commandTag = msg.commandTag
+          if sawCopyOut:
+            info.commandTag = msg.commandTag
+          elif queryError == nil:
+            # Completed without entering copy-out mode (e.g. INSERT/SET).
+            queryError = newException(
+              PgQueryError, "COPY OUT requires a COPY ... TO STDOUT statement"
+            )
         of bmkErrorResponse:
-          queryError = newPgQueryError(msg.errorFields)
+          # Keep an earlier misuse error; after CopyFail the server's
+          # ErrorResponse just echoes the abort.
+          if queryError == nil:
+            queryError = newPgQueryError(msg.errorFields)
         of bmkReadyForQuery:
           conn.txStatus = msg.txStatus
           if conn.state != csClosed:
@@ -607,6 +679,10 @@ proc copyOutStream*(
   ## data before the error surfaces — cancel the query out-of-band (`cancel`)
   ## if you need to abort a large COPY OUT promptly. On timeout the connection
   ## is instead marked csClosed (protocol out of sync).
+  ##
+  ## A statement that is not ``COPY ... TO STDOUT`` raises ``PgQueryError``
+  ## without invoking ``callback``; ``COPY ... FROM STDIN`` is aborted with
+  ## CopyFail instead of deadlocking.
   var info: CopyOutInfo
   withConnTracing(
     conn,

tests/test_e2e_copy.nim

@@ -140,6 +140,113 @@ suite "E2E: COPY Protocol":
 
     waitFor t()
 
+  test "copyOut with COPY FROM STDIN raises PgQueryError":
+    proc t() {.async.} =
+      let conn = await connect(plainConfig())
+      discard await conn.exec("DROP TABLE IF EXISTS test_copyout_wrongdir")
+      discard await conn.exec("CREATE TABLE test_copyout_wrongdir (id int)")
+
+      var raised = false
+      try:
+        # Would deadlock without the CopyFail abort; timeout is a regression guard
+        discard await conn.copyOut(
+          "COPY test_copyout_wrongdir FROM STDIN", timeout = seconds(5)
+        )
+      except PgQueryError as e:
+        raised = true
+        doAssert "copyIn" in e.msg
+      doAssert raised
+      doAssert conn.state == csReady
+
+      # Connection stays usable and no rows were inserted
+      let res = await conn.query("SELECT count(*) FROM test_copyout_wrongdir")
+      doAssert res.rows[0].getStr(0) == "0"
+
+      discard await conn.exec("DROP TABLE test_copyout_wrongdir")
+      await conn.close()
+
+    waitFor t()
+
+  test "copyOut with non-COPY statement raises PgQueryError":
+    proc t() {.async.} =
+      let conn = await connect(plainConfig())
+
+      # Row-returning statement (RowDescription)
+      var raised = false
+      try:
+        discard await conn.copyOut("SELECT 1")
+      except PgQueryError:
+        raised = true
+      doAssert raised
+      doAssert conn.state == csReady
+
+      # Non-row statement (CommandComplete only)
+      raised = false
+      try:
+        discard await conn.copyOut("SET search_path TO public")
+      except PgQueryError:
+        raised = true
+      doAssert raised
+      doAssert conn.state == csReady
+
+      # Connection stays usable
+      let res = await conn.query("SELECT 1")
+      doAssert res.rows[0].getStr(0) == "1"
+
+      await conn.close()
+
+    waitFor t()
+
+  test "copyOutStream with COPY FROM STDIN raises without invoking callback":
+    proc t() {.async.} =
+      let conn = await connect(plainConfig())
+      discard await conn.exec("DROP TABLE IF EXISTS test_copyout_stream_wrongdir")
+      discard await conn.exec("CREATE TABLE test_copyout_stream_wrongdir (id int)")
+
+      var callCount = 0
+      let cb = makeCopyOutCallback:
+        inc callCount
+      var raised = false
+      try:
+        discard await conn.copyOutStream(
+          "COPY test_copyout_stream_wrongdir FROM STDIN", cb, timeout = seconds(5)
+        )
+      except PgQueryError as e:
+        raised = true
+        doAssert "copyInStream" in e.msg
+      doAssert raised
+      doAssert callCount == 0
+      doAssert conn.state == csReady
+
+      discard await conn.exec("DROP TABLE test_copyout_stream_wrongdir")
+      await conn.close()
+
+    waitFor t()
+
+  test "copyOutStream with non-COPY statement raises without invoking callback":
+    proc t() {.async.} =
+      let conn = await connect(plainConfig())
+
+      var callCount = 0
+      let cb = makeCopyOutCallback:
+        inc callCount
+      var raised = false
+      try:
+        discard await conn.copyOutStream("SELECT 1", cb)
+      except PgQueryError:
+        raised = true
+      doAssert raised
+      doAssert callCount == 0
+      doAssert conn.state == csReady
+
+      # Connection stays usable
+      let res = await conn.query("SELECT 1")
+      doAssert res.rows[0].getStr(0) == "1"
+
+      await conn.close()
+
+    waitFor t()
+
   test "copyIn empty data":
     proc t() {.async.} =
       let conn = await connect(plainConfig())