|
14 | 14 | logger = logging.getLogger(__name__) |
15 | 15 |
|
16 | 16 |
|
| 17 | +def _unwrap_durable_context(context: Any) -> LambdaContext: |
| 18 | + """Unwrap Lambda Context from DurableContext if applicable. |
| 19 | +
|
| 20 | + Parameters |
| 21 | + ---------- |
| 22 | + context : object |
| 23 | + Lambda context object or DurableContext |
| 24 | +
|
| 25 | + Returns |
| 26 | + ------- |
| 27 | + LambdaContext |
| 28 | + The unwrapped Lambda context |
| 29 | + """ |
| 30 | + # Check if this is a DurableContext by duck typing |
| 31 | + if hasattr(context, "lambda_context") and hasattr(context, "state"): |
| 32 | + return context.lambda_context |
| 33 | + |
| 34 | + return context |
| 35 | + |
| 36 | + |
17 | 37 | class BaseProvider(ABC): |
18 | 38 | """ |
19 | 39 | Interface to create a metrics provider. |
@@ -177,6 +197,11 @@ def handler(event, context): |
177 | 197 | default_dimensions: dict[str, str], optional |
178 | 198 | metric dimensions as key=value that will always be present |
179 | 199 |
|
| 200 | + Notes |
| 201 | + ----- |
| 202 | + Supports both standard Lambda Context and DurableContext from AWS Durable Execution SDK. |
| 203 | + When DurableContext is passed, it automatically unwraps the underlying Lambda Context. |
| 204 | +
|
180 | 205 | Raises |
181 | 206 | ------ |
182 | 207 | e |
@@ -221,13 +246,15 @@ def _add_cold_start_metric(self, context: Any) -> None: |
221 | 246 | Parameters |
222 | 247 | ---------- |
223 | 248 | context : Any |
224 | | - Lambda context |
| 249 | + Lambda context or DurableContext |
225 | 250 | """ |
226 | 251 | if not cold_start.is_cold_start: |
227 | 252 | return |
228 | 253 |
|
229 | 254 | logger.debug("Adding cold start metric and function_name dimension") |
230 | | - self.add_cold_start_metric(context=context) |
| 255 | + # Unwrap DurableContext if applicable before passing to add_cold_start_metric |
| 256 | + lambda_context = _unwrap_durable_context(context) |
| 257 | + self.add_cold_start_metric(context=lambda_context) |
231 | 258 |
|
232 | 259 | cold_start.is_cold_start = False |
233 | 260 |
|
|
0 commit comments