org.apache.calcite.avatica.Meta.StatementHandle Java Examples

@Override
public DremioJdbc41PreparedStatement newPreparedStatement(AvaticaConnection connection,
                                                     StatementHandle h,
                                                     Meta.Signature signature,
                                                     int resultSetType,
                                                     int resultSetConcurrency,
                                                     int resultSetHoldability)
    throws SQLException {
  DremioConnectionImpl dremioConnection = (DremioConnectionImpl) connection;
  DremioClient client = dremioConnection.getClient();
  if (dremioConnection.getConfig().isServerPreparedStatementDisabled() || !client.getSupportedMethods().contains(ServerMethod.PREPARED_STATEMENT)) {
    // fallback to client side prepared statement
    return new DremioJdbc41PreparedStatement(dremioConnection, h, signature, null, resultSetType, resultSetConcurrency, resultSetHoldability);
  }
  return newServerPreparedStatement(dremioConnection, h, signature, resultSetType,
      resultSetConcurrency, resultSetHoldability);
}
 

@Override protected PreparedStatement newInstance() throws Exception {
  UnregisteredDriver driver = new TestDriver();
  AvaticaConnection connection = new AvaticaConnection(driver, driver.createFactory(),
      "jdbc:avatica", new Properties()) {
  };
  StatementHandle handle = mock(StatementHandle.class);
  Signature signature = new Signature(Collections.emptyList(), "", Collections.emptyList(),
      Collections.emptyMap(), null, Meta.StatementType.SELECT);
  AvaticaPreparedStatement preparedStatement = new AvaticaPreparedStatement(connection, handle,
      signature, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY,
      ResultSet.HOLD_CURSORS_OVER_COMMIT) {

  };

  preparedStatement.close();
  assertTrue("Prepared statement is not closed", preparedStatement.isClosed());
  return preparedStatement;
}
 

@Override protected ResultSet newInstance() throws Exception {
  UnregisteredDriver driver = new TestDriver();
  AvaticaConnection connection =
      new AvaticaConnection(driver, driver.createFactory(), "jdbc:avatica", new Properties()) {
      };
  StatementHandle handle = mock(StatementHandle.class);
  AvaticaStatement statement =
      new AvaticaStatement(connection, handle, ResultSet.TYPE_FORWARD_ONLY,
          ResultSet.CONCUR_READ_ONLY, ResultSet.HOLD_CURSORS_OVER_COMMIT) {
      };
  Signature signature = new Signature(Collections.emptyList(), "", Collections.emptyList(),
      Collections.emptyMap(), null, Meta.StatementType.SELECT);
  AvaticaResultSet resultSet = new AvaticaResultSet(statement, new QueryState(""), signature,
      null, DateTimeUtils.UTC_ZONE, null);
  resultSet.close();
  assertTrue("Resultset is not closed", resultSet.isClosed());

  return resultSet;
}
 

@Test public void testExecuteSerialization() throws Exception {
  Service.ExecuteRequest executeRequest = new Service.ExecuteRequest(
      new StatementHandle("connection", 12345, getSignature()), getTypedValues(), 0);

  Requests.ExecuteRequest pbExecuteRequest = executeRequest.serialize();
  ByteArrayOutputStream baos = new ByteArrayOutputStream(1024);
  pbExecuteRequest.writeTo(baos);

  byte[] serialized = baos.toByteArray();
  baos.reset();
  WireMessage wireMsg = WireMessage.newBuilder().setName(Requests.ExecuteRequest.class.getName())
      .setWrappedMessage(UnsafeByteOperations.unsafeWrap(serialized)).build();
  wireMsg.writeTo(baos);
  serialized = baos.toByteArray();

  ProtobufTranslation translator = new ProtobufTranslationImpl();

  Request newRequest = translator.parseRequest(serialized);

  Assert.assertEquals(executeRequest, newRequest);
}
 

@Test public void testPrepareAndExecuteSetsMaxRows() throws Exception {
  final String id = UUID.randomUUID().toString();
  final int statementId = 12345;
  final String sql = "SELECT * FROM FOO";
  final int maxRows = 500;

  JdbcMeta meta = Mockito.mock(JdbcMeta.class);
  PreparedStatement statement = Mockito.mock(PreparedStatement.class);
  @SuppressWarnings("unchecked")
  Cache<Integer, StatementInfo> statementCache =
      (Cache<Integer, StatementInfo>) Mockito.mock(Cache.class);
  Signature signature = Mockito.mock(Signature.class);

  final StatementInfo statementInfo = new StatementInfo(statement);
  final StatementHandle statementHandle = new StatementHandle(id, statementId, signature);

  Mockito.when(meta.getStatementCache()).thenReturn(statementCache);
  Mockito.when(statementCache.getIfPresent(statementId)).thenReturn(statementInfo);
  Mockito.when(statement.getResultSet()).thenReturn(null);
  // The real methods
  Mockito.when(meta.prepareAndExecute(statementHandle, sql, maxRows, 50, null))
      .thenCallRealMethod();
  Mockito.doCallRealMethod().when(meta).setMaxRows(statement, maxRows);

  // Call our method
  meta.prepareAndExecute(statementHandle, sql, maxRows, 50, null);

  // Verify we called setMaxRows with the right value
  Mockito.verify(statement).setMaxRows(maxRows);
}
 

DremioJdbc41PreparedStatement(DremioConnectionImpl connection,
                             StatementHandle h,
                             Meta.Signature signature,
                             com.dremio.exec.proto.UserProtos.PreparedStatement pstmt,
                             int resultSetType,
                             int resultSetConcurrency,
                             int resultSetHoldability) throws SQLException {
  super(connection, h, signature, pstmt,
        resultSetType, resultSetConcurrency, resultSetHoldability);
}
 

@Override
public DremioStatementImpl newStatement(AvaticaConnection connection,
                                       StatementHandle h,
                                       int resultSetType,
                                       int resultSetConcurrency,
                                       int resultSetHoldability) {
  return new DremioStatementImpl((DremioConnectionImpl) connection,
                                h,
                                resultSetType,
                                resultSetConcurrency,
                                resultSetHoldability);
}
 

@Override protected Statement newInstance() throws Exception {
  UnregisteredDriver driver = new TestDriver();
  AvaticaConnection connection = new AvaticaConnection(driver, driver.createFactory(),
      "jdbc:avatica", new Properties()) {
  };
  StatementHandle handle = mock(StatementHandle.class);
  AvaticaStatement statement = new AvaticaStatement(connection, handle,
      ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY,
      ResultSet.HOLD_CURSORS_OVER_COMMIT) {
  };
  statement.close();
  assertTrue("Statement is not closed", statement.isClosed());

  return statement;
}
 

protected DremioPreparedStatementImpl(DremioConnectionImpl connection,
                                     StatementHandle h,
                                     Meta.Signature signature,
                                     PreparedStatement preparedStatementHandle,
                                     int resultSetType,
                                     int resultSetConcurrency,
                                     int resultSetHoldability) throws SQLException {
  super(connection, h, signature,
        resultSetType, resultSetConcurrency, resultSetHoldability);
  connection.openStatementsRegistry.addStatement(this);
  this.preparedStatementHandle = preparedStatementHandle;
  if (preparedStatementHandle != null) {
    ((DremioColumnMetaDataList) signature.columns).updateColumnMetaData(preparedStatementHandle.getColumnsList());
  }
}
 

protected KylinStatement(AvaticaConnection connection, StatementHandle h, int resultSetType, int resultSetConcurrency, int resultSetHoldability) {
    super(connection, h, resultSetType, resultSetConcurrency, resultSetHoldability);
}
 

protected KylinPreparedStatement(AvaticaConnection connection, StatementHandle h, Signature signature, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
    super(connection, h, signature, resultSetType, resultSetConcurrency, resultSetHoldability);
    if (this.handle.signature == null)
        this.handle.signature = signature;
}
 

@Override
public AvaticaPreparedStatement newPreparedStatement(AvaticaConnection connection, StatementHandle h, Signature signature, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
    return new KylinPreparedStatement((KylinConnection) connection, h, signature, resultSetType, resultSetConcurrency, resultSetHoldability);
}
 

@Override
public AvaticaStatement newStatement(AvaticaConnection connection, StatementHandle h, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
    return new KylinStatement((KylinConnection) connection, h, resultSetType, resultSetConcurrency, resultSetHoldability);
}
 

protected KylinPreparedStatement(AvaticaConnection connection, StatementHandle h, Signature signature, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
    super(connection, h, signature, resultSetType, resultSetConcurrency, resultSetHoldability);
    if (this.handle.signature == null)
        this.handle.signature = signature;
}
 

@Override
protected AvaticaStatement lookupStatement(StatementHandle h) throws SQLException {
  return super.lookupStatement(h);
}
 

DremioStatementImpl(DremioConnectionImpl connection, StatementHandle h, int resultSetType,
                   int resultSetConcurrency, int resultSetHoldability) {
  super(connection, h, resultSetType, resultSetConcurrency, resultSetHoldability);
  this.connection = connection;
  connection.openStatementsRegistry.addStatement(this);
}
 

public StatementHandle getHandle() {
  return handle;
}
 

public MissingResultsException(StatementHandle handle) {
  this.handle = handle;
}
 

public StatementHandle getStatementHandle() {
  return stmtHandle;
}
 

public NoSuchStatementException(StatementHandle stmtHandle) {
  this.stmtHandle = stmtHandle;
}
 

protected KylinStatement(AvaticaConnection connection, StatementHandle h, int resultSetType, int resultSetConcurrency, int resultSetHoldability) {
    super(connection, h, resultSetType, resultSetConcurrency, resultSetHoldability);
}
 

@Override
public AvaticaPreparedStatement newPreparedStatement(AvaticaConnection connection, StatementHandle h, Signature signature, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
    return new KylinPreparedStatement((KylinConnection) connection, h, signature, resultSetType, resultSetConcurrency, resultSetHoldability);
}
 

@Override
public AvaticaStatement newStatement(AvaticaConnection connection, StatementHandle h, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
    return new KylinStatement((KylinConnection) connection, h, resultSetType, resultSetConcurrency, resultSetHoldability);
}