fix(psycopg3): resolve attribute error raised when psycopg3 is patched [backport #5833 to 1.13] (#5866)

Backport of #5833 to 1.13

Resolves: #5821

- Fixes incorrect dbapi_async and psycopg async patches. 
- Adds tests to ensure new patching is correct. Changes test to use
mark_asyncio
- Change psycopg async connection patch to use a factory, and get the
pin object from the psycopg module
- Fix psycop unpatch to correctly rebind the original connect and async
connect class methods.
- Update dbapi_async to check during `TracedAsyncCursor.__init__` if the
cursor is a `FetchTracedAsyncCursor`, in which case we don't enable
analytics sample rate stats. Previously this check was for a
`FetchTracedCursor` type, which was incorrect.

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/contributing.html#Release-Note-Guidelines)
are followed.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] OPTIONAL: PR description includes explicit acknowledgement of the
performance implications of the change as reported in the benchmarks PR
comment.

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.

---------

Co-authored-by: Munir Abdinur <[email protected]>

Author: wconti27

.riot/requirements/119287b.txt

@@ -4,9 +4,9 @@
 #
 #    pip-compile --no-annotate --resolver=backtracking .riot/requirements/1ad8e2d.in
 #
-attrs==22.2.0
-coverage[toml]==7.2.2
-docutils==0.19
+attrs==23.1.0
+coverage[toml]==7.2.5
+docutils==0.20
 exceptiongroup==1.1.1
 flake8==3.8.4
 flake8-blind-except==0.2.1
@@ -19,15 +19,15 @@ hypothesis==6.45.0
 iniconfig==2.0.0
 isort==5.12.0
 mccabe==0.6.1
-mock==5.0.1
+mock==5.0.2
 opentracing==2.4.0
-packaging==23.0
+packaging==23.1
 pluggy==1.0.0
 pycodestyle==2.6.0
 pydocstyle==6.3.0
 pyflakes==2.2.0
-pygments==2.14.0
-pytest==7.2.2
+pygments==2.15.1
+pytest==7.3.1
 pytest-cov==4.0.0
 pytest-mock==3.10.0
 restructuredtext-lint==1.4.0

.riot/requirements/1ad8e2d.txt

@@ -4,19 +4,19 @@
 #
 #    pip-compile --no-annotate --resolver=backtracking .riot/requirements/bc7c1d4.in
 #
-attrs==22.2.0
-coverage[toml]==7.2.2
+attrs==23.1.0
+coverage[toml]==7.2.5
 envier==0.4.0
 exceptiongroup==1.1.1
 hypothesis==6.45.0
 iniconfig==2.0.0
-mock==5.0.1
+mock==5.0.2
 mypy==0.991
 mypy-extensions==1.0.0
 opentracing==2.4.0
-packaging==23.0
+packaging==23.1
 pluggy==1.0.0
-pytest==7.2.2
+pytest==7.3.1
 pytest-cov==4.0.0
 pytest-mock==3.10.0
 sortedcontainers==2.4.0

.riot/requirements/bc7c1d4.txt

@@ -1,3 +1,4 @@
+from ddtrace import config
 from ddtrace.appsec.iast._util import _is_iast_enabled
 from ddtrace.internal.constants import COMPONENT
 
@@ -10,7 +11,6 @@
 from ...internal.utils import ArgumentError
 from ...internal.utils import get_argument_value
 from ...pin import Pin
-from ..dbapi import FetchTracedCursor
 from ..dbapi import TracedConnection
 from ..dbapi import TracedCursor
 from ..trace_utils import ext_service
@@ -21,6 +21,20 @@
 
 
 class TracedAsyncCursor(TracedCursor):
+    async def __aenter__(self):
+        # previous versions of the dbapi didn't support context managers. let's
+        # reference the func that would be called to ensure that error
+        # messages will be the same.
+        await self.__wrapped__.__aenter__()
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        # previous versions of the dbapi didn't support context managers. let's
+        # reference the func that would be called to ensure that error
+        # messages will be the same.
+        return await self.__wrapped__.__aexit__()
+
     async def _trace_method(self, method, name, resource, extra_tags, dbm_propagator, *args, **kwargs):
         """
         Internal function to trace the call to the underlying cursor method
@@ -61,7 +75,7 @@ async def _trace_method(self, method, name, resource, extra_tags, dbm_propagator
                     SqlInjection.report(evidence_value=args[0])
 
             # set analytics sample rate if enabled but only for non-FetchTracedCursor
-            if not isinstance(self, FetchTracedCursor):
+            if not isinstance(self, FetchTracedAsyncCursor):
                 s.set_tag(ANALYTICS_SAMPLE_RATE_KEY, self._self_config.get_analytics_sample_rate())
 
             if dbm_propagator:
@@ -147,8 +161,11 @@ async def fetchmany(self, *args, **kwargs):
 
 
 class TracedAsyncConnection(TracedConnection):
-    traced_cursor_cls = TracedCursor
-    traced_fetch_cursor_cls = FetchTracedCursor
+    def __init__(self, conn, pin=None, cfg=config.dbapi2, cursor_cls=None):
+        if not cursor_cls:
+            # Do not trace `fetch*` methods by default
+            cursor_cls = FetchTracedAsyncCursor if cfg.trace_fetch_methods else TracedAsyncCursor
+        super(TracedAsyncConnection, self).__init__(conn, pin, cfg, cursor_cls)
 
     async def __aenter__(self):
         """Context management is not defined by the dbapi spec.
@@ -165,7 +182,7 @@ async def __aenter__(self):
         - pymysql doesn't implement it.
         - sqlite3 returns the connection.
         """
-        r = self.__wrapped__.__aenter__()
+        r = await self.__wrapped__.__aenter__()
 
         if hasattr(r, "cursor"):
             # r is Connection-like.
@@ -190,7 +207,7 @@ async def __aenter__(self):
                 pin = Pin.get_from(self)
                 if not pin:
                     return r
-                return await self._self_cursor_cls(r, pin, self._self_config)
+                return self._self_cursor_cls(r, pin, self._self_config)
         else:
             # Otherwise r is some other object, so maintain the functionality
             # of the original.
@@ -204,10 +221,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         # previous versions of the dbapi didn't support context managers. let's
         # reference the func that would be called to ensure that errors
         # messages will be the same.
-        self.__wrapped__.__aenter__
-
-        # and finally, yield the traced cursor.
-        return self
+        return await self.__wrapped__.__aexit__()
 
     async def _trace_method(self, method, name, extra_tags, *args, **kwargs):
         pin = Pin.get_from(self)

ddtrace/contrib/dbapi_async/__init__.py

@@ -29,41 +29,38 @@ async def execute(self, *args, **kwargs):
 
         async def patched_execute(*args, **kwargs):
             try:
-                cur = await self.cursor()
+                cur = self.cursor()
                 if kwargs.get("binary", None):
                     cur.format = 1  # set to 1 for binary or 0 if not
-                return cur.execute(*args, **kwargs)
+                return await cur.execute(*args, **kwargs)
             except Exception as ex:
                 raise ex.with_traceback(None)
 
         return await self._trace_method(patched_execute, span_name, {}, *args, **kwargs)
 
 
-async def patched_connect_async(connect_func, _, args, kwargs):
-    traced_conn_cls = Psycopg3TracedAsyncConnection
+def patched_connect_async_factory(psycopg_module):
+    async def patched_connect_async(connect_func, _, args, kwargs):
+        traced_conn_cls = Psycopg3TracedAsyncConnection
 
-    _config = globals()["config"]._config
-    module_name = (
-        connect_func.__module__
-        if len(connect_func.__module__.split(".")) == 1
-        else connect_func.__module__.split(".")[0]
-    )
-    pin = Pin.get_from(_config[module_name].base_module)
+        pin = Pin.get_from(psycopg_module)
 
-    if not pin or not pin.enabled() or not pin._config.trace_connect:
-        conn = await connect_func(*args, **kwargs)
-    else:
-        with pin.tracer.trace(
-            "{}.{}".format(connect_func.__module__, connect_func.__name__),
-            service=ext_service(pin, pin._config),
-            span_type=SpanTypes.SQL,
-        ) as span:
-            span.set_tag_str(SPAN_KIND, SpanKind.CLIENT)
-            span.set_tag_str(COMPONENT, pin._config.integration_name)
-            if span.get_tag(db.SYSTEM) is None:
-                span.set_tag_str(db.SYSTEM, pin._config.dbms_name)
-
-            span.set_tag(SPAN_MEASURED_KEY)
+        if not pin or not pin.enabled() or not pin._config.trace_connect:
             conn = await connect_func(*args, **kwargs)
+        else:
+            with pin.tracer.trace(
+                "{}.{}".format(connect_func.__module__, connect_func.__name__),
+                service=ext_service(pin, pin._config),
+                span_type=SpanTypes.SQL,
+            ) as span:
+                span.set_tag_str(SPAN_KIND, SpanKind.CLIENT)
+                span.set_tag_str(COMPONENT, pin._config.integration_name)
+                if span.get_tag(db.SYSTEM) is None:
+                    span.set_tag_str(db.SYSTEM, pin._config.dbms_name)
+
+                span.set_tag(SPAN_MEASURED_KEY)
+                conn = await connect_func(*args, **kwargs)
+
+        return patch_conn(conn, pin=pin, traced_conn_cls=traced_conn_cls)
 
-    return patch_conn(conn, pin=pin, traced_conn_cls=traced_conn_cls)
+    return patched_connect_async

ddtrace/contrib/psycopg/async_connection.py

@@ -8,7 +8,7 @@
 
 
 try:
-    from ddtrace.contrib.psycopg.async_connection import patched_connect_async
+    from ddtrace.contrib.psycopg.async_connection import patched_connect_async_factory
     from ddtrace.contrib.psycopg.async_cursor import Psycopg3FetchTracedAsyncCursor
     from ddtrace.contrib.psycopg.async_cursor import Psycopg3TracedAsyncCursor
 # catch async function syntax errors when using Python<3.7 with no async support
@@ -28,6 +28,18 @@
 from ...propagation._database_monitoring import _DBM_Propagator
 
 
+try:
+    psycopg_import = import_module("psycopg")
+
+    # must get the original connect class method from the class __dict__ to use later in unpatch
+    # Python 3.11 and wrapt result in the class method being rebinded as an instance method when
+    # using unwrap
+    _original_connect = psycopg_import.Connection.__dict__["connect"]
+    _original_async_connect = psycopg_import.AsyncConnection.__dict__["connect"]
+except ImportError:
+    pass
+
+
 def _psycopg_sql_injector(dbm_comment, sql_statement):
     for psycopg_module in config.psycopg["_patched_modules"]:
         if (
@@ -45,7 +57,6 @@ def _psycopg_sql_injector(dbm_comment, sql_statement):
         _default_service="postgres",
         _dbapi_span_name_prefix="postgres",
         _patched_modules=set(),
-        _patched_functions=dict(),
         trace_fetch_methods=asbool(
             os.getenv("DD_PSYCOPG_TRACE_FETCH_METHODS", default=False)
             or os.getenv("DD_PSYCOPG2_TRACE_FETCH_METHODS", default=False)
@@ -72,7 +83,6 @@ def _psycopg_modules():
             pass
 
 
-# NB: We are patching the default elasticsearch.transport module
 def patch():
     for psycopg_module in _psycopg_modules():
         _patch(psycopg_module)
@@ -89,7 +99,6 @@ def _patch(psycopg_module):
     Pin(_config=config.psycopg).onto(psycopg_module)
 
     if psycopg_module.__name__ == "psycopg2":
-        config.psycopg["_patched_functions"].update({"psycopg2.connect": psycopg_module.connect})
 
         # patch all psycopg2 extensions
         _psycopg2_extensions = get_psycopg2_extensions(psycopg_module)
@@ -100,23 +109,14 @@ def _patch(psycopg_module):
 
         config.psycopg["_patched_modules"].add(psycopg_module)
     else:
-        config.psycopg["_patched_functions"].update(
-            {
-                "psycopg.connect": psycopg_module.connect,
-                "psycopg.Connection": psycopg_module.Connection,
-                "psycopg.Cursor": psycopg_module.Cursor,
-                "psycopg.AsyncConnection": psycopg_module.AsyncConnection,
-                "psycopg.AsyncCursor": psycopg_module.AsyncCursor,
-            }
-        )
 
         _w(psycopg_module, "connect", patched_connect_factory(psycopg_module))
-        _w(psycopg_module.Connection, "connect", patched_connect_factory(psycopg_module))
         _w(psycopg_module, "Cursor", init_cursor_from_connection_factory(psycopg_module))
-
-        _w(psycopg_module.AsyncConnection, "connect", patched_connect_async)
         _w(psycopg_module, "AsyncCursor", init_cursor_from_connection_factory(psycopg_module))
 
+        _w(psycopg_module.Connection, "connect", patched_connect_factory(psycopg_module))
+        _w(psycopg_module.AsyncConnection, "connect", patched_connect_async_factory(psycopg_module))
+
         config.psycopg["_patched_modules"].add(psycopg_module)
 
 
@@ -139,18 +139,10 @@ def _unpatch(psycopg_module):
             _u(psycopg_module, "Cursor")
             _u(psycopg_module, "AsyncCursor")
 
-            try:
-                _u(psycopg_module.Connection, "connect")
-                _u(psycopg_module.AsyncConnection, "connect")
-
-            # _u throws an attribute error for Python 3.11 on method objects because of
-            # no __get__ method on the BoundFunctionWrapper
-            except AttributeError:
-                _original_connection_class = config.psycopg["_patched_functions"]["psycopg.Connection"]
-                _original_asyncconnection_class = config.psycopg["_patched_functions"]["psycopg.AsyncConnection"]
-
-                psycopg_module.Connection = _original_connection_class
-                psycopg_module.AsyncConnection = _original_asyncconnection_class
+            # _u throws an attribute error for Python 3.11, no __get__ on the BoundFunctionWrapper
+            # unlike Python Class Methods which implement __get__
+            setattr(psycopg_module.Connection, "connect", _original_connect)
+            setattr(psycopg_module.AsyncConnection, "connect", _original_async_connect)
 
         pin = Pin.get_from(psycopg_module)
         if pin:
@@ -162,7 +154,14 @@ def init_cursor_from_connection(wrapped_cursor_cls, _, args, kwargs):
         connection = kwargs.pop("connection", None)
         if not connection:
             args = list(args)
-            connection = args.pop(next((i for i, x in enumerate(args) if isinstance(x, dbapi.TracedConnection)), None))
+            index = next((i for i, x in enumerate(args) if isinstance(x, dbapi.TracedConnection)), None)
+            if index is not None:
+                connection = args.pop(index)
+
+            # if we do not have an example of a traced connection, call the original cursor function
+            if not connection:
+                return wrapped_cursor_cls(*args, **kwargs)
+
         pin = Pin.get_from(connection).clone()
         cfg = config.psycopg
 

ddtrace/contrib/psycopg/patch.py

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    psycopg: Resolves an issue where an AttributeError is raised when ``psycopg.AsyncConnection`` is traced.
+

releasenotes/notes/resolve-attribute-error-psycopg3-patch-c2b722b1c9474f09.yaml

@@ -2403,7 +2403,7 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
         Venv(
             name="dbapi_async",
             command="pytest {cmdargs} tests/contrib/dbapi_async",
-            pys=select_pys(min_version="3.5"),
+            pys=select_pys(min_version="3.6"),
             env={
                 "DD_IAST_REQUEST_SAMPLING": "100",  # Override default 30% to analyze all IAST requests
             },
@@ -2412,7 +2412,7 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
             },
             venvs=[
                 Venv(
-                    pys=["3.5", "3.6", "3.8", "3.9", "3.10"],
+                    pys=["3.6", "3.8", "3.9", "3.10"],
                 ),
                 Venv(pys=["3.11"], pkgs={"attrs": latest}),
             ],