1- import asyncio
21import copy
32import functools
43import inspect
@@ -250,10 +249,11 @@ def handler(event, context)
250249 err
251250 Exception raised by method
252251 """
252+ lambda_handler_name = lambda_handler .__name__
253253
254254 @functools .wraps (lambda_handler )
255255 def decorate (event , context ):
256- with self .provider .in_subsegment (name = f"## { lambda_handler . __name__ } " ) as subsegment :
256+ with self .provider .in_subsegment (name = f"## { lambda_handler_name } " ) as subsegment :
257257 global is_cold_start
258258 if is_cold_start :
259259 logger .debug ("Annotating cold start" )
@@ -265,13 +265,12 @@ def decorate(event, context):
265265 response = lambda_handler (event , context )
266266 logger .debug ("Received lambda handler response successfully" )
267267 logger .debug (response )
268- if response :
269- subsegment .put_metadata (
270- key = "lambda handler response" , value = response , namespace = self ._config ["service" ]
271- )
268+ self ._add_response_as_metadata (
269+ function_name = lambda_handler_name , data = response , subsegment = subsegment
270+ )
272271 except Exception as err :
273- logger .exception ("Exception received from lambda handler" , exc_info = True )
274- subsegment . put_metadata ( key = f" { self .service } error" , value = err , namespace = self . _config [ "service" ] )
272+ logger .exception ("Exception received from lambda handler" )
273+ self . _add_full_exception_as_metadata ( function_name = self .service , error = err , subsegment = subsegment )
275274 raise
276275
277276 return response
@@ -392,71 +391,79 @@ async def async_tasks():
392391 """
393392 method_name = f"{ method .__name__ } "
394393
395- async def decorate_logic (
396- decorated_method_with_args : functools .partial = None ,
397- subsegment : aws_xray_sdk .core .models .subsegment = None ,
398- coroutine : bool = False ,
399- ) -> Any :
400- """Decorate logic runs both sync and async decorated methods
401-
402- Parameters
403- ----------
404- decorated_method_with_args : functools.partial
405- Partial decorated method with arguments/keyword arguments
406- subsegment : aws_xray_sdk.core.models.subsegment
407- X-Ray subsegment to reuse
408- coroutine : bool, optional
409- Instruct whether partial decorated method is a wrapped coroutine, by default False
410-
411- Returns
412- -------
413- Any
414- Returns method's response
415- """
416- response = None
417- try :
418- logger .debug (f"Calling method: { method_name } " )
419- if coroutine :
420- response = await decorated_method_with_args ()
421- else :
422- response = decorated_method_with_args ()
423- logger .debug (f"Received { method_name } response successfully" )
424- logger .debug (response )
425- except Exception as err :
426- logger .exception (f"Exception received from '{ method_name } ' method" , exc_info = True )
427- subsegment .put_metadata (key = f"{ method_name } error" , value = err , namespace = self ._config ["service" ])
428- raise
429- finally :
430- if response is not None :
431- subsegment .put_metadata ( # pragma: no cover
432- key = f"{ method_name } response" , value = response , namespace = self ._config ["service" ]
433- )
434-
435- return response
436-
437394 if inspect .iscoroutinefunction (method ):
438395
439396 @functools .wraps (method )
440397 async def decorate (* args , ** kwargs ):
441- decorated_method_with_args = functools .partial (method , * args , ** kwargs )
442398 async with self .provider .in_subsegment_async (name = f"## { method_name } " ) as subsegment :
443- return await decorate_logic (
444- decorated_method_with_args = decorated_method_with_args , subsegment = subsegment , coroutine = True
445- )
399+ try :
400+ logger .debug (f"Calling method: { method_name } " )
401+ response = await method (* args , ** kwargs )
402+ self ._add_response_as_metadata (function_name = method_name , data = response , subsegment = subsegment )
403+ except Exception as err :
404+ logger .exception (f"Exception received from '{ method_name } ' method" )
405+ self ._add_full_exception_as_metadata (
406+ function_name = method_name , error = err , subsegment = subsegment
407+ )
408+ raise
409+
410+ return response
446411
447412 else :
448413
449414 @functools .wraps (method )
450415 def decorate (* args , ** kwargs ):
451- loop = asyncio .get_event_loop ()
452- decorated_method_with_args = functools .partial (method , * args , ** kwargs )
453416 with self .provider .in_subsegment (name = f"## { method_name } " ) as subsegment :
454- return loop .run_until_complete (
455- decorate_logic (decorated_method_with_args = decorated_method_with_args , subsegment = subsegment )
456- )
417+ try :
418+ logger .debug (f"Calling method: { method_name } " )
419+ response = method (* args , ** kwargs )
420+ self ._add_response_as_metadata (function_name = method_name , data = response , subsegment = subsegment )
421+ except Exception as err :
422+ logger .exception (f"Exception received from '{ method_name } ' method" )
423+ self ._add_full_exception_as_metadata (
424+ function_name = method_name , error = err , subsegment = subsegment
425+ )
426+ raise
427+
428+ return response
457429
458430 return decorate
459431
432+ def _add_response_as_metadata (
433+ self , function_name : str = None , data : Any = None , subsegment : aws_xray_sdk .core .models .subsegment = None
434+ ):
435+ """Add response as metadata for given subsegment
436+
437+ Parameters
438+ ----------
439+ function_name : str, optional
440+ function name to add as metadata key, by default None
441+ data : Any, optional
442+ data to add as subsegment metadata, by default None
443+ subsegment : aws_xray_sdk.core.models.subsegment, optional
444+ existing subsegment to add metadata on, by default None
445+ """
446+ if data is None or subsegment is None :
447+ return
448+
449+ subsegment .put_metadata (key = f"{ function_name } response" , value = data , namespace = self ._config ["service" ])
450+
451+ def _add_full_exception_as_metadata (
452+ self , function_name : str = None , error : Exception = None , subsegment : aws_xray_sdk .core .models .subsegment = None
453+ ):
454+ """Add full exception object as metadata for given subsegment
455+
456+ Parameters
457+ ----------
458+ function_name : str, optional
459+ function name to add as metadata key, by default None
460+ error : Exception, optional
461+ error to add as subsegment metadata, by default None
462+ subsegment : aws_xray_sdk.core.models.subsegment, optional
463+ existing subsegment to add metadata on, by default None
464+ """
465+ subsegment .put_metadata (key = f"{ function_name } error" , value = error , namespace = self ._config ["service" ])
466+
460467 def __disable_tracing_provider (self ):
461468 """Forcefully disables tracing"""
462469 logger .debug ("Disabling tracer provider..." )
0 commit comments