nim/server/src/db/mysql_service.cpp

#include "db/mysql_service.h"

#include <cctype>
#include <exception>
#include <string>
#include <vector>

#include "base/exception.h"
#include "common/logger.h"
#include "util/file.h"

namespace im {
namespace {

std::string trim_sql(const std::string& s)
{
    size_t b = 0;
    size_t e = s.size();
    while (b < e && std::isspace(static_cast<unsigned char>(s[b]))) {
        ++b;
    }
    while (e > b && std::isspace(static_cast<unsigned char>(s[e - 1]))) {
        --e;
    }
    return s.substr(b, e - b);
}

bool starts_ci(const std::string& s, const char* prefix)
{
    const size_t n = std::char_traits<char>::length(prefix);
    if (s.size() < n) {
        return false;
    }
    for (size_t i = 0; i < n; ++i) {
        const unsigned char a = static_cast<unsigned char>(s[i]);
        const unsigned char b = static_cast<unsigned char>(prefix[i]);
        if (std::tolower(a) != std::tolower(b)) {
            return false;
        }
    }
    return true;
}

std::vector<std::string> split_sql(const std::string& sql)
{
    std::vector<std::string> out;
    std::string cur;
    bool in_str = false;
    bool in_line = false;
    bool in_block = false;
    for (size_t i = 0; i < sql.size(); ++i) {
        const char c = sql[i];
        const char n = (i + 1 < sql.size()) ? sql[i + 1] : '\0';
        if (in_line) {
            if (c == '\n') {
                in_line = false;
            }
            continue;
        }
        if (in_block) {
            if (c == '*' && n == '/') {
                in_block = false;
                ++i;
            }
            continue;
        }
        if (in_str) {
            cur.push_back(c);
            if (c == '\\' && n != '\0') {
                cur.push_back(n);
                ++i;
                continue;
            }
            if (c == '\'') {
                if (n == '\'') {
                    cur.push_back(n);
                    ++i;
                } else {
                    in_str = false;
                }
            }
            continue;
        }
        if (c == '-' && n == '-') {
            in_line = true;
            ++i;
            continue;
        }
        if (c == '/' && n == '*') {
            in_block = true;
            ++i;
            continue;
        }
        if (c == '\'') {
            in_str = true;
            cur.push_back(c);
            continue;
        }
        if (c == ';') {
            const std::string stmt = trim_sql(cur);
            if (!stmt.empty()) {
                out.push_back(stmt);
            }
            cur.clear();
            continue;
        }
        cur.push_back(c);
    }
    const std::string tail = trim_sql(cur);
    if (!tail.empty()) {
        out.push_back(tail);
    }
    return out;
}

bool exec_sql(ylib::mysql::conn* conn, const std::string& sql)
{
    auto* stmt = conn->setsql(sql);
    if (stmt == nullptr) {
        Logger::error("mysql exec prepare failed: " + sql.substr(0, 96));
        return false;
    }
    stmt->update();
    return true;
}

}  // namespace

MysqlConn::MysqlConn(ylib::mysql::conn* conn) : conn_(conn) {}

MysqlConn::MysqlConn(MysqlConn&& other) noexcept : conn_(other.conn_)
{
    other.conn_ = nullptr;
}

MysqlConn& MysqlConn::operator=(MysqlConn&& other) noexcept
{
    if (this != &other) {
        release();
        conn_ = other.conn_;
        other.conn_ = nullptr;
    }
    return *this;
}

MysqlConn::~MysqlConn()
{
    release();
}

void MysqlConn::release()
{
    if (conn_ == nullptr) {
        return;
    }
    if (conn_->pool() != nullptr) {
        conn_->pool()->recover(conn_);
    }
    conn_ = nullptr;
}

bool MysqlService::start(const MysqlConfig& cfg)
{
    ylib::mysql::mysql_conn_info info;
    info.ipaddress = cfg.host;
    info.port = cfg.port;
    info.username = cfg.user;
    info.password = cfg.password;
    info.database = cfg.database;
    info.charset = cfg.charset;

    if (!pool_.start(info, cfg.pool_size)) {
        Logger::error("mysql pool start failed");
        return false;
    }

    started_.store(true, std::memory_order_release);
    if (!ping()) {
        started_.store(false, std::memory_order_release);
        pool_.close();
        Logger::error("mysql ping failed, check host/user/password/database");
        return false;
    }

    Logger::info("mysql pool ready, size=" + std::to_string(cfg.pool_size) +
                 " db=" + cfg.database + "@" + cfg.host + ":" + std::to_string(cfg.port));
    ensure_schema();
    return true;
}

bool MysqlService::ensure_schema()
{
    const char* stmts[] = {
        "ALTER TABLE im_user ADD COLUMN allow_search_account TINYINT NOT NULL DEFAULT 1",
        "ALTER TABLE im_user ADD COLUMN allow_search_phone TINYINT NOT NULL DEFAULT 1",
        "ALTER TABLE im_user ADD COLUMN allow_add_friend TINYINT NOT NULL DEFAULT 1",
        "ALTER TABLE im_friend_request ADD COLUMN message VARCHAR(64) NOT NULL DEFAULT ''",
        "ALTER TABLE im_friend ADD COLUMN remark VARCHAR(64) NOT NULL DEFAULT ''",
        "ALTER TABLE im_user ADD COLUMN avatar_ver INT NOT NULL DEFAULT 0",
        "ALTER TABLE im_group ADD COLUMN avatar_ver INT NOT NULL DEFAULT 0",
        "ALTER TABLE im_user ADD COLUMN file_max BIGINT NOT NULL DEFAULT 0",
        "ALTER TABLE im_message ADD COLUMN sender_deleted TINYINT NOT NULL DEFAULT 0",
        "ALTER TABLE im_message ADD COLUMN receiver_deleted TINYINT NOT NULL DEFAULT 0",
        "CREATE TABLE IF NOT EXISTS im_friend_request ("
        "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
        "from_uid BIGINT UNSIGNED NOT NULL,"
        "to_uid BIGINT UNSIGNED NOT NULL,"
        "status TINYINT NOT NULL DEFAULT 0,"
        "created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,"
        "updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,"
        "PRIMARY KEY (id), UNIQUE KEY uk_from_to (from_uid, to_uid),"
        "KEY idx_to_status (to_uid, status)"
        ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4",
        "CREATE TABLE IF NOT EXISTS im_sticker_pack ("
        "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
        "name VARCHAR(64) NOT NULL,"
        "version INT NOT NULL DEFAULT 1,"
        "sort_order INT NOT NULL DEFAULT 0,"
        "status TINYINT NOT NULL DEFAULT 1,"
        "PRIMARY KEY (id)"
        ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4",
        "CREATE TABLE IF NOT EXISTS im_sticker ("
        "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
        "pack_id BIGINT UNSIGNED NOT NULL,"
        "code VARCHAR(64) NOT NULL,"
        "name VARCHAR(64) NOT NULL DEFAULT '',"
        "file_path VARCHAR(255) NOT NULL,"
        "sha256 VARCHAR(64) NOT NULL DEFAULT '',"
        "mime VARCHAR(32) NOT NULL DEFAULT 'image/png',"
        "sort_order INT NOT NULL DEFAULT 0,"
        "PRIMARY KEY (id), UNIQUE KEY uk_pack_code (pack_id, code), KEY idx_pack (pack_id)"
        ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4",
        "CREATE TABLE IF NOT EXISTS im_file ("
        "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
        "from_uid BIGINT UNSIGNED NOT NULL,"
        "name VARCHAR(255) NOT NULL,"
        "mime VARCHAR(64) NOT NULL DEFAULT 'application/octet-stream',"
        "size INT NOT NULL DEFAULT 0,"
        "sha256 VARCHAR(64) NOT NULL DEFAULT '',"
        "file_path VARCHAR(255) NOT NULL DEFAULT '',"
        "created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,"
        "PRIMARY KEY (id), KEY idx_from (from_uid)"
        ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4",
        "CREATE TABLE IF NOT EXISTS im_group ("
        "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
        "name VARCHAR(64) NOT NULL DEFAULT '',"
        "owner_uid BIGINT UNSIGNED NOT NULL,"
        "join_mode TINYINT NOT NULL DEFAULT 1,"
        "member_count INT NOT NULL DEFAULT 0,"
        "avatar_ver INT NOT NULL DEFAULT 0,"
        "created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,"
        "updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,"
        "PRIMARY KEY (id), KEY idx_owner (owner_uid)"
        ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4",
        "CREATE TABLE IF NOT EXISTS im_group_member ("
        "group_id BIGINT UNSIGNED NOT NULL,"
        "uid BIGINT UNSIGNED NOT NULL,"
        "role TINYINT NOT NULL DEFAULT 0,"
        "created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,"
        "PRIMARY KEY (group_id, uid), KEY idx_uid (uid)"
        ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4",
        "CREATE TABLE IF NOT EXISTS im_group_request ("
        "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
        "group_id BIGINT UNSIGNED NOT NULL,"
        "from_uid BIGINT UNSIGNED NOT NULL,"
        "status TINYINT NOT NULL DEFAULT 0,"
        "message VARCHAR(64) NOT NULL DEFAULT '',"
        "created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,"
        "updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,"
        "PRIMARY KEY (id), UNIQUE KEY uk_group_from (group_id, from_uid),"
        "KEY idx_group_status (group_id, status)"
        ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4",
        "CREATE TABLE IF NOT EXISTS im_group_message ("
        "id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,"
        "group_id BIGINT UNSIGNED NOT NULL,"
        "from_uid BIGINT UNSIGNED NOT NULL DEFAULT 0,"
        "msg_type TINYINT NOT NULL DEFAULT 1,"
        "content TEXT NOT NULL,"
        "created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,"
        "PRIMARY KEY (id), KEY idx_group_id (group_id, id)"
        ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4",
        "CREATE TABLE IF NOT EXISTS im_group_msg_deleted ("
        "uid BIGINT UNSIGNED NOT NULL,"
        "msg_id BIGINT UNSIGNED NOT NULL,"
        "PRIMARY KEY (uid, msg_id), KEY idx_msg (msg_id)"
        ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4",
        "CREATE TABLE IF NOT EXISTS im_group_inbox ("
        "uid BIGINT UNSIGNED NOT NULL,"
        "group_id BIGINT UNSIGNED NOT NULL,"
        "last_msg_id BIGINT UNSIGNED NOT NULL DEFAULT 0,"
        "last_content VARCHAR(256) NOT NULL DEFAULT '',"
        "last_time DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,"
        "unread INT NOT NULL DEFAULT 0,"
        "PRIMARY KEY (uid, group_id), KEY idx_uid_time (uid, last_time)"
        ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4",
    };
    try {
        auto conn = acquire();
        for (const char* sql : stmts) {
            try {
                exec_sql(conn.get(), sql);
            } catch (const std::exception& e) {
                Logger::info(std::string("ensure_schema skip: ") + e.what());
            }
        }
        return true;
    } catch (const std::exception& e) {
        Logger::error(std::string("ensure_schema: ") + e.what());
        return false;
    }
}

bool MysqlService::init_schema(const MysqlConfig& cfg, const std::string& sql_path)
{
    ylib::buffer raw = ylib::file::read(sql_path);
    std::string sql = raw.to_string();
    if (sql.size() >= 3 && static_cast<unsigned char>(sql[0]) == 0xEF &&
        static_cast<unsigned char>(sql[1]) == 0xBB &&
        static_cast<unsigned char>(sql[2]) == 0xBF) {
        sql.erase(0, 3);
    }
    if (trim_sql(sql).empty()) {
        Logger::error("init sql empty: " + sql_path);
        return false;
    }

    ylib::mysql::mysql_conn_info info;
    info.ipaddress = cfg.host;
    info.port = cfg.port;
    info.username = cfg.user;
    info.password = cfg.password;
    info.charset = cfg.charset;
    info.database = "";

    ylib::mysql::pool tmp;
    if (!tmp.start(info, 1)) {
        tmp.close();
        info.database = cfg.database;
        if (!tmp.start(info, 1)) {
            Logger::error("mysql init connect failed");
            return false;
        }
    }

    bool ok = false;
    try {
        MysqlConn conn(tmp.get());
        if (!conn) {
            Logger::error("mysql init acquire failed");
        } else {
            const std::string create_db =
                "CREATE DATABASE IF NOT EXISTS `" + cfg.database +
                "` DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci";
            try {
                if (!exec_sql(conn.get(), create_db)) {
                    Logger::warn("create database skipped, try existing db=" + cfg.database);
                }
            } catch (const std::exception& e) {
                Logger::warn(std::string("create database skipped: ") + e.what());
            }

            conn->setDatabase(cfg.database);

            ok = true;
            for (const auto& stmt : split_sql(sql)) {
                if (starts_ci(stmt, "USE ") || starts_ci(stmt, "CREATE DATABASE")) {
                    continue;
                }
                if (!exec_sql(conn.get(), stmt)) {
                    ok = false;
                    break;
                }
            }
        }
    } catch (const std::exception& e) {
        Logger::error(std::string("mysql init schema failed: ") + e.what());
        ok = false;
    }

    tmp.close();
    if (!ok) {
        return false;
    }
    Logger::info("mysql schema ready, db=" + cfg.database + " sql=" + sql_path);
    return true;
}

void MysqlService::stop()
{
    if (!started_.load(std::memory_order_acquire)) {
        return;
    }
    pool_.close();
    started_.store(false, std::memory_order_release);
    Logger::info("mysql pool stopped");
}

MysqlConn MysqlService::acquire()
{
    if (!started_.load(std::memory_order_acquire)) {
        throw ylib::exception("mysql pool is not started");
    }
    return MysqlConn(pool_.get());
}

bool MysqlService::ping()
{
    try {
        auto conn = acquire();
        auto* stmt = conn->setsql("SELECT 1");
        if (stmt == nullptr) {
            return false;
        }
        ylib::mysql::result* rs = stmt->query();
        return rs != nullptr;
    } catch (const std::exception& e) {
        Logger::error(std::string("mysql ping exception: ") + e.what());
        return false;
    }
}

}  // namespace im