Java tutorial
/* * Copyright 2018 University of Michigan * * 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 org.verdictdb.connection; import java.lang.reflect.Constructor; import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.sql.Connection; import java.sql.DatabaseMetaData; import java.sql.DriverManager; import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; import java.util.ArrayList; import java.util.List; import java.util.Properties; import org.apache.commons.lang3.tuple.ImmutablePair; import org.apache.commons.lang3.tuple.Pair; import org.verdictdb.commons.StringSplitter; import org.verdictdb.commons.VerdictDBLogger; import org.verdictdb.exception.VerdictDBDbmsException; import org.verdictdb.sqlsyntax.HiveSyntax; import org.verdictdb.sqlsyntax.ImpalaSyntax; import org.verdictdb.sqlsyntax.PostgresqlSyntax; import org.verdictdb.sqlsyntax.PrestoHiveSyntax; import org.verdictdb.sqlsyntax.PrestoMemorySyntax; import org.verdictdb.sqlsyntax.PrestoSyntax; import org.verdictdb.sqlsyntax.RedshiftSyntax; import org.verdictdb.sqlsyntax.SparkSyntax; import org.verdictdb.sqlsyntax.SqlSyntax; import org.verdictdb.sqlsyntax.SqlSyntaxList; public class JdbcConnection extends DbmsConnection { Connection conn; SqlSyntax syntax; String currentSchema = null; JdbcQueryResult jrs = null; protected boolean outputDebugMessage = false; protected Statement runningStatement = null; protected VerdictDBLogger log; protected boolean isAborting = false; public static JdbcConnection create(String jdbcConnectionString, Properties info) throws VerdictDBDbmsException { try { Connection c; if (info == null) { c = DriverManager.getConnection(jdbcConnectionString); } else { c = DriverManager.getConnection(jdbcConnectionString, info); } return JdbcConnection.create(c); } catch (SQLException e) { throw new VerdictDBDbmsException(e); } } public static JdbcConnection create(String jdbcConnectionString) throws VerdictDBDbmsException { try { Connection c = DriverManager.getConnection(jdbcConnectionString); return JdbcConnection.create(c); } catch (SQLException e) { throw new VerdictDBDbmsException(e); } } public static JdbcConnection create(Connection conn) throws VerdictDBDbmsException { String connectionString = null; try { connectionString = conn.getMetaData().getURL(); } catch (SQLException e) { throw new VerdictDBDbmsException(e); } SqlSyntax syntax = SqlSyntaxList.getSyntaxFromConnectionString(connectionString); // This is temporary fix to have 'memory' catalog connection for unit tests try { if (syntax instanceof PrestoHiveSyntax && conn.getCatalog() != null && conn.getCatalog().equalsIgnoreCase("memory")) { syntax = new PrestoMemorySyntax(); } } catch (SQLException e) { e.printStackTrace(); } JdbcConnection jdbcConn = null; if (syntax instanceof PrestoSyntax) { // To handle that Presto's JDBC driver is not compatible with JDK7, // we use Java's reflection-based instantiation. try { Class<?> prestoConnClass = Class.forName("org.verdictdb.connection.PrestoJdbcConnection"); Constructor<?> prestoConnClsConstructor = prestoConnClass.getConstructor(Connection.class, SqlSyntax.class); jdbcConn = (JdbcConnection) prestoConnClsConstructor.newInstance(conn, syntax); Method ensureMethod = prestoConnClass.getMethod("ensureCatalogSet"); ensureMethod.invoke(jdbcConn); } catch (ClassNotFoundException | NoSuchMethodException | SecurityException | InstantiationException | IllegalAccessException | IllegalArgumentException e) { throw new RuntimeException("Instantiating PrestoJdbcConnection failed."); } catch (InvocationTargetException e) { if (e.getTargetException() instanceof VerdictDBDbmsException) { throw new VerdictDBDbmsException(e.getMessage()); } else { throw new RuntimeException("Instantiating PrestoJdbcConnection failed."); } } } else { jdbcConn = new JdbcConnection(conn, syntax); } return jdbcConn; } public JdbcConnection(Connection conn, SqlSyntax syntax) { this.conn = conn; try { this.currentSchema = conn.getSchema(); } catch (SQLException e) { e.printStackTrace(); } // set a default value if an inappropriate value is set. if (currentSchema == null || currentSchema.length() == 0) { currentSchema = syntax.getFallbackDefaultSchema(); } this.syntax = syntax; this.log = VerdictDBLogger.getLogger(this.getClass()); } @Override public void abort() { log.trace("Aborts a statement if running."); isAborting = true; try { synchronized (this) { // having isClosed() check seems to block this statement. if (runningStatement != null) { // if (runningStatement != null && !runningStatement.isClosed()) { log.trace("Aborts a running statement."); runningStatement.cancel(); runningStatement.close(); runningStatement = null; } } isAborting = false; } catch (SQLException e) { e.printStackTrace(); } } @Override public void close() { log.debug("Closes a JDBC connection."); abort(); try { this.conn.close(); } catch (SQLException e) { e.printStackTrace(); } } @Override public DbmsQueryResult execute(String sql) throws VerdictDBDbmsException { String quoteChars = "'\""; List<String> sqls = StringSplitter.splitOnSemicolon(sql, quoteChars); DbmsQueryResult finalResult = null; for (String s : sqls) { finalResult = executeSingle(s); } return finalResult; } protected void setRunningStatement(Statement stmt) { synchronized (this) { runningStatement = stmt; } } protected Statement getRunningStatement() { synchronized (this) { return runningStatement; } } public DbmsQueryResult executeSingle(String sql) throws VerdictDBDbmsException { log.debug("Issues the following query to DBMS: " + sql); try { Statement stmt = conn.createStatement(); setRunningStatement(stmt); JdbcQueryResult jrs = null; boolean doesResultExist = stmt.execute(sql); if (doesResultExist) { ResultSet rs = stmt.getResultSet(); jrs = new JdbcQueryResult(rs); rs.close(); } else { jrs = null; } setRunningStatement(null); stmt.close(); return jrs; } catch (SQLException e) { if (isAborting) { return null; } else { String msg = "Issued the following query: " + sql + "\n" + e.getMessage(); throw new VerdictDBDbmsException(msg); } } } public DbmsQueryResult executeQuery(String sql) throws VerdictDBDbmsException { return execute(sql); } @Override public SqlSyntax getSyntax() { return syntax; } public Connection getConnection() { return conn; } @Override public List<String> getSchemas() throws VerdictDBDbmsException { List<String> schemas = new ArrayList<>(); DbmsQueryResult queryResult = executeQuery(syntax.getSchemaCommand()); while (queryResult.next()) { schemas.add(queryResult.getString(syntax.getSchemaNameColumnIndex())); } return schemas; } @Override public List<String> getTables(String schema) { List<String> tables = new ArrayList<>(); try { DbmsQueryResult queryResult = executeQuery(syntax.getTableCommand(schema)); while (queryResult.next()) { tables.add(queryResult.getString(syntax.getTableNameColumnIndex())); } } catch (VerdictDBDbmsException e) { log.debug(e.getMessage()); } return tables; } @Override public List<Pair<String, String>> getColumns(String schema, String table) throws VerdictDBDbmsException { List<Pair<String, String>> columns = new ArrayList<>(); String sql = syntax.getColumnsCommand(schema, table); try { DbmsQueryResult queryResult = executeQuery(sql); while (queryResult.next()) { String type; if (syntax instanceof PostgresqlSyntax) { type = queryResult.getString(syntax.getColumnTypeColumnIndex()); if (queryResult .getInt(((PostgresqlSyntax) syntax).getCharacterMaximumLengthColumnIndex()) != 0) { type = type + "(" + queryResult .getInt(((PostgresqlSyntax) syntax).getCharacterMaximumLengthColumnIndex()) + ")"; } } else { type = queryResult.getString(syntax.getColumnTypeColumnIndex()); } type = type.toLowerCase(); columns.add(new ImmutablePair<>(queryResult.getString(syntax.getColumnNameColumnIndex()), type)); } } catch (Exception e) { if (syntax instanceof RedshiftSyntax && e.getMessage().matches("(?s).*schema .* does not exist;.*")) { return columns; } else { throw e; } } return columns; } @Override public List<Pair<String, String>> getColumns(String catalog, String schema, String table) throws VerdictDBDbmsException { List<Pair<String, String>> columns = new ArrayList<>(); String sql = syntax.getColumnsCommand(catalog, schema, table); try { DbmsQueryResult queryResult = executeQuery(sql); while (queryResult.next()) { String type; if (syntax instanceof PostgresqlSyntax) { type = queryResult.getString(syntax.getColumnTypeColumnIndex()); if (queryResult .getInt(((PostgresqlSyntax) syntax).getCharacterMaximumLengthColumnIndex()) != 0) { type = type + "(" + queryResult .getInt(((PostgresqlSyntax) syntax).getCharacterMaximumLengthColumnIndex()) + ")"; } } else { type = queryResult.getString(syntax.getColumnTypeColumnIndex()); } type = type.toLowerCase(); columns.add(new ImmutablePair<>(queryResult.getString(syntax.getColumnNameColumnIndex()), type)); } } catch (Exception e) { if (syntax instanceof RedshiftSyntax && e.getMessage().matches("(?s).*schema .* does not exist;.*")) { return columns; } else { throw e; } } return columns; } @Override public List<String> getPartitionColumns(String schema, String table) throws VerdictDBDbmsException { List<String> partition = new ArrayList<>(); if (!syntax.doesSupportTablePartitioning()) { return partition; } DbmsQueryResult queryResult; if (syntax instanceof ImpalaSyntax) { try { queryResult = executeQuery(syntax.getPartitionCommand(schema, table)); for (int i = 0; i < queryResult.getColumnCount(); i++) { String columnName = queryResult.getColumnName(i); if (columnName.equalsIgnoreCase("#rows")) { break; } else { partition.add(columnName); } } return partition; } catch (Exception e) { if (e.getMessage().contains("Table is not partitioned")) { return partition; } else { throw e; } } } else if (syntax instanceof RedshiftSyntax) { try { queryResult = executeQuery(syntax.getPartitionCommand(schema, table)); } catch (Exception e) { if (e.getMessage().matches("(?s).*schema .* does not exist;.*")) { return partition; } else { throw e; } } } else { queryResult = executeQuery(syntax.getPartitionCommand(schema, table)); } // the result of postgresql is a vector of column index if (syntax instanceof PostgresqlSyntax) { if (queryResult.next()) { Object o = queryResult.getValue(0); String[] arr = o.toString().split(" "); List<Pair<String, String>> columns = getColumns(schema, table); for (int i = 0; i < arr.length; i++) { partition.add(columns.get(Integer.valueOf(arr[i]) - 1).getKey()); } } } // Hive and Spark append partition information at the end of the "DESCRIBE TABLE" statement. else if (syntax instanceof HiveSyntax || syntax instanceof SparkSyntax || syntax instanceof PrestoSyntax) { boolean hasPartitionInfoStarted = false; while (queryResult.next()) { String name = queryResult.getString(0); if (hasPartitionInfoStarted && (name.equalsIgnoreCase("# col_name") == false)) { partition.add(name); } else if (name.equalsIgnoreCase("# Partition Information")) { hasPartitionInfoStarted = true; } } } else { while (queryResult.next()) { String column = queryResult.getString(0); if (column != null) { partition.add(queryResult.getString(0)); } } } return partition; } @Override public String getDefaultSchema() { return currentSchema; } @Override public void setDefaultSchema(String schema) throws VerdictDBDbmsException { try { // these database have a different meaning for catalog; thus, does not change. if (syntax instanceof PrestoSyntax || syntax instanceof PostgresqlSyntax || syntax instanceof RedshiftSyntax) { } else { conn.setCatalog(schema); } } catch (SQLException e) { throw new VerdictDBDbmsException(e); } currentSchema = schema; } public DatabaseMetaData getMetadata() throws VerdictDBDbmsException { try { return conn.getMetaData(); } catch (SQLException e) { throw new VerdictDBDbmsException(e); } } public boolean isOutputDebugMessage() { return outputDebugMessage; } public void setOutputDebugMessage(boolean outputDebugMessage) { this.outputDebugMessage = outputDebugMessage; } @Override public DbmsConnection copy() throws VerdictDBDbmsException { JdbcConnection newConn = new JdbcConnection(conn, syntax); newConn.setDefaultSchema(currentSchema); newConn.jrs = this.jrs; newConn.outputDebugMessage = this.outputDebugMessage; return newConn; } /** @return a list of column names of primary key columns. (0-indexed) */ @Override public List<String> getPrimaryKey(String schema, String table) throws VerdictDBDbmsException { List<Integer> primaryKeyIndexList = new ArrayList<>(); List<String> primaryKeyColumnName = new ArrayList<>(); SqlSyntax syntax = getSyntax(); if (syntax.getPrimaryKey(schema, table) != null) { DbmsQueryResult result = execute(syntax.getPrimaryKey(schema, table)); while (result.next()) { primaryKeyIndexList.add(result.getInt(3) - 1); } List<String> columns = new ArrayList<>(); result = execute(syntax.getColumnsCommand(schema, table)); while (result.next()) { columns.add(result.getString(0)); } for (int idx : primaryKeyIndexList) { primaryKeyColumnName.add(columns.get(idx)); } } return primaryKeyColumnName; } }