DbUtil.java

/***************************************************************************
   Copyright 2015 Emily Estes

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
***************************************************************************/
package net.metanotion.sql;


import java.io.File;
import java.io.FileInputStream;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.IOException;
import java.io.PrintWriter;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Iterator;
import javax.sql.DataSource;

import com.zaxxer.hikari.HikariDataSource;

import net.metanotion.functor.Block;
import net.metanotion.util.JDBCConfig;
import net.metanotion.util.JDBCTransaction;

/** This class represents a collection of utility methods for common patterns working with SQL databases via JDBC. */
public final class DbUtil {
	private static final String IF_NOT_EXISTS = "IF NOT EXISTS ";
	private static final String IF_EXISTS = "IF EXISTS ";
	private static final String CASCADE = " CASCADE";
	private static final String CREATE_SCHEMA_STMT = "CREATE SCHEMA ";
	private static final String DROP_SCHEMA_STMT = "DROP SCHEMA ";
	private static final String SET_SEARCH_PATH_STMT = "SET search_path TO ";
	private static final String SEMI = ";";

	/** Create a new schema.
		@param conn Database connection.
		@param schema The name of the schema to create.
		@param force If true, create schema whether it already exists or not (if it already exists, nothing will
			happen). If this is false, and the schema already exists an exception will be thrown.
		@throws SQLException if the schema could not be created or already exists.
	*/
	public static void createSchema(final Connection conn, final String schema, final boolean force) throws SQLException {
		try (final Statement stmt = conn.createStatement()) {
			stmt.executeUpdate(CREATE_SCHEMA_STMT + (force ? IF_NOT_EXISTS : "") + schema + SEMI);
		}
	}

	/** Drop a schema.
		@param conn Database connection.
		@param schema The name of the schema to drop.
		@param force If true, the schema will be dropped even if its not empty. If false and the schema did not exist
			an exception will be thrown.
		@throws SQLException if the schema could not be dropped or did not exist.
	*/
	public static void dropSchema(final Connection conn, final String schema, final boolean force) throws SQLException {
		try (final Statement stmt = conn.createStatement()) {
			stmt.executeUpdate(DROP_SCHEMA_STMT + (force ? IF_EXISTS : "") + schema + (force ? CASCADE : "") + SEMI);
		}
	}

	/** Set the search path to the schema provided.
		@param conn Database connection.
		@param schema The name of the schema to set the search path.
		@throws SQLException if the search path could not be set.
	*/
	public static void setSearchPath(final Connection conn, final String schema) throws SQLException {
		try (final Statement stmt = conn.createStatement()) {
			stmt.execute(SET_SEARCH_PATH_STMT + schema + SEMI);
		}
	}

	/** This method looks for a static method on the class called "schemaFactory" to invoke to provide a schema
	generator instance.
		@param factory The class that will produce the schema factory.
		@return A schema generator from the class.
		@throws NoSuchMethodException if there is no "schemaFactory" method.
		@throws IllegalAccessException if the "schemaFactory" method is not accessible.
		@throws InvocationTargetException if the "schemaFactory" method is not static.
	*/
	public static SchemaGenerator reflectSchemaGenerator(final Class factory)
			throws NoSuchMethodException, IllegalAccessException, InvocationTargetException {
		final Method m = factory.getMethod("schemaFactory", null);
		if(!(SchemaGenerator.class.isAssignableFrom(m.getReturnType()))) {
			throw new
				UnsupportedOperationException("Class " + factory.getName() + " does not provide a schema generator instance.");
		}
		m.setAccessible(true);
		return (SchemaGenerator) m.invoke(null, new Object[0]);
	}

	/** Use the class loader provided to turn a resource in to a schema generator.
		@param classLoader The class loader that will provide the resource.
		@param resource The resource to load.
		@return A schema generator using the contents of the resource.
	*/
	public static SchemaGenerator resourceSchemaGenerator(final ClassLoader classLoader, final String resource) {
		return new SchemaGenerator() {
			@Override public Iterator<String> openSchema() {
				try (final InputStream in = classLoader.getResourceAsStream(resource)) {
					final Iterator<String> it =
						SequenceExecutor.readerSequencer(new InputStreamReader(in, StandardCharsets.UTF_8));
					final ArrayList<String> list = new ArrayList<>();
					while(it.hasNext()) {
						list.add(it.next());
					}
					return list.iterator();
				} catch (final IOException ioe) {
					throw new RuntimeException(ioe);
				}
			}
		};
	}

	/** Read the provided file to create a schema generator.
		@param file The file load
		@return A schema generator using the contents of the file.
	*/
	public static SchemaGenerator fileSchemaGenerator(final File file) {
		return new SchemaGenerator() {
			@Override public Iterator<String> openSchema() {
				try (final InputStream in = new FileInputStream(file)) {
					final Iterator<String> it =
						SequenceExecutor.readerSequencer(new InputStreamReader(in, StandardCharsets.UTF_8));
					final ArrayList<String> list = new ArrayList<>();
					while(it.hasNext()) {
						list.add(it.next());
					}
					return list.iterator();
				} catch (final IOException ioe) {
					throw new RuntimeException(ioe);
				}
			}
		};
	}

	/** Run a schema generator instance against a connection from the data source.
		@param ds The database connection pool to run the schema against.
		@param schema The schema to run against the database.
		@return The number of statements successfully executed.
		@throws SQLException if there is a problem getting the connection.
	*/
	public static int runSchema(final DataSource ds, final SchemaGenerator schema) throws SQLException {
		try (final Connection conn = ds.getConnection()) {
			return runSchema(conn, schema);
		}
	}

	/** Run a schema generator instance against a connection from the data source.
		@param conn The database connection to run the schema against.
		@param schema The schema to run against the database.
		@return The number of statements successfully executed.
	*/
	public static int runSchema(final Connection conn, final SchemaGenerator schema) {
		return JDBCTransaction.doTX(conn, new Block<Connection,Integer>() {
			@Override public Integer eval(final Connection conn) throws Exception {
				return SequenceExecutor.doSequence(conn, schema.openSchema());
			}
		});
	}

	/** Output the statements of a schema generator to an output stream.
		@param out The output stream to write the schema statements into.
		@param schema The schema generator to print.
	*/
	public static void printSchema(final PrintWriter out, final SchemaGenerator schema) {
		printSchema(out, schema, null);
	}

	/** Output the statements of a schema generator to an output stream.
		@param out The output stream to write the schema statements into.
		@param schema The schema generator to print.
		@param path The schema path, null if the default path is used.
	*/
	public static void printSchema(final PrintWriter out, final SchemaGenerator schema, final String path) {
		if(path != null) {
			out.println(CREATE_SCHEMA_STMT + path + SEMI);
			out.println(SET_SEARCH_PATH_STMT + path + SEMI);
		}
		final Iterator<String> statements = schema.openSchema();
		while(statements.hasNext()) {
			out.println(statements.next());
		}
		out.flush();
	}

	/** Create a database.
		@param conn The connection to the server to create the database on.
		@param dbName The name of the database to create.
		@param user The user that will own the database.
		@param dropFirst Attempt to drop the database if it already exists.
		@throws SQLException if there is a problem creating the database.
	*/
	public static void createDatabase(final Connection conn,
			final String dbName,
			final String user,
			boolean dropFirst) throws SQLException {
		try (final Statement stmt = conn.createStatement()) {
			if(dropFirst) {
				try {
					stmt.executeUpdate("DROP DATABASE " + dbName);
				} catch (final SQLException sqle) { }
			}
			stmt.executeUpdate("CREATE DATABASE " + dbName + " WITH OWNER = " + user + " ENCODING = 'UTF8';");
		}
	}

	/** Create a PostgreSQL based database connection pool from a JDBCCoonfig data structure.
		@param config The object containing the connection information for the pool.
		@return a DataSource for grabbing connections from the pool.
	*/
	public static DataSource startDBConnectionPool(final JDBCConfig config) {
		return startDBConnectionPool(config.connect, config.username, config.password);
	}

	/** Create a database connection pool from JDBC connection string, usernmame, and password.
		@param connect the JDBC connection string/uri
		@param username the username to use to connect to the database.
		@param password The password to use to connect to the database.
		@return a DataSource for grabbing connections from the pool.
	*/
	public static DataSource startDBConnectionPool(final String connect, final String username, final String password) {
		final HikariDataSource ds = new HikariDataSource();
		ds.setJdbcUrl(connect);
		ds.setUsername(username);
		ds.setPassword(password);
		return ds;
	}
}