def test_snapshot_loaded(self, create_consumer) -> None:
kafka_consumer = FakeKafkaConsumer()
kafka_consumer.items = [
message(
0,
0,
'{"snapshot-id":"abc123", "product":"somewhere-else", "tables": [], "event":"snapshot-init"}',
False,
),
message(
1,
0,
'{"snapshot-id":"abc123", "product":"snuba", "tables": [], "event":"snapshot-init"}',
False,
),
message(
2,
0,
('{"snapshot-id":"abc123", "event":"snapshot-loaded",'
'"transaction-info": {"xmin":123, "xmax":124, "xip-list": []}'
'}'),
False,
),
message(0, 0, None, True),
]
create_consumer.return_value = kafka_consumer
bootstrap = BootstrapState("cdc_control", "somewhere", "something")
ret = bootstrap.handle(None)
assert ret[0] == ConsumerStateCompletionEvent.SNAPSHOT_READY_RECEIVED
assert kafka_consumer.commit_calls == 2
def test_empty_topic(self, create_consumer) -> None:
kafka_consumer = FakeKafkaConsumer()
kafka_consumer.items = [
message(0, 0, None, True),
]
create_consumer.return_value = kafka_consumer
bootstrap = BootstrapState("cdc_control", "somewhere", "something")
ret = bootstrap.handle(None)
assert ret[0] == ConsumerStateCompletionEvent.NO_SNAPSHOT
assert kafka_consumer.commit_calls == 0
def test_empty_topic(self, create_consumer) -> None:
kafka_consumer = FakeConfluentKafkaConsumer()
kafka_consumer.items = [
build_confluent_kafka_message(0, 0, None, True),
]
create_consumer.return_value = kafka_consumer
bootstrap = BootstrapState(
"cdc_control",
self.broker_config,
"something",
get_cdc_storage(StorageKey.GROUPEDMESSAGES),
)
ret = bootstrap.handle(None)
assert ret[0] == ConsumerStateCompletionEvent.NO_SNAPSHOT
assert kafka_consumer.commit_calls == 0
def test_init_snapshot(self, create_consumer) -> None:
kafka_consumer = FakeKafkaConsumer()
kafka_consumer.items = [
message(
0,
0,
'{"snapshot-id":"abc123", "tables": [], "product":"snuba", "event":"snapshot-init"}',
False,
),
message(0, 0, None, True),
]
create_consumer.return_value = kafka_consumer
bootstrap = BootstrapState("cdc_control", "somewhere", "something")
ret = bootstrap.handle(None)
assert ret[0] == ConsumerStateCompletionEvent.SNAPSHOT_INIT_RECEIVED
assert kafka_consumer.commit_calls == 0
def test_snapshot_loaded(self, create_consumer: Mock) -> None:
kafka_consumer = FakeConfluentKafkaConsumer()
kafka_consumer.items = [
build_confluent_kafka_message(
0,
0,
b'{"snapshot-id":"abc123", "product":"somewhere-else", "tables": [], "event":"snapshot-init"}',
False,
),
build_confluent_kafka_message(
1,
0,
b'{"snapshot-id":"abc123", "product":"snuba", "tables": ["sentry_groupedmessage"], "event":"snapshot-init"}',
False,
),
build_confluent_kafka_message(
2,
0,
(
b'{"snapshot-id":"abc123", "event":"snapshot-loaded",'
b'"transaction-info": {"xmin":123, "xmax":124, "xip-list": []}'
b"}"
),
False,
),
build_confluent_kafka_message(0, 0, None, True),
]
create_consumer.return_value = kafka_consumer
bootstrap = BootstrapState(
"cdc_control",
self.broker_config,
"something",
get_cdc_storage(StorageKey.GROUPEDMESSAGES),
)
ret = bootstrap.handle(None)
assert ret[0] == ConsumerStateCompletionEvent.SNAPSHOT_READY_RECEIVED
assert kafka_consumer.commit_calls == 2
def test_snapshot_for_other_table(self, create_consumer) -> None:
kafka_consumer = FakeConfluentKafkaConsumer()
kafka_consumer.items = [
build_confluent_kafka_message(
0,
0,
b'{"snapshot-id":"abc123", "tables": ["someone_else"], "product":"snuba", "event":"snapshot-init"}',
False,
),
build_confluent_kafka_message(0, 0, None, True),
]
create_consumer.return_value = kafka_consumer
bootstrap = BootstrapState(
"cdc_control",
"somewhere",
"something",
get_dataset("groupedmessage"),
)
ret = bootstrap.handle(None)
assert ret[0] == ConsumerStateCompletionEvent.NO_SNAPSHOT
assert kafka_consumer.commit_calls == 1
def test_init_snapshot(self, create_consumer) -> None:
kafka_consumer = FakeConfluentKafkaConsumer()
kafka_consumer.items = [
build_confluent_kafka_message(
0,
0,
b'{"snapshot-id":"abc123", "tables": ["sentry_groupedmessage"], "product":"snuba", "event":"snapshot-init"}',
False,
),
build_confluent_kafka_message(0, 0, None, True),
]
create_consumer.return_value = kafka_consumer
bootstrap = BootstrapState(
"cdc_control",
self.broker_config,
"something",
get_cdc_storage(StorageKey.GROUPEDMESSAGES),
)
ret = bootstrap.handle(None)
assert ret[0] == ConsumerStateCompletionEvent.SNAPSHOT_INIT_RECEIVED
assert kafka_consumer.commit_calls == 0
def _build_state(
self,
state_class: StateType[ConsumerStateCompletionEvent, Optional[ConsumerStateData]],
) -> State[ConsumerStateCompletionEvent, Optional[ConsumerStateData]]:
if state_class == ConsumingState:
return ConsumingState(self.__consumer_builder)
elif state_class == BootstrapState:
return BootstrapState(
topic=self.__topic,
bootstrap_servers=self.__bootstrap_servers,
group_id=self.__group_id,
dataset=self.__dataset,
)
elif state_class == CatchingUpState:
return CatchingUpState(self.__consumer_builder)
else:
return state_class()
def _build_state(
self,
state_class: StateType[ConsumerStateCompletionEvent,
Optional[ConsumerStateData]],
) -> State[ConsumerStateCompletionEvent, Optional[ConsumerStateData]]:
if state_class == ConsumingState:
return ConsumingState(self.__consumer_builder)
elif state_class == BootstrapState:
return BootstrapState(
topic=self.__topic,
broker_config=self.__consumer_builder.broker_config,
group_id=self.__group_id,
storage=self.__storage,
)
elif state_class == CatchingUpState:
return CatchingUpState(self.__consumer_builder)
else:
return state_class()
