diff --git a/CHANGELOG.asciidoc b/CHANGELOG.asciidoc index c1e848dd5..ea90d81d8 100644 --- a/CHANGELOG.asciidoc +++ b/CHANGELOG.asciidoc @@ -22,6 +22,7 @@ https://github.com/elastic/apm-agent-python/compare/v1.0.0\...master[Check the H * added `transaction_max_spans` setting to limit the amount of spans that are recorded per transaction ({pull}127[#127]) * added configuration options to limit captured local variables to a certain length ({pull}130[#130]) * added options for configuring the amount of context lines that are captured with each frame ({pull}136[#136]) + * added support for tracing queries formatted as http://initd.org/psycopg/docs/sql.html[`psycopg2.sql.SQL`] objects ({pull}148[#148]) * switched to `time.perf_counter` as timing function on Python 3 ({pull}138[#138]) * BREAKING: Several settings and APIs have been renamed ({pull}111[#111], {pull}119[#119], {pull}143[#143]): ** The decorator for custom instrumentation, `elasticapm.trace`, is now `elasticapm.capture_span` diff --git a/elasticapm/instrumentation/packages/dbapi2.py b/elasticapm/instrumentation/packages/dbapi2.py index e13518b3b..d5186648c 100644 --- a/elasticapm/instrumentation/packages/dbapi2.py +++ b/elasticapm/instrumentation/packages/dbapi2.py @@ -121,6 +121,11 @@ def scan(tokens): def extract_signature(sql): + """ + Extracts a minimal signature from a given SQL query + :param sql: the SQL statement + :return: a string representing the signature + """ sql = sql.strip() first_space = sql.find(' ') if first_space < 0: @@ -170,10 +175,18 @@ def executemany(self, sql, param_list): return self._trace_sql(self.__wrapped__.executemany, sql, param_list) + def _bake_sql(self, sql): + """ + Method to turn the "sql" argument into a string. Most database backends simply return + the given object, as it is already a string + """ + return sql + def _trace_sql(self, method, sql, params): - signature = self.extract_signature(sql) + sql_string = self._bake_sql(sql) + signature = self.extract_signature(sql_string) kind = "db.{0}.sql".format(self.provider_name) - with capture_span(signature, kind, {'db': {"type": "sql", "statement": sql}}): + with capture_span(signature, kind, {'db': {"type": "sql", "statement": sql_string}}): if params is None: return method(sql) else: diff --git a/elasticapm/instrumentation/packages/psycopg2.py b/elasticapm/instrumentation/packages/psycopg2.py index ede9eb3b2..c2e79e3ba 100644 --- a/elasticapm/instrumentation/packages/psycopg2.py +++ b/elasticapm/instrumentation/packages/psycopg2.py @@ -9,6 +9,13 @@ class PGCursorProxy(CursorProxy): provider_name = 'postgresql' + def _bake_sql(self, sql): + # if this is a Composable object, use its `as_string` method + # see http://initd.org/psycopg/docs/sql.html + if hasattr(sql, 'as_string'): + return sql.as_string(self.__wrapped__) + return sql + def extract_signature(self, sql): return extract_signature(sql) diff --git a/setup.cfg b/setup.cfg index 2db1bef43..601291d03 100644 --- a/setup.cfg +++ b/setup.cfg @@ -25,7 +25,7 @@ multi_line_output=0 known_standard_library=importlib,types,asyncio known_django=django known_first_party=elasticapm,tests -known_third_party=pytest,flask,aiohttp,urllib3_mock,webob,memcache,pymongo,boto3,logbook,twisted,celery,zope,urllib3,redis,jinja2,requests,certifi,mock,jsonschema,werkzeug,pytest_localserver +known_third_party=pytest,flask,aiohttp,urllib3_mock,webob,memcache,pymongo,boto3,logbook,twisted,celery,zope,urllib3,redis,jinja2,requests,certifi,mock,jsonschema,werkzeug,pytest_localserver,psycopg2 default_section=FIRSTPARTY sections=FUTURE,STDLIB,DJANGO,THIRDPARTY,FIRSTPARTY,LOCALFOLDER diff --git a/tests/instrumentation/psycopg2_tests.py b/tests/instrumentation/psycopg2_tests.py index ee12f1ed0..0de151c91 100644 --- a/tests/instrumentation/psycopg2_tests.py +++ b/tests/instrumentation/psycopg2_tests.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- import os +import psycopg2 import pytest from elasticapm.instrumentation import control @@ -8,10 +9,12 @@ extract_signature) try: - import psycopg2 - has_psycopg2 = True + from psycopg2 import sql + has_sql_module = True except ImportError: - has_psycopg2 = False + # as of Jan 2018, psycopg2cffi doesn't have this module + has_sql_module = False + has_postgres_configured = 'POSTGRES_DB' in os.environ @@ -37,100 +40,100 @@ def postgres_connection(request): def test_insert(): - sql = """INSERT INTO mytable (id, name) VALUE ('2323', 'Ron')""" - actual = extract_signature(sql) + sql_statement = """INSERT INTO mytable (id, name) VALUE ('2323', 'Ron')""" + actual = extract_signature(sql_statement) assert "INSERT INTO mytable" == actual def test_update_with_quotes(): - sql = """UPDATE "my table" set name='Ron' WHERE id = 2323""" - actual = extract_signature(sql) + sql_statement = """UPDATE "my table" set name='Ron' WHERE id = 2323""" + actual = extract_signature(sql_statement) assert "UPDATE my table" == actual def test_update(): - sql = """update mytable set name = 'Ron where id = 'a'""" - actual = extract_signature(sql) + sql_statement = """update mytable set name = 'Ron where id = 'a'""" + actual = extract_signature(sql_statement) assert "UPDATE mytable" == actual def test_delete_simple(): - sql = 'DELETE FROM "mytable"' - actual = extract_signature(sql) + sql_statement = 'DELETE FROM "mytable"' + actual = extract_signature(sql_statement) assert "DELETE FROM mytable" == actual def test_delete(): - sql = """DELETE FROM "my table" WHERE id = 2323""" - actual = extract_signature(sql) + sql_statement = """DELETE FROM "my table" WHERE id = 2323""" + actual = extract_signature(sql_statement) assert "DELETE FROM my table" == actual def test_select_simple(): - sql = """SELECT id, name FROM my_table WHERE id = 2323""" - actual = extract_signature(sql) + sql_statement = """SELECT id, name FROM my_table WHERE id = 2323""" + actual = extract_signature(sql_statement) assert "SELECT FROM my_table" == actual def test_select_with_entity_quotes(): - sql = """SELECT id, name FROM "mytable" WHERE id = 2323""" - actual = extract_signature(sql) + sql_statement = """SELECT id, name FROM "mytable" WHERE id = 2323""" + actual = extract_signature(sql_statement) assert "SELECT FROM mytable" == actual def test_select_with_difficult_values(): - sql = """SELECT id, 'some name' + '" from Denmark' FROM "mytable" WHERE id = 2323""" - actual = extract_signature(sql) + sql_statement = """SELECT id, 'some name' + '" from Denmark' FROM "mytable" WHERE id = 2323""" + actual = extract_signature(sql_statement) assert "SELECT FROM mytable" == actual def test_select_with_dollar_quotes(): - sql = """SELECT id, $$some single doubles ' $$ + '" from Denmark' FROM "mytable" WHERE id = 2323""" - actual = extract_signature(sql) + sql_statement = """SELECT id, $$some single doubles ' $$ + '" from Denmark' FROM "mytable" WHERE id = 2323""" + actual = extract_signature(sql_statement) assert "SELECT FROM mytable" == actual def test_select_with_invalid_dollar_quotes(): - sql = """SELECT id, $fish$some single doubles ' $$ + '" from Denmark' FROM "mytable" WHERE id = 2323""" - actual = extract_signature(sql) + sql_statement = """SELECT id, $fish$some single doubles ' $$ + '" from Denmark' FROM "mytable" WHERE id = 2323""" + actual = extract_signature(sql_statement) assert "SELECT FROM" == actual def test_select_with_dollar_quotes_custom_token(): - sql = """SELECT id, $token $FROM $ FROM $ FROM single doubles ' $token $ + '" from Denmark' FROM "mytable" WHERE id = 2323""" - actual = extract_signature(sql) + sql_statement = """SELECT id, $token $FROM $ FROM $ FROM single doubles ' $token $ + '" from Denmark' FROM "mytable" WHERE id = 2323""" + actual = extract_signature(sql_statement) assert "SELECT FROM mytable" == actual def test_select_with_difficult_table_name(): - sql = "SELECT id FROM \"myta\n-æøåble\" WHERE id = 2323""" - actual = extract_signature(sql) + sql_statement = "SELECT id FROM \"myta\n-æøåble\" WHERE id = 2323""" + actual = extract_signature(sql_statement) assert "SELECT FROM myta\n-æøåble" == actual def test_select_subselect(): - sql = """SELECT id, name FROM ( + sql_statement = """SELECT id, name FROM ( SELECT id, 'not a FROM ''value' FROM mytable WHERE id = 2323 ) LIMIT 20""" - actual = extract_signature(sql) + actual = extract_signature(sql_statement) assert "SELECT FROM mytable" == actual def test_select_subselect_with_alias(): - sql = """ + sql_statement = """ SELECT count(*) FROM ( SELECT count(id) AS some_alias, some_column @@ -139,75 +142,76 @@ def test_select_subselect_with_alias(): HAVING count(id) > 1 ) AS foo """ - actual = extract_signature(sql) + actual = extract_signature(sql_statement) assert "SELECT FROM mytable" == actual def test_select_with_multiple_tables(): - sql = """SELECT count(table2.id) + sql_statement = """SELECT count(table2.id) FROM table1, table2, table2 WHERE table2.id = table1.table2_id """ - actual = extract_signature(sql) + actual = extract_signature(sql_statement) assert "SELECT FROM table1" == actual def test_select_with_invalid_subselect(): - sql = "SELECT id FROM (SELECT * """ - actual = extract_signature(sql) + sql_statement = "SELECT id FROM (SELECT * """ + actual = extract_signature(sql_statement) assert "SELECT FROM" == actual def test_select_with_invalid_literal(): - sql = "SELECT 'neverending literal FROM (SELECT * FROM ...""" - actual = extract_signature(sql) + sql_statement = "SELECT 'neverending literal FROM (SELECT * FROM ...""" + actual = extract_signature(sql_statement) assert "SELECT FROM" == actual def test_savepoint(): - sql = """SAVEPOINT x_asd1234""" - actual = extract_signature(sql) + sql_statement = """SAVEPOINT x_asd1234""" + actual = extract_signature(sql_statement) assert "SAVEPOINT" == actual def test_begin(): - sql = """BEGIN""" - actual = extract_signature(sql) + sql_statement = """BEGIN""" + actual = extract_signature(sql_statement) assert "BEGIN" == actual def test_create_index_with_name(): - sql = """CREATE INDEX myindex ON mytable""" - actual = extract_signature(sql) + sql_statement = """CREATE INDEX myindex ON mytable""" + actual = extract_signature(sql_statement) assert "CREATE INDEX" == actual def test_create_index_without_name(): - sql = """CREATE INDEX ON mytable""" - actual = extract_signature(sql) + sql_statement = """CREATE INDEX ON mytable""" + actual = extract_signature(sql_statement) assert "CREATE INDEX" == actual def test_drop_table(): - sql = """DROP TABLE mytable""" - actual = extract_signature(sql) + sql_statement = """DROP TABLE mytable""" + actual = extract_signature(sql_statement) assert "DROP TABLE" == actual def test_multi_statement_sql(): - sql = """CREATE TABLE mytable; SELECT * FROM mytable; DROP TABLE mytable""" - actual = extract_signature(sql) + sql_statement = """CREATE TABLE mytable; SELECT * FROM mytable; DROP TABLE mytable""" + actual = extract_signature(sql_statement) assert "CREATE TABLE" == actual + @pytest.mark.integrationtest @pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured") def test_psycopg2_register_type(postgres_connection, elasticapm_client): @@ -289,3 +293,35 @@ def test_psycopg2_select_LIKE(postgres_connection, elasticapm_client): assert 'db' in span['context'] assert span['context']['db']['type'] == 'sql' assert span['context']['db']['statement'] == query + + +@pytest.mark.integrationtest +@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured") +@pytest.mark.skipif(not has_sql_module, reason="psycopg2.sql module missing") +def test_psycopg2_composable_query_works(postgres_connection, elasticapm_client): + """ + Check that we parse queries that are psycopg2.sql.Composable correctly + """ + control.instrument() + 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.instrumentation_store.get_all() + spans = transactions[0]['spans'] + span = spans[0] + assert span['name'] == 'SELECT FROM test' + assert 'db' in span['context'] + assert span['context']['db']['type'] == 'sql' + assert span['context']['db']['statement'] == baked_query