AltitudeWeb/database/src/main/java/com/alttd/altitudeweb/setup/Connection.java

package com.alttd.altitudeweb.setup;

import com.alttd.altitudeweb.database.Databases;
import com.alttd.altitudeweb.database.web_db.DatabaseSettings;
import com.alttd.altitudeweb.database.web_db.SettingsMapper;
import com.alttd.altitudeweb.type_handler.UUIDTypeHandler;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.datasource.pooled.PooledDataSource;
import org.apache.ibatis.mapping.Environment;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.apache.ibatis.session.SqlSessionFactoryBuilder;
import org.apache.ibatis.transaction.jdbc.JdbcTransactionFactory;
import org.jetbrains.annotations.NotNull;

import java.util.HashMap;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.function.Consumer;

@Slf4j
public class Connection {

    private static final HashMap<Databases, Connection> connections = new HashMap<>();
    private SqlSessionFactory sqlSessionFactory;
    private final DatabaseSettings settings;
    private final AddMappers addMappers;

    private Connection(DatabaseSettings settings, AddMappers addMappers) {
        this.settings = settings;
        this.addMappers = addMappers;
    }

    public static void initDatabases() {
        InitializeWebDb.init();
        InitializeLiteBans.init();
        InitializeLuckPerms.init();
        InitializeProxyPlaytime.init();
        InitializeDiscord.init();
        InitializeVotingPlugin.init();
    }

    @FunctionalInterface
    public interface AddMappers {
        void apply(Configuration configuration);
    }

    public static Connection getConnection(Databases database) {
        if (connections.containsKey(database)) {
            return connections.get(database);
        }
        throw new RuntimeException("Database " + database + " has not been initialized");
    }

    protected static CompletableFuture<Connection> getConnection(Databases database, AddMappers addMappers) {
        if (connections.containsKey(database)) {
            return CompletableFuture.completedFuture(connections.get(database));
        }
        if (database == Databases.DEFAULT) {
            return loadDefaultDatabase(addMappers);
        }

        log.debug("Loading settings for database {}", database.getInternalName());
        CompletableFuture<DatabaseSettings> settingsFuture = new CompletableFuture<>();

        getConnection(Databases.DEFAULT, (mapper -> mapper.addMapper(SettingsMapper.class))).thenApply(connection -> {
            connection.runQuery(session -> {
                try {
                    log.debug("Running query to load settings for database {}", database.getInternalName());
                    DatabaseSettings loadedSettings = session.getMapper(SettingsMapper.class).getSettings(database.getInternalName());

                    if (loadedSettings == null) {
                        log.error("Failed to load settings for database {}. No settings found in db_connection_settings table.",
                                database.getInternalName());
                        settingsFuture.completeExceptionally(new IllegalStateException(
                                "Database settings for " + database.getInternalName() + " not found in db_connection_settings table"));
                    } else {
                        log.debug("Loaded settings for database {}: host={}, port={}, name={}",
                                database.getInternalName(), loadedSettings.host(), loadedSettings.port(), loadedSettings.name());
                        settingsFuture.complete(loadedSettings);
                    }
                } catch (Exception e) {
                    log.error("Error occurred while loading database settings for {}", database.getInternalName(), e);
                    settingsFuture.completeExceptionally(e);
                }
            });
            return null;
        }).exceptionally(ex -> {
            log.error("Failed to access DEFAULT database to load settings for {}", database.getInternalName(), ex);
            settingsFuture.completeExceptionally(ex);
            return null;
        });

        return settingsFuture.thenApply(loadedSettings -> {
            log.debug("Storing connection for database {}", database.getInternalName());
            Connection connection = new Connection(loadedSettings, addMappers);
            connections.put(database, connection);
            return connection;
        }).exceptionally(ex -> {
            log.error("Failed to create connection for database {}", database.getInternalName(), ex);
            throw new CompletionException("Failed to initialize database connection for " + database.getInternalName(), ex);
        });
    }

    private static CompletableFuture<Connection> loadDefaultDatabase(AddMappers addMappers) {
        DatabaseSettings databaseSettings = new DatabaseSettings(
                System.getenv("DB_HOST"),
                Integer.parseInt(System.getenv("DB_PORT")),
                System.getenv("DB_NAME"),
                System.getenv("DB_USER"),
                System.getenv("DB_PASS")
        );
        log.debug("Loaded default database settings {}", databaseSettings);
        Connection connection = new Connection(databaseSettings, addMappers);
        log.debug("Created default database connection {}", connection);
        connections.put(Databases.DEFAULT, connection);
        return CompletableFuture.completedFuture(connection);
    }

    public void runQuery(Consumer<SqlSession> consumer) {
        new Thread(() -> {
            if (sqlSessionFactory == null) {
                sqlSessionFactory = createSqlSessionFactory(settings, addMappers);
            }

            SqlSession session = null;
            try {
                session = sqlSessionFactory.openSession();
                consumer.accept(session);
                session.commit();
            } catch (Exception e) {
                if (session != null) {
                    session.rollback();
                }
                log.error("Failed to run query", e);
            } finally {
                if (session != null) {
                    session.close();
                }
            }
        }).start();
    }

    private SqlSessionFactory createSqlSessionFactory(DatabaseSettings settings, AddMappers addMappers) {
        try {
            Configuration configuration = getConfiguration(settings);
            configuration.getTypeHandlerRegistry().register(UUID.class, UUIDTypeHandler.class);
            addMappers.apply(configuration);

            return new SqlSessionFactoryBuilder().build(configuration);
        } catch (Exception e) {
            log.error("""
                    Failed to create sql session factory with
                    \thost {}
                    \tport: {}
                    \tname: {}
                    \tusername: {}
                    """, settings.host(), settings.port(), settings.name(), settings.username(), e);
            throw e;
        }
    }

    private static @NotNull Configuration getConfiguration(DatabaseSettings settings) {
        PooledDataSource dataSource = new PooledDataSource();
        dataSource.setDriver("com.mysql.cj.jdbc.Driver");

        String url = String.format(
                "jdbc:mysql://%s:%d/%s?useSSL=true&tcpKeepAlive=true&socketTimeout=60000&connectTimeout=10000&autoReconnect=false&useUnicode=true&characterEncoding=utf8",
                settings.host(),
                settings.port(),
                settings.name()
                                  );
        dataSource.setUrl(url);
        dataSource.setUsername(settings.username());
        dataSource.setPassword(settings.password());

        dataSource.setPoolMaximumActiveConnections(10);
        dataSource.setPoolMaximumIdleConnections(5);
        dataSource.setPoolTimeToWait(20000);
        dataSource.setPoolPingEnabled(true);
        dataSource.setPoolPingQuery("SELECT 1");
        dataSource.setPoolPingConnectionsNotUsedFor(300000); // 5 min

        Environment environment = new Environment("production", new JdbcTransactionFactory(), dataSource);
        return new Configuration(environment);
    }

}