feat: add asyncpg integration (backport #3410) (#3428) (#3434)

* feat: add asyncpg integration (#3410)

# Trace Integration
This PR adds support for [`asyncpg`](https://magicstack.github.io/asyncpg/) starting with `v0.18.0`.

[Connections](https://magicstack.github.io/asyncpg/current/api/index.html#connection), [cursors](https://magicstack.github.io/asyncpg/current/api/index.html#cursors) are traced.

Note that this integration cannot use the `dbapi` patching that we commonly do due to asyncpg not following the api (dbapi in general is not compatible with async usage).

Work for this was started in #394, thanks again to @thehesiod!!

As a follow up, support can be added for [connection pools](https://magicstack.github.io/asyncpg/current/api/index.html#connection-pools).

## Links

- Integration docs: https://magicstack.github.io/asyncpg/
- Corp docs PR: TODO
- https://github.com/MagicStack/asyncpg
- https://cs.github.com/MagicStack/asyncpg/
- https://pypi.org/project/asyncpg
- https://github.com/encode/databases
- [Connection source code](https://cs.github.com/MagicStack/asyncpg/blob/a2f093df6aceec7842709eaf92c5ff9df093efae/asyncpg/connection.py)
- [Cursor source code](https://cs.github.com/MagicStack/asyncpg/blob/a2f093df6aceec7842709eaf92c5ff9df093efae/asyncpg/cursor.py )

## Checklist

- [x] Usage and configuration documentation added in `__init__.py`, `docs/index.rst` and `docs/integrations.rst`.
- [ ] [Corp docs](https://github.com/Datadog/documentation) PR to add new integration to documentation.
- [x] Span metadata
  - [x] Service (use [`int_service`](https://github.com/DataDog/dd-trace-py/blob/90d1d5981c72ea312c21ac04e5be47521d0f0f2e/ddtrace/contrib/trace_utils.py#L55) or [`ext_service`](https://github.com/DataDog/dd-trace-py/blob/90d1d5981c72ea312c21ac04e5be47521d0f0f2e/ddtrace/contrib/trace_utils.py#L87)).
  - [x] Span type should be one of [these](https://github.com/DataDog/dd-trace-py/blob/90d1d5981c72ea312c21ac04e5be47521d0f0f2e/ddtrace/ext/__init__.py#L7).
  - [x] Resource
  - [x] Measured tag
  - [x] ~Sample analytics~
- [x] Global configuration
  - [x] `ddtrace.config` entry is specified.
  - [x] Environment variables are provided for config options.
- [x] Instance configuration
  - [x] Pin overriding.
  - [x] Service name override (if applicable).
- [x] Async
  - [x] Span parenting behaves as expected.
  - [x] Context propagation across async contexts.
- [x] Tests
  - [x] Use `pytest` fixtures found in `tests/conftest.py` or the test helpers on `TracerTestCase` if
    writing `unittest` style test cases.
  - [x] Tests are provided for all the above.
  - [x] Tests are added to CI (`.circleci/config.yml`).
  - [x] Functionality is maintained from original library by adding the library's test suite with the tracer enabled to our CI.
  - [x] Patch test cases are added (see `test_django_patch.py` for an example).
  - [x] All Python versions that the library supports are tested.
  - [x] All significant library versions (including the latest) are tested. This typically includes every minor release going back a few years.

(cherry picked from commit a0515ec)

* Update ddtrace/contrib/asyncpg/__init__.py

Co-authored-by: Kyle Verhoog <[email protected]>
Co-authored-by: Brett Langdon <[email protected]>
(cherry picked from commit 4a18ed0)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: mergify[bot]

.circleci/config.yml

@@ -461,6 +461,14 @@ jobs:
       - run_tox_scenario:
           pattern: '^asyncio_contrib-'
 
+  asyncpg:
+    <<: *machine_executor
+    steps:
+      - run_test:
+          pattern: 'asyncpg'
+          snapshot: true
+          docker_services: 'postgres'
+
   pylons:
     <<: *contrib_job
     steps:
@@ -975,6 +983,7 @@ requires_tests: &requires_tests
     - aiopg
     - aioredis
     - asyncio
+    - asyncpg
     - algoliasearch
     - asgi
     - benchmarks
@@ -1062,6 +1071,7 @@ workflows:
       - aiopg: *requires_base_venvs
       - aioredis: *requires_base_venvs
       - asyncio: *requires_base_venvs
+      - asyncpg: *requires_base_venvs
       - algoliasearch: *requires_base_venvs
       - asgi: *requires_base_venvs
       - benchmarks: *requires_base_venvs

.github/workflows/test_frameworks.yml

@@ -278,3 +278,35 @@ jobs:
         run: "make init"
       - name: Run tests
         run: ddtrace-run pytest -p no:warnings tests
+
+  asyncpg-testsuite-0_25_0:
+    # https://github.com/MagicStack/asyncpg/blob/v0.25.0/.github/workflows/tests.yml#L125
+    runs-on: "ubuntu-latest"
+    env:
+      DD_TESTING_RAISE: true
+      DD_PROFILING_ENABLED: true
+    defaults:
+      run:
+        working-directory: asyncpg
+    steps:
+      - uses: actions/setup-python@v3
+        with:
+          python-version: '3.9'
+      - uses: actions/checkout@v3
+        with:
+          path: ddtrace
+      - uses: actions/checkout@v3
+        with:
+          repository: magicstack/asyncpg
+          ref: v0.25.0
+          path: asyncpg
+          fetch-depth: 50
+          submodules: true
+      - name: Install ddtrace
+        run: pip install ../ddtrace
+      - name: Install dependencies
+        run: |
+          python -m pip install -U pip setuptools wheel
+          python -m pip install -e .[test]
+      - name: Run tests
+        run: ddtrace-run python setup.py test

ddtrace/_monkey.py

@@ -77,6 +77,7 @@
     "dogpile_cache": True,
     "yaaredis": True,
     "aiohttp_jinja2": False,  # disabled as this is handled by aiohttp for now.
+    "asyncpg": True,
 }
 
 _LOCK = threading.Lock()

ddtrace/contrib/asyncpg/__init__.py

@@ -0,0 +1,59 @@
+"""
+The ``asyncpg`` integration traces database requests made using connection
+and cursor objects.
+
+
+Enabling
+~~~~~~~~
+
+The 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::
+
+    from ddtrace import patch
+    patch(asyncpg=True)
+
+
+Global Configuration
+~~~~~~~~~~~~~~~~~~~~
+
+.. py:data:: ddtrace.config.asyncpg['service']
+
+   The service name reported by default for asyncpg connections.
+
+   This option can also be set with the ``DD_ASYNCPG_SERVICE``
+   environment variable.
+
+   Default: ``postgres``
+
+
+Instance Configuration
+~~~~~~~~~~~~~~~~~~~~~~
+
+Service
+^^^^^^^
+
+To configure the service name used by the asyncpg integration on a per-instance
+basis use the ``Pin`` API::
+
+    import asyncpg
+    from ddtrace import Pin
+
+    conn = asyncpg.connect("postgres://localhost:5432")
+    Pin.override(conn, service="custom-service")
+"""
+from ...internal.utils.importlib import require_modules
+
+
+required_modules = ["asyncpg"]
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        from .patch import patch
+        from .patch import unpatch
+
+        __all__ = [
+            "patch",
+            "unpatch",
+        ]

ddtrace/contrib/asyncpg/patch.py

@@ -0,0 +1,135 @@
+from typing import TYPE_CHECKING
+
+from ddtrace import Pin
+from ddtrace import config
+from ddtrace.vendor import wrapt
+
+from ...constants import SPAN_MEASURED_KEY
+from ...ext import SpanTypes
+from ...ext import db
+from ...ext import net
+from ...internal.logger import get_logger
+from ...internal.utils import get_argument_value
+from ..trace_utils import ext_service
+from ..trace_utils import unwrap
+from ..trace_utils import wrap
+from ..trace_utils_async import with_traced_module
+
+
+if TYPE_CHECKING:
+    from types import ModuleType
+    from typing import Dict
+    from typing import Union
+
+    import asyncpg
+    from asyncpg.prepared_stmt import PreparedStatement
+
+
+config._add(
+    "asyncpg",
+    dict(
+        _default_service="postgres",
+    ),
+)
+
+
+log = get_logger(__name__)
+
+
+def _get_connection_tags(conn):
+    # type: (asyncpg.Connection) -> Dict[str, str]
+    addr = conn._addr
+    params = conn._params
+    host = port = ""
+    if isinstance(addr, tuple) and len(addr) == 2:
+        host, port = addr
+    return {
+        net.TARGET_HOST: host,
+        net.TARGET_PORT: port,
+        db.USER: params.user,
+        db.NAME: params.database,
+    }
+
+
+class _TracedConnection(wrapt.ObjectProxy):
+    def __init__(self, conn, pin):
+        super(_TracedConnection, self).__init__(conn)
+        conn_pin = pin.clone(tags=_get_connection_tags(conn))
+        # Keep the pin on the protocol
+        conn_pin.onto(self._protocol)
+
+    def __setddpin__(self, pin):
+        pin.onto(self._protocol)
+
+    def __getddpin__(self):
+        return Pin.get_from(self._protocol)
+
+
+@with_traced_module
+async def _traced_connect(asyncpg, pin, func, instance, args, kwargs):
+    """Traced asyncpg.connect().
+
+    connect() is instrumented and patched to return a connection proxy.
+    """
+    with pin.tracer.trace(
+        "postgres.connect", span_type=SpanTypes.SQL, service=ext_service(pin, config.asyncpg)
+    ) as span:
+        # Need an ObjectProxy since Connection uses slots
+        conn = _TracedConnection(await func(*args, **kwargs), pin)
+        span.set_tags(_get_connection_tags(conn))
+        return conn
+
+
+async def _traced_query(pin, method, query, args, kwargs):
+    with pin.tracer.trace(
+        "postgres.query", resource=query, service=ext_service(pin, config.asyncpg), span_type=SpanTypes.SQL
+    ) as span:
+        span.set_tag(SPAN_MEASURED_KEY)
+        span.set_tags(pin.tags)
+        return await method(*args, **kwargs)
+
+
+@with_traced_module
+async def _traced_protocol_execute(asyncpg, pin, func, instance, args, kwargs):
+    state = get_argument_value(args, kwargs, 0, "state")  # type: Union[str, PreparedStatement]
+    query = state if isinstance(state, str) else state.query
+    return await _traced_query(pin, func, query, args, kwargs)
+
+
+def _patch(asyncpg):
+    # type: (ModuleType) -> None
+    wrap(asyncpg, "connect", _traced_connect(asyncpg))
+    for method in ("execute", "bind_execute", "query", "bind_execute_many"):
+        wrap(asyncpg.protocol, "Protocol.%s" % method, _traced_protocol_execute(asyncpg))
+
+
+def patch():
+    # type: () -> None
+    import asyncpg
+
+    if getattr(asyncpg, "_datadog_patch", False):
+        return
+
+    Pin().onto(asyncpg)
+    _patch(asyncpg)
+
+    setattr(asyncpg, "_datadog_patch", True)
+
+
+def _unpatch(asyncpg):
+    # type: (ModuleType) -> None
+    unwrap(asyncpg, "connect")
+    for method in ("execute", "bind_execute", "query", "bind_execute_many"):
+        unwrap(asyncpg.protocol.Protocol, method)
+
+
+def unpatch():
+    # type: () -> None
+    import asyncpg
+
+    if not getattr(asyncpg, "_datadog_patch", False):
+        return
+
+    _unpatch(asyncpg)
+
+    setattr(asyncpg, "_datadog_patch", False)

docs/index.rst

@@ -59,6 +59,8 @@ contacting support.
 +--------------------------------------------------+---------------+----------------+
 | :ref:`asyncio`                                   | \*            | Yes [4]_       |
 +--------------------------------------------------+---------------+----------------+
+| :ref:`asyncpg`                                   | >= 0.18.0     | Yes            |
++--------------------------------------------------+---------------+----------------+
 | :ref:`boto2`                                     | >= 2.29.0     | Yes            |
 +--------------------------------------------------+---------------+----------------+
 | :ref:`botocore`                                  | >= 1.4.51     | Yes            |

docs/integrations.rst

@@ -64,6 +64,13 @@ asyncio
 .. automodule:: ddtrace.contrib.asyncio
 
 
+.. _asyncpg:
+
+asyncpg
+^^^^^^^
+.. automodule:: ddtrace.contrib.asyncpg
+
+
 .. _botocore:
 
 botocore

docs/spelling_wordlist.txt

@@ -20,6 +20,7 @@ aredis
 args
 asgi
 asyncio
+asyncpg
 attrs
 autodetected
 autopatching

releasenotes/notes/asyncpg-45cdf83efdf9270d.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    asyncpg: add integration supporting v0.18.0 and above. See :ref:`the docs<asyncpg>` for more information.

riotfile.py

@@ -1616,5 +1616,46 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
                 ],
             },
         ),
+        Venv(
+            name="asyncpg",
+            command="pytest {cmdargs} tests/contrib/asyncpg",
+            pkgs={
+                "pytest-asyncio": latest,
+            },
+            venvs=[
+                Venv(
+                    pys=select_pys(min_version="3.6", max_version="3.8"),
+                    pkgs={
+                        "asyncpg": [
+                            "~=0.18.0",
+                            "~=0.20.0",
+                            "~=0.22.0",
+                            "~=0.24.0",
+                            latest,
+                        ],
+                    },
+                ),
+                Venv(
+                    pys=["3.9"],
+                    pkgs={
+                        "asyncpg": [
+                            "~=0.20.0",
+                            "~=0.22.0",
+                            "~=0.24.0",
+                            latest,
+                        ],
+                    },
+                ),
+                Venv(
+                    pys=select_pys(min_version="3.10"),
+                    pkgs={
+                        "asyncpg": [
+                            "~=0.24.0",
+                            latest,
+                        ],
+                    },
+                ),
+            ],
+        ),
     ],
 )