Examples with QueryRunner io.trino.testing.QueryRunner used on opensource projects

Example 1 with QueryRunner

use of io.trino.testing.QueryRunner in project trino by trinodb.

the class TestDeltaLakeAdlsConnectorSmokeTest method createTableFromResources.

@Override
void createTableFromResources(String table, String resourcePath, QueryRunner queryRunner) {
    String targetDirectory = bucketName + "/" + table;
    try {
        List<ClassPath.ResourceInfo> resources = ClassPath.from(TestDeltaLakeAdlsConnectorSmokeTest.class.getClassLoader()).getResources().stream().filter(resourceInfo -> resourceInfo.getResourceName().startsWith(resourcePath + "/")).collect(toImmutableList());
        for (ClassPath.ResourceInfo resourceInfo : resources) {
            String fileName = resourceInfo.getResourceName().replaceFirst("^" + Pattern.quote(resourcePath), quoteReplacement(targetDirectory));
            ByteSource byteSource = resourceInfo.asByteSource();
            azureContainerClient.getBlobClient(fileName).upload(byteSource.openBufferedStream(), byteSource.size());
        }
    } catch (IOException e) {
        throw new UncheckedIOException(e);
    }
    queryRunner.execute(format("CREATE TABLE %s (dummy int) WITH (location = '%s')", table, getLocationForTable(bucketName, table)));
}

Example 2 with QueryRunner

use of io.trino.testing.QueryRunner in project trino by trinodb.

the class TestHivePlans method setUp.

@BeforeClass
public void setUp() {
    QueryRunner queryRunner = getQueryRunner();
    // Use common VALUES for setup so that types are the same and there are no coercions.
    String values = "VALUES ('one', 1), ('two', 2), ('three', 3), ('four', 4), ('five', 5)";
    // partitioned on integer
    queryRunner.execute("CREATE TABLE table_int_partitioned WITH (partitioned_by = ARRAY['int_part']) AS SELECT str_col, int_part FROM (" + values + ") t(str_col, int_part)");
    // partitioned on varchar
    queryRunner.execute("CREATE TABLE table_str_partitioned WITH (partitioned_by = ARRAY['str_part']) AS SELECT int_col, str_part FROM (" + values + ") t(str_part, int_col)");
    // with too many partitions
    queryRunner.execute("CREATE TABLE table_int_with_too_many_partitions WITH (partitioned_by = ARRAY['int_part']) AS SELECT str_col, int_part FROM (" + values + ", ('six', 6)) t(str_col, int_part)");
    // unpartitioned
    queryRunner.execute("CREATE TABLE table_unpartitioned AS SELECT str_col, int_col FROM (" + values + ") t(str_col, int_col)");
}

Example 3 with QueryRunner

use of io.trino.testing.QueryRunner in project trino by trinodb.

the class BaseDeltaLakeConnectorSmokeTest method testDeltaLakeTableLocationChanged.

private void testDeltaLakeTableLocationChanged(boolean fewerEntries, boolean firstPartitioned, boolean secondPartitioned) throws Exception {
    // Create a table with a bunch of transaction log entries
    String tableName = "test_table_location_changed_" + randomTableSuffix();
    String initialLocation = getLocationForTable(bucketName, tableName);
    assertUpdate(format("CREATE TABLE %s (a_number int, a_string varchar) WITH (location = '%s' %s)", tableName, initialLocation, firstPartitioned ? ", partitioned_by = ARRAY['a_number']" : ""));
    BiConsumer<QueryRunner, String> insertABunchOfRows = (queryRunner, prefix) -> {
        queryRunner.execute(format("INSERT INTO %s (a_number, a_string) VALUES (1, '%s one')", tableName, prefix));
        queryRunner.execute(format("INSERT INTO %s (a_number, a_string) VALUES (2, '%s two')", tableName, prefix));
        queryRunner.execute(format("INSERT INTO %s (a_number, a_string) VALUES (3, '%s tree')", tableName, prefix));
        queryRunner.execute(format("INSERT INTO %s (a_number, a_string) VALUES (4, '%s four')", tableName, prefix));
    };
    insertABunchOfRows.accept(getQueryRunner(), "first");
    MaterializedResult initialData = computeActual("SELECT * FROM " + tableName);
    assertThat(initialData.getMaterializedRows()).hasSize(4);
    MaterializedResult expectedDataAfterChange;
    String newLocation;
    try (QueryRunner independentQueryRunner = createDeltaLakeQueryRunner(Map.of())) {
        // Change table's location without main Delta Lake connector (main query runner) knowing about this
        newLocation = getLocationForTable(bucketName, "test_table_location_changed_new_" + randomTableSuffix());
        independentQueryRunner.execute("DROP TABLE " + tableName);
        independentQueryRunner.execute(format("CREATE TABLE %s (a_number int, a_string varchar, another_string varchar) WITH (location = '%s' %s) ", tableName, newLocation, secondPartitioned ? ", partitioned_by = ARRAY['a_number']" : ""));
        if (fewerEntries) {
            // Have fewer transaction log entries so that version mismatch is more apparent (but easier to detect)
            independentQueryRunner.execute(format("INSERT INTO %s VALUES (1, 'second one', 'third column')", tableName));
        } else {
            insertABunchOfRows.accept(independentQueryRunner, "second");
        }
        expectedDataAfterChange = independentQueryRunner.execute("SELECT * FROM " + tableName);
        assertThat(expectedDataAfterChange.getMaterializedRows()).hasSize(fewerEntries ? 1 : 4);
    }
    Stopwatch stopwatch = Stopwatch.createStarted();
    while (true) {
        MaterializedResult currentVisibleData = computeActual("SELECT * FROM " + tableName);
        if (Set.copyOf(currentVisibleData.getMaterializedRows()).equals(Set.copyOf(expectedDataAfterChange.getMaterializedRows()))) {
            // satisfied
            break;
        }
        if (!Set.copyOf(currentVisibleData.getMaterializedRows()).equals(Set.copyOf(initialData.getMaterializedRows()))) {
            throw new AssertionError(format("Unexpected result when reading table: %s,\n expected either initialData: %s\n or expectedDataAfterChange: %s", currentVisibleData, initialData, expectedDataAfterChange));
        }
        if (stopwatch.elapsed(SECONDS) > TEST_METADATA_CACHE_TTL_SECONDS + 10) {
            throw new RuntimeException("Timed out waiting on table to reflect new data from new location");
        }
        SECONDS.sleep(1);
    }
    // Verify table schema gets reflected correctly
    assertThat(computeScalar("SHOW CREATE TABLE " + tableName)).isEqualTo(format("" + "CREATE TABLE %s.%s.%s (\n" + "   a_number integer,\n" + "   a_string varchar,\n" + "   another_string varchar\n" + ")\n" + "WITH (\n" + "   location = '%s',\n" + "   partitioned_by = ARRAY[%s]\n" + ")", getSession().getCatalog().orElseThrow(), getSession().getSchema().orElseThrow(), tableName, newLocation, secondPartitioned ? "'a_number'" : ""));
}

Example 4 with QueryRunner

use of io.trino.testing.QueryRunner in project trino by trinodb.

the class BaseDeltaLakeConnectorSmokeTest method createQueryRunner.

@Override
protected QueryRunner createQueryRunner() throws Exception {
    this.dockerizedDataLake = closeAfterClass(createDockerizedDataLake());
    QueryRunner queryRunner = createDeltaLakeQueryRunner(ImmutableMap.<String, String>builder().put("delta.metadata.cache-ttl", TEST_METADATA_CACHE_TTL_SECONDS + "s").put("hive.metastore-cache-ttl", TEST_METADATA_CACHE_TTL_SECONDS + "s").buildOrThrow());
    queryRunner.execute(format("CREATE SCHEMA %s WITH (location = '%s')", SCHEMA, getLocationForTable(bucketName, SCHEMA)));
    REQUIRED_TPCH_TABLES.forEach(table -> queryRunner.execute(format("CREATE TABLE %s WITH (location = '%s') AS SELECT * FROM tpch.tiny.%1$s", table.getTableName(), getLocationForTable(bucketName, table.getTableName()))));
    /* Data (across 2 files) generated using:
         * INSERT INTO foo VALUES
         *   (1, 100, 'data1'),
         *   (2, 200, 'data2')
         *
         * Data (across 2 files) generated using:
         * INSERT INTO bar VALUES
         *   (100, 'data100'),
         *   (200, 'data200')
         *
         * INSERT INTO old_dates
         * VALUES (DATE '0100-01-01', 1), (DATE '1582-10-15', 2), (DATE '1960-01-01', 3), (DATE '2020-01-01', 4)
         *
         * INSERT INTO test_timestamps VALUES
         * (TIMESTAMP '0100-01-01 01:02:03', 1), (TIMESTAMP '1582-10-15 01:02:03', 2), (TIMESTAMP '1960-01-01 01:02:03', 3), (TIMESTAMP '2020-01-01 01:02:03', 4);
         */
    NON_TPCH_TABLES.forEach(table -> {
        String resourcePath = "databricks/" + table;
        createTableFromResources(table, resourcePath, queryRunner);
    });
    return queryRunner;
}

Example 5 with QueryRunner

use of io.trino.testing.QueryRunner in project trino by trinodb.

the class BaseDeltaLakeMinioConnectorTest method createQueryRunner.

@Override
protected QueryRunner createQueryRunner() throws Exception {
    this.dockerizedMinioDataLake = closeAfterClass(createDockerizedMinioDataLakeForDeltaLake(bucketName, Optional.empty()));
    QueryRunner queryRunner = DeltaLakeQueryRunner.createS3DeltaLakeQueryRunner(DELTA_CATALOG, SCHEMA, ImmutableMap.<String, String>builder().put("delta.enable-non-concurrent-writes", "true").buildOrThrow(), dockerizedMinioDataLake.getMinioAddress(), dockerizedMinioDataLake.getTestingHadoop());
    queryRunner.execute("CREATE SCHEMA " + SCHEMA + " WITH (location = 's3://" + bucketName + "/" + SCHEMA + "')");
    TpchTable.getTables().forEach(table -> {
        String tableName = table.getTableName();
        dockerizedMinioDataLake.copyResources(resourcePath + tableName, SCHEMA + "/" + tableName);
        queryRunner.execute(format("CREATE TABLE %1$s.%2$s.%3$s (dummy int) WITH (location = 's3://%4$s/%2$s/%3$s')", DELTA_CATALOG, SCHEMA, tableName, bucketName));
    });
    return queryRunner;
}