Flush metrics to log

Author: tianchu

.flake8

@@ -0,0 +1,2 @@
+[flake8]
+max-line-length = 100

Dockerfile

@@ -12,7 +12,7 @@ COPY requirements.txt requirements.txt
 RUN pip install -r requirements.txt -t ./python/lib/$runtime/site-packages
 
 # Install datadog_lambda
-COPY datadog_lambda ./python/lib/$runtime/site-packages
+COPY datadog_lambda ./python/lib/$runtime/site-packages/datadog_lambda
 
 # Remove *.pyc files
 RUN find ./python/lib/$runtime/site-packages -name \*.pyc -delete

README.md

@@ -12,10 +12,17 @@ arn:aws:lambda:<AWS_REGION>:464622532012:layer:Datadog-Python37:<VERSION>
 
 Replace `<AWS_REGION>` with the region where your Lambda function lives, and `<VERSION>` with the desired (or the latest) version that can be found from [CHANGELOG](CHANGELOG.md).
 
-The following Datadog environment variables must be defined via [AWS CLI](https://docs.aws.amazon.com/lambda/latest/dg/env_variables.html) or [Serverless Framework](https://serverless-stack.com/chapters/serverless-environment-variables.html):
+The Datadog API must be defined as an environment variable via [AWS CLI](https://docs.aws.amazon.com/lambda/latest/dg/env_variables.html) or [Serverless Framework](https://serverless-stack.com/chapters/serverless-environment-variables.html):
 
-* DATADOG_API_KEY
-* DATADOG_APP_KEY
+* DD_API_KEY or DD_KMS_API_KEY (if encrypted by KMS)
+
+Set the following Datadog environment variable to `datadoghq.eu` to send your data to the Datadog EU site.
+
+* DD_SITE
+
+If your Lambda function powers a performance-critical task (e.g., a consumer-facing API). You can avoid the added latencies of metric-submitting API calls, by setting the following Datadog environment variable to `True`. Custom metrics will be submitted asynchronously through CloudWatch Logs and [the Datadog log forwarder](https://github.com/DataDog/datadog-serverless-functions/tree/master/aws/logs_monitoring).
+
+* DATADOG_FLUSH_TO_LOG
 
 ### The Serverless Framework
 
@@ -40,7 +47,6 @@ functions:
       - arn:aws:lambda:us-east-1:464622532012:layer:Datadog-Python37:1
     environment:
       DATADOG_API_KEY: xxx
-      DATADOG_APP_KEY: yyy
 ```
 
 

datadog_lambda/constants.py

@@ -3,6 +3,7 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 
+
 # Datadog trace sampling priority
 class SamplingPriority(object):
     USER_REJECT = -1

datadog_lambda/metric.py

@@ -5,7 +5,11 @@
 
 import os
 import sys
+import json
+import time
+import base64
 
+import boto3
 from datadog import api
 from datadog.threadstats import ThreadStats
 from datadog_lambda import __version__
@@ -22,26 +26,42 @@ def _format_dd_lambda_layer_tag():
     return "dd_lambda_layer:datadog-{}_{}".format(runtime, __version__)
 
 
-def _tag_dd_lambda_layer(args, kwargs):
+def _tag_dd_lambda_layer(tags):
     """
     Used by lambda_metric to insert the dd_lambda_layer tag
     """
     dd_lambda_layer_tag = _format_dd_lambda_layer_tag()
-    if 'tags' in kwargs:
-        kwargs['tags'].append(dd_lambda_layer_tag)
-    elif len(args) >= 4:
-        args[3].append(dd_lambda_layer_tag)
+    if tags:
+        return tags + [dd_lambda_layer_tag]
     else:
-        kwargs['tags'] = [dd_lambda_layer_tag]
+        return [dd_lambda_layer_tag]
 
 
-def lambda_metric(*args, **kwargs):
+def lambda_metric(metric_name, value, timestamp=None, tags=None):
     """
     Submit a data point to Datadog distribution metrics.
     https://docs.datadoghq.com/graphing/metrics/distributions/
+
+    When DATADOG_LOG_FORWARDER is True, write metric to log, and
+    wait for the Datadog Log Forwarder Lambda function to submit
+    the metrics asynchronously.
+
+    Otherwise, the metrics will be submitted to the Datadog API
+    periodically and at the end of the function execution in a
+    background thread.
     """
-    _tag_dd_lambda_layer(args, kwargs)
-    lambda_stats.distribution(*args, **kwargs)
+    tags = _tag_dd_lambda_layer(tags)
+    if os.environ.get('DATADOG_FLUSH_TO_LOG') == 'True':
+        print(json.dumps({
+            'metric_name': metric_name,
+            'value': value,
+            'timestamp': timestamp or int(time.time()),
+            'tags': tags
+        }))
+    else:
+        lambda_stats.distribution(
+            metric_name, value, timestamp=timestamp, tags=tags
+        )
 
 
 def init_api_client():
@@ -57,8 +77,20 @@ def init_api_client():
     By making the initial request async, we spare a lot of execution
     time in the lambdas.
     """
-    api._api_key = os.environ.get('DATADOG_API_KEY')
-    api._api_host = os.environ.get('DATADOG_HOST', 'https://api.datadoghq.com')
+    if "DD_KMS_API_KEY" in os.environ:
+        DD_KMS_API_KEY = os.environ["DD_KMS_API_KEY"]
+        api._api_key = boto3.client("kms").decrypt(
+            CiphertextBlob=base64.b64decode(DD_KMS_API_KEY)
+        )["Plaintext"]
+    else:
+        api._api_key = os.environ.get(
+            'DATADOG_API_KEY',
+            os.environ.get('DD_API_KEY'),
+        )
+    api._api_host = os.environ.get(
+        'DATADOG_HOST',
+        'https://api.' + os.environ.get('DD_SITE', 'datadoghq.com')
+    )
     try:
         api.api_client.APIClient.submit('GET', 'validate')
     except Exception:

datadog_lambda/tracing.py

@@ -55,8 +55,7 @@ def get_dd_trace_context():
     Return the Datadog trace context to be propogated on the outgoing requests.
 
     If the Lambda function is invoked by a Datadog-traced service, a Datadog
-    trace
-    context may already exist, and it should be used. Otherwise, use the
+    trace context may already exist, and it should be used. Otherwise, use the
     current X-Ray trace entity.
 
     Most of widely-used HTTP clients are patched to inject the context

datadog_lambda/wrapper.py

@@ -3,6 +3,7 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 
+import os
 import traceback
 from threading import Thread
 
@@ -33,12 +34,14 @@ class _LambdaDecorator(object):
 
     def __init__(self, func):
         self.func = func
+        self.flush_to_log = os.environ.get('DATADOG_FLUSH_TO_LOG') == 'True'
 
     def _before(self, event, context):
         try:
             # Async initialization of the TLS connection with Datadog API,
             # and reduces the overhead to the final metric flush at the end.
-            Thread(target=init_api_client).start()
+            if not self.flush_to_log:
+                Thread(target=init_api_client).start()
 
             # Extract Datadog trace context from incoming requests
             extract_dd_trace_context(event)
@@ -50,7 +53,8 @@ def _before(self, event, context):
 
     def _after(self, event, context):
         try:
-            lambda_stats.flush(float("inf"))
+            if not self.flush_to_log:
+                lambda_stats.flush(float("inf"))
         except Exception:
             traceback.print_exc()
 

requirements.txt

@@ -1,3 +1,5 @@
 aws-xray-sdk==2.4.2
 datadog==0.28.0
-wrapt==1.11.1
+wrapt==1.11.1
+setuptools==40.8.0
+boto3==1.9.160

tests/Dockerfile

@@ -3,6 +3,7 @@ FROM python:$python_version
 
 ENV PYTHONDONTWRITEBYTECODE True
 
+COPY .flake8 .flake8
 COPY requirements.txt requirements.txt
 COPY requirements-dev.txt requirements-dev.txt
 RUN pip install -r requirements.txt

tests/test_metric.py

@@ -1,10 +1,12 @@
+import os
 import unittest
 try:
     from unittest.mock import patch, call
 except ImportError:
     from mock import patch, call
 
 from datadog_lambda.metric import (
+    init_api_client,
     lambda_metric,
     _format_dd_lambda_layer_tag,
 )
@@ -18,12 +20,52 @@ def setUp(self):
         self.addCleanup(patcher.stop)
 
     def test_lambda_metric_tagged_with_dd_lambda_layer(self):
-        lambda_metric('test.metric', 1)
-        lambda_metric('test.metric', 1, 123, ['tag1:test'])
-        lambda_metric('test.metric', 1, tags=['tag1:test'])
+        lambda_metric('test', 1)
+        lambda_metric('test', 1, 123, [])
+        lambda_metric('test', 1, tags=['tag1:test'])
         expected_tag = _format_dd_lambda_layer_tag()
         self.mock_metric_lambda_stats.distribution.assert_has_calls([
-            call('test.metric', 1, tags=[expected_tag]),
-            call('test.metric', 1, 123, ['tag1:test', expected_tag]),
-            call('test.metric', 1, tags=['tag1:test', expected_tag]),
+            call('test', 1, timestamp=None, tags=[expected_tag]),
+            call('test', 1, timestamp=123, tags=[expected_tag]),
+            call('test', 1, timestamp=None, tags=['tag1:test', expected_tag]),
         ])
+
+    def test_lambda_metric_flush_to_log(self):
+        os.environ["DATADOG_FLUSH_TO_LOG"] = 'True'
+
+        lambda_metric('test', 1)
+        self.mock_metric_lambda_stats.distribution.assert_not_called()
+
+        del os.environ["DATADOG_FLUSH_TO_LOG"]
+
+
+class TestAPIClient(unittest.TestCase):
+
+    def setUp(self):
+        patcher = patch('datadog_lambda.metric.api')
+        self.mock_api = patcher.start()
+        self.addCleanup(patcher.stop)
+
+    def test_init_api_client(self):
+        os.environ["DD_API_KEY"] = '123456'
+        os.environ["DD_SITE"] = 'dummy-site'
+
+        init_api_client()
+
+        self.assertEqual(self.mock_api._api_key, '123456')
+        self.assertEqual(self.mock_api._api_host, 'https://api.dummy-site')
+        self.mock_api.api_client.APIClient.submit.assert_called()
+
+        os.environ["DATADOG_API_KEY"] = '654321'
+        os.environ["DATADOG_HOST"] = 'dummy-host'
+
+        init_api_client()
+
+        self.assertEqual(self.mock_api._api_key, '654321')
+        self.assertEqual(self.mock_api._api_host, 'dummy-host')
+        self.mock_api.api_client.APIClient.submit.assert_called()
+
+        del os.environ["DATADOG_API_KEY"]
+        del os.environ["DATADOG_HOST"]
+        del os.environ["DD_API_KEY"]
+        del os.environ["DD_SITE"]

tests/test_wrapper.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 try:
     from unittest.mock import patch, call, ANY
@@ -35,9 +36,28 @@ def lambda_handler(event, context):
         lambda_event = {}
         lambda_context = {}
         lambda_handler(lambda_event, lambda_context)
+
         self.mock_metric_lambda_stats.distribution.assert_has_calls([
-            call('test.metric', 100, tags=ANY)
+            call('test.metric', 100, timestamp=None, tags=ANY)
         ])
         self.mock_wrapper_lambda_stats.flush.assert_called()
         self.mock_extract_dd_trace_context.assert_called_with(lambda_event)
         self.mock_patch_all.assert_called()
+
+    def test_datadog_lambda_wrapper_flush_to_log(self):
+        os.environ["DATADOG_FLUSH_TO_LOG"] = 'True'
+
+        @datadog_lambda_wrapper
+        def lambda_handler(event, context):
+            lambda_metric("test.metric", 100)
+
+        lambda_event = {}
+        lambda_context = {}
+        lambda_handler(lambda_event, lambda_context)
+
+        self.mock_metric_lambda_stats.distribution.assert_not_called()
+        self.mock_wrapper_lambda_stats.flush.assert_not_called()
+        self.mock_extract_dd_trace_context.assert_called_with(lambda_event)
+        self.mock_patch_all.assert_called()
+
+        del os.environ["DATADOG_FLUSH_TO_LOG"]