From 34c69a4360648325308063802d93d523bddf9b43 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alja=C5=BE=20Mur=20Er=C5=BEen?=
 <aljazerzen@users.noreply.github.com>
Date: Mon, 26 Aug 2024 22:08:36 +0200
Subject: [PATCH] Fix DML with Execute limit over SQL adapter (#7684)

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.
---
 edb/server/pgcon/pgcon.pyx | 25 +++++++++++++++++--------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/edb/server/pgcon/pgcon.pyx b/edb/server/pgcon/pgcon.pyx
index 992b7c9caa9..8e233999a6b 100644
--- a/edb/server/pgcon/pgcon.pyx
+++ b/edb/server/pgcon/pgcon.pyx
@@ -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)
@@ -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,
@@ -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
@@ -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