TeaSpeakLibrary/src/sql/mysql/MySQL.cpp

#include <memory>

#include "MySQL.h"
#include "mysql_connection.h"
#include "src/log/LogUtils.h"
#include <pipes/misc/http.h>

#include <memory>
#include <utility>

using namespace std;
using namespace sql;
using namespace sql::mysql;

MySQLManager::MySQLManager() : SqlManager(SqlType::TYPE_MYSQL) {}
MySQLManager::~MySQLManager() {}

//Property info: https://dev.mysql.com/doc/connector-j/5.1/en/connector-j-reference-configuration-properties.html
//mysql://[host][:port]/[database][?propertyName1=propertyValue1[&propertyName2=propertyValue2]...]

#define MYSQL_PREFIX "mysql://"
inline result parse_url(const string& url, ConnectOptionsMap& connect_map) {
	string target_url;
	if(url.find(MYSQL_PREFIX) != 0) return {ERROR_MYSQL_INVLID_URL, "Missing mysql:// at begin"};
	auto index_parms = url.find('?');
	if(index_parms == string::npos) {
		target_url = "tcp://" + url.substr(strlen(MYSQL_PREFIX));
	} else {
		target_url = "tcp://" + url.substr(strlen(MYSQL_PREFIX), index_parms - strlen(MYSQL_PREFIX));
		auto parms = url.substr(index_parms + 1);
		size_t index = 0;
		do {
			auto idx = parms.find('&', index);
			auto element = parms.substr(index, idx - index);

			auto key_idx = element.find('=');
			auto key = element.substr(0, key_idx);
			auto value = element.substr(key_idx + 1);
			connect_map[key] = http::decode_url(value);
			logTrace(LOG_GENERAL, "Got mysql property {}. Value: {}", key, value);

			index = idx + 1;
		} while(index != 0);
	}

	//if(!connect_map["hostName"].get<const char*>() || strcmp(*connect_map["hostName"].get<const char*>(), ""))
	connect_map["hostName"] = target_url;
	logTrace(LOG_GENERAL, "Got mysql property {}. Value: {}", "hostName", target_url);
	return result::success;
}

result MySQLManager::connect(const std::string &url) {
	this->disconnecting = false;

	ConnectOptionsMap connect_map;
	connect_map["connections"] = "1";
	auto res = parse_url(url, connect_map);
	if(!res) return res;

    this->driver = get_driver_instance();
    if(!this->driver) return {ERROR_MYSQL_MISSING_DRIVER, "Missing driver!"};

	try {
		auto entry = connect_map["connections"];
		auto connections = std::stoll(connect_map["connections"].get<sql::SQLString>()->asStdString());
		if(connections < 1) return {ERROR_MYSQL_INVLID_PROPERTIES, "Invalid connection count"};

		for(int i = 0; i < connections; i++) {
			auto connection = unique_ptr<Connection>(this->driver->connect(connect_map));
			if(!connection) return {ERROR_MYSQL_INVLID_CONNECT, "Could not spawn new connection"};
			if(!connection->isValid()) return {ERROR_MYSQL_INVLID_CONNECT, "Could not validate connection"};
			if(connection->getSchema().length() == 0) return {ERROR_MYSQL_INVLID_CONNECT, "Missing schema!"};

			this->connections.push_back(shared_ptr<ConnectionEntry>(new ConnectionEntry(std::move(connection), false)));
		}
	} catch (sql::SQLException& ex) {
		return {ERROR_MYSQL_INVLID_CONNECT, ex.what()};
	}
    return result::success;
}

bool MySQLManager::connected() {
	lock_guard<mutex> lock(this->connections_lock);
	for(const auto& conn : this->connections)
		if(conn->used || conn->connection->isValid()) return true;
	return false;
}

result MySQLManager::disconnect() {
	lock_guard<mutex> lock(this->connections_lock);
	this->disconnecting = true;

	for(const auto& entry : this->connections)
		entry->connection->close();

	this->connections.clear();
	this->connections_condition.notify_all();

	this->disconnecting = false;
	return result::success;
}

std::shared_ptr<CommandData> MySQLManager::allocateCommandData() {
	return make_shared<MySQLCommand>();
}

std::shared_ptr<CommandData> MySQLManager::copyCommandData(std::shared_ptr<CommandData> ptr) {
	auto _new = this->allocateCommandData();
	_new->handle = ptr->handle;
	_new->lock = ptr->lock;
	_new->sql_command = ptr->sql_command;
	_new->variables = ptr->variables;

	auto __new = static_pointer_cast<MySQLCommand>(_new);
	auto __ptr = static_pointer_cast<MySQLCommand>(ptr);
	//__new->stmt = __ptr->stmt;
	return __new;
}

void prepared_statement_release(sql::PreparedStatement* stmt) {
	if(stmt) stmt->close();
	delete stmt;
};
typedef unique_ptr<sql::PreparedStatement, decltype(prepared_statement_release)*> PreparedStatementHandle;

void statement_release(sql::Statement* stmt) {
	if(stmt) stmt->close();
	delete stmt;
};
typedef unique_ptr<sql::Statement, decltype(statement_release)*> StatementHandle;

void result_release(sql::ResultSet* set) {
	if(set && !set->isClosed()) set->close();
	delete set;
}
typedef unique_ptr<sql::ResultSet, decltype(result_release)*> ResultHandle;


namespace sql {
	namespace mysql {
		bool evaluate_sql_query(string& sql, const std::vector<variable>& vars, std::vector<variable>& result) {
			char quote = 0;
			for(int index = 0; index < sql.length(); index++) {
				if(sql[index] == '\'' || sql[index] == '"' || sql[index] == '`') {
					if(quote > 0) {
						if(quote == sql[index]) quote = 0;
					} else {
						quote = sql[index];
					}
					continue;
				}
				if(quote > 0) continue;
				if(sql[index] != ':') continue;

				auto index_end = sql.find_first_not_of("0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ_", index + 1);
				if(index_end == string::npos) index_end = sql.length();

				string key = sql.substr(index, index_end - index);
				//Now we can replace it with a ?
				sql.replace(index, index_end - index, "?", 1);

				bool insert = false;
				for(const auto& e : vars)
					if(e.key() == key) {
						result.push_back(e);
						insert = true;
						break;
					}
				if(!insert)
					result.push_back(variable{});
			}
			return true;
		}

		inline bool bind_parms(const PreparedStatementHandle& stmt, const std::vector<variable>& vars) {
			uint32_t index = 1;
			for(const auto& var : vars) {
				switch (var.type()) {
					case VARTYPE_NULL:
						stmt->setNull(index, 0);
						break;
					case VARTYPE_BOOLEAN:
						stmt->setBoolean(index, var.as<bool>());
						break;
					case VARTYPE_INT:
						stmt->setInt(index, var.as<int32_t>());
						break;
					case VARTYPE_LONG:
						stmt->setInt64(index, var.as<int64_t>());
						break;
					case VARTYPE_DOUBLE:
						stmt->setDouble(index, var.as<double>());
						break;
					case VARTYPE_FLOAT:
						stmt->setDouble(index, var.as<float>());
						break;
					case VARTYPE_TEXT:
						stmt->setString(index, var.value());
						break;
					default:
						cerr << "[MySQL] Invalid var type (" << var.type() << ")" << endl;
						break;
				}
				index++;
			}
			return true;
		}
	}
}

LocalConnection::LocalConnection(MySQLManager* mgr, const std::shared_ptr<ConnectionEntry>& entry) : _mgr(mgr), _connection(entry) {
	_mgr->driver->threadInit();
	_connection->used = true;
	//logMessage(LOG_GENERAL, "Allocate local connection {} and thread {}", (void*) _connection.get(), (void*) threads::self::id());
}

LocalConnection::~LocalConnection() {
	//logMessage(LOG_GENERAL, "Deallocate local connection {} and thread {}", (void*) this->_connection.get(), (void*) threads::self::id());
	_mgr->driver->threadEnd();
	_connection->used = false;

	{
		lock_guard<mutex> lock(_mgr->connections_lock);
		_mgr->connections_condition.notify_all();
	}
}

std::unique_ptr<LocalConnection> MySQLManager::next_connection() {
	unique_ptr<LocalConnection> result;
	{
		unique_lock connections_lock(this->connections_lock);

		while(!result) {
			size_t available_connections = 0;
			for(const auto& connection : this->connections) {
				available_connections++;
				if(connection->used) continue;

				result = std::make_unique<LocalConnection>(this, connection);
				break;
			}

			if(!result) {
				if(available_connections == 0) {
					if(this->listener_disconnected)
						this->listener_disconnected(false);
					this->disconnect();
					return nullptr;
				}
				this->connections_condition.wait(connections_lock); /* wait for the next connection */
			}
		}
	}

	if(!result->_connection->connection->isValid()) {
		try {
			logError(0, "MySQL connection is invalid! Closing connection!");
			result->_connection->connection->close();
		} catch(sql::SQLException& ex) {}
	}
	if(result->_connection->connection->isClosed()) {
		logError(0, "MySQL connection was closed! Attempt reconnect!");
		try {
			if(!result->_connection->connection->reconnect()) {
				logError(0, "MySQL connection reconnect attempt failed! Dropping connection!");
				{
					lock_guard<mutex> connections_lock(this->connections_lock);
					auto index = find(this->connections.begin(), this->connections.end(), result->_connection);
					if(index != this->connections.end())
						this->connections.erase(index);
				}
				return this->next_connection();
			}
		} catch (sql::SQLException& ex) {
			logError(0, "Got an exception while reconnecting! Message: " + string(ex.what()));
			logError(0, "Dropping connection!");
			{
				lock_guard<mutex> connections_lock(this->connections_lock);
				auto index = find(this->connections.begin(), this->connections.end(), result->_connection);
				if(index != this->connections.end())
					this->connections.erase(index);
			}
			return this->next_connection();
		}
	}
	return result;
}

result MySQLManager::executeCommand(std::shared_ptr<CommandData> _ptr) {
	auto ptr = static_pointer_cast<MySQLCommand>(_ptr);
	std::lock_guard<threads::Mutex> command_lock(ptr->lock);
	auto command = ptr->sql_command;
	auto variables = ptr->variables;
	vector<variable> mapped_variables;
	if(!sql::mysql::evaluate_sql_query(command, variables, mapped_variables)) return {ptr->sql_command, -1, "Could not map sqlite vars to mysql!"};

	unique_ptr<LocalConnection> connection = this->next_connection();
	if(!connection) return {ptr->sql_command, -1, "Could not get a valid connection!"};
	try {
		PreparedStatementHandle stmt(connection->_connection->connection->prepareStatement(command), prepared_statement_release);
		//logMessage(LOG_GENERAL, "Deleting prepered statement {} and thread {}", (void*) stmt.get(), (void*) threads::self::id());
		if(!stmt) return {ptr->sql_command, -1, "Could not span a prepared statement"};

		if(!sql::mysql::bind_parms(stmt, mapped_variables)) return {ptr->sql_command, -1, "Could not bind variables!"};

		auto update_count = stmt->executeUpdate();
		if(update_count < 0)
			return {ptr->sql_command, -1, "Could not execute update. Code: " + to_string(update_count)};

		stmt.reset();
		return result::success;
	} catch (sql::SQLException& ex) {
		logError(0, "SQL Error: {}", ex.what());
		return {ptr->sql_command, -1, ex.what()};
	}
}

result MySQLManager::queryCommand(shared_ptr<CommandData> _ptr, const QueryCallback &fn) {
	auto ptr = static_pointer_cast<MySQLCommand>(_ptr);
	std::lock_guard<threads::Mutex> lock(ptr->lock);
	auto command = ptr->sql_command;
	auto variables = ptr->variables;
	vector<variable> mapped_variables;
	if(!sql::mysql::evaluate_sql_query(command, variables, mapped_variables)) return {ptr->sql_command, -1, "Could not map sqlite vars to mysql!"};

	unique_ptr<LocalConnection> connection = this->next_connection();
	if(!connection) return {ptr->sql_command, -1, "Could not get a valid connection!"};
	try {
		PreparedStatementHandle stmt(connection->_connection->connection->prepareStatement(command), prepared_statement_release);
		//logMessage(LOG_GENERAL, "Deleting prepered statement {} and thread {}", (void*) stmt.get(), (void*) threads::self::id());
		if(!sql::mysql::bind_parms(stmt, mapped_variables)) return {ptr->sql_command, -1, "Could not bind variables!"};
		ResultHandle result(stmt->executeQuery(), result_release);

		auto column_count = result->getMetaData()->getColumnCount();
		std::string columnNames[column_count];
		std::string columnValues[column_count];

		for(int index = 0; index < column_count; index++)
			columnNames[index] = result->getMetaData()->getColumnName(index + 1);

		bool userQuit = false;
		while(result->next() && !userQuit) {
			for(int index = 0; index < column_count; index++)
				columnValues[index] = result->getString(index + 1);
			if(fn(column_count, columnValues, columnNames) != 0) {
				userQuit = true;
				break;
			}
		}

		stmt.reset();
		return result::success;
	} catch (sql::SQLException& ex) {
		return {ptr->sql_command, -1, ex.what()};
	}
}

result MySQLManager::execute_raw(const std::string &command) {
	auto connection = this->next_connection();
	if(!connection)
		return {command, -1, "no connection available"};

	auto& mysql_connection = connection->_connection->connection;
	mysql_connection->clearWarnings();

	StatementHandle statement(mysql_connection->createStatement(), statement_release);
	statement->clearWarnings();

	try {
		auto result = statement->execute(command);
		if(statement->getWarnings()) {
			cerr << "Got some warnings: " << endl;
			while(statement->getWarnings()) {
				cerr << " - " << statement->getWarnings()->getMessage() << endl;
				statement->getWarnings()->setNextWarning(statement->getWarnings()->getNextWarning());
			}
		}
		if(result)
			return result::success;
		return {command, -1, "return false"};
	} catch (sql::SQLException& ex) {
		return {command, -1, ex.what()};
	}
}