Search in sources :

Example 1 with VerticaStreamWriterBuilder

use of com.github.jcustenborder.vertica.VerticaStreamWriterBuilder in project kafka-connect-vertica by jcustenborder.

the class VerticaSinkTask method put.

@Override
public void put(Collection<SinkRecord> records) {
    Multimap<String, SinkRecord> recordsByTable = HashMultimap.create(this.config.expectedTopics, this.config.expectedRecords);
    Multiset<String> countsByTable = HashMultiset.create(this.config.expectedTopics);
    for (SinkRecord record : records) {
        String table = record.topic();
        countsByTable.add(table);
        recordsByTable.put(table, record);
    }
    for (String table : countsByTable.elementSet()) {
        log.trace("put() - Writing {} record(s) to {}", countsByTable.count(table), table);
    }
    DataSource dataSource = PoolOfPools.get(this.config);
    try (Connection connection = dataSource.getConnection()) {
        VerticaConnection verticaConnection = connection.unwrap(VerticaConnection.class);
        try {
            for (final String tableName : recordsByTable.keys()) {
                log.trace("put() - Processing records for table '{}'", tableName);
                Collection<SinkRecord> tableRecords = recordsByTable.get(tableName);
                VerticaStreamWriterBuilder builder = configureBuilder(verticaConnection, tableName);
                final String statement = new QueryBuilder(builder).toString();
                log.info("put() - Creating VerticaCopyStream with statement:\n{}", statement);
                VerticaCopyStream copyStream = new VerticaCopyStream(verticaConnection, statement);
                copyStream.start();
                final PipedInputStream inputStream = new PipedInputStream(this.config.inputBufferSize);
                final PipedOutputStream outputStream = new PipedOutputStream(inputStream);
                try {
                    Stopwatch stopwatch = Stopwatch.createStarted();
                    Future<?> importFuture = executorService.submit(() -> {
                        try {
                            copyStream.addStream(inputStream);
                            copyStream.execute();
                        } catch (SQLException e) {
                            throw new IllegalStateException(e);
                        }
                    });
                    int count = 0;
                    try (VerticaStreamWriter writer = builder.build(outputStream)) {
                        for (SinkRecord record : tableRecords) {
                            Struct value = (Struct) record.value();
                            int i = 0;
                            Object[] values = new Object[writer.columns().size()];
                            for (VerticaColumnInfo columnInfo : writer.columns()) {
                                values[i] = value.get(columnInfo.name());
                                i++;
                            }
                            log.trace("Writing row");
                            writer.write(values);
                            count++;
                        }
                        log.info("Wrote {} record(s) to stream", count);
                    }
                    outputStream.close();
                    log.info("Waiting for import to complete.");
                    try {
                        importFuture.get(this.config.streamTimeoutMs, TimeUnit.MILLISECONDS);
                    } catch (TimeoutException e) {
                        log.warn("Import exceeded timeout of {} ms. Rolling back", this.config.streamTimeoutMs);
                        connection.rollback();
                    }
                    log.info("put() - Imported {} record(s) in {} millisecond(s).", count, stopwatch.elapsed(TimeUnit.MILLISECONDS));
                    final int rejectedCount = copyStream.getRejects().size();
                    if (rejectedCount > 0) {
                        log.warn("put() - Rejected {} record(s).", copyStream.getRejects().size());
                        for (Long l : copyStream.getRejects()) {
                            log.warn("Rejected row {}", l);
                        }
                    }
                } catch (InterruptedException | ExecutionException e) {
                    log.error("Exception thrown", e);
                } finally {
                    inputStream.close();
                }
            }
        } catch (IOException ex) {
            throw new RetriableException(ex);
        } catch (ExecutionException ex) {
            throw new RetriableException(ex);
        }
        log.trace("put() - committing transaction");
        connection.commit();
    } catch (SQLException ex) {
        throw new RetriableException(ex);
    }
}
Also used : SQLException(java.sql.SQLException) Stopwatch(com.google.common.base.Stopwatch) PipedOutputStream(java.io.PipedOutputStream) QueryBuilder(com.github.jcustenborder.vertica.QueryBuilder) Struct(org.apache.kafka.connect.data.Struct) VerticaConnection(com.vertica.jdbc.VerticaConnection) VerticaCopyStream(com.vertica.jdbc.VerticaCopyStream) ExecutionException(java.util.concurrent.ExecutionException) TimeoutException(java.util.concurrent.TimeoutException) VerticaColumnInfo(com.github.jcustenborder.vertica.VerticaColumnInfo) Connection(java.sql.Connection) VerticaConnection(com.vertica.jdbc.VerticaConnection) PipedInputStream(java.io.PipedInputStream) IOException(java.io.IOException) SinkRecord(org.apache.kafka.connect.sink.SinkRecord) DataSource(javax.sql.DataSource) VerticaStreamWriter(com.github.jcustenborder.vertica.VerticaStreamWriter) VerticaStreamWriterBuilder(com.github.jcustenborder.vertica.VerticaStreamWriterBuilder) RetriableException(org.apache.kafka.connect.errors.RetriableException)

Example 2 with VerticaStreamWriterBuilder

use of com.github.jcustenborder.vertica.VerticaStreamWriterBuilder in project kafka-connect-vertica by jcustenborder.

the class VerticaSinkTask method configureBuilder.

VerticaStreamWriterBuilder configureBuilder(Connection connection, final String tableName) throws ExecutionException {
    return this.builderCache.get(tableName, () -> {
        VerticaStreamWriterBuilder builder = new VerticaStreamWriterBuilder();
        builder.compressionType(this.config.compressionType);
        builder.loadMethod(this.config.loadMethod);
        builder.table(tableName);
        log.trace("configureBuilder() - Finding columns for '{}'", tableName);
        List<String> fields = new ArrayList<>();
        try (PreparedStatement statement = connection.prepareStatement(SQL)) {
            statement.setString(1, tableName);
            try (ResultSet results = statement.executeQuery()) {
                while (results.next()) {
                    String columnName = results.getString(1);
                    fields.add(columnName);
                    String dataType = results.getString(2);
                    Matcher matcher = PATTERN.matcher(dataType);
                    if (!matcher.find()) {
                        log.warn("Could not match '{}' for column '{}'", dataType, columnName);
                        continue;
                    }
                    String stype = matcher.group(0).toUpperCase();
                    VerticaColumnType type = VerticaColumnType.valueOf(stype);
                    log.info("configureBuilder() - '{}' = {}", columnName, type);
                    builder.column(columnName, type);
                }
            }
        }
        return builder;
    });
}
Also used : VerticaStreamWriterBuilder(com.github.jcustenborder.vertica.VerticaStreamWriterBuilder) Matcher(java.util.regex.Matcher) ArrayList(java.util.ArrayList) ResultSet(java.sql.ResultSet) PreparedStatement(java.sql.PreparedStatement) VerticaColumnType(com.github.jcustenborder.vertica.VerticaColumnType)

Aggregations

VerticaStreamWriterBuilder (com.github.jcustenborder.vertica.VerticaStreamWriterBuilder)2 QueryBuilder (com.github.jcustenborder.vertica.QueryBuilder)1 VerticaColumnInfo (com.github.jcustenborder.vertica.VerticaColumnInfo)1 VerticaColumnType (com.github.jcustenborder.vertica.VerticaColumnType)1 VerticaStreamWriter (com.github.jcustenborder.vertica.VerticaStreamWriter)1 Stopwatch (com.google.common.base.Stopwatch)1 VerticaConnection (com.vertica.jdbc.VerticaConnection)1 VerticaCopyStream (com.vertica.jdbc.VerticaCopyStream)1 IOException (java.io.IOException)1 PipedInputStream (java.io.PipedInputStream)1 PipedOutputStream (java.io.PipedOutputStream)1 Connection (java.sql.Connection)1 PreparedStatement (java.sql.PreparedStatement)1 ResultSet (java.sql.ResultSet)1 SQLException (java.sql.SQLException)1 ArrayList (java.util.ArrayList)1 ExecutionException (java.util.concurrent.ExecutionException)1 TimeoutException (java.util.concurrent.TimeoutException)1 Matcher (java.util.regex.Matcher)1 DataSource (javax.sql.DataSource)1