fix(asyncpg): fix _TracedConnection error when using the connect option in create_pool [backport 3.12] (#14450)

Backports #14436 to 3.12
## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

Author: wantsui

ddtrace/contrib/internal/asyncpg/patch.py

@@ -94,6 +94,9 @@ async def _traced_connect(asyncpg, pin, func, instance, args, kwargs):
 
     connect() is instrumented and patched to return a connection proxy.
     """
+    # When using a pool, there's a connection_class args
+    is_pool_context = "connection_class" in kwargs
+
     with pin.tracer.trace(
         "postgres.connect", span_type=SpanTypes.SQL, service=ext_service(pin, config.asyncpg)
     ) as span:
@@ -103,10 +106,23 @@ async def _traced_connect(asyncpg, pin, func, instance, args, kwargs):
         # set span.kind to the type of request being performed
         span.set_tag_str(SPAN_KIND, SpanKind.CLIENT)
 
-        # Need an ObjectProxy since Connection uses slots
-        conn = _TracedConnection(await func(*args, **kwargs), pin)
-        span.set_tags(_get_connection_tags(conn))
-        return conn
+        raw_conn = await func(*args, **kwargs)
+        if is_pool_context:
+            # Return the unwrapped connection to avoid _TracedConnection errors
+            # when using a pool with a custom connect param
+            connection_tags = _get_connection_tags(raw_conn)
+            connection_tags[db.SYSTEM] = DBMS_NAME
+            conn_pin = pin.clone(tags=connection_tags)
+            conn_pin.onto(raw_conn._protocol)
+            span.set_tags(connection_tags)
+            # Returns a asyncpg.connection.Connection object
+            return raw_conn
+        else:
+            # # Need an ObjectProxy when not using pools since Connection uses slots
+            conn = _TracedConnection(raw_conn, pin)
+            span.set_tags(_get_connection_tags(conn))
+            # Returns a _TracedConnection object
+            return conn
 
 
 async def _traced_query(pin, method, query, args, kwargs):

releasenotes/notes/fix-asyncpg-pool-error-97a58b63370428bc.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    asyncpg: Fix the error "Error: expected pool connect callback to return an instance of 'asyncpg.connection.Connection', got 'ddtrace.contrib.internal.asyncpg.patch._TracedConnection'`" when a pool connection" due to using the custom connect option. With this fix, postgres.connect spans will be created when this option is used.

tests/contrib/asyncpg/test_asyncpg.py

@@ -8,6 +8,7 @@
 from ddtrace.contrib.internal.asyncpg.patch import patch
 from ddtrace.contrib.internal.asyncpg.patch import unpatch
 from ddtrace.contrib.internal.trace_utils import iswrapped
+from ddtrace.internal.utils.version import parse_version
 from ddtrace.trace import Pin
 from ddtrace.trace import tracer
 from tests.contrib.asyncio.utils import AsyncioTestCase
@@ -315,6 +316,58 @@ async def test():
     assert err == b""
 
 
+@pytest.mark.skipif(
+    parse_version(getattr(asyncpg, "__version__", "0.0.0")) < (0, 30, 0),
+    reason="the custom connect parameter for create_pool requires asyncpg >= 0.30.0",
+)
+@pytest.mark.snapshot()
+@pytest.mark.asyncio
+async def test_pool_custom_connect():
+    """
+    Test that if someone uses a custom connect parameter when creating a pool,
+    the tracer doesn't cause a _TracedConnection error or throw any other errors
+    The connect option was introduced in 0.30.0
+    """
+    database_url = f"postgresql://{POSTGRES_CONFIG['user']}:{POSTGRES_CONFIG['password']}@{POSTGRES_CONFIG['host']}:{POSTGRES_CONFIG['port']}/{POSTGRES_CONFIG['dbname']}"
+
+    try:
+        # The default is 10 connection pools so the integration will create 10 spans in the snapshot
+        pool = await asyncpg.create_pool(database_url, connect=asyncpg.connect)
+
+        async with pool.acquire() as conn:
+            result = await conn.fetchval("SELECT 1")
+            assert result == 1
+
+        await pool.close()
+    except Exception as err:
+        raise err
+
+    assert True
+
+
+@pytest.mark.snapshot()
+@pytest.mark.asyncio
+async def test_pool_without_custom_connect():
+    """
+    Test that create_pool without the connect option still works
+    """
+    database_url = f"postgresql://{POSTGRES_CONFIG['user']}:{POSTGRES_CONFIG['password']}@{POSTGRES_CONFIG['host']}:{POSTGRES_CONFIG['port']}/{POSTGRES_CONFIG['dbname']}"
+
+    try:
+        # The default is 10 connection pools so the integration will create 10 spans in the snapshot
+        pool = await asyncpg.create_pool(database_url)
+
+        async with pool.acquire() as conn:
+            result = await conn.fetchval("SELECT 1")
+            assert result == 1
+
+        await pool.close()
+    except Exception as err:
+        raise err
+
+    assert True
+
+
 def test_patch_unpatch_asyncpg():
     assert iswrapped(asyncpg.connect)
     assert iswrapped(asyncpg.protocol.Protocol.execute)