From 8a176d64aa94454ef50cbb52dd9719fd577d3037 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Markus=20M=C3=A4kel=C3=A4?= <markus.makela@mariadb.com>
Date: Tue, 9 Jul 2019 14:59:52 +0300
Subject: [PATCH] MXS-2490: Add direct execution support

Certain MariaDB connectors will use the direct execution for batching
COM_STMT_PREPARE and COM_STMT_EXECUTE execution without waiting for the
COM_STMT_PREPARE to complete. In these cases the COM_STMT_EXECUTE (and
other COM_STMT commands as well) will use the special ID 0xffffffff. When
this is detected, it should be substituted with the ID of the latest
statement that was prepared.
---
 include/maxscale/queryclassifier.hh           |  5 +++-
 maxscale-system-test/CMakeLists.txt           |  3 +++
 .../mxs2490_ps_execute_direct.cpp             | 26 +++++++++++++++++++
 server/core/queryclassifier.cc                |  8 ++++++
 4 files changed, 41 insertions(+), 1 deletion(-)
 create mode 100644 maxscale-system-test/mxs2490_ps_execute_direct.cpp

diff --git a/include/maxscale/queryclassifier.hh b/include/maxscale/queryclassifier.hh
index 78ea4fed6..c4a6426b5 100644
--- a/include/maxscale/queryclassifier.hh
+++ b/include/maxscale/queryclassifier.hh
@@ -363,7 +363,7 @@ private:
      */
     bool query_type_is_read_only(uint32_t qtype) const;
 
-    void process_routing_hints(HINT* pHints, uint32_t* target);
+    void     process_routing_hints(HINT* pHints, uint32_t* target);
     uint32_t get_route_target(uint8_t command, uint32_t qtype);
 
     MXS_SESSION* session() const
@@ -415,5 +415,8 @@ private:
     RouteInfo         m_route_info;
     bool              m_trx_is_read_only;
     bool              m_ps_continuation;
+
+    uint32_t m_prev_ps_id = 0;      /**< For direct PS execution, storest latest prepared PS ID.
+                                     * https://mariadb.com/kb/en/library/com_stmt_execute/#statement-id **/
 };
 }
diff --git a/maxscale-system-test/CMakeLists.txt b/maxscale-system-test/CMakeLists.txt
index a8ded287f..e07a028ef 100644
--- a/maxscale-system-test/CMakeLists.txt
+++ b/maxscale-system-test/CMakeLists.txt
@@ -969,6 +969,9 @@ add_test_executable(mxs2563_concurrent_slave_failure.cpp mxs2563_concurrent_slav
 # MXS-2521: COM_STMT_EXECUTE maybe return empty result
 add_test_executable(mxs2521_double_exec.cpp mxs2521_double_exec mxs2521_double_exec LABELS REPL_BACKEND readwritesplit)
 
+# MXS-2490: Direct execution doesn't work with MaxScale
+add_test_executable(mxs2490_ps_execute_direct.cpp mxs2490_ps_execute_direct replication LABELS REPL_BACKEND readwritesplit)
+
 ############################################
 # BEGIN: binlogrouter and avrorouter tests #
 ############################################
diff --git a/maxscale-system-test/mxs2490_ps_execute_direct.cpp b/maxscale-system-test/mxs2490_ps_execute_direct.cpp
new file mode 100644
index 000000000..6e5735ae5
--- /dev/null
+++ b/maxscale-system-test/mxs2490_ps_execute_direct.cpp
@@ -0,0 +1,26 @@
+/**
+ * MXS-2490: Unknown prepared statement handler (0) given to mysqld_stmt_execute
+ *
+ * See:
+ *
+ * https://mariadb.com/kb/en/library/mariadb_stmt_execute_direct/
+ * https://mariadb.com/kb/en/library/com_stmt_execute/#statement-id
+ */
+
+#include "testconnections.h"
+
+int main(int argc, char** argv)
+{
+    TestConnections test(argc, argv);
+
+    test.set_timeout(30);
+    test.maxscales->connect();
+
+    MYSQL_STMT* stmt = mysql_stmt_init(test.maxscales->conn_rwsplit[0]);
+    std::string query = "SELECT user FROM mysql.user";
+    test.expect(mariadb_stmt_execute_direct(stmt, query.c_str(), query.length()) == 0,
+                "execute_direct should work: %s", mysql_stmt_error(stmt));
+    mysql_stmt_close(stmt);
+
+    return test.global_result;
+}
diff --git a/server/core/queryclassifier.cc b/server/core/queryclassifier.cc
index 918db76fa..b8beb4481 100644
--- a/server/core/queryclassifier.cc
+++ b/server/core/queryclassifier.cc
@@ -644,6 +644,13 @@ uint32_t QueryClassifier::ps_id_internal_get(GWBUF* pBuffer)
 
     // All COM_STMT type statements store the ID in the same place
     uint32_t external_id = mysql_extract_ps_id(pBuffer);
+
+    if (external_id == 0xffffffff)
+    {
+        // "Direct execution" that refers to the latest prepared statement
+        external_id = m_prev_ps_id;
+    }
+
     auto it = m_ps_handles.find(external_id);
 
     if (it != m_ps_handles.end())
@@ -663,6 +670,7 @@ uint32_t QueryClassifier::ps_id_internal_get(GWBUF* pBuffer)
 void QueryClassifier::ps_store_response(uint32_t internal_id, GWBUF* buffer)
 {
     auto external_id = qc_mysql_extract_ps_id(buffer);
+    m_prev_ps_id = external_id;
     m_ps_handles[external_id] = internal_id;
 
     if (auto param_count = qc_extract_ps_param_count(buffer))