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()