fix instrumentation of connection when pool.acquire was called multip…

…le times (#381)

CHANGELOG.md

@@ -23,6 +23,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#350](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/350))
 - `opentelemetry-exporter-datadog` Fix warning when DatadogFormat encounters a request with
   no DD_ORIGIN headers ([#368](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/368)).
+- `opentelemetry-instrumentation-aiopg` Fix multiple nested spans when
+  `aiopg.pool` is used
+  ([#336](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/381)).
 - Updated instrumentations to use `opentelemetry.trace.use_span` instead of `Tracer.use_span()`
   ([#364](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/364))
 - `opentelemetry-propagator-ot-trace` Do not throw an exception when headers are not present

instrumentation/opentelemetry-instrumentation-aiopg/src/opentelemetry/instrumentation/aiopg/aiopg_integration.py

@@ -86,9 +86,11 @@ def acquire(self):
         async def _acquire(self):
             # pylint: disable=protected-access
             connection = await self.__wrapped__._acquire()
-            return get_traced_connection_proxy(
-                connection, db_api_integration, *args, **kwargs
-            )
+            if not isinstance(connection, AsyncProxyObject):
+                connection = get_traced_connection_proxy(
+                    connection, db_api_integration, *args, **kwargs
+                )
+            return connection
 
     return TracedPoolProxy(pool, *args, **kwargs)
 

tests/opentelemetry-docker-tests/tests/postgres/test_aiopg_functional.py

@@ -117,6 +117,10 @@ def setUpClass(cls):
         cls._cursor = None
         cls._tracer = cls.tracer_provider.get_tracer(__name__)
         AiopgInstrumentor().instrument(tracer_provider=cls.tracer_provider)
+        cls._dsn = (
+            f"dbname='{POSTGRES_DB_NAME}' user='{POSTGRES_USER}' password='{POSTGRES_PASSWORD}'"
+            f" host='{POSTGRES_HOST}' port='{POSTGRES_PORT}'"
+        )
         cls._pool = async_call(
             aiopg.create_pool(
                 dbname=POSTGRES_DB_NAME,
@@ -185,3 +189,19 @@ def test_callproc(self):
         ):
             async_call(self._cursor.callproc("test", ()))
             self.validate_spans("test")
+
+    def test_instrumented_pool_with_multiple_acquires(self, *_, **__):
+        async def double_acquire():
+            pool = await aiopg.create_pool(dsn=self._dsn)
+            async with pool.acquire() as conn:
+                async with conn.cursor() as cursor:
+                    query = "SELECT 1"
+                    await cursor.execute(query)
+            async with pool.acquire() as conn:
+                async with conn.cursor() as cursor:
+                    query = "SELECT 1"
+                    await cursor.execute(query)
+
+        async_call(double_acquire())
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 2)