diff --git a/.ci/.matrix_framework.yml b/.ci/.matrix_framework.yml index 64d1dc8ef..a47e12259 100644 --- a/.ci/.matrix_framework.yml +++ b/.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 diff --git a/.ci/.matrix_framework_full.yml b/.ci/.matrix_framework_full.yml index 10c8cee73..bf789b7d6 100644 --- a/.ci/.matrix_framework_full.yml +++ b/.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 diff --git a/elasticapm/instrumentation/packages/psycopg.py b/elasticapm/instrumentation/packages/psycopg.py index 0d79cf686..5870f6603 100644 --- a/elasticapm/instrumentation/packages/psycopg.py +++ b/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" diff --git a/tests/instrumentation/psycopg_pool_tests.py b/tests/instrumentation/psycopg_pool_tests.py new file mode 100644 index 000000000..12a73441b --- /dev/null +++ b/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" diff --git a/tests/requirements/reqs-psycopg_pool-newest.txt b/tests/requirements/reqs-psycopg_pool-newest.txt new file mode 100644 index 000000000..408c051c2 --- /dev/null +++ b/tests/requirements/reqs-psycopg_pool-newest.txt @@ -0,0 +1,2 @@ +psycopg[binary,pool] +-r reqs-base.txt diff --git a/tests/scripts/envs/psycopg_pool.sh b/tests/scripts/envs/psycopg_pool.sh new file mode 100644 index 000000000..b4fc1e9c4 --- /dev/null +++ b/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"