org.apache.calcite.avatica.AvaticaStatement Java Examples

protected MetaResultSet createResultSet(
    Map<String, Object> internalParameters, List<ColumnMetaData> columns,
    CursorFactory cursorFactory, final Frame firstFrame) {
  try {
    final CalciteConnectionImpl connection = getConnection();
    final AvaticaStatement statement = connection.createStatement();
    final CalcitePrepare.CalciteSignature<Object> signature =
        new CalcitePrepare.CalciteSignature<Object>("",
            ImmutableList.of(), internalParameters, null,
            columns, cursorFactory, null, ImmutableList.of(), -1,
            null, Meta.StatementType.SELECT) {
          @Override public Enumerable<Object> enumerable(
              DataContext dataContext) {
            return Linq4j.asEnumerable(firstFrame.rows);
          }
        };
    return MetaResultSet.create(connection.id, statement.getId(), true,
        signature, firstFrame);
  } catch (SQLException e) {
    throw new RuntimeException(e);
  }
}
 

@Test public void testUnicodeCharacters() throws Exception {
  ConnectionSpec.getDatabaseLock().lock();
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    final AvaticaStatement statement = conn.createStatement();
    ResultSet rs = statement.executeQuery(
        "select * from (values ('您好', 'こんにちは', '안녕하세요'))");
    assertThat(rs.next(), is(true));
    assertEquals("您好", rs.getString(1));
    assertEquals("こんにちは", rs.getString(2));
    assertEquals("안녕하세요", rs.getString(3));
    rs.close();
    statement.close();
    conn.close();
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

private void checkLargeQuery(int n) throws Exception {
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    final AvaticaStatement statement = conn.createStatement();
    final String frenchDisko = "It said human existence is pointless\n"
        + "As acts of rebellious solidarity\n"
        + "Can bring sense in this world\n"
        + "La resistance!\n";
    final String sql = "select '"
        + longString(frenchDisko, n)
        + "' as s from (values 'x')";
    prepareAndExecuteInternal(conn, statement, sql, -1);
    ResultSet rs = statement.getResultSet();
    int count = 0;
    while (rs.next()) {
      count++;
    }
    assertThat(count, is(1));
    rs.close();
    statement.close();
    conn.close();
  }
}
 

@Test public void testRemoteExecuteMaxRowCount() throws Exception {
  ConnectionSpec.getDatabaseLock().lock();
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    final AvaticaStatement statement = conn.createStatement();
    prepareAndExecuteInternal(conn, statement,
        "select * from (values ('a', 1), ('b', 2))", 0);
    ResultSet rs = statement.getResultSet();
    int count = 0;
    while (rs.next()) {
      count++;
    }
    assertEquals("Check maxRowCount=0 and ResultSets is 0 row", count, 0);
    assertEquals("Check result set meta is still there",
        rs.getMetaData().getColumnCount(), 2);
    rs.close();
    statement.close();
    conn.close();
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

@Test public void testRemoteExecuteMaxRowCount() throws Exception {
  ConnectionSpec.getDatabaseLock().lock();
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    final AvaticaStatement statement = conn.createStatement();
    prepareAndExecuteInternal(conn, statement,
        "select * from (values ('a', 1), ('b', 2))", 0);
    ResultSet rs = statement.getResultSet();
    int count = 0;
    while (rs.next()) {
      count++;
    }
    assertEquals("Check maxRowCount=0 and ResultSets is 0 row", count, 0);
    assertEquals("Check result set meta is still there",
        rs.getMetaData().getColumnCount(), 2);
    rs.close();
    statement.close();
    conn.close();
  } finally {
    ConnectionSpec.getDatabaseLock().unlock();
  }
}
 

public static JdbcResultSet create(String connectionId, int statementId,
    ResultSet resultSet, int maxRowCount, Meta.Signature signature) {
  try {
    final Calendar calendar = DateTimeUtils.calendar();
    final int fetchRowCount;
    if (maxRowCount == JdbcMeta.UNLIMITED_COUNT) {
      fetchRowCount = -1;
    } else if (maxRowCount < 0L) {
      fetchRowCount = AvaticaStatement.DEFAULT_FETCH_SIZE;
    } else if (maxRowCount > AvaticaStatement.DEFAULT_FETCH_SIZE) {
      fetchRowCount = AvaticaStatement.DEFAULT_FETCH_SIZE;
    } else {
      fetchRowCount = maxRowCount;
    }
    final Meta.Frame firstFrame = frame(null, resultSet, 0, fetchRowCount, calendar,
        Optional.of(signature));
    if (firstFrame.done) {
      resultSet.close();
    }
    return new JdbcResultSet(connectionId, statementId, true, signature,
        firstFrame);
  } catch (SQLException e) {
    throw new RuntimeException(e);
  }
}
 

private void checkLargeQuery(int n) throws Exception {
  try (AvaticaConnection conn = (AvaticaConnection) DriverManager.getConnection(url)) {
    final AvaticaStatement statement = conn.createStatement();
    final String frenchDisko = "It said human existence is pointless\n"
        + "As acts of rebellious solidarity\n"
        + "Can bring sense in this world\n"
        + "La resistance!\n";
    final String sql = "select '"
        + longString(frenchDisko, n)
        + "' as s from (values 'x')";
    prepareAndExecuteInternal(conn, statement, sql, -1);
    ResultSet rs = statement.getResultSet();
    int count = 0;
    while (rs.next()) {
      count++;
    }
    assertThat(count, is(1));
    rs.close();
    statement.close();
    conn.close();
  }
}
 

@Test
public void testQueryMetaData() {
    try {
        AvaticaConnection conn = getConnection();
        final AvaticaStatement statement = conn.createStatement();
        ResultSet rs = statement.executeQuery("select * from (values (1, 'a'), (2, 'b'))");
        if (rs.getMetaData() != null) {
            System.out.println(rs.getMetaData().getColumnCount());
            System.out.println(rs.getMetaData().getColumnName(1));
            System.out.println(rs.getMetaData().getColumnName(2));
        }
        close(rs, statement,conn);
    } catch (Exception e) {
        e.printStackTrace();
    }
}
 

@Test
public void testPrepareStatementQuery() {
    try {
        AvaticaConnection conn = getConnection();
        final AvaticaStatement statement = conn.createStatement();
        PreparedStatement preparedStatement = conn.prepareStatement("select * from (values (1, 'a'), (2, 'b'), "
            + "(3, 'c')) "
            + "where expr_col__0 = ? or expr_col__1 = ?");
        preparedStatement.setInt(1,1);
        preparedStatement.setString(2,"b");
        ResultSet rs = preparedStatement.executeQuery();
        while (rs.next()) {
            System.out.println(rs.getString(1));
            System.out.println(rs.getString(2));
        }
        close(rs, statement,conn);
    } catch (Exception e) {
        e.printStackTrace();
    }
}
 

public static QuarkMetaResultSet create(String connectionId, int statementId,
                                        ResultSet resultSet,
                                        long maxRowCount, Meta.Signature signature) {
  try {
    final Calendar calendar = Calendar.getInstance(DateTimeUtils.GMT_ZONE);
    final int fetchRowCount;
    if (maxRowCount == QuarkMetaImpl.UNLIMITED_COUNT) {
      fetchRowCount = -1;
    } else if (maxRowCount < 0L) {
      fetchRowCount = AvaticaStatement.DEFAULT_FETCH_SIZE;
    } else if (maxRowCount > AvaticaStatement.DEFAULT_FETCH_SIZE) {
      fetchRowCount = AvaticaStatement.DEFAULT_FETCH_SIZE;
    } else {
      fetchRowCount = (int) maxRowCount;
    }
    final Meta.Frame firstFrame = frame(resultSet, 0, fetchRowCount, calendar);
    if (firstFrame.done) {
      resultSet.close();
    }
    return new QuarkMetaResultSet(connectionId, statementId, true, signature,
        firstFrame);
  } catch (SQLException e) {
    throw new RuntimeException(e);
  }
}
 

public static QuicksqlServerResultSet create(String connectionId, int statementId,
    ResultSet resultSet, int maxRowCount, Signature signature) {
  try {
    final Calendar calendar = DateTimeUtils.calendar();
    final int fetchRowCount;
    if (maxRowCount == QuicksqlServerMeta.UNLIMITED_COUNT) {
      fetchRowCount = -1;
    } else if (maxRowCount < 0L) {
      fetchRowCount = AvaticaStatement.DEFAULT_FETCH_SIZE;
    } else if (maxRowCount > AvaticaStatement.DEFAULT_FETCH_SIZE) {
      fetchRowCount = AvaticaStatement.DEFAULT_FETCH_SIZE;
    } else {
      fetchRowCount = maxRowCount;
    }
    final Meta.Frame firstFrame = frame(null, resultSet,0, fetchRowCount, calendar,
        Optional.of(signature));
    if (firstFrame.done) {
      resultSet.close();
    }
    return new QuicksqlServerResultSet(connectionId, statementId, true, signature,
        firstFrame);
  } catch (SQLException e) {
    throw new RuntimeException(e);
  }
}
 

protected MetaResultSet createResultSet(
    Map<String, Object> internalParameters, List<ColumnMetaData> columns,
    CursorFactory cursorFactory, final Frame firstFrame) {
  try {
    final QuarkConnectionImpl connection = getConnection();
    final AvaticaStatement statement = connection.createStatement();
    final CalcitePrepare.CalciteSignature<Object> signature =
        new CalcitePrepare.CalciteSignature<Object>("",
            ImmutableList.<AvaticaParameter>of(), internalParameters, null,
            columns, cursorFactory, null, ImmutableList.<RelCollation>of(), -1,
            null, Meta.StatementType.SELECT) {
          @Override public Enumerable<Object> enumerable(
              DataContext dataContext) {
            return Linq4j.asEnumerable(firstFrame.rows);
          }
        };
    return MetaResultSet.create(connection.id, statement.getId(), true,
        signature, firstFrame);
  } catch (SQLException e) {
    throw new RuntimeException(e);
  }
}
 

protected MetaResultSet createResultSet(
    Map<String, Object> internalParameters, List<ColumnMetaData> columns,
    CursorFactory cursorFactory, final Frame firstFrame) {
  try {
    final CalciteConnectionImpl connection = getConnection();
    final AvaticaStatement statement = connection.createStatement();
    final CalcitePrepare.CalciteSignature<Object> signature =
        new CalcitePrepare.CalciteSignature<Object>("",
            ImmutableList.of(), internalParameters, null,
            columns, cursorFactory, null, ImmutableList.of(), -1,
            null, Meta.StatementType.SELECT) {
          @Override public Enumerable<Object> enumerable(
              DataContext dataContext) {
            return Linq4j.asEnumerable(firstFrame.rows);
          }
        };
    return MetaResultSet.create(connection.id, statement.getId(), true,
        signature, firstFrame);
  } catch (SQLException e) {
    throw new RuntimeException(e);
  }
}
 

@Override
public AvaticaResultSet newResultSet(AvaticaStatement statement,
                                     QueryState state,
                                     Meta.Signature signature,
                                     TimeZone timeZone,
                                     Meta.Frame firstFrame) {
  final ResultSetMetaData metaData =
      newResultSetMetaData(statement, signature);
  return new QuarkResultSet(statement, signature, metaData, timeZone,
      firstFrame);
}
 

public static QuarkMetaResultSet create(String connectionId, int statementId,
                                        Iterator iterator,
                                        long maxRowCount, Meta.Signature signature) {
  try {
    final Calendar calendar = Calendar.getInstance(DateTimeUtils.GMT_ZONE);
    final int fetchRowCount;
    if (maxRowCount == QuarkMetaImpl.UNLIMITED_COUNT) {
      fetchRowCount = -1;
    } else if (maxRowCount < 0L) {
      fetchRowCount = AvaticaStatement.DEFAULT_FETCH_SIZE;
    } else if (maxRowCount > AvaticaStatement.DEFAULT_FETCH_SIZE) {
      fetchRowCount = AvaticaStatement.DEFAULT_FETCH_SIZE;
    } else {
      fetchRowCount = (int) maxRowCount;
    }
    final Meta.Frame firstFrame;
    if (!iterator.hasNext()) {
      firstFrame = Meta.Frame.EMPTY;
    } else {
      firstFrame = frame(iterator, signature.columns, 0, fetchRowCount, calendar);
    }
    return new QuarkMetaResultSet(connectionId, statementId, true, signature,
        firstFrame);
  } catch (SQLException e) {
    throw new RuntimeException(e);
  }
}
 

/**
 *
 * @param statement
 * @param signature
 */
DremioCursor(DremioConnectionImpl connection, AvaticaStatement statement, Signature signature) {
  this.connection = connection;
  this.statement = statement;
  this.signature = signature;

  DremioClient client = connection.getClient();
  final int batchQueueThrottlingThreshold =
      client.getConfig().getInt(JDBC_BATCH_QUEUE_THROTTLING_THRESHOLD );
  resultsListener = new ResultsListener(batchQueueThrottlingThreshold);
  currentBatchHolder = new RecordBatchLoader(client.getRecordAllocator());
}
 

@Override
public DremioResultSetImpl newResultSet(AvaticaStatement statement,
                                       QueryState state,
                                       Meta.Signature signature,
                                       TimeZone timeZone,
                                       Meta.Frame firstFrame) throws SQLException {
  final ResultSetMetaData metaData = newResultSetMetaData(statement, signature);
  return new DremioResultSetImpl(statement, state, signature, metaData, timeZone, firstFrame);
}
 

DremioResultSetImpl(AvaticaStatement statement, QueryState state,
                   Meta.Signature signature, ResultSetMetaData resultSetMetaData,
                   TimeZone timeZone, Meta.Frame firstFrame) throws SQLException {
  super(statement, state, signature, resultSetMetaData, timeZone, firstFrame);
  connection = (DremioConnectionImpl) statement.getConnection();
  cursor = new DremioCursor(connection, statement, signature);
}
 

@Override
public ExecuteResult execute(StatementHandle h, List<TypedValue> parameterValues, int maxRowsInFirstFrame)
    throws NoSuchStatementException {
  // Signature might have been zeroed by AvaticaConnection#executeQueryInternal()
  // Get it from the original handle
  final AvaticaStatement stmt;
  try {
     stmt = connection.lookupStatement(h);
  } catch(SQLException e) {
    throw new NoSuchStatementException(h);
  }
    MetaResultSet metaResultSet =
        MetaResultSet.create(h.connectionId, h.id, false, stmt.handle.signature, null);
  return new ExecuteResult(ImmutableList.of(metaResultSet));
}
 

@Override
protected ExecuteResult prepareAndExecuteInternal(AvaticaStatement statement, String sql, long maxRowCount)
    throws SQLException, NoSuchStatementException {
  try {
    return super.prepareAndExecuteInternal(statement, sql, maxRowCount);
  } catch(RuntimeException e) {
    Throwables.propagateIfInstanceOf(e.getCause(), SQLException.class);
    throw e;
  }
}
 

public QuarkMetaResultSet execute(ParserResult parserResult) throws Exception {
  QuarkJdbcStatement stmt = connection.server.getStatement(h);

  QuarkExecutor executor = QuarkExecutorFactory.getQuarkExecutor(parserResult.getKind(),
      connection.parserFactory, connection.getProperties(), connectionCache);
  Object result = executor.execute(parserResult);

  if (result instanceof Integer) {
    // Alter, Create, Drop DDL commands will either return id or 0
    return QuarkMetaResultSet.count(h.connectionId, h.id, ((Integer) result).intValue());
  } else if (result instanceof ArrayList) {
    // Show DDL returns an arraylist
    Class pojoType = getPojoType(parserResult.getParsedSql());

    return getMetaResultSetFromIterator(
        convertToIterator((ArrayList) result, pojoType),
        connection, parserResult, "", connection.server.getStatement(h), h,
        AvaticaStatement.DEFAULT_FETCH_SIZE, pojoType);
  } else if (result instanceof ResultSet) {
    // Querying JdbcDB
    return QuarkMetaResultSet.create(h.connectionId, h.id, (ResultSet) result,
        maxRowCount);
  } else if (result instanceof Iterator) {
    // Querying QuboleDB
    return getMetaResultSetFromIterator((Iterator<Object>) result,
        connection, parserResult, parserResult.getParsedSql(), stmt, h, maxRowCount, null);
  }

  throw new RuntimeException("Cannot handle execution for: " + parserResult.getParsedSql());
}
 

public CalciteResultSet newResultSet(AvaticaStatement statement, QueryState state,
    Meta.Signature signature, TimeZone timeZone, Meta.Frame firstFrame)
    throws SQLException {
  final ResultSetMetaData metaData =
      newResultSetMetaData(statement, signature);
  final CalcitePrepare.CalciteSignature calciteSignature =
      (CalcitePrepare.CalciteSignature) signature;
  return new CalciteResultSet(statement, calciteSignature, metaData, timeZone,
      firstFrame);
}
 

private static Meta.ExecuteResult prepareAndExecuteInternal(AvaticaConnection conn,
    final AvaticaStatement statement, String sql, int maxRowCount) throws Exception {
  Method m =
      AvaticaConnection.class.getDeclaredMethod("prepareAndExecuteInternal",
          AvaticaStatement.class, String.class, long.class);
  m.setAccessible(true);
  return (Meta.ExecuteResult) m.invoke(conn, statement, sql, maxRowCount);
}
 

/** Creates a CalciteResultSet. */
CalciteResultSet(AvaticaStatement statement,
    CalcitePrepare.CalciteSignature calciteSignature,
    ResultSetMetaData resultSetMetaData, TimeZone timeZone,
    Meta.Frame firstFrame) throws SQLException {
  super(statement, null, calciteSignature, resultSetMetaData, timeZone, firstFrame);
}
 

public CalciteResultSet newResultSet(AvaticaStatement statement, QueryState state,
    Meta.Signature signature, TimeZone timeZone, Meta.Frame firstFrame)
    throws SQLException {
  final ResultSetMetaData metaData =
      newResultSetMetaData(statement, signature);
  final CalcitePrepare.CalciteSignature calciteSignature =
      (CalcitePrepare.CalciteSignature) signature;
  return new CalciteResultSet(statement, calciteSignature, metaData, timeZone,
      firstFrame);
}
 

private static Meta.ExecuteResult prepareAndExecuteInternal(AvaticaConnection conn,
    final AvaticaStatement statement, String sql, int maxRowCount) throws Exception {
  Method m =
      AvaticaConnection.class.getDeclaredMethod("prepareAndExecuteInternal",
          AvaticaStatement.class, String.class, long.class);
  m.setAccessible(true);
  return (Meta.ExecuteResult) m.invoke(conn, statement, sql, maxRowCount);
}
 

/** Creates a CalciteResultSet. */
CalciteResultSet(AvaticaStatement statement,
    CalcitePrepare.CalciteSignature calciteSignature,
    ResultSetMetaData resultSetMetaData, TimeZone timeZone,
    Meta.Frame firstFrame) throws SQLException {
  super(statement, null, calciteSignature, resultSetMetaData, timeZone, firstFrame);
}
 

public AvaticaResultSet newResultSet(AvaticaStatement statement, QueryState state,
    Signature signature, TimeZone timeZone, Meta.Frame firstFrame)
    throws SQLException {
  final ResultSetMetaData metaData =
      newResultSetMetaData(statement, signature);
  QuicksqlConnectionImpl connection = (QuicksqlConnectionImpl)statement.getConnection();
  return new QuicksqlResultSet(statement, signature, metaData, timeZone,
      firstFrame);
}
 

/**
 * Creates a QuicksqlResultSet.
 */
public QuicksqlResultSet(AvaticaStatement statement,
    Meta.Signature calciteSignature,
    ResultSetMetaData resultSetMetaData, TimeZone timeZone,
    Meta.Frame firstFrame) throws SQLException {
    super(statement, null, calciteSignature, resultSetMetaData, timeZone, firstFrame);
}
 

protected QuicksqlServerResultSet getResultSet(StatementHandle h, String sql, int maxResNum, QueryResult result)
    throws Exception {
    final List<ColumnMetaData> columnMetaDataList = new ArrayList<>();
    columnMetaDataList.addAll(result.columnMeta);
    final StatementInfo info = getStatementCache().getIfPresent(h.id);
    Signature signature = preparedSignature(result, sql);
    Cursor cursor = MetaImpl.createCursor(signature.cursorFactory, result.iterable);
    QuicksqlResultSet quickSqlResultSet = new QuicksqlResultSet((AvaticaStatement) info.statement, signature,
        new AvaticaResultSetMetaData((AvaticaStatement) info.statement, null, signature), TimeZone.getDefault(),
        null);
    quickSqlResultSet.execute2(cursor, columnMetaDataList);
    info.setResultSet(quickSqlResultSet);
    return QuicksqlServerResultSet.create(h.connectionId, h.id, quickSqlResultSet, signature, maxResNum);
}