"""Utils for generating kafka events for testing.""" import base64 import copy from datetime import datetime from datetime import timedelta from datetime import timezone import json import os from confluent_kafka.serialization import MessageField from confluent_kafka.serialization import SerializationContext from confluent_kafka.serialization import StringSerializer import pkg_resources from kafka_utils.producer.serializer import avro from kafka_utils.producer.serializer.simple_json import SimpleJSONSerializer _SCHEMA_REGISTRY_URL = os.environ.get( 'SCHEMA_REGISTRY_URL', 'https://dev-schema-registry.dev.theorchard.io') _example_cdc_key = pkg_resources.resource_filename( 'kafka_utils', 'testing/integration/_example_cdc_key.json') _example_cdc_value = pkg_resources.resource_filename( 'kafka_utils', 'testing/integration/_example_cdc_value.json') def _build_key_serialization_context(topic, headers=None): return SerializationContext(topic, MessageField.KEY, headers) def _build_value_serialization_context(topic, headers=None): return SerializationContext(topic, MessageField.VALUE, headers) def _base64_encode(key, value): """Base64 encode a key and value.""" try: encoded_key = base64.b64encode(key).decode('ascii') except TypeError: encoded_key = base64.b64encode(key.encode('ascii')).decode('ascii') try: encoded_value = base64.b64encode(value).decode('ascii') except TypeError: encoded_value = base64.b64encode(value.encode('ascii')).decode('ascii') return encoded_key, encoded_value def _serialize_cdc(topic, key, value): """Serialize cdc message.""" simple_json = SimpleJSONSerializer() serialized_key = simple_json( key, _build_key_serialization_context(topic)) serialized_value = simple_json( value, _build_value_serialization_context(topic)) return _base64_encode(serialized_key, serialized_value) def _serialize_avro(topic, schema_version, key, value): """Serialize avro message.""" string_serializer = StringSerializer() avro_serializer = avro.build_avro_serializer( _SCHEMA_REGISTRY_URL, topic, schema_version) serialized_key = string_serializer( key, _build_key_serialization_context(topic)) serialized_value = avro_serializer( value, _build_value_serialization_context(topic)) return _base64_encode(serialized_key, serialized_value) def _group_messages(messages): messages = copy.deepcopy(messages) grouped_messages = {} for message in messages: topic = message['topic'] group_key = f'{topic}-{message.get("partition", 0)}' topic_partition_messages = grouped_messages.get(group_key, []) offset = len(topic_partition_messages) message['offset'] = offset message['timestamp'] = round(datetime.now(timezone.utc).timestamp()) + offset message['timestampType'] = 'CREATE_TIME' message['headers'] = [] topic_partition_messages.append(message) grouped_messages[group_key] = topic_partition_messages return grouped_messages _aws_msk_lambda_event_template = { 'eventSource': 'aws:kafka', 'eventSourceArn': 'arn:aws:kafka:us-east-1:103233932089:cluster/dev-managed-kafka-cdc-destination/194e93a4-1a4e-4e96-85bd-99b4ebce8df0-11', # noqa 'bootstrapServers': 'b-3.dev-managed-kafka-cdc.rk4es0.c11.kafka.us-east-1.amazonaws.com:9094,b-1.dev-managed-kafka-cdc.rk4es0.c11.kafka.us-east-1.amazonaws.com:9094,b-2.dev-managed-kafka-cdc.rk4es0.c11.kafka.us-east-1.amazonaws.com:9094', # noqa 'records': {} } def generate_aws_msk_lambda_event(messages): """Generate aws msk lambda event.""" aws_msk_lambda_event = copy.deepcopy(_aws_msk_lambda_event_template) aws_msk_lambda_event['records'] = _group_messages(messages) return json.dumps(aws_msk_lambda_event) def build_message(topic, key, value, partition=0): """Build message.""" return { 'topic': topic, 'partition': partition, 'key': key, 'value': value, } def build_avro_message(topic, schema_version, key, value, partition=0): """Build avro message.""" serialized_key, serialized_value = _serialize_avro(topic, schema_version, key, value) return { **build_message(topic, serialized_key, serialized_value, partition), 'schema_version': schema_version, } def build_cdc_message(topic, key, value, partition=0): """Build cdc message.""" serialized_key, serialized_value = _serialize_cdc(topic, key, value) return build_message(topic, serialized_key, serialized_value, partition) def example_usage(): """Run example for aws_msk_lambda_utils.""" with open(_example_cdc_key) as file_example_cdc_key,\ open(_example_cdc_value) as file_example_cdc_value: example_cdc_key = json.load(file_example_cdc_key) example_cdc_value = json.load(file_example_cdc_value) aws_msk_lambda_event_messages = [ build_avro_message( topic='event.owsContentReview.review', schema_version=5, key='', value={ 'operation': { 'type': 'create', 'context': 'new', 'timestamp': round(datetime.now(timezone.utc).timestamp()) }, 'payload': { 'product_id': 1, 'review_queue_id': 1, 'review_note': None, 'user_id': None, 'locked_by_user_id': None, 'locked_until_datetime': None, } }, partition=0, ), build_avro_message( topic='event.owsContentReview.review', schema_version=5, key='', value={ 'operation': { 'type': 'patch', 'context': 'reject', 'timestamp': round(datetime.now(timezone.utc).timestamp()) }, 'payload': { 'product_id': 2, 'review_queue_id': 1, 'review_note': 'was not good', 'user_id': '1f35d33c-bcda-45fb-acb3-2e052558cd42', 'locked_by_user_id': None, 'locked_until_datetime': None, } }, partition=0, ), build_avro_message( topic='event.owsContentReview.review', schema_version=5, key='', value={ 'operation': { 'type': 'patch', 'context': 'lock', 'timestamp': round(datetime.now(timezone.utc).timestamp()) }, 'payload': { 'product_id': 3, 'review_queue_id': 3, 'review_note': None, 'user_id': '1f35d33c-bcda-45fb-acb3-2e052558cd42', 'locked_by_user_id': '7e9bba4f-5a38-446c-9499-b5fe8488f4e7', 'locked_until_datetime': int( (datetime.now(timezone.utc) + timedelta(minutes=10)).timestamp() * 1E6), } }, partition=12321, ), build_avro_message( topic='event.owsProductDigital.product', schema_version=1, key='product_id_12321', value={ 'operation': { 'type': 'create', 'context': 'new', 'timestamp': round(datetime.now(timezone.utc).timestamp()) }, 'payload': { 'product_id': 12321, 'upc': '123456789128', 'product_name': 'Applesauce Bananas', 'project_id': 32123 } }, partition=10, ), build_cdc_message( topic='cdc.contentReview.reviewQueue', key=example_cdc_key, value=example_cdc_value, partition=0, ), build_cdc_message( topic='cdc.contentReview.reviewQueue', key=example_cdc_key, value=example_cdc_value, partition=1, ), build_cdc_message( topic='cdc.applesauce.bananas', key=example_cdc_key, value=example_cdc_value, partition=10, ), ] return generate_aws_msk_lambda_event(aws_msk_lambda_event_messages) if __name__ == '__main__': print(example_usage())