forked from DataDog/datadog-lambda-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathwrapper.py
More file actions
90 lines (71 loc) · 2.47 KB
/
wrapper.py
File metadata and controls
90 lines (71 loc) · 2.47 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
# Unless explicitly stated otherwise all files in this repository are licensed
# under the Apache License Version 2.0.
# This product includes software developed at Datadog (https://www.datadoghq.com/).
# Copyright 2019 Datadog, Inc.
import os
import logging
import traceback
from datadog_lambda.cold_start import set_cold_start
from datadog_lambda.metric import (
lambda_stats,
submit_invocations_metric,
submit_errors_metric,
)
from datadog_lambda.patch import patch_all
from datadog_lambda.tracing import (
extract_dd_trace_context,
set_correlation_ids,
inject_correlation_ids,
)
logger = logging.getLogger(__name__)
"""
Usage:
import requests
from datadog_lambda.wrapper import datadog_lambda_wrapper
from datadog_lambda.metric import lambda_metric
@datadog_lambda_wrapper
def my_lambda_handle(event, context):
lambda_metric("my_metric", 10)
requests.get("https://www.datadoghq.com")
"""
class _LambdaDecorator(object):
"""
Decorator to automatically initialize Datadog API client, flush metrics,
and extracts/injects trace context.
"""
def __init__(self, func):
self.func = func
self.flush_to_log = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true"
self.logs_injection = os.environ.get("DD_LOGS_INJECTION", "").lower() == "true"
# Inject trace correlation ids to logs
if self.logs_injection:
inject_correlation_ids()
# Patch HTTP clients to propagate Datadog trace context
patch_all()
logger.debug("datadog_lambda_wrapper initialized")
def _before(self, event, context):
set_cold_start()
try:
submit_invocations_metric(context)
# Extract Datadog trace context from incoming requests
extract_dd_trace_context(event)
# Set log correlation ids using extracted trace context
set_correlation_ids()
except Exception:
traceback.print_exc()
def _after(self, event, context):
try:
if not self.flush_to_log:
lambda_stats.flush(float("inf"))
except Exception:
traceback.print_exc()
def __call__(self, event, context, **kwargs):
self._before(event, context)
try:
return self.func(event, context, **kwargs)
except Exception:
submit_errors_metric(context)
raise
finally:
self._after(event, context)
datadog_lambda_wrapper = _LambdaDecorator