fix: instrument psycopg_pool connections to restore Postgres spans

Signed-off-by: SaJH <[email protected]>

Author: JaeHyuckSa

elasticapm/instrumentation/packages/psycopg.py

@@ -79,7 +79,7 @@ def __enter__(self):
 class PsycopgInstrumentation(DbApi2Instrumentation):
     name = "psycopg"
 
-    instrument_list = [("psycopg", "connect")]
+    instrument_list = [("psycopg", "connect"), ("psycopg", "Connection.connect")]
 
     def call(self, module, method, wrapped, instance, args, kwargs):
         signature = "psycopg.connect"

tests/instrumentation/psycopg_pool_tests.py

@@ -0,0 +1,63 @@
+#  BSD 3-Clause License
+#
+#  Copyright (c) 2019, Elasticsearch BV
+#  All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  * Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+#
+#  * Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+#
+#  * Neither the name of the copyright holder nor the names of its
+#    contributors may be used to endorse or promote products derived from
+#    this software without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import os
+
+import pytest
+
+pytestmark = [pytest.mark.psycopg_pool, pytest.mark.integrationtest]
+
+pytest.importorskip("psycopg")
+pool_mod = pytest.importorskip("psycopg_pool")
+
+
+def test_pool_generates_spans(instrument, elasticapm_client):
+    with pool_mod.ConnectionPool(
+        os.environ.get(
+            "PSYCOPG_TEST_DSN",
+            "postgresql://postgres:[email protected]:5432/postgres",
+        ),
+        min_size=1,
+        max_size=2,
+        open=True,
+    ) as pool:
+        pool.wait()
+
+        elasticapm_client.begin_transaction("request")
+        try:
+            with pool.connection() as connection:
+                with connection.cursor() as cursor:
+                    cursor.execute("SELECT 1")
+                    cursor.fetchone()
+        finally:
+            elasticapm_client.end_transaction("200")
+
+    spans = elasticapm_client.events["span"]
+    assert any(span.get("context", {}).get("db", {}).get("type") == "sql" for span in spans)

tests/requirements/reqs-psycopg-pool-newest.txt

@@ -0,0 +1,2 @@
+psycopg[binary,pool]
+-r reqs-base.txt

tests/scripts/envs/psycopg-pool.sh

@@ -0,0 +1 @@
+export PYTEST_MARKER="-m psycopg_pool"