<commit_before>package water.jdbc;

import water.*;
import water.fvec.*;
import water.parser.BufferedString;
import water.parser.ParseDataset;
import water.util.Log;

import java.math.BigDecimal;
import java.sql.*;
import java.util.Arrays;
import java.util.concurrent.ArrayBlockingQueue;

public class SQLManager {

  private static final String TEMP_TABLE_NAME = "table_for_h2o_import";
  private static final String MAX_USR_CONNECTIONS_KEY = H2O.OptArgs.SYSTEM_PROP_PREFIX + "sql.connections.max";
  private static final String JDBC_DRIVER_CLASS_KEY_PREFIX = H2O.OptArgs.SYSTEM_PROP_PREFIX + "sql.jdbc.driver.";
  //A target upper bound on number of connections to database
  private static final int MAX_CONNECTIONS = 100;
  //A lower bound on number of connections to database per node
  private static final int MIN_CONNECTIONS_PER_NODE = 1;
  private static final String NETEZZA_DB_TYPE = "netezza";
  private static final String HIVE_DB_TYPE = "hive2";
  private static final String ORACLE_DB_TYPE = "oracle";
  private static final String SQL_SERVER_DB_TYPE = "sqlserver";
  private static final String TERADATA_DB_TYPE = "teradata";

  private static final String NETEZZA_JDBC_DRIVER_CLASS = "org.netezza.Driver";
  private static final String HIVE_JDBC_DRIVER_CLASS = "org.apache.hive.jdbc.HiveDriver";

  private static final String TMP_TABLE_ENABLED = H2O.OptArgs.SYSTEM_PROP_PREFIX + "sql.tmp_table.enabled";

  public static Job<Frame> importSqlTable(final String connection_url, String table, final String select_query,
                                          final String username, final String password, final String columns,
                                          final boolean optimize) {
    return importSqlTable(connection_url, table, select_query, username, password, columns, optimize, false);
  }

  /**
   * @param connection_url (Input) 
   * @param table (Input)
   * @param select_query (Input)
   * @param username (Input)
   * @param password (Input)
   * @param columns (Input)
   * @param optimize (Input)
   * @param streaming (Input)
   */
  public static Job<Frame> importSqlTable(final String connection_url, String table, final String select_query,
                                          final String username, final String password, final String columns,
                                          final boolean optimize, final boolean streaming) {
    Connection conn = null;
    Statement stmt = null;
    ResultSet rs = null;
    final String databaseType = connection_url.split(":", 3)[1];
    initializeDatabaseDriver(databaseType);
    int catcols = 0, intcols = 0, bincols = 0, realcols = 0, timecols = 0, stringcols = 0;
    final int numCol;
    long numRow = 0;
    final String[] columnNames;
    final byte[] columnH2OTypes;
    try {
      conn = DriverManager.getConnection(connection_url, username, password);
      stmt = conn.createStatement();
      //set fetch size for improved performance
      stmt.setFetchSize(1);
      //if select_query has been specified instead of table
      if (table.equals("")) {
        if (!select_query.toLowerCase().startsWith("select")) {
          throw new IllegalArgumentException("The select_query must start with `SELECT`, but instead is: " + select_query);
        }

        //if tmp table disabled, we use sub-select instead, which outperforms tmp table for very large queries/tables
        // the main drawback of sub-selects is that we lose isolation, but as we're only reading data
        // and counting the rows from the beginning, it should not an issue (at least when using hive...)
        final boolean createTmpTable = Boolean.parseBoolean(System.getProperty(TMP_TABLE_ENABLED, "true")); //default to true to keep old behaviour
        if (createTmpTable) {
          table = SQLManager.TEMP_TABLE_NAME;
          //returns number of rows, but as an int, not long. if int max value is exceeded, result is negative
          numRow = stmt.executeUpdate("CREATE TABLE " + table + " AS " + select_query);
        } else {
          table = "(" + select_query + ") sub_h2o_import";
        }
      } else if (table.equals(SQLManager.TEMP_TABLE_NAME)) {
        //tables with this name are assumed to be created here temporarily and are dropped
        throw new IllegalArgumentException("The specified table cannot be named: " + SQLManager.TEMP_TABLE_NAME);
      }
      //get number of rows. check for negative row count
      if (numRow <= 0) {
        rs = stmt.executeQuery("SELECT COUNT(1) FROM " + table);
        rs.next();
        numRow = rs.getLong(1);
        rs.close();
      }
      //get H2O column names and types
      if (streaming) {
        stmt.setMaxRows(1);
        rs = stmt.executeQuery("SELECT " + columns + " FROM " + table);
      } else {
        rs = stmt.executeQuery(buildSelectSingleRowSql(databaseType, table, columns));
      }
      ResultSetMetaData rsmd = rs.getMetaData();
      numCol = rsmd.getColumnCount();

      columnNames = new String[numCol];
      columnH2OTypes = new byte[numCol];

      rs.next();
      for (int i = 0; i < numCol; i++) {
        columnNames[i] = rsmd.getColumnName(i + 1);
        //must iterate through sql types instead of getObject bc object could be null
        switch (rsmd.getColumnType(i + 1)) {
          case Types.NUMERIC:
          case Types.REAL:
          case Types.DOUBLE:
          case Types.FLOAT:
          case Types.DECIMAL:
            columnH2OTypes[i] = Vec.T_NUM;
            realcols += 1;
            break;
          case Types.INTEGER:
          case Types.TINYINT:
          case Types.SMALLINT:
          case Types.BIGINT:
            columnH2OTypes[i] = Vec.T_NUM;
            intcols += 1;
            break;
          case Types.BIT:
          case Types.BOOLEAN:
            columnH2OTypes[i] = Vec.T_NUM;
            bincols += 1;
            break;
          case Types.VARCHAR:
          case Types.NVARCHAR:
          case Types.CHAR:
          case Types.NCHAR:
          case Types.LONGVARCHAR:
          case Types.LONGNVARCHAR:
            columnH2OTypes[i] = Vec.T_STR;
            stringcols += 1;
            break;
          case Types.DATE:
          case Types.TIME:
          case Types.TIMESTAMP:
            columnH2OTypes[i] = Vec.T_TIME;
            timecols += 1;
            break;
          default:
            Log.warn("Unsupported column type: " + rsmd.getColumnTypeName(i + 1));
            columnH2OTypes[i] = Vec.T_BAD;
        }
      }

    } catch (SQLException ex) {
      throw new RuntimeException("SQLException: " + ex.getMessage() + "\nFailed to connect and read from SQL database with connection_url: " + connection_url);
    } finally {
      // release resources in a finally{} block in reverse-order of their creation
      if (rs != null) {
        try {
          rs.close();
        } catch (SQLException sqlEx) {} // ignore
        rs = null;
      }

      if (stmt != null) {
        try {
          stmt.close();
        } catch (SQLException sqlEx) {} // ignore
        stmt = null;
      }

      if (conn != null) {
        try {
          conn.close();
        } catch (SQLException sqlEx) {} // ignore
        conn = null;
      }
    }

    double binary_ones_fraction = 0.5; //estimate

    //create template vectors in advance and run MR
    final long totSize =
            (long)((float)(catcols+intcols)*numRow*4 //4 bytes for categoricals and integers
                    +(float)bincols          *numRow*1*binary_ones_fraction //sparse uses a fraction of one byte (or even less)
                    +(float)(realcols+timecols+stringcols) *numRow*8); //8 bytes for real and time (long) values

    final Vec vec;
    final int chunk_size = FileVec.calcOptimalChunkSize(totSize, numCol, numCol * 4,
            H2O.ARGS.nthreads, H2O.getCloudSize(), false, false);
    final double rows_per_chunk = chunk_size; //why not numRow * chunk_size / totSize; it's supposed to be rows per chunk, not the byte size
    final int num_chunks = Vec.nChunksFor(numRow, (int) Math.ceil(Math.log1p(rows_per_chunk)), false);

    if (streaming && optimize) {
      Log.warn("Parameter 'optimize' will be ignored in streaming ingest.");
    }
    final boolean estimateConcurrentConnections = optimize && (! streaming);
    if (estimateConcurrentConnections) {
      assert ! streaming;
      final int num_retrieval_chunks = ConnectionPoolProvider.estimateConcurrentConnections(H2O.getCloudSize(), H2O.ARGS.nthreads);
      vec = num_retrieval_chunks >= num_chunks
              ? Vec.makeConN(numRow, num_chunks)
              : Vec.makeConN(numRow, num_retrieval_chunks);
    } else {
      vec = Vec.makeConN(numRow, num_chunks);
    }
    Log.info("Number of chunks for data retrieval: " + vec.nChunks() + ", number of rows:" + numRow);
    //create frame
    final Key<Frame> destination_key = Key.make((table + "_sql_to_hex").replaceAll("\\W", "_"));
    final Job<Frame> j = new Job(destination_key, Frame.class.getName(), "Import SQL Table");

    final String finalTable = table;
    H2O.H2OCountedCompleter work = new H2O.H2OCountedCompleter() {
      @Override
      public void compute2() {
        final ConnectionPoolProvider provider = new ConnectionPoolProvider(connection_url, username, password, vec.nChunks());
        final Frame fr;
<commit_msg>Removed an assertion from SQLManager that is in vain.<commit_after>package water.jdbc;

import water.*;
import water.fvec.*;
import water.parser.BufferedString;
import water.parser.ParseDataset;
import water.util.Log;

import java.math.BigDecimal;
import java.sql.*;
import java.util.Arrays;
import java.util.concurrent.ArrayBlockingQueue;

public class SQLManager {

  private static final String TEMP_TABLE_NAME = "table_for_h2o_import";
  private static final String MAX_USR_CONNECTIONS_KEY = H2O.OptArgs.SYSTEM_PROP_PREFIX + "sql.connections.max";
  private static final String JDBC_DRIVER_CLASS_KEY_PREFIX = H2O.OptArgs.SYSTEM_PROP_PREFIX + "sql.jdbc.driver.";
  //A target upper bound on number of connections to database
  private static final int MAX_CONNECTIONS = 100;
  //A lower bound on number of connections to database per node
  private static final int MIN_CONNECTIONS_PER_NODE = 1;
  private static final String NETEZZA_DB_TYPE = "netezza";
  private static final String HIVE_DB_TYPE = "hive2";
  private static final String ORACLE_DB_TYPE = "oracle";
  private static final String SQL_SERVER_DB_TYPE = "sqlserver";
  private static final String TERADATA_DB_TYPE = "teradata";

  private static final String NETEZZA_JDBC_DRIVER_CLASS = "org.netezza.Driver";
  private static final String HIVE_JDBC_DRIVER_CLASS = "org.apache.hive.jdbc.HiveDriver";

  private static final String TMP_TABLE_ENABLED = H2O.OptArgs.SYSTEM_PROP_PREFIX + "sql.tmp_table.enabled";

  public static Job<Frame> importSqlTable(final String connection_url, String table, final String select_query,
                                          final String username, final String password, final String columns,
                                          final boolean optimize) {
    return importSqlTable(connection_url, table, select_query, username, password, columns, optimize, false);
  }

  /**
   * @param connection_url (Input) 
   * @param table (Input)
   * @param select_query (Input)
   * @param username (Input)
   * @param password (Input)
   * @param columns (Input)
   * @param optimize (Input)
   * @param streaming (Input)
   */
  public static Job<Frame> importSqlTable(final String connection_url, String table, final String select_query,
                                          final String username, final String password, final String columns,
                                          final boolean optimize, final boolean streaming) {
    Connection conn = null;
    Statement stmt = null;
    ResultSet rs = null;
    final String databaseType = connection_url.split(":", 3)[1];
    initializeDatabaseDriver(databaseType);
    int catcols = 0, intcols = 0, bincols = 0, realcols = 0, timecols = 0, stringcols = 0;
    final int numCol;
    long numRow = 0;
    final String[] columnNames;
    final byte[] columnH2OTypes;
    try {
      conn = DriverManager.getConnection(connection_url, username, password);
      stmt = conn.createStatement();
      //set fetch size for improved performance
      stmt.setFetchSize(1);
      //if select_query has been specified instead of table
      if (table.equals("")) {
        if (!select_query.toLowerCase().startsWith("select")) {
          throw new IllegalArgumentException("The select_query must start with `SELECT`, but instead is: " + select_query);
        }

        //if tmp table disabled, we use sub-select instead, which outperforms tmp table for very large queries/tables
        // the main drawback of sub-selects is that we lose isolation, but as we're only reading data
        // and counting the rows from the beginning, it should not an issue (at least when using hive...)
        final boolean createTmpTable = Boolean.parseBoolean(System.getProperty(TMP_TABLE_ENABLED, "true")); //default to true to keep old behaviour
        if (createTmpTable) {
          table = SQLManager.TEMP_TABLE_NAME;
          //returns number of rows, but as an int, not long. if int max value is exceeded, result is negative
          numRow = stmt.executeUpdate("CREATE TABLE " + table + " AS " + select_query);
        } else {
          table = "(" + select_query + ") sub_h2o_import";
        }
      } else if (table.equals(SQLManager.TEMP_TABLE_NAME)) {
        //tables with this name are assumed to be created here temporarily and are dropped
        throw new IllegalArgumentException("The specified table cannot be named: " + SQLManager.TEMP_TABLE_NAME);
      }
      //get number of rows. check for negative row count
      if (numRow <= 0) {
        rs = stmt.executeQuery("SELECT COUNT(1) FROM " + table);
        rs.next();
        numRow = rs.getLong(1);
        rs.close();
      }
      //get H2O column names and types
      if (streaming) {
        stmt.setMaxRows(1);
        rs = stmt.executeQuery("SELECT " + columns + " FROM " + table);
      } else {
        rs = stmt.executeQuery(buildSelectSingleRowSql(databaseType, table, columns));
      }
      ResultSetMetaData rsmd = rs.getMetaData();
      numCol = rsmd.getColumnCount();

      columnNames = new String[numCol];
      columnH2OTypes = new byte[numCol];

      rs.next();
      for (int i = 0; i < numCol; i++) {
        columnNames[i] = rsmd.getColumnName(i + 1);
        //must iterate through sql types instead of getObject bc object could be null
        switch (rsmd.getColumnType(i + 1)) {
          case Types.NUMERIC:
          case Types.REAL:
          case Types.DOUBLE:
          case Types.FLOAT:
          case Types.DECIMAL:
            columnH2OTypes[i] = Vec.T_NUM;
            realcols += 1;
            break;
          case Types.INTEGER:
          case Types.TINYINT:
          case Types.SMALLINT:
          case Types.BIGINT:
            columnH2OTypes[i] = Vec.T_NUM;
            intcols += 1;
            break;
          case Types.BIT:
          case Types.BOOLEAN:
            columnH2OTypes[i] = Vec.T_NUM;
            bincols += 1;
            break;
          case Types.VARCHAR:
          case Types.NVARCHAR:
          case Types.CHAR:
          case Types.NCHAR:
          case Types.LONGVARCHAR:
          case Types.LONGNVARCHAR:
            columnH2OTypes[i] = Vec.T_STR;
            stringcols += 1;
            break;
          case Types.DATE:
          case Types.TIME:
          case Types.TIMESTAMP:
            columnH2OTypes[i] = Vec.T_TIME;
            timecols += 1;
            break;
          default:
            Log.warn("Unsupported column type: " + rsmd.getColumnTypeName(i + 1));
            columnH2OTypes[i] = Vec.T_BAD;
        }
      }

    } catch (SQLException ex) {
      throw new RuntimeException("SQLException: " + ex.getMessage() + "\nFailed to connect and read from SQL database with connection_url: " + connection_url);
    } finally {
      // release resources in a finally{} block in reverse-order of their creation
      if (rs != null) {
        try {
          rs.close();
        } catch (SQLException sqlEx) {} // ignore
        rs = null;
      }

      if (stmt != null) {
        try {
          stmt.close();
        } catch (SQLException sqlEx) {} // ignore
        stmt = null;
      }

      if (conn != null) {
        try {
          conn.close();
        } catch (SQLException sqlEx) {} // ignore
        conn = null;
      }
    }

    double binary_ones_fraction = 0.5; //estimate

    //create template vectors in advance and run MR
    final long totSize =
            (long)((float)(catcols+intcols)*numRow*4 //4 bytes for categoricals and integers
                    +(float)bincols          *numRow*1*binary_ones_fraction //sparse uses a fraction of one byte (or even less)
                    +(float)(realcols+timecols+stringcols) *numRow*8); //8 bytes for real and time (long) values

    final Vec vec;
    final int chunk_size = FileVec.calcOptimalChunkSize(totSize, numCol, numCol * 4,
            H2O.ARGS.nthreads, H2O.getCloudSize(), false, false);
    final double rows_per_chunk = chunk_size; //why not numRow * chunk_size / totSize; it's supposed to be rows per chunk, not the byte size
    final int num_chunks = Vec.nChunksFor(numRow, (int) Math.ceil(Math.log1p(rows_per_chunk)), false);

    if (streaming && optimize) {
      Log.warn("Parameter 'optimize' will be ignored in streaming ingest.");
    }
    final boolean estimateConcurrentConnections = optimize && (! streaming);
    if (estimateConcurrentConnections) {
      final int num_retrieval_chunks = ConnectionPoolProvider.estimateConcurrentConnections(H2O.getCloudSize(), H2O.ARGS.nthreads);
      vec = num_retrieval_chunks >= num_chunks
              ? Vec.makeConN(numRow, num_chunks)
              : Vec.makeConN(numRow, num_retrieval_chunks);
    } else {
      vec = Vec.makeConN(numRow, num_chunks);
    }
    Log.info("Number of chunks for data retrieval: " + vec.nChunks() + ", number of rows:" + numRow);
    //create frame
    final Key<Frame> destination_key = Key.make((table + "_sql_to_hex").replaceAll("\\W", "_"));
    final Job<Frame> j = new Job(destination_key, Frame.class.getName(), "Import SQL Table");

    final String finalTable = table;
    H2O.H2OCountedCompleter work = new H2O.H2OCountedCompleter() {
      @Override
      public void compute2() {
        final ConnectionPoolProvider provider = new ConnectionPoolProvider(connection_url, username, password, vec.nChunks());
        final Frame fr;

{'fraction_non_alphanumeric': 0.09567405503801263, 'fraction_numerical': 0.006424670735624799, 'mean_word_length': 2.5383595377912482, 'pattern_counts': {'":': 2, '<': 15, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2361317', 'n_tokens_mistral': 5392, 'n_tokens_neox': 5306, 'n_words': 1588}