aws_lambda.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499
  1. import functools
  2. import json
  3. import re
  4. import sys
  5. from copy import deepcopy
  6. from datetime import datetime, timedelta, timezone
  7. from os import environ
  8. import sentry_sdk
  9. from sentry_sdk.api import continue_trace
  10. from sentry_sdk.consts import OP
  11. from sentry_sdk.scope import should_send_default_pii
  12. from sentry_sdk.tracing import TransactionSource
  13. from sentry_sdk.utils import (
  14. AnnotatedValue,
  15. capture_internal_exceptions,
  16. ensure_integration_enabled,
  17. event_from_exception,
  18. logger,
  19. TimeoutThread,
  20. reraise,
  21. )
  22. from sentry_sdk.integrations import Integration
  23. from sentry_sdk.integrations._wsgi_common import _filter_headers
  24. from typing import TYPE_CHECKING
  25. if TYPE_CHECKING:
  26. from typing import Any
  27. from typing import TypeVar
  28. from typing import Callable
  29. from typing import Optional
  30. from sentry_sdk._types import EventProcessor, Event, Hint
  31. F = TypeVar("F", bound=Callable[..., Any])
  32. # Constants
  33. TIMEOUT_WARNING_BUFFER = 1500 # Buffer time required to send timeout warning to Sentry
  34. MILLIS_TO_SECONDS = 1000.0
  35. def _wrap_init_error(init_error):
  36. # type: (F) -> F
  37. @ensure_integration_enabled(AwsLambdaIntegration, init_error)
  38. def sentry_init_error(*args, **kwargs):
  39. # type: (*Any, **Any) -> Any
  40. client = sentry_sdk.get_client()
  41. with capture_internal_exceptions():
  42. sentry_sdk.get_isolation_scope().clear_breadcrumbs()
  43. exc_info = sys.exc_info()
  44. if exc_info and all(exc_info):
  45. sentry_event, hint = event_from_exception(
  46. exc_info,
  47. client_options=client.options,
  48. mechanism={"type": "aws_lambda", "handled": False},
  49. )
  50. sentry_sdk.capture_event(sentry_event, hint=hint)
  51. else:
  52. # Fall back to AWS lambdas JSON representation of the error
  53. error_info = args[1]
  54. if isinstance(error_info, str):
  55. error_info = json.loads(error_info)
  56. sentry_event = _event_from_error_json(error_info)
  57. sentry_sdk.capture_event(sentry_event)
  58. return init_error(*args, **kwargs)
  59. return sentry_init_error # type: ignore
  60. def _wrap_handler(handler):
  61. # type: (F) -> F
  62. @functools.wraps(handler)
  63. def sentry_handler(aws_event, aws_context, *args, **kwargs):
  64. # type: (Any, Any, *Any, **Any) -> Any
  65. # Per https://docs.aws.amazon.com/lambda/latest/dg/python-handler.html,
  66. # `event` here is *likely* a dictionary, but also might be a number of
  67. # other types (str, int, float, None).
  68. #
  69. # In some cases, it is a list (if the user is batch-invoking their
  70. # function, for example), in which case we'll use the first entry as a
  71. # representative from which to try pulling request data. (Presumably it
  72. # will be the same for all events in the list, since they're all hitting
  73. # the lambda in the same request.)
  74. client = sentry_sdk.get_client()
  75. integration = client.get_integration(AwsLambdaIntegration)
  76. if integration is None:
  77. return handler(aws_event, aws_context, *args, **kwargs)
  78. if isinstance(aws_event, list) and len(aws_event) >= 1:
  79. request_data = aws_event[0]
  80. batch_size = len(aws_event)
  81. else:
  82. request_data = aws_event
  83. batch_size = 1
  84. if not isinstance(request_data, dict):
  85. # If we're not dealing with a dictionary, we won't be able to get
  86. # headers, path, http method, etc in any case, so it's fine that
  87. # this is empty
  88. request_data = {}
  89. configured_time = aws_context.get_remaining_time_in_millis()
  90. with sentry_sdk.isolation_scope() as scope:
  91. timeout_thread = None
  92. with capture_internal_exceptions():
  93. scope.clear_breadcrumbs()
  94. scope.add_event_processor(
  95. _make_request_event_processor(
  96. request_data, aws_context, configured_time
  97. )
  98. )
  99. scope.set_tag(
  100. "aws_region", aws_context.invoked_function_arn.split(":")[3]
  101. )
  102. if batch_size > 1:
  103. scope.set_tag("batch_request", True)
  104. scope.set_tag("batch_size", batch_size)
  105. # Starting the Timeout thread only if the configured time is greater than Timeout warning
  106. # buffer and timeout_warning parameter is set True.
  107. if (
  108. integration.timeout_warning
  109. and configured_time > TIMEOUT_WARNING_BUFFER
  110. ):
  111. waiting_time = (
  112. configured_time - TIMEOUT_WARNING_BUFFER
  113. ) / MILLIS_TO_SECONDS
  114. timeout_thread = TimeoutThread(
  115. waiting_time,
  116. configured_time / MILLIS_TO_SECONDS,
  117. )
  118. # Starting the thread to raise timeout warning exception
  119. timeout_thread.start()
  120. headers = request_data.get("headers", {})
  121. # Some AWS Services (ie. EventBridge) set headers as a list
  122. # or None, so we must ensure it is a dict
  123. if not isinstance(headers, dict):
  124. headers = {}
  125. transaction = continue_trace(
  126. headers,
  127. op=OP.FUNCTION_AWS,
  128. name=aws_context.function_name,
  129. source=TransactionSource.COMPONENT,
  130. origin=AwsLambdaIntegration.origin,
  131. )
  132. with sentry_sdk.start_transaction(
  133. transaction,
  134. custom_sampling_context={
  135. "aws_event": aws_event,
  136. "aws_context": aws_context,
  137. },
  138. ):
  139. try:
  140. return handler(aws_event, aws_context, *args, **kwargs)
  141. except Exception:
  142. exc_info = sys.exc_info()
  143. sentry_event, hint = event_from_exception(
  144. exc_info,
  145. client_options=client.options,
  146. mechanism={"type": "aws_lambda", "handled": False},
  147. )
  148. sentry_sdk.capture_event(sentry_event, hint=hint)
  149. reraise(*exc_info)
  150. finally:
  151. if timeout_thread:
  152. timeout_thread.stop()
  153. return sentry_handler # type: ignore
  154. def _drain_queue():
  155. # type: () -> None
  156. with capture_internal_exceptions():
  157. client = sentry_sdk.get_client()
  158. integration = client.get_integration(AwsLambdaIntegration)
  159. if integration is not None:
  160. # Flush out the event queue before AWS kills the
  161. # process.
  162. client.flush()
  163. class AwsLambdaIntegration(Integration):
  164. identifier = "aws_lambda"
  165. origin = f"auto.function.{identifier}"
  166. def __init__(self, timeout_warning=False):
  167. # type: (bool) -> None
  168. self.timeout_warning = timeout_warning
  169. @staticmethod
  170. def setup_once():
  171. # type: () -> None
  172. lambda_bootstrap = get_lambda_bootstrap()
  173. if not lambda_bootstrap:
  174. logger.warning(
  175. "Not running in AWS Lambda environment, "
  176. "AwsLambdaIntegration disabled (could not find bootstrap module)"
  177. )
  178. return
  179. if not hasattr(lambda_bootstrap, "handle_event_request"):
  180. logger.warning(
  181. "Not running in AWS Lambda environment, "
  182. "AwsLambdaIntegration disabled (could not find handle_event_request)"
  183. )
  184. return
  185. pre_37 = hasattr(lambda_bootstrap, "handle_http_request") # Python 3.6
  186. if pre_37:
  187. old_handle_event_request = lambda_bootstrap.handle_event_request
  188. def sentry_handle_event_request(request_handler, *args, **kwargs):
  189. # type: (Any, *Any, **Any) -> Any
  190. request_handler = _wrap_handler(request_handler)
  191. return old_handle_event_request(request_handler, *args, **kwargs)
  192. lambda_bootstrap.handle_event_request = sentry_handle_event_request
  193. old_handle_http_request = lambda_bootstrap.handle_http_request
  194. def sentry_handle_http_request(request_handler, *args, **kwargs):
  195. # type: (Any, *Any, **Any) -> Any
  196. request_handler = _wrap_handler(request_handler)
  197. return old_handle_http_request(request_handler, *args, **kwargs)
  198. lambda_bootstrap.handle_http_request = sentry_handle_http_request
  199. # Patch to_json to drain the queue. This should work even when the
  200. # SDK is initialized inside of the handler
  201. old_to_json = lambda_bootstrap.to_json
  202. def sentry_to_json(*args, **kwargs):
  203. # type: (*Any, **Any) -> Any
  204. _drain_queue()
  205. return old_to_json(*args, **kwargs)
  206. lambda_bootstrap.to_json = sentry_to_json
  207. else:
  208. lambda_bootstrap.LambdaRuntimeClient.post_init_error = _wrap_init_error(
  209. lambda_bootstrap.LambdaRuntimeClient.post_init_error
  210. )
  211. old_handle_event_request = lambda_bootstrap.handle_event_request
  212. def sentry_handle_event_request( # type: ignore
  213. lambda_runtime_client, request_handler, *args, **kwargs
  214. ):
  215. request_handler = _wrap_handler(request_handler)
  216. return old_handle_event_request(
  217. lambda_runtime_client, request_handler, *args, **kwargs
  218. )
  219. lambda_bootstrap.handle_event_request = sentry_handle_event_request
  220. # Patch the runtime client to drain the queue. This should work
  221. # even when the SDK is initialized inside of the handler
  222. def _wrap_post_function(f):
  223. # type: (F) -> F
  224. def inner(*args, **kwargs):
  225. # type: (*Any, **Any) -> Any
  226. _drain_queue()
  227. return f(*args, **kwargs)
  228. return inner # type: ignore
  229. lambda_bootstrap.LambdaRuntimeClient.post_invocation_result = (
  230. _wrap_post_function(
  231. lambda_bootstrap.LambdaRuntimeClient.post_invocation_result
  232. )
  233. )
  234. lambda_bootstrap.LambdaRuntimeClient.post_invocation_error = (
  235. _wrap_post_function(
  236. lambda_bootstrap.LambdaRuntimeClient.post_invocation_error
  237. )
  238. )
  239. def get_lambda_bootstrap():
  240. # type: () -> Optional[Any]
  241. # Python 3.7: If the bootstrap module is *already imported*, it is the
  242. # one we actually want to use (no idea what's in __main__)
  243. #
  244. # Python 3.8: bootstrap is also importable, but will be the same file
  245. # as __main__ imported under a different name:
  246. #
  247. # sys.modules['__main__'].__file__ == sys.modules['bootstrap'].__file__
  248. # sys.modules['__main__'] is not sys.modules['bootstrap']
  249. #
  250. # Python 3.9: bootstrap is in __main__.awslambdaricmain
  251. #
  252. # On container builds using the `aws-lambda-python-runtime-interface-client`
  253. # (awslamdaric) module, bootstrap is located in sys.modules['__main__'].bootstrap
  254. #
  255. # Such a setup would then make all monkeypatches useless.
  256. if "bootstrap" in sys.modules:
  257. return sys.modules["bootstrap"]
  258. elif "__main__" in sys.modules:
  259. module = sys.modules["__main__"]
  260. # python3.9 runtime
  261. if hasattr(module, "awslambdaricmain") and hasattr(
  262. module.awslambdaricmain, "bootstrap"
  263. ):
  264. return module.awslambdaricmain.bootstrap
  265. elif hasattr(module, "bootstrap"):
  266. # awslambdaric python module in container builds
  267. return module.bootstrap
  268. # python3.8 runtime
  269. return module
  270. else:
  271. return None
  272. def _make_request_event_processor(aws_event, aws_context, configured_timeout):
  273. # type: (Any, Any, Any) -> EventProcessor
  274. start_time = datetime.now(timezone.utc)
  275. def event_processor(sentry_event, hint, start_time=start_time):
  276. # type: (Event, Hint, datetime) -> Optional[Event]
  277. remaining_time_in_milis = aws_context.get_remaining_time_in_millis()
  278. exec_duration = configured_timeout - remaining_time_in_milis
  279. extra = sentry_event.setdefault("extra", {})
  280. extra["lambda"] = {
  281. "function_name": aws_context.function_name,
  282. "function_version": aws_context.function_version,
  283. "invoked_function_arn": aws_context.invoked_function_arn,
  284. "aws_request_id": aws_context.aws_request_id,
  285. "execution_duration_in_millis": exec_duration,
  286. "remaining_time_in_millis": remaining_time_in_milis,
  287. }
  288. extra["cloudwatch logs"] = {
  289. "url": _get_cloudwatch_logs_url(aws_context, start_time),
  290. "log_group": aws_context.log_group_name,
  291. "log_stream": aws_context.log_stream_name,
  292. }
  293. request = sentry_event.get("request", {})
  294. if "httpMethod" in aws_event:
  295. request["method"] = aws_event["httpMethod"]
  296. request["url"] = _get_url(aws_event, aws_context)
  297. if "queryStringParameters" in aws_event:
  298. request["query_string"] = aws_event["queryStringParameters"]
  299. if "headers" in aws_event:
  300. request["headers"] = _filter_headers(aws_event["headers"])
  301. if should_send_default_pii():
  302. user_info = sentry_event.setdefault("user", {})
  303. identity = aws_event.get("identity")
  304. if identity is None:
  305. identity = {}
  306. id = identity.get("userArn")
  307. if id is not None:
  308. user_info.setdefault("id", id)
  309. ip = identity.get("sourceIp")
  310. if ip is not None:
  311. user_info.setdefault("ip_address", ip)
  312. if "body" in aws_event:
  313. request["data"] = aws_event.get("body", "")
  314. else:
  315. if aws_event.get("body", None):
  316. # Unfortunately couldn't find a way to get structured body from AWS
  317. # event. Meaning every body is unstructured to us.
  318. request["data"] = AnnotatedValue.removed_because_raw_data()
  319. sentry_event["request"] = deepcopy(request)
  320. return sentry_event
  321. return event_processor
  322. def _get_url(aws_event, aws_context):
  323. # type: (Any, Any) -> str
  324. path = aws_event.get("path", None)
  325. headers = aws_event.get("headers")
  326. if headers is None:
  327. headers = {}
  328. host = headers.get("Host", None)
  329. proto = headers.get("X-Forwarded-Proto", None)
  330. if proto and host and path:
  331. return "{}://{}{}".format(proto, host, path)
  332. return "awslambda:///{}".format(aws_context.function_name)
  333. def _get_cloudwatch_logs_url(aws_context, start_time):
  334. # type: (Any, datetime) -> str
  335. """
  336. Generates a CloudWatchLogs console URL based on the context object
  337. Arguments:
  338. aws_context {Any} -- context from lambda handler
  339. Returns:
  340. str -- AWS Console URL to logs.
  341. """
  342. formatstring = "%Y-%m-%dT%H:%M:%SZ"
  343. region = environ.get("AWS_REGION", "")
  344. url = (
  345. "https://console.{domain}/cloudwatch/home?region={region}"
  346. "#logEventViewer:group={log_group};stream={log_stream}"
  347. ";start={start_time};end={end_time}"
  348. ).format(
  349. domain="amazonaws.cn" if region.startswith("cn-") else "aws.amazon.com",
  350. region=region,
  351. log_group=aws_context.log_group_name,
  352. log_stream=aws_context.log_stream_name,
  353. start_time=(start_time - timedelta(seconds=1)).strftime(formatstring),
  354. end_time=(datetime.now(timezone.utc) + timedelta(seconds=2)).strftime(
  355. formatstring
  356. ),
  357. )
  358. return url
  359. def _parse_formatted_traceback(formatted_tb):
  360. # type: (list[str]) -> list[dict[str, Any]]
  361. frames = []
  362. for frame in formatted_tb:
  363. match = re.match(r'File "(.+)", line (\d+), in (.+)', frame.strip())
  364. if match:
  365. file_name, line_number, func_name = match.groups()
  366. line_number = int(line_number)
  367. frames.append(
  368. {
  369. "filename": file_name,
  370. "function": func_name,
  371. "lineno": line_number,
  372. "vars": None,
  373. "pre_context": None,
  374. "context_line": None,
  375. "post_context": None,
  376. }
  377. )
  378. return frames
  379. def _event_from_error_json(error_json):
  380. # type: (dict[str, Any]) -> Event
  381. """
  382. Converts the error JSON from AWS Lambda into a Sentry error event.
  383. This is not a full fletched event, but better than nothing.
  384. This is an example of where AWS creates the error JSON:
  385. https://github.com/aws/aws-lambda-python-runtime-interface-client/blob/2.2.1/awslambdaric/bootstrap.py#L479
  386. """
  387. event = {
  388. "level": "error",
  389. "exception": {
  390. "values": [
  391. {
  392. "type": error_json.get("errorType"),
  393. "value": error_json.get("errorMessage"),
  394. "stacktrace": {
  395. "frames": _parse_formatted_traceback(
  396. error_json.get("stackTrace", [])
  397. ),
  398. },
  399. "mechanism": {
  400. "type": "aws_lambda",
  401. "handled": False,
  402. },
  403. }
  404. ],
  405. },
  406. } # type: Event
  407. return event