Skip to content

Commit

Permalink
Fix DML with Execute limit over SQL adapter (#7684)
Browse files Browse the repository at this point in the history
This PR sets row limit in `Execute` command over SQL adapter to 0 when
we are injecting `CommandComplete` tags.

### Background

For queries like this one:
```
INSERT INTO X VALUES ... RETURNING ...
```
... we count number of rows returned and add a `CommandComplete("INSERT
0 n")` where `n` is the number of rows in response.

Because of this, we must not set result row limit in `Execute` command
message, since that would change the number of reported modified rows.
  • Loading branch information
aljazerzen authored Aug 26, 2024
1 parent 3c2bc42 commit 34c69a4
Showing 1 changed file with 17 additions and 8 deletions.
25 changes: 17 additions & 8 deletions edb/server/pgcon/pgcon.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -1792,7 +1792,19 @@ cdef class PGConnection:
# on this connection, so there's nothing to DEALLOCATE.
action.frontend_only = True

if not action.is_frontend_only():
if action.is_frontend_only():
pass
elif isinstance(
action.query_unit.command_complete_tag,
(dbstate.TagCountMessages, dbstate.TagUnpackRow),
):
# when executing TagUnpackRow, don't pass the limit through
msg_buf = WriteBuffer.new_message(b'E')
msg_buf.write_bytestring(action.portal_name)
msg_buf.write_int32(0)
buf.write_buffer(msg_buf.end_message())
else:
# base case
msg_buf = WriteBuffer.new_message(b'E')
msg_buf.write_bytestring(action.portal_name)
msg_buf.write_int32(action.args)
Expand Down Expand Up @@ -2054,8 +2066,7 @@ cdef class PGConnection:
if self.debug:
self.debug_print('END OF DESCRIBE', mtype)
if (
mtype == b'T' and
action.action == PGAction.DESCRIBE_STMT_ROWS and
mtype == b'T' and
isinstance(
action.query_unit.command_complete_tag,
dbstate.TagUnpackRow,
Expand All @@ -2068,10 +2079,9 @@ cdef class PGConnection:
# that col.
is_binary_format = data[-1] == 1

if is_binary_format:
# TagUnpackRow converts RowDescription into NoData
msg_buf = WriteBuffer.new_message(b'n')
buf.write_buffer(msg_buf.end_message())
# TagUnpackRow converts RowDescription into NoData
msg_buf = WriteBuffer.new_message(b'n')
buf.write_buffer(msg_buf.end_message())

elif not action.is_injected() and not (
mtype == b'n' and
Expand Down Expand Up @@ -2153,7 +2163,6 @@ cdef class PGConnection:

if (
not action.is_injected()
and mtype == b'C'
and action.query_unit.command_complete_tag
):
tag = action.query_unit.command_complete_tag
Expand Down

0 comments on commit 34c69a4

Please sign in to comment.