Merge pull request #4 from DataDog/tian.chu/flush-metric-to-log

Submit custom metrics via CloudWatch Logs

Author: tianchu

.flake8

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

CHANGELOG.md

@@ -1,8 +1,14 @@
 CHANGELOG
 =========
+# v2 / 2019-06-10
+
+* Support submitting metrics through CloudWatch Logs
+* Support submitting metrics to `datadoghq.eu`
+* Support KMS-encrypted DD API Key
+* Fix a few bugs
 
 # v1 / 2019-05-06
 
 * First release
 * Support submitting distribution metrics from Lambda functions
-* Supports distributed tracing between serverful and serverless services
+* Support distributed tracing between serverful and serverless services

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/__init__.py

@@ -1 +1 @@
-__version__ = "1"
+__version__ = "2"

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,44 +26,58 @@ 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/
-    """
-    _tag_dd_lambda_layer(args, kwargs)
-    lambda_stats.distribution(*args, **kwargs)
 
+    When DATADOG_LOG_FORWARDER is True, write metric to log, and
+    wait for the Datadog Log Forwarder Lambda function to submit
+    the metrics asynchronously.
 
-def init_api_client():
+    Otherwise, the metrics will be submitted to the Datadog API
+    periodically and at the end of the function execution in a
+    background thread.
     """
-    No-op GET to initialize the requests connection with DD's endpoints,
-    to make the final flush faster.
+    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
+        )
 
-    We keep alive the Requests session, this means that we can re-use
-    the connection. The consequence is that the HTTP Handshake, which
-    can take hundreds of ms, is now made at the beginning of a lambda
-    instead of at the end.
 
-    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')
-    try:
-        api.api_client.APIClient.submit('GET', 'validate')
-    except Exception:
-        pass
+# Decrypt code should run once and variables stored outside of the function
+# handler so that these are decrypted once per container
+DD_KMS_API_KEY = os.environ.get("DD_KMS_API_KEY")
+if DD_KMS_API_KEY:
+    DD_KMS_API_KEY = boto3.client("kms").decrypt(
+        CiphertextBlob=base64.b64decode(DD_KMS_API_KEY)
+    )["Plaintext"]
+
+# Set API Key and Host in the module, so they only set once per container
+api._api_key = os.environ.get(
+    'DATADOG_API_KEY',
+    os.environ.get('DD_API_KEY', DD_KMS_API_KEY),
+)
+api._api_host = os.environ.get(
+    'DATADOG_HOST',
+    'https://api.' + os.environ.get('DD_SITE', 'datadoghq.com')
+)

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,10 +3,10 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 
+import os
 import traceback
-from threading import Thread
 
-from datadog_lambda.metric import init_api_client, lambda_stats
+from datadog_lambda.metric import lambda_stats
 from datadog_lambda.tracing import extract_dd_trace_context
 from datadog_lambda.patch import patch_all
 
@@ -33,13 +33,10 @@ 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()
-
             # Extract Datadog trace context from incoming requests
             extract_dd_trace_context(event)
 
@@ -50,7 +47,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

setup.py

@@ -0,0 +1,27 @@
+from setuptools import setup
+from os import path
+from io import open
+
+here = path.abspath(path.dirname(__file__))
+
+with open(path.join(here, 'README.md'), encoding='utf-8') as f:
+    long_description = f.read()
+
+setup(
+    name='ddlambda',
+    version='0.2.0',
+    description='The Datadog AWS Lambda Layer',
+    long_description=long_description,
+    long_description_content_type='text/markdown',
+    url='https://github.com/DataDog/datadog-lambda-layer-python',
+    author='Datadog, Inc.',
+    author_email='dev@datadoghq.com',
+    classifiers=[
+        'Programming Language :: Python :: 2.7',
+        'Programming Language :: Python :: 3.6',
+        'Programming Language :: Python :: 3.7',
+    ],
+    keywords='datadog aws lambda layer',
+    packages=['datadog_lambda'],
+    python_requires='>=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*, <4',
+)

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,3 +1,4 @@
+import os
 import unittest
 try:
     from unittest.mock import patch, call
@@ -18,12 +19,20 @@ 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"]

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"]