Feature request: Add the ability to access the LambdaContext when using batch processing

[mew1033]

Use case

I'm looking to use the get_remaining_time_in_millis function specifically. In this case, the reason is that my batch size sometimes includes very large items that might time out if they are at the end of the batch. I want to add something to the top of the record_handler that checks the remaining time, and if it's below a threshold it just raises an exception for that record. Then the partial failure mechanism can put the failed items back on the queue.

As it is right now, if the function just times out, everything is put back, not just the failed items.

Solution/User Experience

Passing the LambdaContext object into the record_handler would be great.

Alternative solutions

No response

Acknowledgment

• This feature request meets Lambda Powertools Tenets
• Should this be considered in other Lambda Powertools languages? i.e. Java, TypeScript

[boring-cyborg]

Thanks for opening your first issue here! We'll come back to you as soon as we can.

[heitorlessa]

hey @mew1033 thank you so much for creating the feature request we asked. Until we implement it, you can use functools.partial to create a new Callable with lambda_context populated, and instead of using the decorator you can use the Context Manager we provide for situations like these.

Here's a working snippet I tested. Please do let us know if this doesn't help you in the meantime - we will implement it as soon as our pause ends (end of July).

from functools import partial
from typing import Optional

from aws_lambda_powertools import Logger, Tracer
from aws_lambda_powertools.utilities.batch import BatchProcessor, EventType
from aws_lambda_powertools.utilities.data_classes.sqs_event import SQSRecord
from aws_lambda_powertools.utilities.typing import LambdaContext

processor = BatchProcessor(event_type=EventType.SQS)
tracer = Tracer()
logger = Logger()


@tracer.capture_method
def record_handler(record: SQSRecord, lambda_context: Optional[LambdaContext] = None):
    ...

@logger.inject_lambda_context
@tracer.capture_lambda_handler
def lambda_handler(event, context: LambdaContext):
    
    # partial creates a callable with default args (lambda_context in this case)
    # BatchProcessor will then call the record_handler and pass the data as it normally would
    # as the parameter is optional, it won't interfere in how BatchProcessor calls the callable 
    # For example, record_handler(record=data))
    with processor(records=batch, handler=partial(record_handler, lambda_context=context)) as batch:
        batch.process()

    return batch.response()

[mew1033]

Cool, this should work pretty well. Thank you! And thanks for looking into the feature request.

[heitorlessa]

Long delay but we've finally managed to prioritize this now - WIP PR if you're interested in the details: #1561

Once this is merged, all you'd need to do is to update your record_handler to include a new optional parameter lambda_context. When using the decorator @batch_processor, we'll know we can automatically inject the Lambda context into your record_handler

Preview of the experience

from typing import Optional

from aws_lambda_powertools import Logger, Tracer
from aws_lambda_powertools.utilities.batch import (BatchProcessor, EventType,
                                                   batch_processor)
from aws_lambda_powertools.utilities.data_classes.sqs_event import SQSRecord
from aws_lambda_powertools.utilities.typing import LambdaContext

processor = BatchProcessor(event_type=EventType.SQS)
tracer = Tracer()
logger = Logger()


@tracer.capture_method
def record_handler(record: SQSRecord, lambda_context: Optional[LambdaContext] = None):
    if lambda_context is not None:
        remaining_time = lambda_context.get_remaining_time_in_millis()
    ...


@logger.inject_lambda_context
@tracer.capture_lambda_handler
@batch_processor(record_handler=record_handler, processor=processor)
def lambda_handler(event, context: LambdaContext):
    return processor.response()

When you're using the context manager, you'll need to pass on the Lambda context so we can inject it successfully, otherwise it'll simply be None as your signature expects:

from typing import Optional

from aws_lambda_powertools import Logger, Tracer
from aws_lambda_powertools.utilities.batch import BatchProcessor, EventType
from aws_lambda_powertools.utilities.data_classes.sqs_event import SQSRecord
from aws_lambda_powertools.utilities.typing import LambdaContext

processor = BatchProcessor(event_type=EventType.SQS)
tracer = Tracer()
logger = Logger()


@tracer.capture_method
def record_handler(record: SQSRecord, lambda_context: Optional[LambdaContext] = None):
    if lambda_context is not None:
        remaining_time = lambda_context.get_remaining_time_in_millis()
    ...

@logger.inject_lambda_context
@tracer.capture_lambda_handler
def lambda_handler(event, context: LambdaContext):
    batch = event["Records"]
    with processor(records=batch, handler=record_handler, lambda_context=context):
        result = processor.process()

    return result

Hope you like it :-)