Fix psycopg_pool instrumentation to restore Postgres spans (#2460)

* fix: instrument psycopg_pool connections to restore Postgres spans

Signed-off-by: SaJH <wogur981208@gmail.com>

* fix: add PostgreSQL dependencies for psycopg_pool tests

Signed-off-by: JaeHyuck Sa <wogur981208@gmail.com>

* Fix tests

Assisted-by: Cursor

* Fix test name

---------

Signed-off-by: SaJH <wogur981208@gmail.com>
Signed-off-by: JaeHyuck Sa <wogur981208@gmail.com>
Co-authored-by: Riccardo Magliocchetti <riccardo.magliocchetti@gmail.com>

Author: JaeHyuckSa

.ci/.matrix_framework.yml

@@ -25,6 +25,7 @@ FRAMEWORK:
   #- aioredis-2  # not supported yet
   - psycopg-newest
   - psycopg2-newest
+  - psycopg_pool-newest
   - pymssql-newest
   - pyodbc-newest
   - memcached-newest

.ci/.matrix_framework_full.yml

@@ -52,6 +52,7 @@ FRAMEWORK:
   - redis-newest
   - psycopg-newest
   - psycopg2-newest
+  - psycopg_pool-newest
   - pymssql-newest
   - memcached-newest
   - pylibmc-1.4

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,97 @@
+#  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
+
+from elasticapm.conf.constants import SPAN
+
+psycopg = pytest.importorskip("psycopg")
+pool_mod = pytest.importorskip("psycopg_pool")
+
+pytestmark = [pytest.mark.psycopg_pool, pytest.mark.integrationtest]
+
+has_postgres_configured = "POSTGRES_DB" in os.environ
+
+
+def connect_kwargs():
+    return {
+        "dbname": os.environ.get("POSTGRES_DB", "elasticapm_test"),
+        "user": os.environ.get("POSTGRES_USER", "postgres"),
+        "password": os.environ.get("POSTGRES_PASSWORD", "postgres"),
+        "host": os.environ.get("POSTGRES_HOST", None),
+        "port": os.environ.get("POSTGRES_PORT", None),
+    }
+
+
+def make_conninfo():
+    kw = connect_kwargs()
+    host = kw["host"] or "localhost"
+    port = kw["port"] or "5432"
+    return f"postgresql://{kw['user']}:{kw['password']}@{host}:{port}/{kw['dbname']}"
+
+
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgreSQL not configured")
+def test_connection_connect_generates_span(instrument, elasticapm_client):
+    elasticapm_client.begin_transaction("request")
+    try:
+        connection = psycopg.Connection.connect(**connect_kwargs())
+        connection.close()
+    finally:
+        elasticapm_client.end_transaction("200")
+
+    spans = elasticapm_client.events[SPAN]
+    assert len(spans) == 1
+    assert spans[0]["action"] == "connect"
+
+
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgreSQL not configured")
+def test_pool_generates_query_span(instrument, elasticapm_client):
+    with pool_mod.ConnectionPool(
+        make_conninfo(),
+        min_size=1,
+        max_size=2,
+    ) 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]
+    span = next(span for span in spans if span["action"] == "query")
+    assert span["action"] == "query"
+    assert span["context"]["db"]["type"] == "sql"

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,7 @@
+export PYTEST_MARKER="-m psycopg_pool"
+export DOCKER_DEPS="postgres"
+export POSTGRES_HOST="postgres"
+export POSTGRES_USER="postgres"
+export POSTGRES_PASSWORD="postgres"
+export POSTGRES_DB="elasticapm_test"
+export POSTGRES_PORT="5432"