From f59dfed6d0ae7230594178cbe216df1d65b78bf9 Mon Sep 17 00:00:00 2001 From: fox0430 Date: Fri, 3 Jul 2026 17:52:33 +0900 Subject: [PATCH] Abort with CopyFail and raise when copyOut is given a non-COPY-TO-STDOUT statement --- async_postgres/pg_client/copy.nim | 84 +++++++++++++++++++++-- tests/test_e2e_copy.nim | 107 ++++++++++++++++++++++++++++++ 2 files changed, 187 insertions(+), 4 deletions(-) diff --git a/async_postgres/pg_client/copy.nim b/async_postgres/pg_client/copy.nim index 8154017..fc8c753 100644 --- a/async_postgres/pg_client/copy.nim +++ b/async_postgres/pg_client/copy.nim @@ -478,6 +478,7 @@ proc copyOutImpl*(conn: PgConnection, sql: string): Future[CopyResult] {.async.} var cr = CopyResult() var queryError: ref PgQueryError + var sawCopyOut = false block recvLoop: while true: @@ -485,16 +486,47 @@ proc copyOutImpl*(conn: PgConnection, sql: string): Future[CopyResult] {.async.} 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,6 +577,7 @@ proc copyOutStreamImpl*( var info = CopyOutInfo() var queryError: ref PgQueryError + var sawCopyOut = false block recvLoop: while true: @@ -548,9 +585,29 @@ proc copyOutStreamImpl*( 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, diff --git a/tests/test_e2e_copy.nim b/tests/test_e2e_copy.nim index 4447fff..fa619a5 100644 --- a/tests/test_e2e_copy.nim +++ b/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())