Examples with ResultSet com.datastax.oss.driver.api.core.cql.ResultSet used on opensource projects

Example 1 with ResultSet

use of com.datastax.oss.driver.api.core.cql.ResultSet in project janusgraph by JanusGraph.

the class CQLResultSetKeyIteratorTest method testIterator.

@Test
public void testIterator() throws IOException {
    final Array<Row> rows = Array.rangeClosed(1, 100).map(idx -> {
        final Row row = mock(Row.class);
        when(row.getByteBuffer("key")).thenReturn(ByteBuffer.wrap(Integer.toString(idx / 5).getBytes()));
        when(row.getByteBuffer("column1")).thenReturn(ByteBuffer.wrap(Integer.toString(idx % 5).getBytes()));
        when(row.getByteBuffer("value")).thenReturn(ByteBuffer.wrap(Integer.toString(idx).getBytes()));
        return row;
    });
    final ResultSet resultSet = mock(ResultSet.class);
    when(resultSet.iterator()).thenReturn(rows.iterator());
    final CQLColValGetter getter = new CQLColValGetter(new EntryMetaData[0]);
    try (final CQLResultSetKeyIterator resultSetKeyIterator = new CQLResultSetKeyIterator(ALL_COLUMNS, getter, resultSet)) {
        int i = 0;
        while (resultSetKeyIterator.hasNext()) {
            final StaticBuffer next = resultSetKeyIterator.next();
            final RecordIterator<Entry> entries = resultSetKeyIterator.getEntries();
            while (entries.hasNext()) {
                final Row row = rows.get(i++);
                final Entry entry = entries.next();
                assertEquals(row.getByteBuffer("key"), next.asByteBuffer());
                assertEquals(row.getByteBuffer("column1"), entry.getColumn().asByteBuffer());
                assertEquals(row.getByteBuffer("value"), entry.getValue().asByteBuffer());
            }
        }
    }
}

Example 2 with ResultSet

use of com.datastax.oss.driver.api.core.cql.ResultSet in project janusgraph by JanusGraph.

the class CQLResultSetKeyIteratorTest method testPartialIterateColumns.

@Test
public void testPartialIterateColumns() throws IOException {
    final Random random = new Random();
    final Array<Tuple2<ByteBuffer, Array<Tuple2<ByteBuffer, ByteBuffer>>>> keysMap = generateRandomKeysMap();
    final ResultSet resultSet = generateMockedResultSet(keysMap);
    final CQLColValGetter getter = new CQLColValGetter(new EntryMetaData[0]);
    try (final CQLResultSetKeyIterator resultSetKeyIterator = new CQLResultSetKeyIterator(ALL_COLUMNS, getter, resultSet)) {
        final Iterator<Tuple2<ByteBuffer, Array<Tuple2<ByteBuffer, ByteBuffer>>>> iterator = keysMap.iterator();
        while (resultSetKeyIterator.hasNext()) {
            final StaticBuffer next = resultSetKeyIterator.next();
            try (final RecordIterator<Entry> entries = resultSetKeyIterator.getEntries()) {
                final Tuple2<ByteBuffer, Array<Tuple2<ByteBuffer, ByteBuffer>>> current = iterator.next();
                final ByteBuffer currentKey = current._1;
                final Array<Tuple2<ByteBuffer, ByteBuffer>> columnValues = current._2;
                final Iterator<Tuple2<ByteBuffer, ByteBuffer>> columnIterator = columnValues.iterator();
                while (entries.hasNext()) {
                    final Entry entry = entries.next();
                    final Tuple2<ByteBuffer, ByteBuffer> columnAndValue = columnIterator.next();
                    assertEquals(currentKey, next.asByteBuffer());
                    assertEquals(columnAndValue._1, entry.getColumn().asByteBuffer());
                    assertEquals(columnAndValue._2, entry.getValue().asByteBuffer());
                    assertEquals(columnIterator.hasNext(), entries.hasNext());
                    // 10% of the time, don't complete the iteration
                    if (random.nextInt(10) == 0) {
                        break;
                    }
                }
            }
        }
    }
}

Example 3 with ResultSet

use of com.datastax.oss.driver.api.core.cql.ResultSet in project rocketmq-externals by apache.

the class Querier method poll.

public void poll() {
    try {
        LinkedList<Table> tableLinkedList = new LinkedList<>();
        for (Map.Entry<String, Database> entry : schema.getDbMap().entrySet()) {
            String dbName = entry.getKey();
            Iterator<Map.Entry<String, Table>> iterator = entry.getValue().getTableMap().entrySet().iterator();
            while (iterator.hasNext()) {
                Map.Entry<String, Table> tableEntry = iterator.next();
                String tableName = tableEntry.getKey();
                Table table = tableEntry.getValue();
                Map<String, String> tableFilterMap = table.getFilterMap();
                Select selectFrom = QueryBuilder.selectFrom(dbName, tableName).all();
                if (tableFilterMap != null && !tableFilterMap.keySet().contains("NO-FILTER")) {
                    for (String key : tableFilterMap.keySet()) {
                        String value = tableFilterMap.get(key);
                        selectFrom.whereColumn(key).isEqualTo(QueryBuilder.literal(value));
                    }
                }
                SimpleStatement stmt;
                boolean finishUpdate = false;
                log.info("trying to execute sql query,{}", selectFrom.asCql());
                ResultSet result = null;
                while (!cqlSession.isClosed() && !finishUpdate) {
                    stmt = selectFrom.build();
                    result = cqlSession.execute(stmt);
                    if (result.wasApplied()) {
                        log.info("query columns success, executed cql query {}", selectFrom.asCql());
                    }
                    finishUpdate = true;
                }
                List<String> colList = tableEntry.getValue().getColList();
                List<String> dataTypeList = tableEntry.getValue().getRawDataTypeList();
                List<ColumnParser> parserList = tableEntry.getValue().getParserList();
                for (Row row : result) {
                    Table tableWithData = new Table(dbName, tableName);
                    tableWithData.setColList(colList);
                    tableWithData.setRawDataTypeList(dataTypeList);
                    tableWithData.setParserList(parserList);
                    for (String col : colList) {
                        tableWithData.getDataList().add(row.getObject(col));
                    }
                    tableLinkedList.add(tableWithData);
                }
            }
        }
        list = tableLinkedList;
    } catch (Exception e) {
        log.error("fail to poll data, {}", e);
    }
}

Example 4 with ResultSet

use of com.datastax.oss.driver.api.core.cql.ResultSet in project rocketmq-externals by apache.

the class Schema method load.

public void load() throws Exception {
    dbMap = new HashMap<>();
    Select selectFrom = QueryBuilder.selectFrom(SCHEMA_SYSTEM_SCHEMA, TABLE_KEYSPACES).column(COLUMN_KEYSPACE_NAME);
    SimpleStatement stmt;
    boolean finishUpdate = false;
    LOGGER.info("trying to execute sql query,{}", selectFrom.asCql());
    ResultSet result = null;
    try {
        while (!cqlSession.isClosed() && !finishUpdate) {
            stmt = selectFrom.build();
            result = cqlSession.execute(stmt);
            if (result.wasApplied()) {
                LOGGER.info("update table success, executed cql query {}", selectFrom.asCql());
            }
            finishUpdate = true;
        }
        for (Row row : result) {
            String dbName = row.getString(COLUMN_KEYSPACE_NAME);
            if (!IGNORED_DATABASES.contains(dbName) && dbTableMap.keySet().contains(dbName)) {
                Database database = new Database(dbName, cqlSession, dbTableMap.get(dbName), tableFilterMap);
                dbMap.put(dbName, database);
            }
        }
    } catch (Exception e) {
        LOGGER.error("init cassandra Schema failure,{}", e);
    }
    for (Database db : dbMap.values()) {
        db.init();
    }
}

Example 5 with ResultSet

use of com.datastax.oss.driver.api.core.cql.ResultSet in project rocketmq-externals by apache.

the class Updater method deleteRow.

private boolean deleteRow(String dbName, String tableName, Map<Field, Object[]> fieldMap) {
    DeleteSelection deleteSelection = QueryBuilder.deleteFrom(dbName, tableName);
    Delete delete = null;
    int count = 0;
    for (Map.Entry<Field, Object[]> entry : fieldMap.entrySet()) {
        count++;
        String fieldName = entry.getKey().getName();
        FieldType fieldType = entry.getKey().getType();
        Object fieldValue = entry.getValue()[1];
        if (count == 1) {
            delete = deleteSelection.whereColumn(fieldName).isEqualTo(buildTerm(fieldType, fieldValue));
        } else {
            delete = delete.whereColumn(fieldName).isEqualTo(buildTerm(fieldType, fieldValue));
        }
    }
    boolean finishDelete = false;
    SimpleStatement stmt = delete.build();
    try {
        while (!cqlSession.isClosed() && !finishDelete) {
            ResultSet result = cqlSession.execute(stmt);
            if (result.wasApplied()) {
                log.info("delete from table success, executed query {}", delete);
                return true;
            }
            finishDelete = true;
        }
    } catch (Exception e) {
        log.error("delete from table error,{}", e);
    }
    return false;
}