diff --git a/elasticapm/instrumentation/packages/psycopg2.py b/elasticapm/instrumentation/packages/psycopg2.py index a058597da..850d849e3 100644 --- a/elasticapm/instrumentation/packages/psycopg2.py +++ b/elasticapm/instrumentation/packages/psycopg2.py @@ -59,6 +59,12 @@ def _bake_sql(self, sql): def extract_signature(self, sql): return extract_signature(sql) + def execute(self, query, vars=None): + return self._trace_sql(self.__wrapped__.execute, query, vars) + + def executemany(self, query, vars_list): + return self._trace_sql(self.__wrapped__.executemany, query, vars_list) + def __enter__(self): return PGCursorProxy(self.__wrapped__.__enter__(), destination_info=self._self_destination_info) diff --git a/tests/instrumentation/psycopg2_tests.py b/tests/instrumentation/psycopg2_tests.py index 70c0d6329..1cee5b269 100644 --- a/tests/instrumentation/psycopg2_tests.py +++ b/tests/instrumentation/psycopg2_tests.py @@ -266,6 +266,27 @@ def test_fully_qualified_table_name(): assert "SELECT FROM db.schema.mytable" == actual +@pytest.mark.integrationtest +@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured") +def test_cursor_execute_signature(instrument, postgres_connection, elasticapm_client): + cursor = postgres_connection.cursor() + cursor.execute(query="SELECT 1", vars=None) + row = cursor.fetchone() + + assert row + + +@pytest.mark.integrationtest +@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured") +def test_cursor_executemany_signature(instrument, postgres_connection, elasticapm_client): + cursor = postgres_connection.cursor() + res = cursor.executemany( + query="INSERT INTO test VALUES (%s, %s)", + vars_list=((4, "four"),), + ) + assert res is None + + @pytest.mark.integrationtest @pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured") def test_destination(instrument, postgres_connection, elasticapm_client):