DBConnector.java

package utils;

import java.io.FileReader;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;

import javatools.database.Database;
import javatools.database.PostgresDatabase;

import org.postgresql.copy.CopyManager;
import org.postgresql.core.BaseConnection;

/**
 * DID NOT WORK USING LOAD. Created ~6 files of 200MB each, to \copy from psql.
 * 
 * @author ntandon
 * 
 */
public class DBConnector {

	private static Database db;
	private static CopyManager cm;

	public static ResultSet q(String sql, boolean setAutoCommit)
		throws SQLException {
		Statement st = DBConnector.getDB().getConnection().createStatement();
		st.setFetchSize(1000);
		db.getConnection().setAutoCommit(setAutoCommit);

		try {
			return st.executeQuery(sql);
		} catch (Error | Exception e) {
			System.out.println("Exception (" + e.getMessage() + ") "
				+ "while initializing the DB. \n "
				+ "Trying to automatically resolve...");

			// returning empty result set for problematic query
			return q("SELECT NULL LIMIT 0;", true);
		}
	}

	public static ResultSet qNoTryCatch(String sql, boolean setAutoCommit)
		throws SQLException {
		Statement st = DBConnector.getDB().getConnection().createStatement();
		st.setFetchSize(1000);
		db.getConnection().setAutoCommit(setAutoCommit);
		return st.executeQuery(sql);
	}

	public static ResultSet q(String sql) throws SQLException {

		// Most of our queries are select queries (in batch size of 1000,
		// these select queries are considered as batch queries.
		return q(sql, false);
	}

	public static void resetDb() throws SQLException {
		q("SELECT NULL LIMIT 0;", true);
	}

	public static Database getDB() {
		return getDB("sreyasi", "123ertghjm,.", "wiki",
			"postgres2.d5.mpi-inf.mpg.de", null);
	}

	public static Database getDB(String user, String password, String database,
		String host, String port) {
		if (db != null)
			return db;
		try {
			db = new PostgresDatabase(user, password, database, host, port);

			// Turn off autocommit mode
			db.getConnection().setAutoCommit(false);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return db;
	}

	public static void closeConnections() {
		try {
			if (!db.getConnection().isClosed())
				db.close();
			db = null;
			cm = null;
		} catch (Exception e) {
			db = null;
			cm = null;
		}
	}

	/**
	 * cm.copyIn("COPY mytable FROM STDIN WITH DELIMITER '\t'", new
	 * FileReader(fileToLoadFrom.tsv));
	 */
	public static void load(String tbName, String fileToCopyFrom,
		boolean clearTable, char delim) throws Exception {
		// TODO set autocommit true to have an effect.
		if (cm == null)
			cm = new CopyManager((BaseConnection) getDB().getConnection());
		if (clearTable)
			getDB().executeUpdate("DELETE FROM " + tbName);
		cm.copyIn("COPY " + tbName + " FROM STDIN WITH DELIMITER '" + delim
			+ "'", new FileReader(fileToCopyFrom));
	}

	public static List<String> singletonColumn(String sql, int MAX_OUTPUT)
		throws SQLException {
		List<String> l = new ArrayList<>();
		ResultSet rs = q(sql);
		while (rs.next())
			if (l.size() <= MAX_OUTPUT)
				l.add(rs.getString(1));
			else
				break;

		return l;
	}

	public static List<String> singletonColumn(String sql) throws SQLException {
		List<String> l = new ArrayList<>();
		ResultSet rs = q(sql);
		while (rs.next())
			l.add(rs.getString(1));

		return l;

	}

}