From b45e0e80b578ed1c60487c85ff0e0d5d66cc796a Mon Sep 17 00:00:00 2001
From: Chris Miles <akkadius1@gmail.com>
Date: Sat, 17 Jun 2023 18:20:13 -0500
Subject: [PATCH] [Database] Add query multi statement execution support
 (#3414)

---
 common/dbcore.cpp | 90 +++++++++++++++++++++++++++++++++++++++++++++++
 common/dbcore.h   | 15 +++++++-
 2 files changed, 104 insertions(+), 1 deletion(-)

diff --git a/common/dbcore.cpp b/common/dbcore.cpp
index 773fb882f..8358e123c 100644
--- a/common/dbcore.cpp
+++ b/common/dbcore.cpp
@@ -13,6 +13,7 @@
 #include <iostream>
 #include <mysqld_error.h>
 #include <string.h>
+#include "strings.h"
 
 #ifdef _WINDOWS
 #define snprintf	_snprintf
@@ -305,3 +306,92 @@ void DBcore::SetMutex(Mutex *mutex)
 
 	DBcore::m_mutex = mutex;
 }
+
+// executes multiple statements in one query
+// do not use this in application logic
+// this was built and maintained for database migrations only
+MySQLRequestResult DBcore::QueryDatabaseMulti(const std::string &query)
+{
+	SetMultiStatementsOn();
+
+	BenchTimer timer;
+	timer.reset();
+
+	LockMutex lock(m_mutex);
+
+	// Reconnect if we are not connected before hand.
+	if (pStatus != Connected) {
+		Open();
+	}
+
+	int status = mysql_real_query(mysql, query.c_str(), query.length());
+	if (status != 0) {
+		unsigned int error_number = mysql_errno(mysql);
+
+		if (error_number == CR_SERVER_GONE_ERROR) {
+			pStatus = Error;
+		}
+
+		// error logging
+		if (mysql_errno(mysql) > 0 && query.length() > 0 && mysql_errno(mysql) != 1065) {
+			std::string error_raw   = fmt::format("{}", mysql_error(mysql));
+			std::string mysql_err   = Strings::Trim(error_raw);
+			std::string clean_query = Strings::Replace(query, "\n", "");
+			LogMySQLQuery("[{}] ({}) query [{}]", mysql_err, mysql_errno(mysql), clean_query);
+		}
+	}
+
+	auto result = MySQLRequestResult{};
+
+	int index = 0;
+
+	// there could be a query with a semicolon in the actual data, this is best effort for
+	// logging / display purposes
+	// rare that we see this when this is only used in DDL statements
+	auto pieces = Strings::Split(query, ";");
+
+	// process each statement result
+	do {
+		uint32    row_count = 0;
+		MYSQL_RES *res      = mysql_store_result(mysql);
+
+		result = MySQLRequestResult(
+			res,
+			(uint32) mysql_affected_rows(mysql),
+			row_count,
+			(uint32) mysql_field_count(mysql),
+			(uint32) mysql_insert_id(mysql)
+		);
+
+		if (pieces.size() >= index) {
+			auto piece = pieces[index];
+			LogMySQLQuery(
+				"{} -- ({} row{} affected) ({}s)",
+				piece,
+				result.RowsAffected(),
+				result.RowsAffected() == 1 ? "" : "s",
+				std::to_string(timer.elapsed())
+			);
+		}
+
+		if (res) {
+			row_count = (uint32) mysql_num_rows(res);
+			mysql_free_result(res);
+		}
+
+		// more results? -1 = no, >0 = error, 0 = yes (keep looping)
+		if ((status = mysql_next_result(mysql)) > 0) {
+			if (mysql_errno(mysql) > 0) {
+				LogMySQLError("[{}] [{}]", mysql_errno(mysql), mysql_error(mysql));
+			}
+
+			// we handle errors elsewhere
+			return result;
+		}
+		index++;
+	} while (status == 0);
+
+	SetMultiStatementsOff();
+
+	return result;
+}
diff --git a/common/dbcore.h b/common/dbcore.h
index 6da00cb62..30f5d1304 100644
--- a/common/dbcore.h
+++ b/common/dbcore.h
@@ -25,6 +25,7 @@ public:
 	eStatus GetStatus() { return pStatus; }
 	MySQLRequestResult QueryDatabase(const char *query, uint32 querylen, bool retryOnFailureOnce = true);
 	MySQLRequestResult QueryDatabase(const std::string& query, bool retryOnFailureOnce = true);
+	MySQLRequestResult QueryDatabaseMulti(const std::string &query);
 	void TransactionBegin();
 	void TransactionCommit();
 	void TransactionRollback();
@@ -77,8 +78,20 @@ private:
 	uint32 pPort;
 	bool   pSSL;
 
+	// allows multiple queries to be executed within the same query
+	// do not use this under normal operation
+	// we use this during database migrations only currently
+	void SetMultiStatementsOn()
+	{
+		mysql_set_server_option(mysql, MYSQL_OPTION_MULTI_STATEMENTS_ON);
+	}
+
+	// disables multiple statements to be executed in one query
+	void SetMultiStatementsOff()
+	{
+		mysql_set_server_option(mysql, MYSQL_OPTION_MULTI_STATEMENTS_OFF);
+	}
 };
 
 
 #endif
-