kiafumi/src/main/java/moe/oko/Kiafumi/model/db/KiafumiDB.java

package moe.oko.Kiafumi.model.db;

import com.google.common.base.Strings;
import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import moe.oko.Kiafumi.Kiafumi;
import moe.oko.Kiafumi.model.Server;
import net.dv8tion.jda.api.entities.Guild;

import java.sql.*;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

import static moe.oko.Kiafumi.util.LoggingManager.*;

/**
 * Kiafumi DB Class
 * Basically our helpful MySQL functions that pertain to data persistence.
 * @author Kay, with moral support from Yodabird
 */
public class KiafumiDB {

    private HikariDataSource dataSource;
    private Connection connection;

    // The prepared statement strings
    private final String CREATE_SERVERINFO_TABLE = "CREATE TABLE IF NOT EXISTS `serverInfo`" +
            "(`id` LONGTEXT NOT NULL, `welcomeChannel` LONGTEXT NULL, `modChannel` LONGTEXT NULL, joinRole LONGTEXT NULL);";
    private final String INSERT_SERVERINFO_DEFAULT = "INSERT INTO serverInfo(id, welcomeChannel, modChannel, joinRole) values (?,?,?,?)";
    private final String SERVERINFO_EXISTS = "select * from serverInfo where id = ?";
    private final String SERVER_INFO_LOAD = "select * from serverInfo";
    // 1 = welcomeChannel, 2 = modChannel , 3 = joinRole, 4 = Guild_ID (immutable)
    private final String SERVER_INFO_MODIFY = "UPDATE `serverInfo` SET welcomeChannel=?, modChannel=?, joinRole=? WHERE id = ?";

    /**
     * KiafumiDB Connection Constructor
     */
    public KiafumiDB(DBCredentials credentials) {
        var config = new HikariConfig();
        config.setJdbcUrl("jdbc:mysql://%s:%s/%s".formatted(credentials.host(), credentials.port(), credentials.db()));

        config.setConnectionTimeout(credentials.connTimeout());
        config.setIdleTimeout(credentials.idleTimeout());
        config.setMaxLifetime(credentials.maxLifetime());
        config.setMaximumPoolSize(credentials.poolSize());
        config.setUsername(credentials.username());

        if (!Strings.isNullOrEmpty(credentials.password()))
            config.setPassword(credentials.password());
        this.dataSource = new HikariDataSource(config);

        try {
            connection = dataSource.getConnection();
            info("Connected to database.");
        } catch (SQLException e) {
            error("Unable to connect to database.");
            e.printStackTrace();
            this.dataSource = null;
        }

        initializeTables();
    }

    /**
     * Table initialization function, ran on every startup.
     */
    private void initializeTables() {
        try {
            connection.prepareStatement(CREATE_SERVERINFO_TABLE).execute();
        } catch (Exception ex) {
            ex.printStackTrace();
            error("Failed to initialize SQL tables. They may need to be created manually.");
        }
    }

    /**
     * Creates the default server information when the bot joins a discord.
     * @param guild - the guild that the bot joined
     * @return whether the function succeeded.
     */
    public boolean createServerInformation(Guild guild) {
        try {
            PreparedStatement prep = connection.prepareStatement(SERVERINFO_EXISTS);
            prep.setInt(1, (int) guild.getIdLong());
            ResultSet rsCheck = prep.executeQuery();
            while (rsCheck.next()) {
                //Server already exists, no need to initialize a section for it.
                if(rsCheck.getInt(1) != 0) {
                    info("Server already existed. Skipping initialization.");
                    return true;
                }
            }
            //Proceed with making defaults.
            PreparedStatement ps = connection.prepareStatement(INSERT_SERVERINFO_DEFAULT);
            ps.setString(1, guild.getId());
            ps.setNull(2, Types.LONGVARCHAR);
            ps.setNull(3, Types.LONGVARCHAR);
            ps.setNull(4, Types.LONGVARCHAR);
            ps.execute();
            return true;
        } catch (Exception ex) {
            error("Failed to create default server info for guild " + guild.getId());
            ex.printStackTrace();
            return false;
        }
    }

    /**
     * Loads all the server information from MySQL into memory.
     * @return - a list of all servers loaded from MySQL.
     */
    public List<Server> loadServerInformation() {
        List<Server> servers = new ArrayList<>();
        try {
            PreparedStatement ps = connection.prepareStatement(SERVER_INFO_LOAD);
            ResultSet rs = ps.executeQuery();
            while(rs.next()) {
                info("Starting new load for server: " + rs.getString(1));
                String id = rs.getString(1);
                String welcomeChannel = rs.getString(2);
                String modChannel = rs.getString(3);
                String joinRole = rs.getString(4);
                Server server = new Server(id, welcomeChannel, modChannel, joinRole);
                debug("Loaded " + server + "from database.");
                servers.add(server);
            }
            return servers;
        } catch (Exception ex) {
            ex.printStackTrace();
            error("Failed to load server information, check stack.");
            return null;
        }
    }

    /**
     * Saves modified servers into persistent MySQL server.
     * @return - whether the method succeeded.
     */
    public boolean saveServerInformation() {
        Collection<Server> servers = Kiafumi.instance.getServerManager().getServers();
        info("Starting save on " + servers.size() + " servers.");
        try {
            // 1 = welcomeChannel, 2 = modChannel, 3 = joinRole, 4 = Guild_ID (immutable)
            PreparedStatement ps = connection.prepareStatement(SERVER_INFO_MODIFY);
            int i = 0;
            for (Server server : servers) {
                if(!server.isModified()) { continue; } // Skip, unmodified server.
                info("Starting save on modified " + server);
                if(server.getWelcomeChannel() != null) {
                    ps.setString(1, server.getWelcomeChannel());
                } else {
                    ps.setNull(1, Types.LONGVARCHAR);
                }
                if(server.getModChannel() != null) {
                    ps.setString(2, server.getModChannel());
                } else {
                    ps.setNull(2, Types.LONGVARCHAR);
                }
                if(server.getJoinRole() != null) {
                    ps.setString(3, server.getJoinRole());
                } else {
                    ps.setNull(3, Types.LONGVARCHAR);
                }
                ps.setString(4, server.getId());
                ps.addBatch();
                i++;
            }
            info("Total Batches: " + i + ". Starting SQL save.");
            ps.executeBatch();
            return true;
        } catch (Exception ex) {
            error("Failed to persist server data. Check stack.");
            ex.printStackTrace();
            return false;
        }
    }
}