package tech.easyflow.datacenter.connector.impl;

import com.zaxxer.hikari.HikariDataSource;
import org.springframework.stereotype.Component;
import tech.easyflow.datacenter.connector.dialect.GaussdbSqlDialect;
import tech.easyflow.datacenter.connector.support.AbstractJdbcConnector;
import tech.easyflow.datacenter.connector.support.DatacenterDatasourceManager;
import tech.easyflow.datacenter.meta.entity.DatacenterSource;
import tech.easyflow.datacenter.meta.enums.DatacenterCapability;
import tech.easyflow.datacenter.meta.enums.DatacenterSourceType;

import java.sql.Connection;
import java.util.EnumSet;

@Component
public class GaussdbNativeConnector extends AbstractJdbcConnector {

    private final DatacenterDatasourceManager datasourceManager;

    public GaussdbNativeConnector(DatacenterDatasourceManager datasourceManager) {
        super(DatacenterSourceType.GAUSSDB_NATIVE, new GaussdbSqlDialect(), EnumSet.of(
                DatacenterCapability.TEST_CONNECTION,
                DatacenterCapability.BROWSE_METADATA,
                DatacenterCapability.READ_QUERY
        ));
        this.datasourceManager = datasourceManager;
    }

    @Override
    protected <T> T withConnection(DatacenterSource source, boolean cacheable, JdbcCallback<T> callback) throws Exception {
        HikariDataSource dataSource = cacheable ? datasourceManager.getOrCreateExternalDatasource(source) : datasourceManager.createExternalDatasource(source);
        try (Connection connection = dataSource.getConnection()) {
            return callback.apply(connection);
        } finally {
            if (!cacheable || source.getId() == null) {
                dataSource.close();
            }
        }
    }
}
