fix extract_signature to work with sql.Composable objects (#148)

see http://initd.org/psycopg/docs/sql.html

closes #148

Author: beniwohli

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`

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:

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)
 

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
 

tests/instrumentation/psycopg2_tests.py

@@ -1,17 +1,20 @@
 # -*- coding: utf-8 -*-
 import os
 
+import psycopg2
 import pytest
 
 from elasticapm.instrumentation import control
 from elasticapm.instrumentation.packages.psycopg2 import (PGCursorProxy,
                                                           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