Examples with RecordSet io.trino.spi.connector.RecordSet used on opensource projects

Example 6 with RecordSet

use of io.trino.spi.connector.RecordSet in project trino by trinodb.

the class TestPrometheusRecordSet method testCursorSimple.

@Test
public void testCursorSimple() {
    RecordSet recordSet = new PrometheusRecordSet(new PrometheusClient(new PrometheusConnectorConfig(), METRIC_CODEC, TESTING_TYPE_MANAGER), new PrometheusSplit(dataUri), ImmutableList.of(new PrometheusColumnHandle("labels", varcharMapType, 0), new PrometheusColumnHandle("timestamp", TIMESTAMP_COLUMN_TYPE, 1), new PrometheusColumnHandle("value", DoubleType.DOUBLE, 2)));
    RecordCursor cursor = recordSet.cursor();
    assertEquals(cursor.getType(0), varcharMapType);
    assertEquals(cursor.getType(1), TIMESTAMP_COLUMN_TYPE);
    assertEquals(cursor.getType(2), DoubleType.DOUBLE);
    List<PrometheusStandardizedRow> actual = new ArrayList<>();
    while (cursor.advanceNextPosition()) {
        actual.add(new PrometheusStandardizedRow((Block) cursor.getObject(0), (Instant) cursor.getObject(1), cursor.getDouble(2)));
        assertFalse(cursor.isNull(0));
        assertFalse(cursor.isNull(1));
        assertFalse(cursor.isNull(2));
    }
    List<PrometheusStandardizedRow> expected = ImmutableList.<PrometheusStandardizedRow>builder().add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType, ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), ofEpochMilli(1565962969044L), 1.0)).add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType, ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), ofEpochMilli(1565962984045L), 1.0)).add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType, ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), ofEpochMilli(1565962999044L), 1.0)).add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType, ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), ofEpochMilli(1565963014044L), 1.0)).build();
    assertThat(actual).as("actual").hasSize(expected.size());
    for (int i = 0; i < actual.size(); i++) {
        PrometheusStandardizedRow actualRow = actual.get(i);
        PrometheusStandardizedRow expectedRow = expected.get(i);
        assertEquals(getMapFromBlock(varcharMapType, actualRow.getLabels()), getMapFromBlock(varcharMapType, expectedRow.getLabels()));
        assertEquals(actualRow.getTimestamp(), expectedRow.getTimestamp());
        assertEquals(actualRow.getValue(), expectedRow.getValue());
    }
}

Example 7 with RecordSet

use of io.trino.spi.connector.RecordSet in project trino by trinodb.

the class TestPrometheusRecordSetProvider method testGetRecordSet.

@Test
public void testGetRecordSet() {
    ConnectorTableHandle tableHandle = new PrometheusTableHandle("schema", "table");
    PrometheusRecordSetProvider recordSetProvider = new PrometheusRecordSetProvider(client);
    RecordSet recordSet = recordSetProvider.getRecordSet(PrometheusTransactionHandle.INSTANCE, SESSION, new PrometheusSplit(dataUri), tableHandle, ImmutableList.of(new PrometheusColumnHandle("labels", varcharMapType, 0), new PrometheusColumnHandle("timestamp", TIMESTAMP_COLUMN_TYPE, 1), new PrometheusColumnHandle("value", DoubleType.DOUBLE, 2)));
    assertNotNull(recordSet, "recordSet is null");
    RecordCursor cursor = recordSet.cursor();
    assertNotNull(cursor, "cursor is null");
    Map<Instant, Map<?, ?>> actual = new LinkedHashMap<>();
    while (cursor.advanceNextPosition()) {
        actual.put((Instant) cursor.getObject(1), getMapFromBlock(varcharMapType, (Block) cursor.getObject(0)));
    }
    Map<Instant, Map<String, String>> expected = ImmutableMap.<Instant, Map<String, String>>builder().put(ofEpochMilli(1565962969044L), ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")).put(ofEpochMilli(1565962984045L), ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")).put(ofEpochMilli(1565962999044L), ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")).put(ofEpochMilli(1565963014044L), ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")).buildOrThrow();
    assertEquals(actual, expected);
}

Example 8 with RecordSet

use of io.trino.spi.connector.RecordSet in project trino by trinodb.

the class TestPrometheusIntegrationMetrics method testGetColumnTypes.

@Test
public void testGetColumnTypes() {
    String dataUri = server.getUri().toString();
    RecordSet recordSet = new PrometheusRecordSet(client, new PrometheusSplit(dataUri), ImmutableList.of(new PrometheusColumnHandle("labels", createUnboundedVarcharType(), 0), new PrometheusColumnHandle("value", DoubleType.DOUBLE, 1), new PrometheusColumnHandle("timestamp", TIMESTAMP_COLUMN_TYPE, 2)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(createUnboundedVarcharType(), DoubleType.DOUBLE, TIMESTAMP_COLUMN_TYPE));
    recordSet = new PrometheusRecordSet(client, new PrometheusSplit(dataUri), ImmutableList.of(new PrometheusColumnHandle("value", BIGINT, 1), new PrometheusColumnHandle("text", createUnboundedVarcharType(), 0)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, createUnboundedVarcharType()));
    recordSet = new PrometheusRecordSet(client, new PrometheusSplit(dataUri), ImmutableList.of(new PrometheusColumnHandle("value", BIGINT, 1), new PrometheusColumnHandle("value", BIGINT, 1), new PrometheusColumnHandle("text", createUnboundedVarcharType(), 0)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, BIGINT, createUnboundedVarcharType()));
    recordSet = new PrometheusRecordSet(client, new PrometheusSplit(dataUri), ImmutableList.of());
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of());
}

Example 9 with RecordSet

use of io.trino.spi.connector.RecordSet in project trino by trinodb.

the class H2QueryRunner method insertRows.

private static void insertRows(ConnectorTableMetadata tableMetadata, Handle handle, RecordSet data) {
    List<ColumnMetadata> columns = tableMetadata.getColumns().stream().filter(columnMetadata -> !columnMetadata.isHidden()).collect(toImmutableList());
    String vars = Joiner.on(',').join(nCopies(columns.size(), "?"));
    String sql = format("INSERT INTO %s VALUES (%s)", tableMetadata.getTable().getTableName(), vars);
    RecordCursor cursor = data.cursor();
    while (true) {
        // insert 1000 rows at a time
        PreparedBatch batch = handle.prepareBatch(sql);
        for (int row = 0; row < 1000; row++) {
            if (!cursor.advanceNextPosition()) {
                if (batch.size() > 0) {
                    batch.execute();
                }
                return;
            }
            for (int column = 0; column < columns.size(); column++) {
                Type type = columns.get(column).getType();
                if (BOOLEAN.equals(type)) {
                    batch.bind(column, cursor.getBoolean(column));
                } else if (BIGINT.equals(type)) {
                    batch.bind(column, cursor.getLong(column));
                } else if (INTEGER.equals(type)) {
                    batch.bind(column, toIntExact(cursor.getLong(column)));
                } else if (DOUBLE.equals(type)) {
                    batch.bind(column, cursor.getDouble(column));
                } else if (type instanceof VarcharType) {
                    batch.bind(column, cursor.getSlice(column).toStringUtf8());
                } else if (DATE.equals(type)) {
                    long millisUtc = TimeUnit.DAYS.toMillis(cursor.getLong(column));
                    // H2 expects dates in to be millis at midnight in the JVM timezone
                    long localMillis = DateTimeZone.UTC.getMillisKeepLocal(DateTimeZone.getDefault(), millisUtc);
                    batch.bind(column, new Date(localMillis));
                } else {
                    throw new IllegalArgumentException("Unsupported type " + type);
                }
            }
            batch.add();
        }
        batch.execute();
    }
}

Example 10 with RecordSet

use of io.trino.spi.connector.RecordSet in project trino by trinodb.

the class TestJdbcRecordSet method testGetColumnTypes.

@Test
public void testGetColumnTypes() {
    RecordSet recordSet = createRecordSet(ImmutableList.of(new JdbcColumnHandle("text", JDBC_VARCHAR, VARCHAR), new JdbcColumnHandle("text_short", JDBC_VARCHAR, createVarcharType(32)), new JdbcColumnHandle("value", JDBC_BIGINT, BIGINT)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(VARCHAR, createVarcharType(32), BIGINT));
    recordSet = createRecordSet(ImmutableList.of(new JdbcColumnHandle("value", JDBC_BIGINT, BIGINT), new JdbcColumnHandle("text", JDBC_VARCHAR, VARCHAR)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, VARCHAR));
    recordSet = createRecordSet(ImmutableList.of(new JdbcColumnHandle("value", JDBC_BIGINT, BIGINT), new JdbcColumnHandle("value", JDBC_BIGINT, BIGINT), new JdbcColumnHandle("text", JDBC_VARCHAR, VARCHAR)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, BIGINT, VARCHAR));
    recordSet = createRecordSet(ImmutableList.of());
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of());
}