Repository: elastic/apm-agent-python

Issue:
DB interactions not traced when using context manager with psycopg2 connections or cursors
When using a context manager with psycopg2 connections or cursors, db interactions are not captured in spans.

The code below generates a span for `psycopg2.connect`, but not the query:
```
with psycopg2.connect(DSN) as conn:
    with conn.cursor() as curs:
        curs.execute("SELECT * FROM data.portfolio;")
        portfolios = curs.fetchall()
```

whereas the following captures both spans as expected:
```
conn = psycopg2.connect(DSN)
curs = conn.cursor()
curs.execute("SELECT * FROM data.portfolio;")
portfolios = curs.fetchall()
curs.close()
conn.close()
```


Discussion:

Resolving patch:
diff --git a/elasticapm/instrumentation/packages/psycopg2.py b/elasticapm/instrumentation/packages/psycopg2.py
--- a/elasticapm/instrumentation/packages/psycopg2.py
+++ b/elasticapm/instrumentation/packages/psycopg2.py
@@ -51,10 +51,16 @@ def _bake_sql(self, sql):
     def extract_signature(self, sql):
         return extract_signature(sql)
 
+    def __enter__(self):
+        return PGCursorProxy(self.__wrapped__.__enter__())
+
 
 class PGConnectionProxy(ConnectionProxy):
     cursor_proxy = PGCursorProxy
 
+    def __enter__(self):
+        return PGConnectionProxy(self.__wrapped__.__enter__())
+
 
 class Psycopg2Instrumentation(DbApi2Instrumentation):
     name = "psycopg2"


Test patch:
diff --git a/tests/instrumentation/psycopg2_tests.py b/tests/instrumentation/psycopg2_tests.py
--- a/tests/instrumentation/psycopg2_tests.py
+++ b/tests/instrumentation/psycopg2_tests.py
@@ -54,14 +54,18 @@
 has_postgres_configured = "POSTGRES_DB" in os.environ
 
 
+def connect_kwargs():
+    return {
+        "database": os.environ.get("POSTGRES_DB", "elasticapm_test"),
+        "user": os.environ.get("POSTGRES_USER", "postgres"),
+        "host": os.environ.get("POSTGRES_HOST", None),
+        "port": os.environ.get("POSTGRES_PORT", None),
+    }
+
+
 @pytest.yield_fixture(scope="function")
 def postgres_connection(request):
-    conn = psycopg2.connect(
-        database=os.environ.get("POSTGRES_DB", "elasticapm_test"),
-        user=os.environ.get("POSTGRES_USER", "postgres"),
-        host=os.environ.get("POSTGRES_HOST", None),
-        port=os.environ.get("POSTGRES_PORT", None),
-    )
+    conn = psycopg2.connect(**connect_kwargs())
     cursor = conn.cursor()
     cursor.execute(
         "CREATE TABLE test(id int, name VARCHAR(5) NOT NULL);"
@@ -376,3 +380,22 @@ def test_psycopg2_call_stored_procedure(instrument, postgres_connection, elastic
     span = elasticapm_client.spans_for_transaction(transactions[0])[0]
     assert span["name"] == "squareme()"
     assert span["action"] == "exec"
+
+
+@pytest.mark.integrationtest
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
+def test_psycopg_context_manager(instrument, elasticapm_client):
+    elasticapm_client.begin_transaction("test")
+    with psycopg2.connect(**connect_kwargs()) as conn:
+        with conn.cursor() as curs:
+            curs.execute("SELECT 1;")
+            curs.fetchall()
+    elasticapm_client.end_transaction("test", "OK")
+    transactions = elasticapm_client.events[TRANSACTION]
+    spans = elasticapm_client.spans_for_transaction(transactions[0])
+    assert len(spans) == 2
+    assert spans[0]["subtype"] == "postgresql"
+    assert spans[0]["action"] == "connect"
+
+    assert spans[1]["subtype"] == "postgresql"
+    assert spans[1]["action"] == "query"