Repository: DataDog/dd-trace-py

Issue:
Support trace for pynamodb
Hello,

I’m using `Pynamodb` library in order to use AWS Dynamodb. `Pynamodb` uses `botocore` internally.
There is `botocore` in supported library list so I was sure that this creates spans but it seems not.

Is there any solution to create spans with `Pynamodb`?



Discussion:

Resolving patch:
diff --git a/ddtrace/contrib/pynamodb/__init__.py b/ddtrace/contrib/pynamodb/__init__.py
new file mode 100644
--- /dev/null
+++ b/ddtrace/contrib/pynamodb/__init__.py
@@ -0,0 +1,37 @@
+"""
+The PynamoDB integration traces all db calls made with the pynamodb
+library through the connection API.
+
+Enabling
+~~~~~~~~
+
+The PynamoDB integration is enabled automatically when using
+:ref:`ddtrace-run<ddtracerun>` or :ref:`patch_all()<patch_all>`.
+
+Or use :ref:`patch()<patch>` to manually enable the integration::
+
+    import pynamodb
+    from ddtrace import patch, config
+    patch(pynamodb=True)
+
+Global Configuration
+~~~~~~~~~~~~~~~~~~~~
+
+.. py:data:: ddtrace.config.pynamodb["service"]
+   The service name reported by default for the PynamoDB instance.
+   This option can also be set with the ``DD_PYNAMODB_SERVICE`` environment
+   variable.
+   Default: ``"pynamodb"``
+
+"""
+
+
+from ...utils.importlib import require_modules
+
+required_modules = ["pynamodb.connection.base"]
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        from .patch import patch
+
+        __all__ = ["patch"]
diff --git a/ddtrace/contrib/pynamodb/patch.py b/ddtrace/contrib/pynamodb/patch.py
new file mode 100644
--- /dev/null
+++ b/ddtrace/contrib/pynamodb/patch.py
@@ -0,0 +1,80 @@
+"""
+Trace queries to botocore api done via a pynamodb client
+"""
+
+from ddtrace.vendor import wrapt
+from ddtrace import config
+import pynamodb.connection.base
+
+
+from ...constants import ANALYTICS_SAMPLE_RATE_KEY, SPAN_MEASURED_KEY
+from ...pin import Pin
+from ...ext import SpanTypes
+from ...utils.formats import deep_getattr
+from ...utils.wrappers import unwrap
+from .. import trace_utils
+
+# Pynamodb connection class
+_PynamoDB_client = pynamodb.connection.base.Connection
+
+config._add("pynamodb", {"_default_service": "pynamodb",})
+
+
+def patch():
+    if getattr(pynamodb.connection.base, "_datadog_patch", False):
+        return
+    setattr(pynamodb.connection.base, "_datadog_patch", True)
+
+    wrapt.wrap_function_wrapper("pynamodb.connection.base", "Connection._make_api_call", patched_api_call)
+    Pin(service=None).onto(pynamodb.connection.base.Connection)
+
+
+def unpatch():
+    if getattr(pynamodb.connection.base, "_datadog_patch", False):
+        setattr(pynamodb.connection.base, "_datadog_patch", False)
+        unwrap(pynamodb.connection.base.Connection, "_make_api_call")
+
+
+def patched_api_call(original_func, instance, args, kwargs):
+
+    pin = Pin.get_from(instance)
+    if not pin or not pin.enabled():
+        return original_func(*args, **kwargs)
+
+    with pin.tracer.trace(
+        "pynamodb.command", service=trace_utils.ext_service(pin, config.pynamodb, "pynamodb"), span_type=SpanTypes.HTTP
+    ) as span:
+
+        span.set_tag(SPAN_MEASURED_KEY)
+
+        if args and args[0]:
+            operation = args[0]
+            span.resource = operation
+
+            if args[1] and "TableName" in args[1]:
+                table_name = args[1]["TableName"]
+                span.set_tag("table_name", table_name)
+                span.resource = span.resource + " " + table_name
+
+        else:
+            span.resource = "Unknown"
+            operation = None
+
+        region_name = deep_getattr(instance, "client.meta.region_name")
+
+        meta = {
+            "aws.agent": "pynamodb",
+            "aws.operation": operation,
+            "aws.region": region_name,
+        }
+        span.set_tags(meta)
+
+        # set analytics sample rate
+        sample_rate = config.pynamodb.get_analytics_sample_rate(use_global_config=True)
+
+        if sample_rate is not None:
+            span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, sample_rate)
+
+        result = original_func(*args, **kwargs)
+
+        return result
diff --git a/ddtrace/monkey.py b/ddtrace/monkey.py
--- a/ddtrace/monkey.py
+++ b/ddtrace/monkey.py
@@ -64,6 +64,7 @@
     "pyramid": False,
     # Auto-enable logging if the environment variable DD_LOGS_INJECTION is true
     "logging": config.logs_injection,
+    "pynamodb": True,
 }
 
 _LOCK = threading.Lock()


Test patch:
diff --git a/tests/contrib/pynamodb/__init__.py b/tests/contrib/pynamodb/__init__.py
new file mode 100644
diff --git a/tests/contrib/pynamodb/test_pynamodb.py b/tests/contrib/pynamodb/test_pynamodb.py
new file mode 100644
--- /dev/null
+++ b/tests/contrib/pynamodb/test_pynamodb.py
@@ -0,0 +1,195 @@
+import pynamodb.connection.base
+from pynamodb.connection.base import Connection
+from moto import mock_dynamodb
+
+# project
+from ddtrace import Pin
+from ddtrace.contrib.pynamodb.patch import patch, unpatch
+
+# testing
+import pytest
+from ... import TracerTestCase, assert_is_measured
+from moto.dynamodb import dynamodb_backend
+
+
+class PynamodbTest(TracerTestCase):
+
+    TEST_SERVICE = "pynamodb"
+
+    def setUp(self):
+        patch()
+
+        self.conn = Connection(region="us-east-1")
+        self.conn.session.set_credentials("aws-access-key", "aws-secret-access-key", "session-token")
+
+        super(PynamodbTest, self).setUp()
+        Pin.override(self.conn, tracer=self.tracer)
+
+    def tearDown(self):
+        super(PynamodbTest, self).tearDown()
+        unpatch()
+
+    @mock_dynamodb
+    def test_list_tables(self):
+        dynamodb_backend.create_table("Test", hash_key_attr="content", hash_key_type="S")
+        list_result = self.conn.list_tables()
+
+        spans = self.get_spans()
+
+        assert spans
+        span = spans[0]
+
+        assert span.name == "pynamodb.command"
+        assert span.service == "pynamodb"
+        assert span.resource == "ListTables"
+        assert len(spans) == 1
+        assert_is_measured(span)
+        assert span.span_type == "http"
+        assert span.get_tag("aws.operation") == "ListTables"
+        assert span.get_tag("aws.region") == "us-east-1"
+        assert span.get_tag("aws.agent") == "pynamodb"
+        assert span.duration >= 0
+        assert span.error == 0
+
+        assert len(list_result["TableNames"]) == 1
+        assert list_result["TableNames"][0] == "Test"
+
+    @mock_dynamodb
+    def test_delete_table(self):
+        dynamodb_backend.create_table("Test", hash_key_attr="content", hash_key_type="S")
+        delete_result = self.conn.delete_table("Test")
+        spans = self.get_spans()
+
+        assert spans
+        span = spans[0]
+
+        assert span.name == "pynamodb.command"
+        assert span.service == "pynamodb"
+        assert span.resource == "DeleteTable Test"
+        assert len(spans) == 1
+        assert_is_measured(span)
+        assert span.span_type == "http"
+        assert span.get_tag("aws.operation") == "DeleteTable"
+        assert span.get_tag("aws.region") == "us-east-1"
+        assert span.get_tag("aws.agent") == "pynamodb"
+        assert span.duration >= 0
+        assert span.error == 0
+
+        assert delete_result["Table"]["TableName"] == "Test"
+        assert len(self.conn.list_tables()["TableNames"]) == 0
+
+    @mock_dynamodb
+    def test_scan(self):
+        dynamodb_backend.create_table("Test", hash_key_attr="content", hash_key_type="S")
+        scan_result = self.conn.scan("Test")
+        spans = self.get_spans()
+
+        assert spans
+        span = spans[0]
+
+        assert span.name == "pynamodb.command"
+        assert span.service == "pynamodb"
+        assert span.resource == "Scan Test"
+        assert len(spans) == 1
+        assert_is_measured(span)
+        assert span.span_type == "http"
+        assert span.get_tag("aws.operation") == "Scan"
+        assert span.get_tag("aws.region") == "us-east-1"
+        assert span.get_tag("aws.agent") == "pynamodb"
+        assert span.duration >= 0
+        assert span.error == 0
+
+        assert scan_result["ScannedCount"] == 0
+        assert len(scan_result["Items"]) == 0
+
+    @mock_dynamodb
+    def test_scan_on_error(self):
+        with pytest.raises(pynamodb.exceptions.ScanError):
+            self.conn.scan("OtherTable")
+
+        spans = self.get_spans()
+        assert spans
+        span = spans[0]
+        assert span.name == "pynamodb.command"
+        assert span.service == "pynamodb"
+        assert span.resource == "Scan OtherTable"
+        assert len(spans) == 1
+        assert_is_measured(span)
+        assert span.span_type == "http"
+        assert span.get_tag("aws.operation") == "Scan"
+        assert span.get_tag("aws.region") == "us-east-1"
+        assert span.get_tag("aws.agent") == "pynamodb"
+        assert span.duration >= 0
+        assert span.error == 1
+        assert span.meta["error.type"] != ""
+
+    @TracerTestCase.run_in_subprocess(env_overrides=dict(DD_SERVICE="mysvc"))
+    @mock_dynamodb
+    def test_user_specified_service(self):
+        from ddtrace import config
+
+        assert config.service == "mysvc"
+
+        dynamodb_backend.create_table("Test", hash_key_attr="content", hash_key_type="S")
+        list_result = self.conn.list_tables()
+
+        span = self.get_spans()[0]
+        assert span.service == "pynamodb"
+        assert len(list_result["TableNames"]) == 1
+        assert list_result["TableNames"][0] == "Test"
+
+    @TracerTestCase.run_in_subprocess(env_overrides=dict(DD_PYNAMODB_SERVICE="mypynamodb"))
+    @mock_dynamodb
+    def test_env_user_specified_pynamodb_service(self):
+        dynamodb_backend.create_table("Test", hash_key_attr="content", hash_key_type="S")
+        list_result = self.conn.list_tables()
+
+        span = self.get_spans()[0]
+
+        assert span.service == "mypynamodb", span.service
+        assert len(list_result["TableNames"]) == 1
+        assert list_result["TableNames"][0] == "Test"
+
+        self.reset()
+
+        # Global config
+        with self.override_config("pynamodb", dict(service="cfg-pynamodb")):
+            dynamodb_backend.create_table("Test", hash_key_attr="content", hash_key_type="S")
+            list_result = self.conn.list_tables()
+            span = self.get_spans()[0]
+
+            assert span.service == "cfg-pynamodb", span.service
+            assert len(list_result["TableNames"]) == 1
+            assert list_result["TableNames"][0] == "Test"
+
+        self.reset()
+
+        # Manual override
+        dynamodb_backend.create_table("Test", hash_key_attr="content", hash_key_type="S")
+        Pin.override(self.conn, service="mypynamodb", tracer=self.tracer)
+        list_result = self.conn.list_tables()
+        span = self.get_spans()[0]
+        assert span.service == "mypynamodb", span.service
+        assert len(list_result["TableNames"]) == 1
+        assert list_result["TableNames"][0] == "Test"
+
+    @TracerTestCase.run_in_subprocess(env_overrides=dict(DD_SERVICE="app-svc", DD_PYNAMODB_SERVICE="env-pynamodb"))
+    @mock_dynamodb
+    def test_service_precedence(self):
+        dynamodb_backend.create_table("Test", hash_key_attr="content", hash_key_type="S")
+        list_result = self.conn.list_tables()
+        span = self.get_spans()[0]
+        assert span.service == "env-pynamodb", span.service
+        assert len(list_result["TableNames"]) == 1
+        assert list_result["TableNames"][0] == "Test"
+
+        self.reset()
+
+        # Manual overide
+        dynamodb_backend.create_table("Test", hash_key_attr="content", hash_key_type="S")
+        Pin.override(self.conn, service="override-pynamodb", tracer=self.tracer)
+        list_result = self.conn.list_tables()
+        span = self.get_spans()[0]
+        assert span.service == "override-pynamodb", span.service
+        assert len(list_result["TableNames"]) == 1
+        assert list_result["TableNames"][0] == "Test"