feature: Kafka Parser support

Author: Ran Isenberg

.github/scripts/label_related_issue.js

@@ -1,29 +1,53 @@
-module.exports = async ({github, context, core}) => {
-    const prBody = context.payload.body;
-    const prNumber = context.payload.number;
-    const releaseLabel = process.env.RELEASE_LABEL;
-    const maintainersTeam = process.env.MAINTAINERS_TEAM
+const {
+  PR_AUTHOR,
+  PR_BODY,
+  PR_NUMBER,
+  IGNORE_AUTHORS,
+  LABEL_PENDING_RELEASE,
+  HANDLE_MAINTAINERS_TEAM,
+  PR_IS_MERGED,
+} = require("./constants")
 
-    const RELATED_ISSUE_REGEX = /Issue number:[^\d\r\n]+(?<issue>\d+)/;
+module.exports = async ({ github, context, core }) => {
+  if (IGNORE_AUTHORS.includes(PR_AUTHOR)) {
+    return core.notice("Author in IGNORE_AUTHORS list; skipping...")
+  }
 
-    const isMatch = RELATED_ISSUE_REGEX.exec(prBody);
+  if (PR_IS_MERGED == "false") {
+    return core.notice("Only merged PRs to avoid spam; skipping")
+  }
+
+  const RELATED_ISSUE_REGEX = /Issue number:[^\d\r\n]+(?<issue>\d+)/;
+
+  const isMatch = RELATED_ISSUE_REGEX.exec(PR_BODY);
+
+  try {
     if (!isMatch) {
-      core.setFailed(`Unable to find related issue for PR number ${prNumber}.\n\n Body details: ${prBody}`);
+      core.setFailed(`Unable to find related issue for PR number ${PR_NUMBER}.\n\n Body details: ${PR_BODY}`);
       return await github.rest.issues.createComment({
         owner: context.repo.owner,
         repo: context.repo.repo,
-        body: `${maintainersTeam} No related issues found. Please ensure '${releaseLabel}' label is applied before releasing.`,
-        issue_number: prNumber,
+        body: `${HANDLE_MAINTAINERS_TEAM} No related issues found. Please ensure '${LABEL_PENDING_RELEASE}' label is applied before releasing.`,
+        issue_number: PR_NUMBER,
       });
     }
+  } catch (error) {
+    core.setFailed(`Unable to create comment on PR number ${PR_NUMBER}.\n\n Error details: ${error}`);
+    throw new Error(error);
+  }
 
-    const { groups: {relatedIssueNumber} } = isMatch
+  const { groups: { issue } } = isMatch
 
-    core.info(`Auto-labeling related issue ${relatedIssueNumber} for release`)
+  try {
+    core.info(`Auto-labeling related issue ${issue} for release`)
     return await github.rest.issues.addLabels({
-      issue_number: relatedIssueNumber,
+      issue_number: issue,
       owner: context.repo.owner,
       repo: context.repo.repo,
-      labels: [releaseLabel]
+      labels: [LABEL_PENDING_RELEASE]
     })
-}
+  } catch (error) {
+    core.setFailed(`Is this issue number (${issue}) valid? Perhaps a discussion?`);
+    throw new Error(error);
+  }
+}

CHANGELOG.md

@@ -4,6 +4,7 @@
 from .cloudwatch import CloudWatchLogsEnvelope
 from .dynamodb import DynamoDBStreamEnvelope
 from .event_bridge import EventBridgeEnvelope
+from .kafka import KafkaEnvelope
 from .kinesis import KinesisDataStreamEnvelope
 from .lambda_function_url import LambdaFunctionUrlEnvelope
 from .sns import SnsEnvelope, SnsSqsEnvelope
@@ -20,5 +21,6 @@
     "SnsEnvelope",
     "SnsSqsEnvelope",
     "SqsEnvelope",
+    "KafkaEnvelope",
     "BaseEnvelope",
 ]

aws_lambda_powertools/utilities/parser/envelopes/__init__.py

@@ -0,0 +1,41 @@
+import logging
+from typing import Any, Dict, List, Optional, Type, Union
+
+from ..models import KafkaEventModel
+from ..types import Model
+from .base import BaseEnvelope
+
+logger = logging.getLogger(__name__)
+
+
+class KafkaEnvelope(BaseEnvelope):
+    """Kafka event envelope to extract data within body key
+    The record's body parameter is a string, though it can also be a JSON encoded string.
+    Regardless of its type it'll be parsed into a BaseModel object.
+
+    Note: Records will be parsed the same way so if model is str,
+    all items in the list will be parsed as str and npt as JSON (and vice versa)
+    """
+
+    def parse(self, data: Optional[Union[Dict[str, Any], Any]], model: Type[Model]) -> List[Optional[Model]]:
+        """Parses data found with model provided
+
+        Parameters
+        ----------
+        data : Dict
+            Lambda event to be parsed
+        model : Type[Model]
+            Data model provided to parse after extracting data using envelope
+
+        Returns
+        -------
+        Any
+            Parsed detail payload with model provided
+        """
+        logger.debug(f"Parsing incoming data with Kafka event model {KafkaEventModel}")
+        parsed_envelope: KafkaEventModel = KafkaEventModel.parse_obj(data)
+        logger.debug(f"Parsing Kafka event records in `value` with {model}")
+        ret_list = []
+        for records in parsed_envelope.records.values():
+            ret_list += [self._parse(data=record.value, model=model) for record in records]
+        return ret_list

aws_lambda_powertools/utilities/parser/envelopes/kafka.py

@@ -17,6 +17,7 @@
 from .cloudwatch import CloudWatchLogsData, CloudWatchLogsDecode, CloudWatchLogsLogEvent, CloudWatchLogsModel
 from .dynamodb import DynamoDBStreamChangedRecordModel, DynamoDBStreamModel, DynamoDBStreamRecordModel
 from .event_bridge import EventBridgeModel
+from .kafka import KafkaEventModel, KafkaRecordModel
 from .kinesis import KinesisDataStreamModel, KinesisDataStreamRecord, KinesisDataStreamRecordPayload
 from .lambda_function_url import LambdaFunctionUrlModel
 from .s3 import S3Model, S3RecordModel
@@ -98,4 +99,6 @@
     "APIGatewayEventRequestContext",
     "APIGatewayEventAuthorizer",
     "APIGatewayEventIdentity",
+    "KafkaEventModel",
+    "KafkaRecordModel",
 ]

aws_lambda_powertools/utilities/parser/models/__init__.py

@@ -0,0 +1,70 @@
+import base64
+import logging
+from binascii import Error as BinAsciiError
+from datetime import datetime
+from typing import Dict, List, Optional, Type, Union
+
+from pydantic import BaseModel, validator
+
+from aws_lambda_powertools.utilities.parser.types import Literal
+
+SERVERS_DELIMITER = ","
+
+logger = logging.getLogger(__name__)
+
+
+def _base64_decode(value: str) -> bytes:
+    try:
+        logger.debug("Decoding base64 Kafka record item before parsing")
+        return base64.b64decode(value)
+    except (BinAsciiError, TypeError):
+        raise ValueError("base64 decode failed")
+
+
+def _bytes_to_string(value: bytes) -> str:
+    try:
+        return value.decode("utf-8")
+    except (BinAsciiError, TypeError):
+        raise ValueError("base64 UTF-8 decode failed")
+
+
+class KafkaRecordModel(BaseModel):
+    topic: str
+    partition: int
+    offset: int
+    timestamp: datetime
+    timestampType: str
+    key: bytes
+    value: Union[str, Type[BaseModel]]
+    headers: List[Dict[str, bytes]]
+
+    # validators
+    _decode_key = validator("key", allow_reuse=True)(_base64_decode)
+
+    @validator("value", pre=True, allow_reuse=True)
+    def data_base64_decode(cls, value):
+        as_bytes = _base64_decode(value)
+        return _bytes_to_string(as_bytes)
+
+    @validator("headers", pre=True, allow_reuse=True)
+    def decode_headers_list(cls, value):
+        for header in value:
+            for key, values in header.items():
+                header[key] = bytes(values)
+        return value
+
+
+class KafkaEventModel(BaseModel):
+    """Self-managed Apache Kafka event trigger
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/lambda/latest/dg/with-kafka.html
+    """
+
+    eventSource: Literal["aws:SelfManagedKafka"]
+    bootstrapServers: Optional[List[str]]
+    records: Dict[str, List[KafkaRecordModel]]
+
+    @validator("bootstrapServers", pre=True, allow_reuse=True)
+    def split_servers(cls, value):
+        return None if not value else value.split(SERVERS_DELIMITER)

aws_lambda_powertools/utilities/parser/models/kafka.py

@@ -168,6 +168,7 @@ Parser comes with the following built-in models:
 | **APIGatewayProxyEventModel**   | Lambda Event Source payload for Amazon API Gateway                 |
 | **APIGatewayProxyEventV2Model** | Lambda Event Source payload for Amazon API Gateway v2 payload      |
 | **LambdaFunctionUrlModel**      | Lambda Event Source payload for Lambda Function URL payload        |
+| **KafkaModel**                  | Lambda Event Source payload for self managed Kafka payload         |
 
 ### extending built-in models
 
@@ -308,6 +309,7 @@ Parser comes with the following built-in envelopes, where `Model` in the return
 | **ApiGatewayEnvelope**        | 1. Parses data using `APIGatewayProxyEventModel`. <br/> 2. Parses `body` key using your model and returns it.                                                                                               | `Model`                            |
 | **ApiGatewayV2Envelope**      | 1. Parses data using `APIGatewayProxyEventV2Model`. <br/> 2. Parses `body` key using your model and returns it.                                                                                             | `Model`                            |
 | **LambdaFunctionUrlEnvelope** | 1. Parses data using `LambdaFunctionUrlModel`. <br/> 2. Parses `body` key using your model and returns it.                                                                                                  | `Model`                            |
+| **KafkaEnvelope**             | 1. Parses data using `KafkaRecordModel`. <br/> 2. Parses `value` key using your model and returns it.                                                                                                       | `Model`                            |
 
 ### Bringing your own envelope
 

docs/utilities/parser.md

@@ -0,0 +1,34 @@
+{
+    "eventSource": "aws:SelfManagedKafka",
+    "bootstrapServers": "b-2.demo-cluster-1.a1bcde.c1.kafka.us-east-1.amazonaws.com:9092,b-1.demo-cluster-1.a1bcde.c1.kafka.us-east-1.amazonaws.com:9092",
+    "records": {
+        "mytopic-0": [
+            {
+                "topic": "mytopic",
+                "partition": 0,
+                "offset": 15,
+                "timestamp": 1545084650987,
+                "timestampType": "CREATE_TIME",
+                "key": "cmVjb3JkS2V5",
+                "value": "eyJrZXkiOiJ2YWx1ZSJ9",
+                "headers": [
+                    {
+                        "headerKey": [
+                            104,
+                            101,
+                            97,
+                            100,
+                            101,
+                            114,
+                            86,
+                            97,
+                            108,
+                            117,
+                            101
+                        ]
+                    }
+                ]
+            }
+        ]
+    }
+}

tests/events/kafkaEvent.json

@@ -91,3 +91,7 @@ class MyApiGatewayBusiness(BaseModel):
 class MyALambdaFuncUrlBusiness(BaseModel):
     message: str
     username: str
+
+
+class MyALambdaKafkaBusiness(BaseModel):
+    key: str

tests/functional/parser/schemas.py

@@ -0,0 +1,49 @@
+from typing import List
+
+from aws_lambda_powertools.utilities.parser import envelopes, event_parser
+from aws_lambda_powertools.utilities.parser.models import KafkaEventModel, KafkaRecordModel
+from aws_lambda_powertools.utilities.typing import LambdaContext
+from tests.functional.parser.schemas import MyALambdaKafkaBusiness
+from tests.functional.utils import load_event
+
+
+@event_parser(model=MyALambdaKafkaBusiness, envelope=envelopes.KafkaEnvelope)
+def handle_lambda_kafka_with_envelope(event: List[MyALambdaKafkaBusiness], _: LambdaContext):
+    assert event[0].key == "value"
+    assert len(event) == 1
+
+
+@event_parser(model=KafkaEventModel)
+def handle_kafka_event(event: KafkaEventModel, _: LambdaContext):
+    return event
+
+
+def test_kafka_event_with_envelope():
+    event = load_event("kafkaEvent.json")
+    handle_lambda_kafka_with_envelope(event, LambdaContext())
+
+
+def test_kafka_event():
+    json_event = load_event("kafkaEvent.json")
+    event: KafkaEventModel = handle_kafka_event(json_event, LambdaContext())
+    assert event.eventSource == "aws:SelfManagedKafka"
+    bootstrap_servers = [
+        "b-2.demo-cluster-1.a1bcde.c1.kafka.us-east-1.amazonaws.com:9092",
+        "b-1.demo-cluster-1.a1bcde.c1.kafka.us-east-1.amazonaws.com:9092",
+    ]
+    assert event.bootstrapServers == bootstrap_servers
+
+    records = list(event.records["mytopic-0"])
+    assert len(records) == 1
+    record: KafkaRecordModel = records[0]
+    assert record.topic == "mytopic"
+    assert record.partition == 0
+    assert record.offset == 15
+    assert record.timestamp is not None
+    convert_time = int(round(record.timestamp.timestamp() * 1000))
+    assert convert_time == 1545084650987
+    assert record.timestampType == "CREATE_TIME"
+    assert record.key == b"recordKey"
+    assert record.value == '{"key":"value"}'
+    assert len(record.headers) == 1
+    assert record.headers[0]["headerKey"] == b"headerValue"