From 9ad710fb94d917c1ce74e46fb10ec27e3f31907b Mon Sep 17 00:00:00 2001 From: fox0430 Date: Sun, 5 Jul 2026 19:09:46 +0900 Subject: [PATCH] Add pumpUntilReady template and refactor ~15 pump loops --- async_postgres/pg_client/copy.nim | 440 ++++++++---------- async_postgres/pg_client/core.nim | 186 ++++---- async_postgres/pg_client/cursor.nim | 43 +- async_postgres/pg_client/prepared.nim | 106 ++--- .../pg_client/transaction_helpers.nim | 97 ++-- async_postgres/pg_connection/buffer_io.nim | 68 +++ async_postgres/pg_connection/simple_query.nim | 106 ++--- 7 files changed, 452 insertions(+), 594 deletions(-) diff --git a/async_postgres/pg_client/copy.nim b/async_postgres/pg_client/copy.nim index fc8c753..1008b13 100644 --- a/async_postgres/pg_client/copy.nim +++ b/async_postgres/pg_client/copy.nim @@ -38,9 +38,9 @@ proc drainToReady(conn: PgConnection) {.async.} = block drainLoop: while true: while (let opt = conn.nextMessage(); opt.isSome): - let msg = opt.get - if msg.kind == bmkReadyForQuery: - conn.txStatus = msg.txStatus + let pumpMsg = opt.get + if pumpMsg.kind == bmkReadyForQuery: + conn.txStatus = pumpMsg.txStatus if conn.state != csClosed: conn.state = csReady break drainLoop @@ -77,40 +77,32 @@ proc copyInRawImpl*( await conn.sendMsg(encodeQuery(sql)) var commandTag = "" - var queryError: ref PgQueryError + # Server-abort error detected by pollCopyInError during the CopyData send + # loop; kept separate from the pumps' own injected `queryError`. + var abortError: ref PgQueryError # Wait for CopyInResponse (or error) - block recvLoop: - while true: - while (let opt = conn.nextMessage(); opt.isSome): - let msg = opt.get - case msg.kind - of bmkCopyInResponse: - break recvLoop - of bmkCopyOutResponse: - # Wrong direction; keep draining to ReadyForQuery, then raise there. - if queryError == nil: - queryError = newException( - PgQueryError, "COPY IN got a COPY ... TO STDOUT statement; use copyOut" - ) - of bmkRowDescription, bmkCommandComplete, bmkEmptyQueryResponse: - # Not a COPY statement; keep draining to ReadyForQuery, then raise there. - if queryError == nil: - queryError = newException( - PgQueryError, "COPY IN requires a COPY ... FROM STDIN statement" - ) - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - raise queryError - return commandTag - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady: + case pumpMsg.kind + of bmkCopyInResponse: + break pumpLoop + of bmkCopyOutResponse: + # Wrong direction; keep draining to ReadyForQuery, then raise there. + if queryError == nil: + queryError = newException( + PgQueryError, "COPY IN got a COPY ... TO STDOUT statement; use copyOut" + ) + of bmkRowDescription, bmkCommandComplete, bmkEmptyQueryResponse: + # Not a COPY statement; keep draining to ReadyForQuery, then raise there. + if queryError == nil: + queryError = + newException(PgQueryError, "COPY IN requires a COPY ... FROM STDIN statement") + else: + discard + do: + if queryError != nil: + raise queryError + return commandTag # Send CopyData in batches, slicing from the input buffer, while watching for # an early server ErrorResponse so a doomed COPY stops streaming instead of @@ -128,11 +120,11 @@ proc copyInRawImpl*( if conn.sendBuf.len >= copyBatchSize: await conn.sendBufMsg() conn.sendBuf.setLen(0) - queryError = await conn.pollCopyInError(watch) - if queryError != nil: + abortError = await conn.pollCopyInError(watch) + if abortError != nil: conn.sendBuf.setLen(0) break - if queryError == nil: + if abortError == nil: # Flush remaining data + CopyDone in one send conn.sendBuf.addCopyDone() await conn.sendBufMsg() @@ -143,6 +135,14 @@ proc copyInRawImpl*( conn.abortCopyWatch(watch) raise e + if abortError != nil: + # Server aborted the COPY mid-stream (already consumed by pollCopyInError): + # the backend has left copy-in mode, so drain to ReadyForQuery and raise, + # mirroring copyInStreamImpl. (The second pump cannot surface it.) + conn.sendBuf.setLen(0) + await conn.drainToReadyBestEffort() + raise abortError + # Settle the in-flight watch read before parsing: on normal completion it # carries the CommandComplete/ReadyForQuery response; on an early error it was # already consumed (`watch.pending` is false) and the remaining bytes drain @@ -151,25 +151,14 @@ proc copyInRawImpl*( await watch.take() # Wait for CommandComplete + ReadyForQuery - block recvLoop2: - while true: - while (let opt = conn.nextMessage(); opt.isSome): - let msg = opt.get - case msg.kind - of bmkCommandComplete: - commandTag = msg.commandTag - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - raise queryError - break recvLoop2 - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady: + case pumpMsg.kind + of bmkCommandComplete: + commandTag = pumpMsg.commandTag + else: + discard + do: + discard return commandTag @@ -253,43 +242,34 @@ proc copyInStreamImpl*( await conn.sendMsg(encodeQuery(sql)) var info = CopyInInfo() - var queryError: ref PgQueryError + # Server-abort error detected by pollCopyInError during the CopyData send + # loop; kept separate from the pumps' own injected `queryError`. + var abortError: ref PgQueryError # Wait for CopyInResponse (or error) - block recvLoop: - while true: - while (let opt = conn.nextMessage(); opt.isSome): - let msg = opt.get - case msg.kind - of bmkCopyInResponse: - info.format = msg.copyFormat - info.columnFormats = msg.copyColumnFormats - break recvLoop - of bmkCopyOutResponse: - # Wrong direction; keep draining to ReadyForQuery, then raise there. - if queryError == nil: - queryError = newException( - PgQueryError, - "COPY IN got a COPY ... TO STDOUT statement; use copyOutStream", - ) - of bmkRowDescription, bmkCommandComplete, bmkEmptyQueryResponse: - # Not a COPY statement; keep draining to ReadyForQuery, then raise there. - if queryError == nil: - queryError = newException( - PgQueryError, "COPY IN requires a COPY ... FROM STDIN statement" - ) - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - raise queryError - return info - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady: + case pumpMsg.kind + of bmkCopyInResponse: + info.format = pumpMsg.copyFormat + info.columnFormats = pumpMsg.copyColumnFormats + break pumpLoop + of bmkCopyOutResponse: + # Wrong direction; keep draining to ReadyForQuery, then raise there. + if queryError == nil: + queryError = newException( + PgQueryError, "COPY IN got a COPY ... TO STDOUT statement; use copyOutStream" + ) + of bmkRowDescription, bmkCommandComplete, bmkEmptyQueryResponse: + # Not a COPY statement; keep draining to ReadyForQuery, then raise there. + if queryError == nil: + queryError = + newException(PgQueryError, "COPY IN requires a COPY ... FROM STDIN statement") + else: + discard + do: + if queryError != nil: + raise queryError + return info # Pull data from the callback and send as CopyData in batches, watching for an # early server ErrorResponse so a doomed COPY stops pulling/streaming instead @@ -323,8 +303,8 @@ proc copyInStreamImpl*( if conn.sendBuf.len >= batchThreshold: await conn.sendBufMsg() conn.sendBuf.setLen(0) - queryError = await conn.pollCopyInError(watch) - if queryError != nil: + abortError = await conn.pollCopyInError(watch) + if abortError != nil: break except CatchableError as e: # Transport failure on the send or on the background watch read @@ -335,7 +315,7 @@ proc copyInStreamImpl*( conn.abortCopyWatch(watch) raise e - if queryError != nil: + if abortError != nil: # Server aborted the COPY mid-stream (already detected; watch consumed). In # the simple-query protocol the backend has left copy-in mode and will emit # ReadyForQuery, so neither CopyDone nor CopyFail is needed — drain and raise. @@ -343,7 +323,7 @@ proc copyInStreamImpl*( # is already csClosed; surface the server's error, not that secondary failure. conn.sendBuf.setLen(0) await conn.drainToReadyBestEffort() - raise queryError + raise abortError elif callbackError != nil: # The callback raised. Flushing pending data is pointless, but before # sending CopyFail check whether the server already aborted the COPY in the @@ -353,7 +333,7 @@ proc copyInStreamImpl*( # instead of the callback's. conn.sendBuf.setLen(0) try: - queryError = await conn.pollCopyInError(watch) + abortError = await conn.pollCopyInError(watch) except CatchableError as e: # `pollCopyInError` re-raises a transport failure on the watch read. The # recv side is gone, so (like the outer handler and copyInRawImpl) abandon @@ -361,11 +341,11 @@ proc copyInStreamImpl*( # failure rather than attempting CopyFail. conn.abortCopyWatch(watch) raise e - if queryError != nil: + if abortError != nil: # Backend already left copy-in mode; drain its abort and surface the # server's error, preferring it over a secondary transport failure here. await conn.drainToReadyBestEffort() - raise queryError + raise abortError # Transport healthy and the backend still in copy-in mode: abort cleanly. try: await conn.sendMsg(encodeCopyFail(callbackError.msg)) @@ -410,25 +390,14 @@ proc copyInStreamImpl*( await watch.take() # Wait for CommandComplete + ReadyForQuery - block recvLoop2: - while true: - while (let opt = conn.nextMessage(); opt.isSome): - let msg = opt.get - case msg.kind - of bmkCommandComplete: - info.commandTag = msg.commandTag - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - raise queryError - break recvLoop2 - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady: + case pumpMsg.kind + of bmkCommandComplete: + info.commandTag = pumpMsg.commandTag + else: + discard + do: + discard return info @@ -477,66 +446,47 @@ proc copyOutImpl*(conn: PgConnection, sql: string): Future[CopyResult] {.async.} await conn.sendMsg(encodeQuery(sql)) 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: - 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: - # 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: - conn.state = csReady - if queryError != nil: - raise queryError - break recvLoop - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady: + case pumpMsg.kind + of bmkCopyOutResponse: + sawCopyOut = true + cr.format = pumpMsg.copyFormat + cr.columnFormats = pumpMsg.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(pumpMsg.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: + if sawCopyOut: + cr.commandTag = pumpMsg.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") + else: + discard + do: + discard return cr @@ -576,87 +526,67 @@ proc copyOutStreamImpl*( await conn.sendMsg(encodeQuery(sql)) 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: - # Cancellation tears the operation down; do not run the recovery - # drain (more I/O on a future being unwound). The COPY OUT is left - # mid-stream so the protocol is out of sync — invalidate the - # connection (the next caller must reconnect rather than see a - # misleading busy state) and propagate the cancellation as-is. - conn.state = csClosed - raise e - except CatchableError as e: - # The callback raised. Drain remaining messages until ReadyForQuery - # to keep the protocol in sync, then re-raise the callback error. If - # the transport dies during the drain, `fillRecvBuf` has already - # invalidated the connection (csClosed); surface the original - # callback error rather than that secondary transport failure. - await conn.drainToReadyBestEffort() - 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: - 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: - # 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: - conn.state = csReady - if queryError != nil: - raise queryError - break recvLoop - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady: + case pumpMsg.kind + of bmkCopyOutResponse: + sawCopyOut = true + info.format = pumpMsg.copyFormat + info.columnFormats = pumpMsg.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(pumpMsg.copyData) + except CancelledError as e: + # Cancellation tears the operation down; do not run the recovery + # drain (more I/O on a future being unwound). The COPY OUT is left + # mid-stream so the protocol is out of sync — invalidate the + # connection (the next caller must reconnect rather than see a + # misleading busy state) and propagate the cancellation as-is. + conn.state = csClosed + raise e + except CatchableError as e: + # The callback raised. Drain remaining messages until ReadyForQuery + # to keep the protocol in sync, then re-raise the callback error. If + # the transport dies during the drain, `fillRecvBuf` has already + # invalidated the connection (csClosed); surface the original + # callback error rather than that secondary transport failure. + await conn.drainToReadyBestEffort() + 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: + if sawCopyOut: + info.commandTag = pumpMsg.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") + else: + discard + do: + discard return info diff --git a/async_postgres/pg_client/core.nim b/async_postgres/pg_client/core.nim index 651ffc2..5d5e1ea 100644 --- a/async_postgres/pg_client/core.nim +++ b/async_postgres/pg_client/core.nim @@ -442,7 +442,6 @@ template queryRecvLoop*( cachedColOids: seq[int32], qr: var QueryResult, ) = - var queryError: ref PgQueryError var cachedParamOids: seq[int32] if cacheHit: @@ -460,60 +459,45 @@ template queryRecvLoop*( qr.data = newRowData(int16(qr.fields.len), colFmts, cachedColOids) qr.data.fields = qr.fields - block recvLoop: - while true: - while (let opt = conn.nextMessage(qr.data, addr qr.rowCount); opt.isSome): - let msg = opt.get - case msg.kind - of bmkParseComplete, bmkBindComplete, bmkCloseComplete: - discard - of bmkParameterDescription: - if cacheMiss: - cachedParamOids = msg.paramTypeOids - of bmkRowDescription: - var fields = msg.fields - var cf: seq[int16] - var co: seq[int32] - if cacheMiss: - cachedFields = msg.fields - if resultFormats.len > 0: - cf = deriveColFmts(resultFormats, fields.len) - co = newSeq[int32](fields.len) - for i in 0 ..< fields.len: - co[i] = fields[i].typeOid - fields[i].formatCode = cf[i] - qr.fields = fields - qr.data = newRowData(int16(qr.fields.len), cf, co) - qr.data.fields = qr.fields - of bmkNoData: - discard - of bmkCommandComplete: - qr.commandTag = msg.commandTag - of bmkEmptyQueryResponse: - discard - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - # csClosed is terminal — an orphan loop after `invalidateOnTimeout` - # must not resurrect the connection. - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - if cacheHit and queryError.sqlState in StmtCacheInvalidatingStates: - conn.removeStmtCache(sql) - raise queryError - if cacheMiss: - conn.addStmtCache( - sql, - CachedStmt( - name: stmtName, fields: cachedFields, paramOids: cachedParamOids - ), - ) - break recvLoop - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady(qr.data, addr qr.rowCount): + case pumpMsg.kind + of bmkParseComplete, bmkBindComplete, bmkCloseComplete: + discard + of bmkParameterDescription: + if cacheMiss: + cachedParamOids = pumpMsg.paramTypeOids + of bmkRowDescription: + var fields = pumpMsg.fields + var cf: seq[int16] + var co: seq[int32] + if cacheMiss: + cachedFields = pumpMsg.fields + if resultFormats.len > 0: + cf = deriveColFmts(resultFormats, fields.len) + co = newSeq[int32](fields.len) + for i in 0 ..< fields.len: + co[i] = fields[i].typeOid + fields[i].formatCode = cf[i] + qr.fields = fields + qr.data = newRowData(int16(qr.fields.len), cf, co) + qr.data.fields = qr.fields + of bmkNoData: + discard + of bmkCommandComplete: + qr.commandTag = pumpMsg.commandTag + of bmkEmptyQueryResponse: + discard + else: + discard + do: + if queryError != nil: + if cacheHit and queryError.sqlState in StmtCacheInvalidatingStates: + conn.removeStmtCache(sql) + elif cacheMiss: + conn.addStmtCache( + sql, + CachedStmt(name: stmtName, fields: cachedFields, paramOids: cachedParamOids), + ) template queryEachRecvLoop*( conn: PgConnection, @@ -578,30 +562,30 @@ template queryEachRecvLoop*( rd.buf.setLen(0) rd.cellIndex.setLen(0) continue - let msg = res.message - case msg.kind + let pumpMsg = res.message + case pumpMsg.kind of bmkNotificationResponse: - conn.dispatchNotification(msg) + conn.dispatchNotification(pumpMsg) continue of bmkNoticeResponse: - conn.dispatchNotice(msg) + conn.dispatchNotice(pumpMsg) continue of bmkParameterStatus: # Parsed directly from recvBuf, so not recorded by `nextMessage`; # mirror buffer_io to keep serverParams current (e.g. SET, promotion). - conn.serverParams[msg.paramName] = msg.paramValue + conn.serverParams[pumpMsg.paramName] = pumpMsg.paramValue continue of bmkParseComplete, bmkBindComplete, bmkCloseComplete: discard of bmkParameterDescription: if cacheMiss: - cachedParamOids = msg.paramTypeOids + cachedParamOids = pumpMsg.paramTypeOids of bmkRowDescription: - var fields = msg.fields + var fields = pumpMsg.fields var cf: seq[int16] var co: seq[int32] if cacheMiss: - cachedFields = msg.fields + cachedFields = pumpMsg.fields if resultFormats.len > 0: cf = deriveColFmts(resultFormats, fields.len) co = newSeq[int32](fields.len) @@ -617,10 +601,10 @@ template queryEachRecvLoop*( of bmkEmptyQueryResponse: discard of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) + queryError = newPgQueryError(pumpMsg.errorFields) of bmkReadyForQuery: conn.recvBufStart = pos - conn.txStatus = msg.txStatus + conn.txStatus = pumpMsg.txStatus if conn.state != csClosed: conn.state = csReady if callbackError != nil: @@ -653,49 +637,35 @@ template execRecvLoop*( ## path: discards `DataRow`s (exec callers don't need rows) and exposes only ## the `CommandComplete` tag via the `commandTag` out-parameter. Shared by ## `execImpl` (both overloads), `execInlineImpl`, and `execDirectRunImpl`. - var queryError: ref PgQueryError var cachedFields: seq[FieldDescription] var cachedParamOids: seq[int32] - block recvLoop: - while true: - while (let opt = conn.nextMessage(); opt.isSome): - let msg = opt.get - case msg.kind - of bmkParseComplete, bmkBindComplete, bmkCloseComplete: - discard - of bmkParameterDescription: - if cacheMiss: - cachedParamOids = msg.paramTypeOids - of bmkRowDescription: - if cacheMiss: - cachedFields = msg.fields - of bmkNoData: - discard - of bmkDataRow: - discard - of bmkCommandComplete: - commandTag = msg.commandTag - of bmkEmptyQueryResponse: - discard - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - if cacheHit and queryError.sqlState in StmtCacheInvalidatingStates: - conn.removeStmtCache(sql) - raise queryError - if cacheMiss: - conn.addStmtCache( - sql, - CachedStmt( - name: stmtName, fields: cachedFields, paramOids: cachedParamOids - ), - ) - break recvLoop - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady: + case pumpMsg.kind + of bmkParseComplete, bmkBindComplete, bmkCloseComplete: + discard + of bmkParameterDescription: + if cacheMiss: + cachedParamOids = pumpMsg.paramTypeOids + of bmkRowDescription: + if cacheMiss: + cachedFields = pumpMsg.fields + of bmkNoData: + discard + of bmkDataRow: + discard + of bmkCommandComplete: + commandTag = pumpMsg.commandTag + of bmkEmptyQueryResponse: + discard + else: + discard + do: + if queryError != nil: + if cacheHit and queryError.sqlState in StmtCacheInvalidatingStates: + conn.removeStmtCache(sql) + elif cacheMiss: + conn.addStmtCache( + sql, + CachedStmt(name: stmtName, fields: cachedFields, paramOids: cachedParamOids), + ) diff --git a/async_postgres/pg_client/cursor.nim b/async_postgres/pg_client/cursor.nim index 6bc7e3b..83a874b 100644 --- a/async_postgres/pg_client/cursor.nim +++ b/async_postgres/pg_client/cursor.nim @@ -63,12 +63,12 @@ proc openCursorImpl( opt.isSome ) : - let msg = opt.get - case msg.kind + let pumpMsg = opt.get + case pumpMsg.kind of bmkParseComplete, bmkBindComplete: discard of bmkRowDescription: - cursor.fields = msg.fields + cursor.fields = pumpMsg.fields # Describe(Portal) runs after Bind, so the server reports the bound # result formats. Mirror the query path: derive per-column format # codes and type OIDs so binary DataRows are decoded as binary @@ -85,7 +85,7 @@ proc openCursorImpl( cursor.fields[i].formatCode = resultFormats[i] cursor.colFormats[i] = resultFormats[i] cursor.bufferedData = - newRowData(int16(msg.fields.len), cursor.colFormats, cursor.colTypeOids) + newRowData(int16(pumpMsg.fields.len), cursor.colFormats, cursor.colTypeOids) cursor.bufferedData.fields = cursor.fields of bmkNoData: discard @@ -110,7 +110,7 @@ proc openCursorImpl( await conn.fillRecvBuf() break recvLoop of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) + queryError = newPgQueryError(pumpMsg.errorFields) # Drain until ReadyForQuery await conn.sendMsg(encodeSync()) block errDrain: @@ -144,8 +144,8 @@ proc fetchNextImpl(cursor: Cursor): Future[seq[Row]] {.async.} = block recvLoop: while true: while (let opt = conn.nextMessage(rd, addr rowCount); opt.isSome): - let msg = opt.get - case msg.kind + let pumpMsg = opt.get + case pumpMsg.kind of bmkPortalSuspended: break recvLoop of bmkCommandComplete: @@ -172,7 +172,7 @@ proc fetchNextImpl(cursor: Cursor): Future[seq[Row]] {.async.} = await conn.fillRecvBuf() break recvLoop of bmkErrorResponse: - let queryError = newPgQueryError(msg.errorFields) + let queryError = newPgQueryError(pumpMsg.errorFields) await conn.sendMsg(encodeSync()) block errDrain: while true: @@ -243,27 +243,12 @@ proc closeCursorImpl(cursor: Cursor): Future[void] {.async.} = batch.addSync() await conn.sendMsg(batch) - var queryError: ref PgQueryError - - block recvLoop: - while true: - while (let opt = conn.nextMessage(); opt.isSome): - let msg = opt.get - case msg.kind - of bmkCloseComplete: - discard - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - raise queryError - break recvLoop - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady: + case pumpMsg.kind + of bmkCloseComplete: discard + else: discard + do: + discard cursor.exhausted = true diff --git a/async_postgres/pg_client/prepared.nim b/async_postgres/pg_client/prepared.nim index c12549d..22e2d1a 100644 --- a/async_postgres/pg_client/prepared.nim +++ b/async_postgres/pg_client/prepared.nim @@ -38,33 +38,21 @@ proc prepareImpl*( await conn.sendMsg(batch) var stmt = PreparedStatement(conn: conn, name: name, sql: sql) - var queryError: ref PgQueryError - - block recvLoop: - while true: - while (let opt = conn.nextMessage(); opt.isSome): - let msg = opt.get - case msg.kind - of bmkParseComplete: - discard - of bmkParameterDescription: - stmt.paramOids = msg.paramTypeOids - of bmkRowDescription: - stmt.fields = msg.fields - of bmkNoData: - discard - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - raise queryError - break recvLoop - else: - discard - await conn.fillRecvBuf() + + conn.pumpUntilReady: + case pumpMsg.kind + of bmkParseComplete: + discard + of bmkParameterDescription: + stmt.paramOids = pumpMsg.paramTypeOids + of bmkRowDescription: + stmt.fields = pumpMsg.fields + of bmkNoData: + discard + else: + discard + do: + discard return stmt @@ -133,31 +121,18 @@ proc executeImpl*( colOids[i] = qr.fields[i].typeOid qr.data = newRowData(int16(qr.fields.len), colFmts, colOids) qr.data.fields = qr.fields - var queryError: ref PgQueryError - - block recvLoop: - while true: - while (let opt = conn.nextMessage(qr.data, addr qr.rowCount); opt.isSome): - let msg = opt.get - case msg.kind - of bmkBindComplete: - discard - of bmkCommandComplete: - qr.commandTag = msg.commandTag - of bmkEmptyQueryResponse: - discard - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - raise queryError - break recvLoop - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady(qr.data, addr qr.rowCount): + case pumpMsg.kind + of bmkBindComplete: + discard + of bmkCommandComplete: + qr.commandTag = pumpMsg.commandTag + of bmkEmptyQueryResponse: + discard + else: + discard + do: + discard return qr @@ -204,27 +179,12 @@ proc closeImpl*(stmt: PreparedStatement): Future[void] {.async.} = batch.addSync() await conn.sendMsg(batch) - var queryError: ref PgQueryError - - block recvLoop: - while true: - while (let opt = conn.nextMessage(); opt.isSome): - let msg = opt.get - case msg.kind - of bmkCloseComplete: - discard - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - raise queryError - break recvLoop - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady: + case pumpMsg.kind + of bmkCloseComplete: discard + else: discard + do: + discard proc close*( stmt: PreparedStatement, timeout: Duration = ZeroDuration diff --git a/async_postgres/pg_client/transaction_helpers.nim b/async_postgres/pg_client/transaction_helpers.nim index 4b630a2..18540b4 100644 --- a/async_postgres/pg_client/transaction_helpers.nim +++ b/async_postgres/pg_client/transaction_helpers.nim @@ -46,63 +46,50 @@ proc queryInTransactionImpl( var qr = QueryResult() var phase = 0 - var queryError: ref PgQueryError - block recvLoop: - while true: - while (let opt = conn.nextMessage(qr.data, addr qr.rowCount); opt.isSome): - let msg = opt.get - case msg.kind - of bmkParseComplete, bmkBindComplete: - discard - of bmkRowDescription: - var fields = msg.fields - var cf: seq[int16] - var co: seq[int32] - if resultFormats.len > 0: - cf = deriveColFmts(resultFormats, fields.len) - co = newSeq[int32](fields.len) - for i in 0 ..< fields.len: - co[i] = fields[i].typeOid - fields[i].formatCode = cf[i] - qr.fields = fields - qr.data = newRowData(int16(qr.fields.len), cf, co) - qr.data.fields = qr.fields - of bmkNoData: - discard - of bmkEmptyQueryResponse: - # Empty/comment-only user SQL yields EmptyQueryResponse instead of - # CommandComplete; advance the phase anyway so the trailing COMMIT's - # CommandComplete isn't captured as the user statement's tag. - inc phase - of bmkCommandComplete: - if phase == 1: - qr.commandTag = msg.commandTag - inc phase - of bmkErrorResponse: - if queryError == nil: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - # ROLLBACK a failed transaction, swallowing any failure so it cannot - # mask the query error; the outer wait(timeout) bounds the cleanup. - if msg.txStatus == tsInFailedTransaction: - try: - discard await conn.simpleExec("ROLLBACK") - except CancelledError as e: - # Don't swallow cancellation (e.g. the outer wait(timeout) - # cancelling this future under chronos) — propagate it. - raise e - except CatchableError: - discard - raise queryError - break recvLoop - else: + conn.pumpUntilReady(qr.data, addr qr.rowCount): + case pumpMsg.kind + of bmkParseComplete, bmkBindComplete: + discard + of bmkRowDescription: + var fields = pumpMsg.fields + var cf: seq[int16] + var co: seq[int32] + if resultFormats.len > 0: + cf = deriveColFmts(resultFormats, fields.len) + co = newSeq[int32](fields.len) + for i in 0 ..< fields.len: + co[i] = fields[i].typeOid + fields[i].formatCode = cf[i] + qr.fields = fields + qr.data = newRowData(int16(qr.fields.len), cf, co) + qr.data.fields = qr.fields + of bmkNoData: + discard + of bmkEmptyQueryResponse: + # Empty/comment-only user SQL yields EmptyQueryResponse instead of + # CommandComplete; advance the phase anyway so the trailing COMMIT's + # CommandComplete isn't captured as the user statement's tag. + inc phase + of bmkCommandComplete: + if phase == 1: + qr.commandTag = pumpMsg.commandTag + inc phase + else: + discard + do: + if queryError != nil: + # ROLLBACK a failed transaction, swallowing any failure so it cannot + # mask the query error; the outer wait(timeout) bounds the cleanup. + if conn.txStatus == tsInFailedTransaction: + try: + discard await conn.simpleExec("ROLLBACK") + except CancelledError as e: + # Don't swallow cancellation (e.g. the outer wait(timeout) + # cancelling this future under chronos) — propagate it. + raise e + except CatchableError: discard - await conn.fillRecvBuf() return qr diff --git a/async_postgres/pg_connection/buffer_io.nim b/async_postgres/pg_connection/buffer_io.nim index 65a2479..b446585 100644 --- a/async_postgres/pg_connection/buffer_io.nim +++ b/async_postgres/pg_connection/buffer_io.nim @@ -361,6 +361,74 @@ proc recvMessage*( return opt.get await conn.fillRecvBuf(timeout) +template pumpUntilReady*( + conn: PgConnection, + resultData: untyped, + rowCountPtr: untyped, + body: untyped, + readyBody: untyped, +) {.dirty.} = + ## Generic protocol pump loop. `pumpMsg` and `queryError` are accessible + ## in `body` and `readyBody`. `DataRow` messages are decoded in-place into + ## `resultData` and counted through `rowCountPtr` by `nextMessage`, so they + ## never surface in `body`. + ## + ## The loop is spelled out in both overloads rather than one forwarding to + ## the other: `{.dirty.}` injection only reaches `body`/`readyBody` across a + ## single template boundary, so forwarding would leave their references to + ## `pumpMsg`/`queryError` undeclared. A single template with defaulted + ## `resultData`/`rowCountPtr` fails the same way — typed params ahead of the + ## untyped bodies suppress the injection (Nim 2.2.x). + block pumpLoop: + # Declared inside the block so two pumps in one proc scope (e.g. copy.nim's + # main loop plus its recvLoop2) don't collide on these dirty-injected names. + var queryError: ref PgQueryError + var pumpMsg: BackendMessage + while true: + while (let opt = conn.nextMessage(resultData, rowCountPtr); opt.isSome): + pumpMsg = opt.get + if pumpMsg.kind == bmkErrorResponse: + if queryError == nil: + queryError = newPgQueryError(pumpMsg.errorFields) + elif pumpMsg.kind == bmkReadyForQuery: + conn.txStatus = pumpMsg.txStatus + if conn.state != csClosed: + conn.state = csReady + readyBody + if queryError != nil: + raise queryError + break pumpLoop + else: + body + await conn.fillRecvBuf() + +template pumpUntilReady*( + conn: PgConnection, body: untyped, readyBody: untyped +) {.dirty.} = + ## Bare overload for callers that decode `DataRow` inside `body` (or ignore + ## it) instead of accumulating into a `RowData`. Body mirrors the data + ## overload above with `nextMessage()` in place of the accumulating call. + block pumpLoop: + var queryError: ref PgQueryError + var pumpMsg: BackendMessage + while true: + while (let opt = conn.nextMessage(); opt.isSome): + pumpMsg = opt.get + if pumpMsg.kind == bmkErrorResponse: + if queryError == nil: + queryError = newPgQueryError(pumpMsg.errorFields) + elif pumpMsg.kind == bmkReadyForQuery: + conn.txStatus = pumpMsg.txStatus + if conn.state != csClosed: + conn.state = csReady + readyBody + if queryError != nil: + raise queryError + break pumpLoop + else: + body + await conn.fillRecvBuf() + # Non-blocking receive watch # # During an otherwise send-only phase (COPY IN) the client must keep streaming diff --git a/async_postgres/pg_connection/simple_query.nim b/async_postgres/pg_connection/simple_query.nim index 281a610..619f844 100644 --- a/async_postgres/pg_connection/simple_query.nim +++ b/async_postgres/pg_connection/simple_query.nim @@ -96,38 +96,22 @@ proc simpleQueryImpl*( var results: seq[QueryResult] var current = QueryResult() - var queryError: ref PgQueryError - - block recvLoop: - while true: - while ( - let opt = conn.nextMessage(current.data, addr current.rowCount) - opt.isSome - ) - : - let msg = opt.get - case msg.kind - of bmkRowDescription: - current = - QueryResult(fields: msg.fields, data: newRowData(int16(msg.fields.len))) - of bmkCommandComplete: - current.commandTag = msg.commandTag - results.add(current) - current = QueryResult() - of bmkEmptyQueryResponse: - results.add(QueryResult()) - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - raise queryError - break recvLoop - else: - discard - await conn.fillRecvBuf() + + conn.pumpUntilReady(current.data, addr current.rowCount): + case pumpMsg.kind + of bmkRowDescription: + current = + QueryResult(fields: pumpMsg.fields, data: newRowData(int16(pumpMsg.fields.len))) + of bmkCommandComplete: + current.commandTag = pumpMsg.commandTag + results.add(current) + current = QueryResult() + of bmkEmptyQueryResponse: + results.add(QueryResult()) + else: + discard + do: + discard return results @@ -136,28 +120,16 @@ proc simpleExecImpl*(conn: PgConnection, sql: string): Future[string] {.async.} conn.state = csBusy await conn.sendMsg(encodeQuery(sql)) var commandTag = "" - var queryError: ref PgQueryError - block recvLoop: - while true: - while (let opt = conn.nextMessage(); opt.isSome): - let msg = opt.get - case msg.kind - of bmkCommandComplete: - commandTag = msg.commandTag - of bmkRowDescription, bmkDataRow, bmkEmptyQueryResponse: - discard - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - raise queryError - break recvLoop - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady: + case pumpMsg.kind + of bmkCommandComplete: + commandTag = pumpMsg.commandTag + of bmkRowDescription, bmkDataRow, bmkEmptyQueryResponse: + discard + else: + discard + do: + discard return commandTag # Cancellation (out-of-band CancelRequest over a separate socket) @@ -336,26 +308,12 @@ proc ping*(conn: PgConnection, timeout = ZeroDuration): Future[void] = conn.state = csBusy await conn.sendMsg(encodeQuery("")) - var queryError: ref PgQueryError - block recvLoop: - while true: - while (let opt = conn.nextMessage(); opt.isSome): - let msg = opt.get - case msg.kind - of bmkEmptyQueryResponse: - discard - of bmkErrorResponse: - queryError = newPgQueryError(msg.errorFields) - of bmkReadyForQuery: - conn.txStatus = msg.txStatus - if conn.state != csClosed: - conn.state = csReady - if queryError != nil: - raise queryError - break recvLoop - else: - discard - await conn.fillRecvBuf() + conn.pumpUntilReady: + case pumpMsg.kind + of bmkEmptyQueryResponse: discard + else: discard + do: + discard if timeout > ZeroDuration: proc withTimeout(): Future[void] {.async.} =