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

Example 6 with ConnectorInsertTableHandle

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

the class IcebergMetadata method finishInsert.

@Override
public Optional<ConnectorOutputMetadata> finishInsert(ConnectorSession session, ConnectorInsertTableHandle insertHandle, Collection<Slice> fragments, Collection<ComputedStatistics> computedStatistics) {
    IcebergWritableTableHandle table = (IcebergWritableTableHandle) insertHandle;
    Table icebergTable = transaction.table();
    List<CommitTaskData> commitTasks = fragments.stream().map(slice -> commitTaskCodec.fromJson(slice.getBytes())).collect(toImmutableList());
    Type[] partitionColumnTypes = icebergTable.spec().fields().stream().map(field -> field.transform().getResultType(icebergTable.schema().findType(field.sourceId()))).toArray(Type[]::new);
    AppendFiles appendFiles = transaction.newAppend();
    ImmutableSet.Builder<String> writtenFiles = ImmutableSet.builder();
    for (CommitTaskData task : commitTasks) {
        DataFiles.Builder builder = DataFiles.builder(icebergTable.spec()).withPath(task.getPath()).withFileSizeInBytes(task.getFileSizeInBytes()).withFormat(table.getFileFormat().toIceberg()).withMetrics(task.getMetrics().metrics());
        if (!icebergTable.spec().fields().isEmpty()) {
            String partitionDataJson = task.getPartitionDataJson().orElseThrow(() -> new VerifyException("No partition data for partitioned table"));
            builder.withPartition(PartitionData.fromJson(partitionDataJson, partitionColumnTypes));
        }
        appendFiles.appendFile(builder.build());
        writtenFiles.add(task.getPath());
    }
    // try to leave as little garbage as possible behind
    if (table.getRetryMode() != NO_RETRIES) {
        cleanExtraOutputFiles(session, writtenFiles.build());
    }
    appendFiles.commit();
    transaction.commitTransaction();
    transaction = null;
    return Optional.of(new HiveWrittenPartitions(commitTasks.stream().map(CommitTaskData::getPath).collect(toImmutableList())));
}

Example 7 with ConnectorInsertTableHandle

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

the class TestRaptorMetadata method testTransactionInsert.

@Test
public void testTransactionInsert() {
    // creating a table allocates a transaction
    long transactionId = 1;
    metadata.createTable(SESSION, getOrdersTable(), false);
    assertTrue(transactionSuccessful(transactionId));
    // start insert
    transactionId++;
    ConnectorTableHandle tableHandle = metadata.getTableHandle(SESSION, DEFAULT_TEST_ORDERS);
    ConnectorInsertTableHandle insertHandle = metadata.beginInsert(SESSION, tableHandle);
    // transaction is in progress
    assertTrue(transactionExists(transactionId));
    assertNull(transactionSuccessful(transactionId));
    // commit insert
    metadata.finishInsert(SESSION, insertHandle, ImmutableList.of(), ImmutableList.of());
    assertTrue(transactionExists(transactionId));
    assertTrue(transactionSuccessful(transactionId));
}

Example 8 with ConnectorInsertTableHandle

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

the class AbstractTestHive method doInsertIntoExistingPartition.

private void doInsertIntoExistingPartition(HiveStorageFormat storageFormat, SchemaTableName tableName) throws Exception {
    // creating the table
    doCreateEmptyTable(tableName, storageFormat, CREATE_TABLE_COLUMNS_PARTITIONED);
    MaterializedResult.Builder resultBuilder = MaterializedResult.resultBuilder(SESSION, CREATE_TABLE_PARTITIONED_DATA.getTypes());
    for (int i = 0; i < 3; i++) {
        // insert the data
        insertData(tableName, CREATE_TABLE_PARTITIONED_DATA);
        try (Transaction transaction = newTransaction()) {
            ConnectorSession session = newSession();
            ConnectorMetadata metadata = transaction.getMetadata();
            metadata.beginQuery(session);
            ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName);
            // verify partitions were created
            List<String> partitionNames = transaction.getMetastore().getPartitionNames(tableName.getSchemaName(), tableName.getTableName()).orElseThrow(() -> new AssertionError("Table does not exist: " + tableName));
            assertEqualsIgnoreOrder(partitionNames, CREATE_TABLE_PARTITIONED_DATA.getMaterializedRows().stream().map(row -> "ds=" + row.getField(CREATE_TABLE_PARTITIONED_DATA.getTypes().size() - 1)).collect(toImmutableList()));
            // load the new table
            List<ColumnHandle> columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values());
            // verify the data
            resultBuilder.rows(CREATE_TABLE_PARTITIONED_DATA.getMaterializedRows());
            MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat));
            assertEqualsIgnoreOrder(result.getMaterializedRows(), resultBuilder.build().getMaterializedRows());
            // test statistics
            for (String partitionName : partitionNames) {
                HiveBasicStatistics statistics = getBasicStatisticsForPartition(transaction, tableName, partitionName);
                assertEquals(statistics.getRowCount().getAsLong(), i + 1L);
                assertEquals(statistics.getFileCount().getAsLong(), i + 1L);
                assertGreaterThan(statistics.getInMemoryDataSizeInBytes().getAsLong(), 0L);
                assertGreaterThan(statistics.getOnDiskDataSizeInBytes().getAsLong(), 0L);
            }
        }
    }
    // test rollback
    Set<String> existingFiles;
    Path stagingPathRoot;
    try (Transaction transaction = newTransaction()) {
        ConnectorMetadata metadata = transaction.getMetadata();
        ConnectorSession session = newSession();
        existingFiles = listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName());
        assertFalse(existingFiles.isEmpty());
        ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName);
        // "stage" insert data
        ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES);
        stagingPathRoot = getStagingPathRoot(insertTableHandle);
        ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle);
        sink.appendPage(CREATE_TABLE_PARTITIONED_DATA.toPage());
        sink.appendPage(CREATE_TABLE_PARTITIONED_DATA.toPage());
        Collection<Slice> fragments = getFutureValue(sink.finish());
        metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of());
        // verify all temp files start with the unique prefix
        HdfsContext context = new HdfsContext(session);
        Set<String> tempFiles = listAllDataFiles(context, getStagingPathRoot(insertTableHandle));
        assertTrue(!tempFiles.isEmpty());
        for (String filePath : tempFiles) {
            assertThat(new Path(filePath).getName()).startsWith(session.getQueryId());
        }
        // verify statistics are visible from within of the current transaction
        List<String> partitionNames = transaction.getMetastore().getPartitionNames(tableName.getSchemaName(), tableName.getTableName()).orElseThrow(() -> new AssertionError("Table does not exist: " + tableName));
        for (String partitionName : partitionNames) {
            HiveBasicStatistics partitionStatistics = getBasicStatisticsForPartition(transaction, tableName, partitionName);
            assertEquals(partitionStatistics.getRowCount().getAsLong(), 5L);
        }
        // rollback insert
        transaction.rollback();
    }
    try (Transaction transaction = newTransaction()) {
        ConnectorMetadata metadata = transaction.getMetadata();
        ConnectorSession session = newSession();
        metadata.beginQuery(session);
        ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName);
        List<ColumnHandle> columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values());
        // verify the data is unchanged
        MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty());
        assertEqualsIgnoreOrder(result.getMaterializedRows(), resultBuilder.build().getMaterializedRows());
        // verify we did not modify the table directory
        assertEquals(listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()), existingFiles);
        // verify temp directory is empty
        HdfsContext hdfsContext = new HdfsContext(session);
        assertTrue(listAllDataFiles(hdfsContext, stagingPathRoot).isEmpty());
        // verify statistics have been rolled back
        List<String> partitionNames = transaction.getMetastore().getPartitionNames(tableName.getSchemaName(), tableName.getTableName()).orElseThrow(() -> new AssertionError("Table does not exist: " + tableName));
        for (String partitionName : partitionNames) {
            HiveBasicStatistics partitionStatistics = getBasicStatisticsForPartition(transaction, tableName, partitionName);
            assertEquals(partitionStatistics.getRowCount().getAsLong(), 3L);
        }
    }
}

Example 9 with ConnectorInsertTableHandle

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

the class AbstractTestHive method doTestTransactionDeleteInsert.

private void doTestTransactionDeleteInsert(HiveStorageFormat storageFormat, SchemaTableName tableName, Domain domainToDrop, MaterializedResult insertData, MaterializedResult expectedData, TransactionDeleteInsertTestTag tag, boolean expectQuerySucceed, Optional<ConflictTrigger> conflictTrigger) throws Exception {
    Path writePath = null;
    Path targetPath = null;
    try (Transaction transaction = newTransaction()) {
        try {
            ConnectorMetadata metadata = transaction.getMetadata();
            ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName);
            ConnectorSession session;
            rollbackIfEquals(tag, ROLLBACK_RIGHT_AWAY);
            // Query 1: delete
            session = newSession();
            HiveColumnHandle dsColumnHandle = (HiveColumnHandle) metadata.getColumnHandles(session, tableHandle).get("pk2");
            TupleDomain<ColumnHandle> tupleDomain = TupleDomain.withColumnDomains(ImmutableMap.of(dsColumnHandle, domainToDrop));
            Constraint constraint = new Constraint(tupleDomain, tupleDomain.asPredicate(), tupleDomain.getDomains().orElseThrow().keySet());
            tableHandle = applyFilter(metadata, tableHandle, constraint);
            tableHandle = metadata.applyDelete(session, tableHandle).get();
            metadata.executeDelete(session, tableHandle);
            rollbackIfEquals(tag, ROLLBACK_AFTER_DELETE);
            // Query 2: insert
            session = newSession();
            ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES);
            rollbackIfEquals(tag, ROLLBACK_AFTER_BEGIN_INSERT);
            writePath = getStagingPathRoot(insertTableHandle);
            targetPath = getTargetPathRoot(insertTableHandle);
            ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle);
            sink.appendPage(insertData.toPage());
            rollbackIfEquals(tag, ROLLBACK_AFTER_APPEND_PAGE);
            Collection<Slice> fragments = getFutureValue(sink.finish());
            rollbackIfEquals(tag, ROLLBACK_AFTER_SINK_FINISH);
            metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of());
            rollbackIfEquals(tag, ROLLBACK_AFTER_FINISH_INSERT);
            assertEquals(tag, COMMIT);
            if (conflictTrigger.isPresent()) {
                JsonCodec<PartitionUpdate> partitionUpdateCodec = JsonCodec.jsonCodec(PartitionUpdate.class);
                List<PartitionUpdate> partitionUpdates = fragments.stream().map(Slice::getBytes).map(partitionUpdateCodec::fromJson).collect(toList());
                conflictTrigger.get().triggerConflict(session, tableName, insertTableHandle, partitionUpdates);
            }
            transaction.commit();
            if (conflictTrigger.isPresent()) {
                assertTrue(expectQuerySucceed);
                conflictTrigger.get().verifyAndCleanup(session, tableName);
            }
        } catch (TestingRollbackException e) {
            transaction.rollback();
        } catch (TrinoException e) {
            assertFalse(expectQuerySucceed);
            if (conflictTrigger.isPresent()) {
                conflictTrigger.get().verifyAndCleanup(newSession(), tableName);
            }
        }
    }
    // check that temporary files are removed
    if (writePath != null && !writePath.equals(targetPath)) {
        HdfsContext context = new HdfsContext(newSession());
        FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, writePath);
        assertFalse(fileSystem.exists(writePath));
    }
    try (Transaction transaction = newTransaction()) {
        // verify partitions
        List<String> partitionNames = transaction.getMetastore().getPartitionNames(tableName.getSchemaName(), tableName.getTableName()).orElseThrow(() -> new AssertionError("Table does not exist: " + tableName));
        assertEqualsIgnoreOrder(partitionNames, expectedData.getMaterializedRows().stream().map(row -> format("pk1=%s/pk2=%s", row.getField(1), row.getField(2))).distinct().collect(toImmutableList()));
        // load the new table
        ConnectorSession session = newSession();
        ConnectorMetadata metadata = transaction.getMetadata();
        metadata.beginQuery(session);
        ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName);
        List<ColumnHandle> columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values());
        // verify the data
        MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat));
        assertEqualsIgnoreOrder(result.getMaterializedRows(), expectedData.getMaterializedRows());
    }
}

Example 10 with ConnectorInsertTableHandle

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

the class AbstractTestHive method doTestMismatchSchemaTable.

protected void doTestMismatchSchemaTable(SchemaTableName schemaTableName, HiveStorageFormat storageFormat, List<ColumnMetadata> tableBefore, MaterializedResult dataBefore, List<ColumnMetadata> tableAfter, MaterializedResult dataAfter) throws Exception {
    String schemaName = schemaTableName.getSchemaName();
    String tableName = schemaTableName.getTableName();
    doCreateEmptyTable(schemaTableName, storageFormat, tableBefore);
    // insert the data
    try (Transaction transaction = newTransaction()) {
        ConnectorSession session = newSession();
        ConnectorMetadata metadata = transaction.getMetadata();
        ConnectorTableHandle tableHandle = getTableHandle(metadata, schemaTableName);
        ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES);
        ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle);
        sink.appendPage(dataBefore.toPage());
        Collection<Slice> fragments = getFutureValue(sink.finish());
        metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of());
        transaction.commit();
    }
    // load the table and verify the data
    try (Transaction transaction = newTransaction()) {
        ConnectorSession session = newSession();
        ConnectorMetadata metadata = transaction.getMetadata();
        metadata.beginQuery(session);
        ConnectorTableHandle tableHandle = getTableHandle(metadata, schemaTableName);
        List<ColumnHandle> columnHandles = metadata.getColumnHandles(session, tableHandle).values().stream().filter(columnHandle -> !((HiveColumnHandle) columnHandle).isHidden()).collect(toList());
        MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty());
        assertEqualsIgnoreOrder(result.getMaterializedRows(), dataBefore.getMaterializedRows());
        transaction.commit();
    }
    // alter the table schema
    try (Transaction transaction = newTransaction()) {
        ConnectorSession session = newSession();
        PrincipalPrivileges principalPrivileges = testingPrincipalPrivilege(session);
        Table oldTable = transaction.getMetastore().getTable(schemaName, tableName).get();
        List<Column> dataColumns = tableAfter.stream().filter(columnMetadata -> !columnMetadata.getName().equals("ds")).map(columnMetadata -> new Column(columnMetadata.getName(), toHiveType(columnMetadata.getType()), Optional.empty())).collect(toList());
        Table.Builder newTable = Table.builder(oldTable).setDataColumns(dataColumns);
        transaction.getMetastore().replaceTable(schemaName, tableName, newTable.build(), principalPrivileges);
        transaction.commit();
    }
    // load the altered table and verify the data
    try (Transaction transaction = newTransaction()) {
        ConnectorSession session = newSession();
        ConnectorMetadata metadata = transaction.getMetadata();
        metadata.beginQuery(session);
        ConnectorTableHandle tableHandle = getTableHandle(metadata, schemaTableName);
        List<ColumnHandle> columnHandles = metadata.getColumnHandles(session, tableHandle).values().stream().filter(columnHandle -> !((HiveColumnHandle) columnHandle).isHidden()).collect(toList());
        MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty());
        assertEqualsIgnoreOrder(result.getMaterializedRows(), dataAfter.getMaterializedRows());
        transaction.commit();
    }
    // insertions to the partitions with type mismatches should fail
    try (Transaction transaction = newTransaction()) {
        ConnectorSession session = newSession();
        ConnectorMetadata metadata = transaction.getMetadata();
        ConnectorTableHandle tableHandle = getTableHandle(metadata, schemaTableName);
        ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES);
        ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle);
        sink.appendPage(dataAfter.toPage());
        Collection<Slice> fragments = getFutureValue(sink.finish());
        metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of());
        transaction.commit();
        fail("expected exception");
    } catch (TrinoException e) {
        // expected
        assertEquals(e.getErrorCode(), HIVE_PARTITION_SCHEMA_MISMATCH.toErrorCode());
    }
}