Repository: elastic/apm-agent-python

Issue:
Instrument psycopg3
psycopg3 was released some time ago. We currently only support psycopg2. While psycopg3 is quite similar from psycopg2, the differences are large enough to probably require its own instrumentation (e.g., extensions are gone in psycopg3 AFAICT).

Also, there is an asyncio interface in psycopg3, which requires its own instrumentation.

https://www.psycopg.org/psycopg3/docs/index.html
https://www.psycopg.org/psycopg3/docs/basic/from_pg2.html



Discussion:
Hello, is there an approximate release date for this?
We don't currently have a release date. However, it's good to know someone is in need of this instrumentation, it helps us prioritize!
It would be useful for us as well.
+1
Hi @basepi @beniwohli! We're using Django (synchronous mode) and would love to upgrade to Psycopg3. Is there anything I could do to speed up the Elastic APM support for Psycopg3? I could open a PR adding support for Psycopg3 synchronous operations which is pretty much similar to Psycopg2, adopting [`Psycopg2Instrumentation`](https://github.com/elastic/apm-agent-python/blob/main/elasticapm/instrumentation/packages/psycopg2.py#L77) class.
@minhtule that would be a welcome contribution! I assumed this would be pretty easy work, we just haven't gotten to it yet.
Hi @minhtule, we need it too as a team, if you haven't started yet we can do it! otherwise tell me how we can help you ❤️

Resolving patch:
diff --git a/elasticapm/instrumentation/packages/psycopg.py b/elasticapm/instrumentation/packages/psycopg.py
new file mode 100644
--- /dev/null
+++ b/elasticapm/instrumentation/packages/psycopg.py
@@ -0,0 +1,96 @@
+#  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.
+from __future__ import absolute_import
+
+from elasticapm.instrumentation.packages.dbapi2 import (
+    ConnectionProxy,
+    CursorProxy,
+    DbApi2Instrumentation,
+    extract_signature,
+)
+from elasticapm.instrumentation.packages.psycopg2 import get_destination_info
+from elasticapm.traces import capture_span
+
+
+class PGCursorProxy(CursorProxy):
+    provider_name = "postgresql"
+
+    def _bake_sql(self, sql):
+        # If this is a Composable object, use its `as_string` method.
+        # See https://www.psycopg.org/psycopg3/docs/api/sql.html
+        if hasattr(sql, "as_string"):
+            sql = sql.as_string(self.__wrapped__)
+        # If the sql string is already a byte string, we need to decode it using the connection encoding
+        if isinstance(sql, bytes):
+            sql = sql.decode(self.connection.info.encoding)
+        return sql
+
+    def extract_signature(self, sql):
+        return extract_signature(sql)
+
+    def __enter__(self):
+        return PGCursorProxy(self.__wrapped__.__enter__(), destination_info=self._self_destination_info)
+
+    @property
+    def _self_database(self):
+        return self.connection.info.dbname or ""
+
+
+class PGConnectionProxy(ConnectionProxy):
+    cursor_proxy = PGCursorProxy
+
+    def __enter__(self):
+        return PGConnectionProxy(self.__wrapped__.__enter__(), destination_info=self._self_destination_info)
+
+
+class PsycopgInstrumentation(DbApi2Instrumentation):
+    name = "psycopg"
+
+    instrument_list = [("psycopg", "connect")]
+
+    def call(self, module, method, wrapped, instance, args, kwargs):
+        signature = "psycopg.connect"
+
+        host, port = get_destination_info(kwargs.get("host"), kwargs.get("port"))
+        database = kwargs.get("dbname")
+        signature = f"{signature} {host}:{port}"
+        destination_info = {
+            "address": host,
+            "port": port,
+        }
+        with capture_span(
+            signature,
+            span_type="db",
+            span_subtype="postgresql",
+            span_action="connect",
+            leaf=True,
+            extra={"destination": destination_info, "db": {"type": "sql", "instance": database}},
+        ):
+            return PGConnectionProxy(wrapped(*args, **kwargs), destination_info=destination_info)
diff --git a/elasticapm/instrumentation/register.py b/elasticapm/instrumentation/register.py
--- a/elasticapm/instrumentation/register.py
+++ b/elasticapm/instrumentation/register.py
@@ -36,6 +36,7 @@
     "elasticapm.instrumentation.packages.botocore.BotocoreInstrumentation",
     "elasticapm.instrumentation.packages.httpx.sync.httpx.HttpxClientInstrumentation",
     "elasticapm.instrumentation.packages.jinja2.Jinja2Instrumentation",
+    "elasticapm.instrumentation.packages.psycopg.PsycopgInstrumentation",
     "elasticapm.instrumentation.packages.psycopg2.Psycopg2Instrumentation",
     "elasticapm.instrumentation.packages.psycopg2.Psycopg2ExtensionsInstrumentation",
     "elasticapm.instrumentation.packages.mysql.MySQLInstrumentation",


Test patch:
diff --git a/tests/instrumentation/psycopg_tests.py b/tests/instrumentation/psycopg_tests.py
new file mode 100644
--- /dev/null
+++ b/tests/instrumentation/psycopg_tests.py
@@ -0,0 +1,254 @@
+# -*- coding: utf-8 -*-
+
+#  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
+from typing import cast
+
+import pytest
+
+from elasticapm import get_client
+from elasticapm.conf.constants import SPAN, TRANSACTION
+from elasticapm.instrumentation.packages.psycopg import PGCursorProxy
+from elasticapm.utils import default_ports
+from tests.fixtures import TempStoreClient
+
+psycopg = pytest.importorskip("psycopg")
+
+pytestmark = pytest.mark.psycopg
+
+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),
+    }
+
+
+@pytest.fixture(scope="function")
+def postgres_connection(request):
+    conn = psycopg.connect(**connect_kwargs())
+    cursor = conn.cursor()
+    cursor.execute(
+        "CREATE TABLE test(id int, name VARCHAR(5) NOT NULL);"
+        "INSERT INTO test VALUES (1, 'one'), (2, 'two'), (3, 'three');"
+    )
+
+    yield conn
+
+    # cleanup
+    cursor.execute("ROLLBACK")
+
+
+@pytest.mark.integrationtest
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
+def test_destination(instrument, postgres_connection, elasticapm_client):
+    elasticapm_client.begin_transaction("test")
+    cursor = postgres_connection.cursor()
+    cursor.execute("SELECT 1")
+    elasticapm_client.end_transaction("test")
+    transaction = elasticapm_client.events[TRANSACTION][0]
+    span = elasticapm_client.spans_for_transaction(transaction)[0]
+    assert span["context"]["destination"] == {
+        "address": os.environ.get("POSTGRES_HOST", None),
+        "port": default_ports["postgresql"],
+        "service": {"name": "", "resource": "postgresql/elasticapm_test", "type": ""},
+    }
+
+
+@pytest.mark.integrationtest
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
+def test_psycopg_tracing_outside_of_elasticapm_transaction(instrument, postgres_connection, elasticapm_client):
+    cursor = postgres_connection.cursor()
+    # check that the cursor is a proxy, even though we're not in an elasticapm
+    # transaction
+    assert isinstance(cursor, PGCursorProxy)
+    cursor.execute("SELECT 1")
+    transactions = elasticapm_client.events[TRANSACTION]
+    assert not transactions
+
+
+@pytest.mark.integrationtest
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
+def test_psycopg_select_LIKE(instrument, postgres_connection, elasticapm_client):
+    """
+    Check that we pass queries with %-notation but without parameters
+    properly to the dbapi backend
+    """
+    cursor = postgres_connection.cursor()
+    query = "SELECT * FROM test WHERE name LIKE 't%'"
+
+    try:
+        elasticapm_client.begin_transaction("web.django")
+        cursor.execute(query)
+        cursor.fetchall()
+        elasticapm_client.end_transaction(None, "test-transaction")
+    finally:
+        # make sure we've cleared out the spans for the other tests.
+        transactions = elasticapm_client.events[TRANSACTION]
+        spans = elasticapm_client.spans_for_transaction(transactions[0])
+        span = spans[0]
+        assert span["name"] == "SELECT FROM test"
+        assert span["type"] == "db"
+        assert span["subtype"] == "postgresql"
+        assert span["action"] == "query"
+        assert "db" in span["context"]
+        assert span["context"]["db"]["instance"] == "elasticapm_test"
+        assert span["context"]["db"]["type"] == "sql"
+        assert span["context"]["db"]["statement"] == query
+        assert span["context"]["service"]["target"]["type"] == "postgresql"
+        assert span["context"]["service"]["target"]["name"] == "elasticapm_test"
+
+
+@pytest.mark.integrationtest
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
+def test_psycopg_composable_query_works(instrument, postgres_connection, elasticapm_client):
+    """
+    Check that we parse queries that are psycopg.sql.Composable correctly
+    """
+    from psycopg import sql
+
+    cursor = postgres_connection.cursor()
+    query = sql.SQL("SELECT * FROM {table} WHERE {row} LIKE 't%' ORDER BY {row} DESC").format(
+        table=sql.Identifier("test"), row=sql.Identifier("name")
+    )
+    baked_query = query.as_string(cursor.__wrapped__)
+    result = None
+    try:
+        elasticapm_client.begin_transaction("web.django")
+        cursor.execute(query)
+        result = cursor.fetchall()
+        elasticapm_client.end_transaction(None, "test-transaction")
+    finally:
+        # make sure we've cleared out the spans for the other tests.
+        assert [(2, "two"), (3, "three")] == result
+        transactions = elasticapm_client.events[TRANSACTION]
+        spans = elasticapm_client.spans_for_transaction(transactions[0])
+        span = spans[0]
+        assert span["name"] == "SELECT FROM test"
+        assert "db" in span["context"]
+        assert span["context"]["db"]["instance"] == "elasticapm_test"
+        assert span["context"]["db"]["type"] == "sql"
+        assert span["context"]["db"]["statement"] == baked_query
+
+
+@pytest.mark.integrationtest
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
+def test_psycopg_binary_query_works(instrument, postgres_connection, elasticapm_client):
+    """
+    Check that we pass queries with %-notation but without parameters
+    properly to the dbapi backend
+    """
+    cursor = postgres_connection.cursor()
+    query = b"SELECT * FROM test WHERE name LIKE 't%'"
+
+    baked_query = query.decode()
+    try:
+        elasticapm_client.begin_transaction("web.django")
+        cursor.execute(query)
+        result = cursor.fetchall()
+        elasticapm_client.end_transaction(None, "test-transaction")
+    finally:
+        # make sure we've cleared out the spans for the other tests.
+        assert [(2, "two"), (3, "three")] == result
+        transactions = elasticapm_client.events[TRANSACTION]
+        spans = elasticapm_client.spans_for_transaction(transactions[0])
+        span = spans[0]
+        assert span["name"] == "SELECT FROM test"
+        assert "db" in span["context"]
+        assert span["context"]["db"]["instance"] == "elasticapm_test"
+        assert span["context"]["db"]["type"] == "sql"
+        assert span["context"]["db"]["statement"] == baked_query
+
+
+@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 psycopg.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[0]["context"]["service"]["target"]["type"] == "postgresql"
+    assert spans[0]["context"]["service"]["target"]["name"] == "elasticapm_test"
+
+    assert spans[1]["subtype"] == "postgresql"
+    assert spans[1]["action"] == "query"
+
+
+@pytest.mark.integrationtest
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
+def test_psycopg_rows_affected(instrument, postgres_connection, elasticapm_client):
+    cursor = postgres_connection.cursor()
+    try:
+        elasticapm_client.begin_transaction("web.django")
+        cursor.execute("INSERT INTO test VALUES (4, 'four')")
+        cursor.execute("SELECT * FROM test")
+        cursor.execute("UPDATE test SET name = 'five' WHERE  id = 4")
+        cursor.execute("DELETE FROM test WHERE  id = 4")
+        elasticapm_client.end_transaction(None, "test-transaction")
+    finally:
+        transactions = elasticapm_client.events[TRANSACTION]
+        spans = elasticapm_client.spans_for_transaction(transactions[0])
+
+        assert spans[0]["name"] == "INSERT INTO test"
+        assert spans[0]["context"]["db"]["rows_affected"] == 1
+
+        assert spans[1]["name"] == "SELECT FROM test"
+        assert "rows_affected" not in spans[1]["context"]["db"]
+
+        assert spans[2]["name"] == "UPDATE test"
+        assert spans[2]["context"]["db"]["rows_affected"] == 1
+
+        assert spans[3]["name"] == "DELETE FROM test"
+        assert spans[3]["context"]["db"]["rows_affected"] == 1
+
+
+@pytest.mark.integrationtest
+def test_psycopg_connection(instrument, elasticapm_transaction, postgres_connection):
+    # elastciapm_client.events is only available on `TempStoreClient`, this keeps the type checkers happy
+    elasticapm_client = cast(TempStoreClient, get_client())
+    elasticapm_client.end_transaction("test", "success")
+    span = elasticapm_client.events[SPAN][0]
+    host = os.environ.get("POSTGRES_HOST", "localhost")
+    assert span["name"] == f"psycopg.connect {host}:5432"
+    assert span["action"] == "connect"
diff --git a/tests/requirements/reqs-psycopg-newest.txt b/tests/requirements/reqs-psycopg-newest.txt
new file mode 100644
--- /dev/null
+++ b/tests/requirements/reqs-psycopg-newest.txt
@@ -0,0 +1,2 @@
+psycopg
+-r reqs-base.txt
diff --git a/tests/scripts/envs/psycopg.sh b/tests/scripts/envs/psycopg.sh
new file mode 100644
--- /dev/null
+++ b/tests/scripts/envs/psycopg.sh
@@ -0,0 +1,8 @@
+export PYTEST_MARKER="-m psycopg"
+export DOCKER_DEPS="postgres"
+export POSTGRES_HOST="postgres"
+export POSTGRES_USER="postgres"
+export POSTGRES_PASSWORD="postgres"
+export POSTGRES_DB="elasticapm_test"
+export POSTGRES_HOST="postgres"
+export POSTGRES_PORT="5432"