package example.datasource;

import java.sql.*;
import javax.sql.*;

public class HiveConnectionPoolDataSource extends org.apache.hive.jdbc.HiveDataSource implements ConnectionPoolDataSource {
    public PooledConnection getPooledConnection() throws SQLException {
        return new HivePooledConnection(null, null);
    }

    public PooledConnection getPooledConnection(String user, String password) throws SQLException {
        return new HivePooledConnection(user, password);
    }

    public boolean isWrapperFor(Class<?> iface) throws SQLException {
        return ConnectionPoolDataSource.class.equals(iface) || super.isWrapperFor(iface);
    }

    public <T> T unwrap(Class<T> iface) throws SQLException {
        return ConnectionPoolDataSource.class.equals(iface) ? (T) this : super.unwrap(iface);
    }

    class HivePooledConnection implements PooledConnection {
        private Connection con;
        private final String user;
        private final String password;

        HivePooledConnection(String user, String password) {
            this.user = user;
            this.password = password;
        }

        public void addConnectionEventListener(ConnectionEventListener listener) {}

        public void addStatementEventListener(StatementEventListener listener) {}

        public void close() throws SQLException {
            if (con != null) {
                con.close();
                con = null;
            }
        }

        public Connection getConnection() throws SQLException {
            if (con == null) {
                con = user == null
                        ? HiveConnectionPoolDataSource.this.getConnection()
                        : HiveConnectionPoolDataSource.this.getConnection(user, password);
                return con;
            } else
                throw new IllegalStateException();
        }

        public void removeConnectionEventListener(ConnectionEventListener listener) {}

        public void removeStatementEventListener(StatementEventListener listener) {}
    }
}