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

@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;
}
 

@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;
}
 

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);
  }
}
 

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);
}
 

Signature mockPreparedSignature(String sql) {
    List<AvaticaParameter> params = new ArrayList<AvaticaParameter>();
    int startIndex = 0;
    while (sql.indexOf("?", startIndex) >= 0) {
        AvaticaParameter param = new AvaticaParameter(false, 0, 0, 0, null, null, null);
        params.add(param);
        startIndex = sql.indexOf("?", startIndex) + 1;
    }

    ArrayList<ColumnMetaData> columns = new ArrayList<ColumnMetaData>();
    Map<String, Object> internalParams = Collections.<String, Object> emptyMap();

    return new Meta.Signature(columns, sql, params, internalParams, CursorFactory.ARRAY, Meta.StatementType.SELECT);
}
 

/**
 *
 * @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());
}
 

public ResultSet executeQuery() throws SQLException {
  checkOpen();
  this.updateCount = -1;
  final Signature sig = getSignature();
  return getConnection().executeQueryInternal(this, sig, null,
      new QueryState(sig.sql), false);
}
 

/**
 * Creates an AvaticaPreparedStatement.
 *
 * @param connection Connection
 * @param h Statement handle
 * @param signature Result of preparing statement
 * @param resultSetType Result set type
 * @param resultSetConcurrency Result set concurrency
 * @param resultSetHoldability Result set holdability
 * @throws SQLException If fails due to underlying implementation reasons.
 */
protected AvaticaPreparedStatement(AvaticaConnection connection,
    Meta.StatementHandle h,
    Meta.Signature signature,
    int resultSetType,
    int resultSetConcurrency,
    int resultSetHoldability) throws SQLException {
  super(connection, h, resultSetType, resultSetConcurrency,
      resultSetHoldability, signature);
  this.slots = new TypedValue[signature.parameters.size()];
  this.resultSetMetaData =
      connection.factory.newResultSetMetaData(this, signature);
  this.parameterValueBatch = new ArrayList<>();
}
 

@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);
}
 

Signature mockPreparedSignature(String sql) {
    List<AvaticaParameter> params = new ArrayList<AvaticaParameter>();
    int startIndex = 0;
    while (sql.indexOf("?", startIndex) >= 0) {
        AvaticaParameter param = new AvaticaParameter(false, 0, 0, 0, null, null, null);
        params.add(param);
        startIndex = sql.indexOf("?", startIndex) + 1;
    }

    ArrayList<ColumnMetaData> columns = new ArrayList<ColumnMetaData>();
    Map<String, Object> internalParams = Collections.<String, Object> emptyMap();

    return new Meta.Signature(columns, sql, params, internalParams, CursorFactory.ARRAY, Meta.StatementType.SELECT);
}
 

public Signature mockPreparedSignature(String sql) {
    List<AvaticaParameter> params = new ArrayList<AvaticaParameter>();
    int startIndex = 0;
    while (sql.indexOf("?", startIndex) >= 0) {
        AvaticaParameter param = new AvaticaParameter(false, 0, 0, 0, null, null, null);
        params.add(param);
        startIndex = sql.indexOf("?", startIndex) + 1;
    }

    ArrayList<ColumnMetaData> columns = new ArrayList<ColumnMetaData>();
    Map<String, Object> internalParams = Collections.<String, Object> emptyMap();

    return new Meta.Signature(columns, sql, params, internalParams, CursorFactory.ARRAY, Meta.StatementType.SELECT);
}
 

public AvaticaPreparedStatement newPreparedStatement(
    AvaticaConnection connection,
    Meta.StatementHandle h,
    Signature signature,
    int resultSetType,
    int resultSetConcurrency,
    int resultSetHoldability) throws SQLException {
  return new QuickSqlJdbc41PreparedStatement(
      (QuicksqlConnectionImpl) connection, h,
       signature, resultSetType,
      resultSetConcurrency, resultSetHoldability);
}
 

QuickSqlJdbc41PreparedStatement(QuicksqlConnectionImpl connection,
    Meta.StatementHandle h, Signature signature,
    int resultSetType, int resultSetConcurrency, int resultSetHoldability)
    throws SQLException {
  super(connection, h, signature, resultSetType, resultSetConcurrency,
      resultSetHoldability);
}
 

/** Creates a result set with maxRowCount.
 *
 * <p>If {@code maxRowCount} is -2 ({@link QuicksqlServerMeta#UNLIMITED_COUNT}),
 * returns an unlimited number of rows in a single frame; any other
 * negative value (typically -1) returns an unlimited number of rows
 * in frames of the default frame size. */
public static QuicksqlServerResultSet create(String connectionId, int statementId,
    ResultSet resultSet, int maxRowCount) {
  try {
    Signature sig = QuicksqlServerMeta.signature(resultSet.getMetaData());
    return create(connectionId, statementId, resultSet, maxRowCount, sig);
  } catch (SQLException e) {
    throw new RuntimeException(e);
  }
}
 

/**
 * Generates a collection of Requests whose serialization will be tested.
 */
private static List<Request> getRequests() {
  LinkedList<Request> requests = new LinkedList<>();

  requests.add(new CatalogsRequest());
  requests.add(new DatabasePropertyRequest());
  requests.add(new SchemasRequest("connectionId", "catalog", "schemaPattern"));
  requests.add(
      new TablesRequest("connectionId", "catalog", "schemaPattern", "tableNamePattern",
          Arrays.asList("STRING", "BOOLEAN", "INT")));
  requests.add(new TableTypesRequest());
  requests.add(
      new ColumnsRequest("connectionId", "catalog", "schemaPattern", "tableNamePattern",
          "columnNamePattern"));
  requests.add(new TypeInfoRequest());
  requests.add(
      new PrepareAndExecuteRequest("connectionId", Integer.MAX_VALUE, "sql",
          Long.MAX_VALUE));
  requests.add(new PrepareRequest("connectionId", "sql", Long.MAX_VALUE));

  List<TypedValue> paramValues =
      Arrays.asList(TypedValue.create(Rep.BOOLEAN.name(), Boolean.TRUE),
          TypedValue.create(Rep.STRING.name(), "string"));
  FetchRequest fetchRequest = new FetchRequest("connectionId", Integer.MAX_VALUE,
      Long.MAX_VALUE, Integer.MAX_VALUE);
  requests.add(fetchRequest);

  requests.add(new CreateStatementRequest("connectionId"));
  requests.add(new CloseStatementRequest("connectionId", Integer.MAX_VALUE));
  Map<String, String> info = new HashMap<>();
  info.put("param1", "value1");
  info.put("param2", "value2");
  requests.add(new OpenConnectionRequest("connectionId", info));
  requests.add(new CloseConnectionRequest("connectionId"));
  requests.add(
      new ConnectionSyncRequest("connectionId",
          new ConnectionPropertiesImpl(Boolean.FALSE, Boolean.FALSE,
              Integer.MAX_VALUE, "catalog", "schema")));

  requests.add(new SyncResultsRequest("connectionId", 12345, getSqlQueryState(), 150));
  requests.add(new SyncResultsRequest("connectionId2", 54321, getMetadataQueryState1(), 0));
  requests.add(new SyncResultsRequest("connectionId3", 5, getMetadataQueryState2(), 10));

  requests.add(new CommitRequest("connectionId"));
  requests.add(new RollbackRequest("connectionId"));

  // ExecuteBatchRequest omitted because of the special protobuf conversion it does

  List<String> commands = Arrays.asList("command1", "command2", "command3");
  requests.add(new PrepareAndExecuteBatchRequest("connectionId", 12345, commands));


  List<ColumnMetaData> columns = Collections.emptyList();
  List<AvaticaParameter> params = Collections.emptyList();
  Meta.CursorFactory cursorFactory = Meta.CursorFactory.create(Style.LIST, Object.class,
      Collections.<String>emptyList());
  Signature signature = Signature.create(columns, "sql", params, cursorFactory,
      Meta.StatementType.SELECT);
  Meta.StatementHandle handle = new Meta.StatementHandle("1234", 1, signature);
  requests.add(new ExecuteRequest(handle, Arrays.<TypedValue>asList((TypedValue) null), 10));
  requests.add(new ExecuteRequest(handle, Arrays.asList(TypedValue.EXPLICIT_NULL), 10));

  return requests;
}
 

/** Creates a frame containing a given number or unlimited number of rows
 * from a result set. */
static Meta.Frame frame(StatementInfo info, ResultSet resultSet, long offset,
    int fetchMaxRowCount, Calendar calendar, Optional<Signature> sig) throws SQLException {
  final ResultSetMetaData metaData = resultSet.getMetaData();
  final int columnCount = metaData.getColumnCount();
  final int[] types = new int[columnCount];
  Set<Integer> arrayOffsets = new HashSet<>();
  for (int i = 0; i < types.length; i++) {
    types[i] = metaData.getColumnType(i + 1);
    if (Types.ARRAY == types[i]) {
      arrayOffsets.add(i);
    }
  }
  final List<Object> rows = new ArrayList<>();
  // Meta prepare/prepareAndExecute 0 return 0 row and done
  boolean done = fetchMaxRowCount == 0;
  for (int i = 0; fetchMaxRowCount < 0 || i < fetchMaxRowCount; i++) {
    final boolean hasRow;
    if (null != info) {
      hasRow = info.next();
    } else {
      hasRow = resultSet.next();
    }
    if (!hasRow) {
      done = true;
      resultSet.close();
      break;
    }
    Object[] columns = new Object[columnCount];
    for (int j = 0; j < columnCount; j++) {
      columns[j] = getValue(resultSet, types[j], j, calendar);
      if (arrayOffsets.contains(j)) {
        // If we have an Array type, our Signature is lacking precision. We can't extract the
        // component type of an Array from metadata, we have to update it as we're serializing
        // the ResultSet.
        final Array array = resultSet.getArray(j + 1);
        // Only attempt to determine the component type for the array when non-null
        if (null != array && sig.isPresent()) {
          ColumnMetaData columnMetaData = sig.get().columns.get(j);
          ArrayType arrayType = (ArrayType) columnMetaData.type;
          SqlType componentSqlType = SqlType.valueOf(array.getBaseType());

          // Avatica Server will always return non-primitives to ensure nullable is guaranteed.
          ColumnMetaData.Rep rep = ColumnMetaData.Rep.serialRepOf(componentSqlType);
          AvaticaType componentType = ColumnMetaData.scalar(array.getBaseType(),
              array.getBaseTypeName(), rep);
          // Update the ArrayType from the Signature
          arrayType.updateComponentType(componentType);

          // We only need to update the array's type once.
          arrayOffsets.remove(j);
        }
      }
    }
    rows.add(columns);
  }
  // final List<Object> rows = new ArrayList<>();
  // rows.addAll(signature.columns);
  return new Meta.Frame(offset, done, rows);
}
 

@Override
public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
    Meta.Signature sig = mockPreparedSignature(sql);
    return factory().newPreparedStatement(this, null, sig, resultSetType, resultSetConcurrency, resultSetHoldability);
}
 

public KylinResultSet(AvaticaStatement statement, QueryState state, Signature signature, ResultSetMetaData resultSetMetaData, TimeZone timeZone, Frame firstFrame) throws SQLException {
    super(statement, state, signature, resultSetMetaData, timeZone, firstFrame);
}
 

@Override
public ResultSetMetaData newResultSetMetaData(AvaticaStatement statement, Signature signature) throws SQLException {
    return new AvaticaResultSetMetaData(statement, null, signature);
}
 

@Override
public AvaticaResultSet newResultSet(AvaticaStatement statement, QueryState state, Signature signature, TimeZone timeZone, Frame firstFrame) throws SQLException {
    AvaticaResultSetMetaData resultSetMetaData = new AvaticaResultSetMetaData(statement, null, signature);
    return new KylinResultSet(statement, state, signature, resultSetMetaData, timeZone, firstFrame);
}
 

@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);
}
 

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;
}
 

public ResultSetMetaData newResultSetMetaData(AvaticaStatement statement,
    Signature signature) {
  return new AvaticaResultSetMetaData(statement, null, signature);
}
 

private Signature getSignature() {
  return null;
}
 

/**
 * Generates a collection of Responses whose serialization will be tested.
 */
private static List<Response> getResponses() {
  final RpcMetadataResponse rpcMetadata = new RpcMetadataResponse("localhost:8765");
  LinkedList<Response> responses = new LinkedList<>();

  // Nested classes (Signature, ColumnMetaData, CursorFactory, etc) are implicitly getting tested)

  // Stub out the metadata for a row
  ScalarType arrayComponentType = ColumnMetaData.scalar(Types.INTEGER, "integer", Rep.INTEGER);
  ColumnMetaData arrayColumnMetaData = getArrayColumnMetaData(arrayComponentType, 2, "counts");
  List<ColumnMetaData> columns =
      Arrays.asList(MetaImpl.columnMetaData("str", 0, String.class, true),
          MetaImpl.columnMetaData("count", 1, Integer.class, true),
          arrayColumnMetaData);
  List<AvaticaParameter> params =
      Arrays.asList(
          new AvaticaParameter(false, 10, 0, Types.VARCHAR, "VARCHAR",
              String.class.getName(), "str"));
  Meta.CursorFactory cursorFactory = Meta.CursorFactory.create(Style.LIST, Object.class,
      Arrays.asList("str", "count", "counts"));
  // The row values
  List<Object> rows = new ArrayList<>();
  rows.add(new Object[] {"str_value1", 50, Arrays.asList(1, 2, 3)});
  rows.add(new Object[] {"str_value2", 100, Arrays.asList(1)});

  // Create the signature and frame using the metadata and values
  Signature signature = Signature.create(columns, "sql", params, cursorFactory,
      Meta.StatementType.SELECT);
  Frame frame = Frame.create(Integer.MAX_VALUE, true, rows);

  // And then create a ResultSetResponse
  ResultSetResponse results1 = new ResultSetResponse("connectionId", Integer.MAX_VALUE, true,
      signature, frame, Long.MAX_VALUE, rpcMetadata);
  responses.add(results1);

  responses.add(new CloseStatementResponse(rpcMetadata));

  ConnectionPropertiesImpl connProps = new ConnectionPropertiesImpl(false, true,
      Integer.MAX_VALUE, "catalog", "schema");
  responses.add(new ConnectionSyncResponse(connProps, rpcMetadata));

  responses.add(new OpenConnectionResponse(rpcMetadata));
  responses.add(new CloseConnectionResponse(rpcMetadata));

  responses.add(new CreateStatementResponse("connectionId", Integer.MAX_VALUE, rpcMetadata));

  Map<Meta.DatabaseProperty, Object> propertyMap = new HashMap<>();
  for (Meta.DatabaseProperty prop : Meta.DatabaseProperty.values()) {
    propertyMap.put(prop, prop.defaultValue);
  }
  responses.add(new DatabasePropertyResponse(propertyMap, rpcMetadata));

  responses.add(
      new ExecuteResponse(Arrays.asList(results1, results1, results1), false, rpcMetadata));
  responses.add(new FetchResponse(frame, false, false, rpcMetadata));
  responses.add(new FetchResponse(frame, true, true, rpcMetadata));
  responses.add(new FetchResponse(frame, false, true, rpcMetadata));
  responses.add(
      new PrepareResponse(
          new Meta.StatementHandle("connectionId", Integer.MAX_VALUE, signature),
          rpcMetadata));

  StringWriter sw = new StringWriter();
  new Exception().printStackTrace(new PrintWriter(sw));
  responses.add(
      new ErrorResponse(Collections.singletonList(sw.toString()), "Test Error Message",
          ErrorResponse.UNKNOWN_ERROR_CODE, ErrorResponse.UNKNOWN_SQL_STATE,
          AvaticaSeverity.WARNING, rpcMetadata));

  // No more results, statement not missing
  responses.add(new SyncResultsResponse(false, false, rpcMetadata));
  // Missing statement, no results
  responses.add(new SyncResultsResponse(false, true, rpcMetadata));
  // More results, no missing statement
  responses.add(new SyncResultsResponse(true, false, rpcMetadata));

  // Some tests to make sure ErrorResponse doesn't fail.
  responses.add(new ErrorResponse((List<String>) null, null, 0, null, null, null));
  responses.add(
      new ErrorResponse(Arrays.asList("stacktrace1", "stacktrace2"), null, 0, null, null, null));

  responses.add(new CommitResponse());
  responses.add(new RollbackResponse());

  long[] updateCounts = new long[]{1, 0, 1, 1};
  responses.add(
      new ExecuteBatchResponse("connectionId", 12345, updateCounts, false, rpcMetadata));

  return responses;
}
 

@Override
public AvaticaResultSet newResultSet(AvaticaStatement statement, QueryState state, Signature signature, TimeZone timeZone, Frame firstFrame) throws SQLException {
    AvaticaResultSetMetaData resultSetMetaData = new AvaticaResultSetMetaData(statement, null, signature);
    return new KylinResultSet(statement, state, signature, resultSetMetaData, timeZone, firstFrame);
}
 

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;
}
 

/** Creates a empty result set with empty frame */
public static QuicksqlServerResultSet empty(String connectionId, int statementId,
    Signature signature) {
  return new QuicksqlServerResultSet(connectionId, statementId, true, signature,
      Meta.Frame.EMPTY);
}
 

protected QuicksqlServerResultSet(String connectionId, int statementId,
    boolean ownStatement, Signature signature, Meta.Frame firstFrame) {
  this(connectionId, statementId, ownStatement, signature, firstFrame, -1L);
}
 

protected QuicksqlServerResultSet(String connectionId, int statementId,
    boolean ownStatement, Signature signature, Meta.Frame firstFrame,
    long updateCount) {
  super(connectionId, statementId, ownStatement, signature, firstFrame, updateCount);
}