/etlunit-core/src/main/java/org/bitbucket/bradleysmithllc/etlunit/util/jdbc/JdbcVisitors.java

package org.bitbucket.bradleysmithllc.etlunit.util.jdbc;

/*
 * #%L
 * etlunit-core
 * %%
 * Copyright (C) 2010 - 2021 bradleysmithllc
 * %%
 * 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.
 * #L%
 */

import java.sql.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

public class JdbcVisitors {
	public static JdbcVisitorCommands withConnection(ExceptionalSupplier<Connection> withConnnection, ExceptionalConsumer<Connection> returnTo) throws Exception {
		Connection c = withConnnection.get();

		return new JdbcVisitorCommands(c, returnTo);
	}

	public static JdbcVisitorCommands withConnection(Connection withConnnection, ExceptionalConsumer<Connection> returnTo) throws Exception {
		return new JdbcVisitorCommands(withConnnection, returnTo);
	}

	public static JdbcVisitorCommands withPersistentConnection(ExceptionalSupplier<Connection> withConnnection) throws Exception {
		Connection c = withConnnection.get();

		return new JdbcVisitorCommands(c, (connection) -> {});
	}

	public static JdbcVisitorCommands withPersistentConnection(Connection withConnnection) throws Exception {
		return new JdbcVisitorCommands(withConnnection, (connection) -> {});
	}

	public static JdbcVisitorCommands withDisposableConnection(ExceptionalSupplier<Connection> withConnnection) throws Exception {
		Connection c = withConnnection.get();

		return new JdbcVisitorCommands(c, (connection) -> {connection.close();});
	}

	public static JdbcVisitorCommands withDisposableConnection(Connection withConnnection) throws Exception {
		return new JdbcVisitorCommands(withConnnection, (connection) -> {connection.close();});
	}

	public static class JdbcVisitorCommands {
		private final Connection connection;
		private final ExceptionalConsumer<Connection> returnTo;

		public JdbcVisitorCommands(Connection connection, ExceptionalConsumer<Connection> returnTo) {
			this.connection = connection;
			this.returnTo = returnTo;
		}

		public JdbcVisitorCommands withConnection(ExceptionalConsumer<Connection> consumer) throws Exception {
			consumer.accept(connection);

			return this;
		}

		public JdbcVisitorCommands withCatalogs(ExceptionalConsumer<String> consumer) throws Exception {
			return withDatabaseMetadata((meta) -> {
				scanResultSet(meta.getCatalogs(), (resultset) -> {
					String catalog = resultset.getString("TABLE_CAT");

					consumer.accept(catalog);
				});
			});
		}

		public JdbcVisitorCommands withSchemas(ExceptionalConsumer<MetadataSchema> consumer) throws Exception {
			return withSchemas(null, null, consumer);
		}

		public JdbcVisitorCommands withSchemas(String catalog, String schema, ExceptionalConsumer<MetadataSchema> consumer) throws Exception {
			return withDatabaseMetadata((meta) -> {
				scanResultSet(meta.getSchemas(catalog, schema), (resultset) -> {
					String catalogName = resultset.getString("TABLE_CATALOG");
					String schemaName = resultset.getString("TABLE_SCHEM");

					consumer.accept(new MetadataSchema(catalogName, schemaName));
				});
			});
		}

		public JdbcVisitorCommands withColumns(String catalog, String schema, String table, String column, ExceptionalConsumer<MetadataColumn> consumer) throws Exception {
			return withDatabaseMetadata((meta) -> {
				scanResultSet(meta.getColumns(catalog, schema, table, column), (resultset) -> {
					consumer.accept(
							new MetadataColumn(
									resultset.getString("TABLE_CAT"),
									resultset.getString("TABLE_SCHEM"),
									resultset.getString("TABLE_NAME"),
									resultset.getString("COLUMN_NAME")
							)
							.withDataType(resultset.getInt("DATA_TYPE"))
							.withColumnSize(resultset.getInt("COLUMN_SIZE"))
							.withDecimalDigits(resultset.getInt("DECIMAL_DIGITS"))
							.withNullable(resultset.getInt("NULLABLE"))
							.withSqlDataType(resultset.getInt("SQL_DATA_TYPE"))
							.withSqlDatetimeSub(resultset.getInt("SQL_DATETIME_SUB"))
							.withCharOctetLength(resultset.getInt("CHAR_OCTET_LENGTH"))
							.withOrdinalPosition(resultset.getInt("ORDINAL_POSITION"))
							.withSourceDataType(resultset.getShort("SOURCE_DATA_TYPE"))

							.withTypeName(resultset.getString("TYPE_NAME"))
							.withRemarks(resultset.getString("REMARKS"))
							.withColumnDef(resultset.getString("COLUMN_DEF"))
							.withIsNullable(resultset.getString("IS_NULLABLE"))
							.withScopeCatalog(resultset.getString("SCOPE_CATALOG"))
							.withScopeSchema(resultset.getString("SCOPE_SCHEMA"))
							.withScopeTable(resultset.getString("SCOPE_TABLE"))
							.withIsAutoincrement(resultset.getString("IS_AUTOINCREMENT"))
							.withIsGeneratedColumn(resultset.getString("IS_GENERATEDCOLUMN"))
					);
				});
			});
		}

		public JdbcVisitorCommands withTables(String catalog, String schema, String table, ExceptionalConsumer<MetadataTable> consumer) throws Exception {
			return withDatabaseMetadata((meta) -> {
				scanResultSet(meta.getTables(catalog, schema, table, null), (resultset) -> {
					consumer.accept(
							new MetadataTable(resultset.getString("TABLE_CAT"), resultset.getString("TABLE_SCHEM"), resultset.getString("TABLE_NAME"))
							.withRemarks(resultset.getString("REMARKS"))
							.withTableType(resultset.getString("TABLE_TYPE"))
							.withSelfReferencingColName(resultset.getString("SELF_REFERENCING_COL_NAME"))
							.withRefGeneration(resultset.getString("REF_GENERATION"))
							.withTypeCatalog(resultset.getString("TYPE_CAT"))
							.withTypeSchema(resultset.getString("TYPE_SCHEM"))
							.withTypeName(resultset.getString("TYPE_NAME"))
					);
				});
			});
		}

		public JdbcVisitorCommands withPrimaryKeys(String catalog, String schema, String table, ExceptionalConsumer<MetadataPrimaryKey> consumer) throws Exception {
			return withDatabaseMetadata((meta) -> {
				scanResultSet(meta.getPrimaryKeys(catalog, schema, table), (resultSet) -> {
					consumer.accept(new MetadataPrimaryKey(
							resultSet.getString("TABLE_CAT"),
							resultSet.getString("TABLE_SCHEM"),
							resultSet.getString("TABLE_NAME"),
							resultSet.getString("PK_NAME")
					)
							.withColumnName(resultSet.getString("COLUMN_NAME"))
							.withKeySequence(resultSet.getShort("KEY_SEQ"))
					);
				});
			});
		}

		public JdbcVisitorCommands withDatabaseMetadata(ExceptionalConsumer<DatabaseMetaData> st) throws Exception {
			st.accept(connection.getMetaData());

			return this;
		}

		public JdbcVisitorCommands withStatement(ExceptionalConsumer<Statement> st) throws Exception {
			Statement statement = connection.createStatement();

			try {
				st.accept(statement);
			} finally {
				statement.close();
			}

			return this;
		}

		public void withResultSet(ResultSet st, ExceptionalConsumer<ResultSet> consumer) throws Exception {
				try {
					consumer.accept(st);
				} finally {
					st.close();
				}
		}

		/* This version calls the consumer once for every single row. */
		public void scanResultSet(ResultSet st, ExceptionalConsumer<ResultSet> consumer) throws Exception {
			scanResultSet(st, consumer, (rowCount) -> {});
		}

		/* This version calls the consumer once for every single row. */
		public void scanResultSet(
				ResultSet st,
				ExceptionalConsumer<ResultSet> consumer,
				ExceptionalConsumer<Integer> allRowsComplete
		) throws Exception {
			AtomicInteger rowCountRef = new AtomicInteger();

			try {
				withResultSet(st, (resultset) -> {
					while (resultset.next()) {
						rowCountRef.incrementAndGet();
						consumer.accept(resultset);
					}
				});
			} finally {
				st.close();
			}

			allRowsComplete.accept(rowCountRef.get());
		}

		public JdbcVisitorCommands scanQuery(String sql, ExceptionalConsumer<ResultSet> st) throws Exception {
			return scanQuery(sql, st, null, (rowCount) -> {});
		}

		public JdbcVisitorCommands withQuery(String sql, ExceptionalConsumer<ResultSet> st) throws Exception {
			return withQuery(sql, st, null);
		}

		/* Inject arbitrary code into the pipeline */
		public JdbcVisitorCommands withStep(Runnable step) throws Exception {
			step.run();
			return this;
		}

		public JdbcVisitorCommands withPreparedStatement(String sql, ExceptionalConsumer<PreparedStatement> st) throws Exception {
			PreparedStatement statement = connection.prepareStatement(sql);

			try {
				st.accept(statement);
			} finally {
				statement.close();
			}

			return this;
		}

		public JdbcVisitorCommands withUpdate(String sql, ExceptionalConsumer<Integer> st) throws Exception {
			withStatement((statement) -> {
				int resultSet = statement.executeUpdate(sql);

				st.accept(resultSet);
			});

			return this;
		}

		public JdbcVisitorCommands scanQuery(
				String sql,
				ExceptionalConsumer<ResultSet> st,
				ExceptionalConsumer<ResultSetMetaData> metadataConsumer
		) throws Exception {
			return scanQuery(sql, st, metadataConsumer, (rowCount) -> {});
		}

		public JdbcVisitorCommands scanQuery(
				String sql,
				ExceptionalConsumer<ResultSet> st,
				ExceptionalConsumer<ResultSetMetaData> metadataConsumer,
				ExceptionalConsumer<Integer> completionConsumer
		) throws Exception {
			withStatement((statement) -> {
				ResultSet resultSet = statement.executeQuery(sql);

				if (metadataConsumer != null) {
					metadataConsumer.accept(resultSet.getMetaData());
				}

				scanResultSet(resultSet, st, completionConsumer);
			});

			return this;
		}

		public JdbcVisitorCommands withQuery(String sql, ExceptionalConsumer<ResultSet> st, ExceptionalConsumer<ResultSetMetaData> metadataConsumer) throws Exception {
			withStatement((statement) -> {
				ResultSet resultSet = statement.executeQuery(sql);

				if (metadataConsumer != null) {
					metadataConsumer.accept(resultSet.getMetaData());
				}

				withResultSet(resultSet, st);
			});

			return this;
		}

		public void dispose() throws Exception {
			returnTo.accept(connection);
		}
	}
}