com.taosdata.jdbc.tmq.JNIConsumer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of taos-jdbcdriver Show documentation
TDengine JDBC Driver
There is a newer version: 3.4.0
package com.taosdata.jdbc.tmq;

import com.taosdata.jdbc.TSDBConstants;
import com.taosdata.jdbc.common.Consumer;

import java.sql.SQLException;
import java.time.Duration;
import java.util.*;
import java.util.stream.Collectors;

import static com.taosdata.jdbc.TSDBErrorNumbers.ERROR_TMQ_CONSUMER_NULL;

public class JNIConsumer implements  Consumer {

    private final TMQConnector connector;

    public JNIConsumer() {
        connector = new TMQConnector();
    }

    @Override
    public void create(Properties properties) throws SQLException {
        long config = connector.createConfig(properties);
        try {
            connector.createConsumer(config);
        } finally {
            connector.destroyConf(config);
        }
    }

    @Override
    public void subscribe(Collection topics) throws SQLException {
        long topicPointer = 0L;
        try {
            topicPointer = connector.createTopic(topics);
            connector.subscribe(topicPointer);
        } finally {
            if (topicPointer != TSDBConstants.JNI_NULL_POINTER) {
                connector.destroyTopic(topicPointer);
            }
        }
    }

    @Override
    public void unsubscribe() throws SQLException {
        connector.unsubscribe();
    }

    @Override
    public Set subscription() throws SQLException {
        return connector.subscription();
    }

    @Override
    public ConsumerRecords poll(Duration timeout, Deserializer deserializer) throws SQLException {
        long resultSet = connector.poll(timeout.toMillis());
        // when tmq pointer is null or result set is null
        if (resultSet == 0 || resultSet == ERROR_TMQ_CONSUMER_NULL) {
            return ConsumerRecords.emptyRecord();
        }

        int timestampPrecision = connector.getResultTimePrecision(resultSet);

        ConsumerRecords records = new ConsumerRecords<>();
        String topic = connector.getTopicName(resultSet);
        String dbName = connector.getDbName(resultSet);
        int vGroupId = connector.getVgroupId(resultSet);
        long offset = connector.getOffset(resultSet);
        String tableName = connector.getTableName(resultSet);

        TopicPartition tp = new TopicPartition(topic, vGroupId);

        try (TMQResultSet rs = new TMQResultSet(connector, resultSet, timestampPrecision, dbName, tableName)) {
            while (rs.next()) {
                V v = deserializer.deserialize(rs, topic, dbName);
                ConsumerRecord r = new ConsumerRecord<>(topic, dbName, vGroupId, offset, v);
                records.put(tp, r);
            }
        }

        return records;
    }

    @Override
    public void commitAsync(OffsetCommitCallback callback) throws SQLException {
        OffsetWaitCallback call = new OffsetWaitCallback<>(getAllConsumed(), this, callback);

        connector.asyncCommit(call);
    }

    private Map getAllConsumed() throws SQLException {
        Map offsets = new HashMap<>();
        this.subscription().forEach(topic -> {
            List topicAssignment = connector.getTopicAssignment(topic);
            topicAssignment.forEach(assignment -> {
                TopicPartition tp = new TopicPartition(topic, assignment.getVgId());
                OffsetAndMetadata metadata = new OffsetAndMetadata(assignment.getCurrentOffset());
                offsets.put(tp, metadata);
            });
        });
        return offsets;
    }

    public void commitAsync(final Map offsets, OffsetCommitCallback callback) {
        for (Map.Entry entry : offsets.entrySet()) {
            Map offset = new HashMap<>();
            offset.put(entry.getKey(), entry.getValue());
            connector.asyncCommit(entry.getKey().getTopic(), entry.getKey().getVGroupId(), entry.getValue().offset(),
                    new OffsetWaitCallback<>(offset, this, callback));
        }
    }

    @Override
    public void commitSync() throws SQLException {
        connector.syncCommit();
    }

    @Override
    public void commitSync(Map offsets) throws SQLException {
        for (Map.Entry entry : offsets.entrySet()) {
            connector.commitOffsetSync(entry.getKey().getTopic(), entry.getKey().getVGroupId(), entry.getValue().offset());
        }
    }

    @Override
    public void seek(TopicPartition partition, long offset) {
        connector.seek(partition.getTopic(), partition.getVGroupId(), offset);
    }

    @Override
    public long position(TopicPartition partition) throws SQLException {
        return connector.position(partition.getTopic(), partition.getVGroupId());
    }

    @Override
    public Map position(String topic) throws SQLException {
        List collect = connector.getTopicAssignment(topic).stream()
                .map(a -> new TopicPartition(topic, a.getVgId())).collect(Collectors.toList());
        Map map = new HashMap<>();
        for (TopicPartition topicPartition : collect) {
            map.put(topicPartition, position(topicPartition));
        }
        return map;
    }

    @Override
    public Map beginningOffsets(String topic) {
        return connector.getTopicAssignment(topic).stream()
                .collect(HashMap::new
                        , (m, a) -> m.put(new TopicPartition(topic, a.getVgId()), a.getBegin())
                        , HashMap::putAll
                );
    }

    @Override
    public Map endOffsets(String topic) {
        return connector.getTopicAssignment(topic).stream()
                .collect(HashMap::new
                        , (m, a) -> m.put(new TopicPartition(topic, a.getVgId()), a.getEnd())
                        , HashMap::putAll
                );
    }

    @Override
    public void seekToBeginning(Collection partitions) throws SQLException {
        Map beginningOffsets = new HashMap<>();
        for (TopicPartition partition : partitions) {
            if (beginningOffsets.containsKey(partition)) {
                Long aLong = beginningOffsets.get(partition);
                seek(partition, aLong);
            } else {
                beginningOffsets(partition.getTopic()).forEach((tp, offset) -> {
                    if (tp.getVGroupId() == partition.getVGroupId()) {
                        seek(tp, offset);
                    } else {
                        beginningOffsets.put(tp, offset);
                    }
                });
            }
        }
    }

    @Override
    public void seekToEnd(Collection partitions) throws SQLException {
        Map endOffsets = new HashMap<>();
        for (TopicPartition partition : partitions) {
            if (endOffsets.containsKey(partition)) {
                Long aLong = endOffsets.get(partition);
                seek(partition, aLong);
            } else {
                endOffsets(partition.getTopic()).forEach((tp, offset) -> {
                    if (tp.getVGroupId() == partition.getVGroupId()) {
                        seek(tp, offset);
                    } else {
                        endOffsets.put(tp, offset);
                    }
                });
            }
        }
    }

    @Override
    public Set assignment() throws SQLException {
        return subscription().stream().map(topic -> {
            List topicAssignment = connector.getTopicAssignment(topic);
            return topicAssignment.stream().map(a -> new TopicPartition(topic, a.getVgId())).collect(Collectors.toList());
        }).flatMap(Collection::stream).collect(Collectors.toSet());
    }

    @Override
    public OffsetAndMetadata committed(TopicPartition partition) throws SQLException {
        long l = connector.committed(partition.getTopic(), partition.getVGroupId());
        return new OffsetAndMetadata(l, null);
    }

    @Override
    public Map committed(Set partitions) throws SQLException {
        Map map = new HashMap<>();
        for (TopicPartition partition : partitions) {
            map.put(partition, committed(partition));
        }
        return map;
    }

    @Override
    public void close() throws SQLException {
        connector.closeConsumer();
    }

    public String getErrMsg(int code) {
        return connector.getErrMsg(code);
    }
}