BUG26376334: Fix prepared statements in MySQL 8.0

In MySQL 8.0, all prepared statements hang on calling `fetchall` on a cursor.
This happens because the server expects a CMD_STMT_FETCH command from the
client to send the results of the prepared query.

This patch adds new flags to the `MySQLCursorPrepared` class to check for
`SERVER_STATUS_CURSOR_EXISTS` and `SERVER_STATUS_LAST_ROW_SENT` flags in the
EOF/OK packets and then send `CMD_STMT_FETCH` command accordingly.

Tests failing with MySQL 8.0 will now pass with this patch.

Author: amitab

lib/mysql/connector/connection.py

@@ -933,6 +933,17 @@ def _handle_binary_result(self, packet):
         eof = self._handle_eof(self._socket.recv())
         return (column_count, columns, eof)
 
+    def cmd_stmt_fetch(self, statement_id, rows=1):
+        """Fetch a MySQL statement Result Set
+
+        This method will send the FETCH command to MySQL together with the
+        given statement id and the number of rows to fetch.
+        """
+        packet = self._protocol.make_stmt_fetch(statement_id, rows)
+        self.unread_result = False
+        self._send_cmd(ServerCmd.STMT_FETCH, packet, expect_response=False)
+        self.unread_result = True
+
     def cmd_stmt_prepare(self, statement):
         """Prepare a MySQL statement
 

lib/mysql/connector/constants.py

@@ -472,6 +472,11 @@ class ServerFlag(_Flags):
     STATUS_LAST_ROW_SENT = 1 << 7
     STATUS_DB_DROPPED = 1 << 8
     STATUS_NO_BACKSLASH_ESCAPES = 1 << 9
+    SERVER_STATUS_METADATA_CHANGED = 1 << 10
+    SERVER_QUERY_WAS_SLOW = 1 << 11
+    SERVER_PS_OUT_PARAMS = 1 << 12
+    SERVER_STATUS_IN_TRANS_READONLY = 1 << 13
+    SERVER_SESSION_STATE_CHANGED = 1 << 14
 
     desc = {
         'SERVER_STATUS_IN_TRANS': (1 << 0,
@@ -483,10 +488,31 @@ class ServerFlag(_Flags):
                                        'next query exists'),
         'SERVER_QUERY_NO_GOOD_INDEX_USED': (1 << 4, ''),
         'SERVER_QUERY_NO_INDEX_USED': (1 << 5, ''),
-        'SERVER_STATUS_CURSOR_EXISTS': (1 << 6, ''),
-        'SERVER_STATUS_LAST_ROW_SENT': (1 << 7, ''),
+        'SERVER_STATUS_CURSOR_EXISTS': (1 << 6,
+                                        'Set when server opened a read-only '
+                                        'non-scrollable cursor for a query.'),
+        'SERVER_STATUS_LAST_ROW_SENT': (1 << 7,
+                                        'Set when a read-only cursor is '
+                                        'exhausted'),
         'SERVER_STATUS_DB_DROPPED': (1 << 8, 'A database was dropped'),
         'SERVER_STATUS_NO_BACKSLASH_ESCAPES': (1 << 9, ''),
+        'SERVER_STATUS_METADATA_CHANGED': (1024,
+                                           'Set if after a prepared statement '
+                                           'reprepare we discovered that the '
+                                           'new statement returns a different '
+                                           'number of result set columns.'),
+        'SERVER_QUERY_WAS_SLOW': (2048, ''),
+        'SERVER_PS_OUT_PARAMS': (4096,
+                                 'To mark ResultSet containing output '
+                                 'parameter values.'),
+        'SERVER_STATUS_IN_TRANS_READONLY': (8192,
+                                            'Set if multi-statement '
+                                            'transaction is a read-only '
+                                            'transaction.'),
+        'SERVER_SESSION_STATE_CHANGED': (1 << 14,
+                                         'Session state has changed on the '
+                                         'server because of the execution of '
+                                         'the last statement'),
     }
 
 

lib/mysql/connector/cursor.py

@@ -31,6 +31,7 @@
 from . import errors
 from .abstracts import MySQLCursorAbstract, NAMED_TUPLE_CACHE
 from .catch23 import PY2
+from .constants import ServerFlag
 
 SQL_COMMENT = r"\/\*.*?\*\/"
 RE_SQL_COMMENT = re.compile(
@@ -60,6 +61,7 @@
 
 ERR_NO_RESULT_TO_FETCH = "No result set to fetch from"
 
+MAX_RESULTS = 4294967295
 
 class _ParamSubstitutor(object):
     """
@@ -1080,6 +1082,36 @@ def __init__(self, connection=None):
         self._prepared = None
         self._binary = True
         self._have_result = None
+        self._last_row_sent = False
+        self._cursor_exists = False
+
+    def reset(self, free=True):
+        if self._prepared:
+            try:
+                self._connection.cmd_stmt_close(self._prepared['statement_id'])
+            except errors.Error:
+                # We tried to deallocate, but it's OK when we fail.
+                pass
+            self._prepared = None
+        self._last_row_sent = False
+        self._cursor_exists = False
+
+    def _handle_noresultset(self, res):
+        self._handle_server_status(res.get('status_flag',
+                                           res.get('server_status', 0)))
+        super(MySQLCursorPrepared, self)._handle_noresultset(res)
+
+    def _handle_server_status(self, flags):
+        """Check for SERVER_STATUS_CURSOR_EXISTS and
+           SERVER_STATUS_LAST_ROW_SENT flags set by the server.
+        """
+        self._cursor_exists = flags & ServerFlag.STATUS_CURSOR_EXISTS != 0
+        self._last_row_sent = flags & ServerFlag.STATUS_LAST_ROW_SENT != 0
+
+    def _handle_eof(self, eof):
+        self._handle_server_status(eof.get('status_flag',
+                                           eof.get('server_status', 0)))
+        super(MySQLCursorPrepared, self)._handle_eof(eof)
 
     def callproc(self, *args, **kwargs):
         """Calls a stored procedue
@@ -1094,13 +1126,7 @@ def close(self):
         This method will try to deallocate the prepared statement and close
         the cursor.
         """
-        if self._prepared:
-            try:
-                self._connection.cmd_stmt_close(self._prepared['statement_id'])
-            except errors.Error:
-                # We tried to deallocate, but it's OK when we fail.
-                pass
-            self._prepared = None
+        self.reset()
         super(MySQLCursorPrepared, self).close()
 
     def _row_to_python(self, rowdata, desc=None):
@@ -1122,6 +1148,11 @@ def _handle_result(self, res):
             self._connection.unread_result = True
             self._have_result = True
 
+            if 'status_flag' in res[2]:
+                self._handle_server_status(res[2]['status_flag'])
+            elif 'server_status' in res[2]:
+                self._handle_server_status(res[2]['server_status'])
+
     def execute(self, operation, params=(), multi=False):  # multi is unused
         """Prepare and execute a MySQL Prepared Statement
 
@@ -1199,6 +1230,8 @@ def fetchone(self):
 
         Returns a tuple or None.
         """
+        if self._cursor_exists:
+            self._connection.cmd_stmt_fetch(self._prepared['statement_id'])
         return self._fetch_row() or None
 
     def fetchmany(self, size=None):
@@ -1214,10 +1247,18 @@ def fetchmany(self, size=None):
     def fetchall(self):
         if not self._have_unread_result():
             raise errors.InterfaceError("No result set to fetch from.")
-        (rows, eof) = self._connection.get_rows(
-            binary=self._binary, columns=self.description)
+        rows = []
+        if self._nextrow[0]:
+            rows.append(self._nextrow[0])
+        while self._have_unread_result():
+            if self._cursor_exists:
+                self._connection.cmd_stmt_fetch(
+                    self._prepared['statement_id'], MAX_RESULTS)
+            (tmp, eof) = self._connection.get_rows(
+                binary=self._binary, columns=self.description)
+            rows.extend(tmp)
+            self._handle_eof(eof)
         self._rowcount = len(rows)
-        self._handle_eof(eof)
         return rows
 
 

lib/mysql/connector/protocol.py

@@ -123,6 +123,10 @@ def make_command(self, command, argument=None):
             data += argument
         return data
 
+    def make_stmt_fetch(self, statement_id, rows=1):
+        """Make a MySQL packet with Fetch Statement command"""
+        return utils.int4store(statement_id) + utils.int4store(rows)
+
     def make_change_user(self, handshake, username=None, password=None,
                          database=None, charset=33, client_flags=0,
                          ssl_enabled=False, auth_plugin=None):