com.datastax.driver.core.ColumnDefinitions.Definition Java Examples

/**
 * This method creates a Beam {@link Schema schema} based on the list of Definition supplied as
 * input. If a field name matches a key from the keyOrder map (supplied as a input parameter) this
 * denotes that the field was previously part of the Cassandra primary-key. To preserve the order
 * of these primary keys metadata is added to the field on the format "key_order": X - where X is
 * an integer indicating the order of the key in Cassandra.
 *
 * @param definitions a list of field definitions.
 * @param keyOrder the order of the fields that make up the Cassandra row key. key is field name,
 *     value denotes the order from the row-key.
 * @return a {@link Schema schema} schema for all fields supplied as input.
 */
private Schema createBeamRowSchema(List<Definition> definitions, Map<String, Integer> keyOrder) {

  List<Field> fields =
      definitions.stream()
          .map(def -> Field.of(def.getName(), toBeamRowType(def.getType())).withNullable(true))
          .map(
              fld ->
                  keyOrder.containsKey(fld.getName())
                      ? fld.withType(
                          fld.getType()
                              .withMetadata(KEY_ORDER_METADATA_KEY, keyOrder.get(fld.getName()).toString()))
                      : fld)
          .collect(Collectors.toList());

  return Schema.builder().addFields(fields).build();
}
 

/**
 * Returns a map with column name as key and column date type as value.
 *
 * The value might be as simple as "Boolean" or more complex like
 *  - "Set|Boolean"
 *  - "List|String"
 *  - "Map|String|Integer"
 *  these are cases when the data type is a container of primitive data types.
 *
 * @param tableName
 * @return
 * @throws DbException
 */
public Map<String, String> getColumnInfo(
                                          String tableName ) throws DbException {

    connect();

    ResultSet results = session.execute("SELECT * FROM " + this.dbName + "." + tableName + " LIMIT 1");

    Map<String, String> columnInfo = new HashMap<String, String>();
    for (Definition columnDefinition : results.getColumnDefinitions()) {
        DataType dataType = columnDefinition.getType();
        String dataTypeName = dataType.getName().name();
        if ("Set".equalsIgnoreCase(dataTypeName)) {
            dataTypeName = dataTypeName + "|" + dataType.getTypeArguments().get(0);
        } else if ("List".equalsIgnoreCase(dataTypeName)) {
            dataTypeName = dataTypeName + "|" + dataType.getTypeArguments().get(0);
        } else if ("Map".equalsIgnoreCase(dataTypeName)) {
            dataTypeName = dataTypeName + "|" + dataType.getTypeArguments().get(0) + "|"
                           + dataType.getTypeArguments().get(1);
        }
        columnInfo.put(columnDefinition.getName(), dataTypeName);
    }

    return columnInfo;
}
 

public RawNumericMetric map(Row row) {
    RawNumericMetric metricRow =  new RawNumericMetric(row.getString(0), row.getString(1), row.getDate(2).getTime(), row.getInt(3));
    ColumnDefinitions columeDef = row.getColumnDefinitions();
    List<Definition> columeDefList = columeDef.asList();
    Map<String, String> tagMap = new HashMap<String, String>();
    for(Definition def: columeDefList){
        if(def.getName().startsWith("tag_")){
            tagMap.put(def.getName(), row.getString(def.getName()));
        }
    }
    
    if(tagMap.size()>0){
        metricRow.setTagMap(tagMap);
    }
    return metricRow;
}
 

@VisibleForTesting
Map<String, DataType> getColumnDefinitions(Row row) {
  return row.getColumnDefinitions()
      .asList()
      .stream()
      .collect(
          Collectors.toMap(
              Definition::getName,
              Definition::getType,
              (d1, d2) -> {
                return d1;
              }));
}
 

/**
 * @param resultSet the {@link ResultSet} to map.
 * @return an iterator containing the mapped rows.
 * @see ResultSet
 * @see Row
 */
@Override
public Iterator<Row> map(ResultSet resultSet) {
  List<Definition> columnDefinitions = resultSet.getColumnDefinitions().asList();

  // If there is no schema then generate one. We assume all rows have the same schema.
  if (schema == null) {
    // Get the order of the primary keys. We store the order of the primary keys in the field
    // metadata.
    Map<String, Integer> keyOrder =
        CassandraKeyUtils.primaryKeyOrder(session, keyspace.get(), table.get());
    schema = createBeamRowSchema(columnDefinitions, keyOrder);
  }

  List<Row> rows = new ArrayList<>();
  // Loop over each cassandra row and covert this to a Beam row.
  for (com.datastax.driver.core.Row cassandraRow : resultSet) {
    List<Object> values = new ArrayList<>();
    for (int i = 0; i < columnDefinitions.size(); i++) {
      DataType columnType = columnDefinitions.get(i).getType();
      Object columnValue = mapValue(cassandraRow.getObject(i), columnType);
      values.add(columnValue);
    }
    Row row = Row.withSchema(schema).addValues(values).build();

    rows.add(row);
  }
  return rows.iterator();
}
 

protected CassandraResultSet(BaseCassandraStatement statement, CassandraCqlStatement parsedStmt, ResultSet rs) {
    super(statement, parsedStmt);

    if (rs != null) {
        for (Definition def : rs.getColumnDefinitions()) {
            CassandraColumnDefinition d = new CassandraColumnDefinition(
                    def.getKeyspace(), def.getTable(), def.getName(), def
                    .getType().getName().toString(), false);
            metadata.addColumnDefinition(d);
        }
    }

    _resultSet = rs;
}
 

@SuppressWarnings("rawtypes")
public int getColumnDisplaySize(int column) throws SQLException
      {
          //checkIndex(column);
          Definition col = null;
          if(currentRow!=null){
          	col = currentRow.getColumnDefinitions().asList().get(column-1);
          }else{
          	col = driverResultSet.getColumnDefinitions().asList().get(column-1);
          }
          try{
          	
          	int length = -1;
           AbstractJdbcType jtype = TypesMap.getTypeForComparator(col.getType().toString());
           if (jtype instanceof JdbcBytes) length = Integer.MAX_VALUE / 2;
           if (jtype instanceof JdbcAscii || jtype instanceof JdbcUTF8) length = Integer.MAX_VALUE;
           if (jtype instanceof JdbcUUID) length = 36;
           if (jtype instanceof JdbcInt32) length = 4;
           if (jtype instanceof JdbcLong) length = 8;
          	// String stringValue = getObject(column).toString();
          	//return (stringValue == null ? -1 : stringValue.length());
           
           return length;
          }catch(Exception e){
          	return -1;
          }
          //return -1;
      }
 

@Override
public DbRecordValuesList[] select(
                                    DbQuery dbQuery,
                                    DbReturnModes dbReturnMode ) throws DbException {

    connect();

    ArrayList<DbRecordValuesList> dbRecords = new ArrayList<DbRecordValuesList>();

    String sqlQuery = dbQuery.getQuery();
    if (allowFiltering) {
        sqlQuery += " ALLOW FILTERING";
    }

    if (log.isDebugEnabled()) {
        log.debug(sqlQuery);
    }

    ResultSet results = session.execute(sqlQuery);

    int currentRow = 0;
    Iterator<Row> it = results.iterator();
    while (it.hasNext()) {
        Row row = it.next();

        currentRow++;
        if (log.isDebugEnabled()) {
            log.debug("Result row number: " + currentRow);
        }

        DbRecordValuesList recordList = new DbRecordValuesList();

        for (Definition columnDefinition : row.getColumnDefinitions()) {
            DbColumn dbColumn = new DbColumn(columnDefinition.getTable(), columnDefinition.getName());
            dbColumn.setColumnType(columnDefinition.getType().getName().toString());

            Object value = extractObjectFromResultSet(row, columnDefinition);

            DbRecordValue recordValue = new DbRecordValue(dbColumn, value);
            recordList.add(recordValue);
        }
        dbRecords.add(recordList);
    }

    return dbRecords.toArray(new DbRecordValuesList[]{});
}