def test_offset_manager_from_file(tmpdir):
offsets = [Offset(consumer_group='group_1', topic='topic_1', partition=0, value=1),
Offset(consumer_group='group_1', topic='topic_1', partition=1, value=3),
Offset(consumer_group='group_1', topic='topic_1', partition=2, value=4)]
offset_manager = OffsetManager(offsets)
file_path = os.path.join(tmpdir, 'offset.sqlite3')
offset_manager.to_file(file_path)
from_file_offset_manager = OffsetManager.from_file(file_path)
assert from_file_offset_manager.offsets == offset_manager.offsets
def test_offset_manager_to_file(tmpdir):
offsets = [Offset(consumer_group='group_1', topic='topic_1', partition=0, value=1),
Offset(consumer_group='group_1', topic='topic_1', partition=1, value=3),
Offset(consumer_group='group_1', topic='topic_1', partition=2, value=4)]
offset_manager = OffsetManager(offsets)
file_path = os.path.join(tmpdir, 'offset.sqlite3')
offset_manager.to_file(file_path)
conn = sqlite3.connect(file_path)
cursor = conn.cursor()
cursor.execute('SELECT * FROM offset')
rows = cursor.fetchall()
assert rows
assert len(rows) == 3
assert [Offset(*r) for r in rows] == offsets
def test_generate_new_offsets():
offsets = [Offset('group_1', 'topic_1', 0, 100),
Offset('group_1', 'topic_1', 1, 200),
Offset('group_1', 'topic_1', 2, 500),
Offset('group_1', 'topic_1', 3, 505)]
partitions = [
Partition('topic_1', 0, [Message(1, 'key_1', 'val_1', 123456789, None),
Message(99, 'key_2', 'val_2', 123456789, None),
Message(101, 'key_3', 'val_3', 123456789, None),
Message(199, 'key_4', 'val_4', 123456789, None),
Message(499, 'key_5', 'val_5', 123456789, None),]),
Partition('topic_1', 1, [Message(1, 'key_1', 'val_1', 123456789, None),
Message(99, 'key_2', 'val_2', 123456789, None),
Message(101, 'key_3', 'val_3', 123456789, None),
Message(199, 'key_4', 'val_4', 123456789, None),
Message(499, 'key_5', 'val_5', 123456789, None),]),
Partition('topic_1', 2, [Message(1, 'key_1', 'val_1', 123456789, None),
Message(99, 'key_2', 'val_2', 123456789, None),
Message(101, 'key_3', 'val_3', 123456789, None),
Message(199, 'key_4', 'val_4', 123456789, None),
Message(499, 'key_5', 'val_5', 123456789, None),]),
Partition('topic_1', 3, [Message(1, 'key_1', 'val_1', 123456789, None),
Message(99, 'key_2', 'val_2', 123456789, None),
Message(101, 'key_3', 'val_3', 123456789, None),
Message(199, 'key_4', 'val_4', 123456789, None),
Message(499, 'key_5', 'val_5', 123456789, None),])
]
new_offsets = generate_new_offsets(offsets, partitions)
assert new_offsets == [Offset(consumer_group='group_1', topic='topic_1', partition=0, value=2),
Offset(consumer_group='group_1', topic='topic_1', partition=1, value=4),
Offset(consumer_group='group_1', topic='topic_1', partition=2, value=5),
Offset(consumer_group='group_1', topic='topic_1', partition=3, value=5)]
def test_data_flow_manager_write(tmpdir):
data_flow_manager = DataFlowManager(tmpdir)
offsets = [
Offset('group_1', 'topic_1', 0, 10),
Offset('group_1', 'topic_2', 0, 1)
]
partitions = [
Partition('topic_1', 0,
[Message(0, 'key_1', 'val_1', 123456789, None)]),
Partition('topic_1', 1,
[Message(0, 'key_1', 'val_1', 123456789, None)]),
Partition('topic_2', 0,
[Message(0, 'key_1', 'val_1', 123456789, None)]),
]
data_flow_manager.write(offsets, partitions)
assert os.path.isdir(os.path.join(tmpdir, 'partitions'))
assert len(os.listdir(os.path.join(tmpdir, 'partitions'))) == 3
assert os.path.isfile(os.path.join(tmpdir, 'offsets.sqlite3'))
def test_data_flow_manager_read(tmpdir):
data_flow_manager = DataFlowManager(tmpdir)
offsets = [
Offset('group_1', 'topic_1', 0, 10),
Offset('group_1', 'topic_2', 0, 1)
]
partitions = [
Partition('topic_1', 0,
[Message(0, 'key_1', 'val_1', 123456789, None)]),
Partition('topic_1', 1,
[Message(0, 'key_1', 'val_1', 123456789, None)]),
Partition('topic_2', 0,
[Message(0, 'key_1', 'val_1', 123456789, None)]),
]
data_flow_manager.write(offsets, partitions)
from_disk_offsets, from_disk_partitions = data_flow_manager.read()
assert from_disk_offsets == offsets
assert from_disk_partitions == partitions
def _get_offsets(consumer_group, partitions, config,
timeout=CONSUMER_OFFSET_TIMEOUT) -> List[Offset]:
offsets = []
consumer = ck.Consumer({**config, 'group.id': consumer_group})
for tp in consumer.committed(partitions, timeout=timeout):
if tp.offset == -1001:
continue
offset = Offset(consumer_group, tp.topic,
tp.partition, tp.offset)
offsets.append(offset)
consumer.close()
return offsets
def get_consumer_offsets(
self, topics: List[str], ignore_group_regex: str = IGNORE_GROUP_REGEX
) -> List[Offset]:
broker_topics = self.client.list_topics().topics
partitions = []
for topic_name in topics:
partitions.extend([TopicPartition(topic_name, k)
for k in broker_topics[topic_name].partitions])
offsets = []
for consumer_group in self.get_consumer_groups():
if re.findall(ignore_group_regex, consumer_group):
logger.debug(f'Ignoring consumer group: {consumer_group}')
continue
consumer = Consumer({**self.config, 'group.id': consumer_group})
for tp in consumer.committed(partitions, timeout=10):
if tp.offset == -1001:
continue
offset = Offset(consumer_group, tp.topic,
tp.partition, tp.offset)
offsets.append(offset)
return offsets