From 675971bd3c5bf37d082f33b47c13839563d581dc Mon Sep 17 00:00:00 2001
From: Johan Wikman <johan.wikman@mariadb.com>
Date: Fri, 22 Feb 2019 11:22:32 +0200
Subject: [PATCH] MXS-2345 Query classifier collects db of USE stmt

qc_get_database_names() will now return the database used
in a USE statement.
---
 .../qc_mysqlembedded/qc_mysqlembedded.cc      | 96 ++++++++++++-------
 query_classifier/qc_sqlite/qc_sqlite.cc       |  7 ++
 2 files changed, 67 insertions(+), 36 deletions(-)

diff --git a/query_classifier/qc_mysqlembedded/qc_mysqlembedded.cc b/query_classifier/qc_mysqlembedded/qc_mysqlembedded.cc
index e8e73dfc3..e83c4d2c2 100644
--- a/query_classifier/qc_mysqlembedded/qc_mysqlembedded.cc
+++ b/query_classifier/qc_mysqlembedded/qc_mysqlembedded.cc
@@ -1853,58 +1853,82 @@ int32_t qc_mysql_get_database_names(GWBUF* querybuf, char*** databasesp, int* si
         goto retblock;
     }
 
-    lex->current_select = lex->all_selects_list;
-
-    while (lex->current_select)
+    if (lex->sql_command == SQLCOM_CHANGE_DB)
     {
-        tbl = lex->current_select->table_list.first;
-
-        while (tbl)
+        if (lex->select_lex.db)
         {
-            if (lex->sql_command == SQLCOM_SHOW_FIELDS)
+            if (i >= currsz)
             {
-                // If we are describing, we want the actual table, not the information_schema.
-                if (tbl->schema_select_lex)
+                tmp = (char**) realloc(databases,
+                                       sizeof(char*) * (currsz * 2 + 1));
+
+                if (tmp == NULL)
                 {
-                    tbl = tbl->schema_select_lex->table_list.first;
+                    goto retblock;
                 }
+
+                databases = tmp;
+                currsz = currsz * 2 + 1;
             }
 
-            // The database is sometimes an empty string. So as not to return
-            // an array of empty strings, we need to check for that possibility.
-            if ((strcmp(tbl->db, "skygw_virtual") != 0) && (*tbl->db != 0))
-            {
-                if (i >= currsz)
-                {
-                    tmp = (char**) realloc(databases,
-                                           sizeof(char*) * (currsz * 2 + 1));
+            databases[i++] = strdup(lex->select_lex.db);
+        }
+    }
+    else
+    {
+        lex->current_select = lex->all_selects_list;
 
-                    if (tmp == NULL)
+        while (lex->current_select)
+        {
+            tbl = lex->current_select->table_list.first;
+
+            while (tbl)
+            {
+                if (lex->sql_command == SQLCOM_SHOW_FIELDS)
+                {
+                    // If we are describing, we want the actual table, not the information_schema.
+                    if (tbl->schema_select_lex)
                     {
-                        goto retblock;
+                        tbl = tbl->schema_select_lex->table_list.first;
+                    }
+                }
+
+                // The database is sometimes an empty string. So as not to return
+                // an array of empty strings, we need to check for that possibility.
+                if ((strcmp(tbl->db, "skygw_virtual") != 0) && (*tbl->db != 0))
+                {
+                    if (i >= currsz)
+                    {
+                        tmp = (char**) realloc(databases,
+                                               sizeof(char*) * (currsz * 2 + 1));
+
+                        if (tmp == NULL)
+                        {
+                            goto retblock;
+                        }
+
+                        databases = tmp;
+                        currsz = currsz * 2 + 1;
                     }
 
-                    databases = tmp;
-                    currsz = currsz * 2 + 1;
+                    int j = 0;
+
+                    while ((j < i) && (strcmp(tbl->db, databases[j]) != 0))
+                    {
+                        ++j;
+                    }
+
+                    if (j == i)     // Not found
+                    {
+                        databases[i++] = strdup(tbl->db);
+                    }
                 }
 
-                int j = 0;
-
-                while ((j < i) && (strcmp(tbl->db, databases[j]) != 0))
-                {
-                    ++j;
-                }
-
-                if (j == i)     // Not found
-                {
-                    databases[i++] = strdup(tbl->db);
-                }
+                tbl = tbl->next_local;
             }
 
-            tbl = tbl->next_local;
+            lex->current_select = lex->current_select->next_select_in_list();
         }
-
-        lex->current_select = lex->current_select->next_select_in_list();
     }
 
 retblock:
diff --git a/query_classifier/qc_sqlite/qc_sqlite.cc b/query_classifier/qc_sqlite/qc_sqlite.cc
index ec23c64bc..4193e824f 100644
--- a/query_classifier/qc_sqlite/qc_sqlite.cc
+++ b/query_classifier/qc_sqlite/qc_sqlite.cc
@@ -3047,6 +3047,13 @@ public:
         m_status = QC_QUERY_PARSED;
         m_type_mask = QUERY_TYPE_SESSION_WRITE;
         m_operation = QUERY_OP_CHANGE_DB;
+
+        if (should_collect(QC_COLLECT_DATABASES))
+        {
+            char* zCopy = MXS_STRNDUP_A(pToken->z, pToken->n);
+
+            m_database_names.push_back(zCopy);
+        }
     }
 
     void set_type_mask(uint32_t type_mask)