add instrumentation for sqlite3's Connection.execute (#271)

Python's sqlite3 module has non-standard shortcuts for `execute`
and `executemany`, so we need wrappers for those as well.

closes #271

Signed-off-by: Benjamin Wohlwend <beni@elastic.co>

CHANGELOG.md

@@ -5,6 +5,7 @@
 [Check the diff](https://github.com/elastic/apm-agent-python/compare/v3.0.0...master)
 
  * added sanitization for `Set-Cookie` response headers (#264)
+ * added instrumentation for the non-standard `Connection.execute()` method for SQLite3 (#271)
 
 ## v3.0.0
 

elasticapm/instrumentation/packages/sqlite.py

@@ -17,6 +17,24 @@ def extract_signature(self, sql):
 class SQLiteConnectionProxy(ConnectionProxy):
     cursor_proxy = SQLiteCursorProxy
 
+    # we need to implement wrappers for the non-standard Connection.execute and
+    # Connection.executemany methods
+
+    def _trace_sql(self, method, sql, params):
+        signature = extract_signature(sql)
+        kind = "db.sqlite.sql"
+        with capture_span(signature, kind, {"db": {"type": "sql", "statement": sql}}):
+            if params is None:
+                return method(sql)
+            else:
+                return method(sql, params)
+
+    def execute(self, sql, params=None):
+        return self._trace_sql(self.__wrapped__.execute, sql, params)
+
+    def executemany(self, sql, params=None):
+        return self._trace_sql(self.__wrapped__.executemany, sql, params)
+
 
 class SQLiteInstrumentation(DbApi2Instrumentation):
     name = "sqlite"

tests/contrib/django/django_tests.py

@@ -1067,7 +1067,8 @@ def test_perf_database_render(benchmark, client, django_elasticapm_client):
 
         assert len(transactions) == len(responses)
         for transaction in transactions:
-            assert len(transaction["spans"]) in (102, 103)
+            # number of spans per transaction can vary slightly
+            assert 102 <= len(transaction["spans"]) < 105
 
 
 @pytest.mark.django_db

tests/instrumentation/sqlite_tests.py

@@ -4,26 +4,68 @@
 def test_connect(instrument, elasticapm_client):
     elasticapm_client.begin_transaction("transaction.test")
 
+    sqlite3.connect(":memory:")
+    elasticapm_client.end_transaction("MyView")
+
+    transactions = elasticapm_client.transaction_store.get_all()
+    spans = transactions[0]["spans"]
+
+    assert spans[0]["name"] == "sqlite3.connect :memory:"
+    assert spans[0]["type"] == "db.sqlite.connect"
+
+
+def test_cursor(instrument, elasticapm_client):
     conn = sqlite3.connect(":memory:")
+
+    elasticapm_client.begin_transaction("transaction.test")
     cursor = conn.cursor()
+    cursor.execute("CREATE TABLE testdb (id integer, username text)")
+    cursor.execute('INSERT INTO testdb VALUES (1, "Ron")')
+    cursor.executemany("INSERT INTO testdb VALUES (?, ?)", ((2, "Rasmus"), (3, "Shay")))
+    cursor.execute("DROP TABLE testdb")
+    elasticapm_client.end_transaction("MyView")
+
+    transactions = elasticapm_client.transaction_store.get_all()
+    spans = transactions[0]["spans"]
+    expected_signatures = {"CREATE TABLE", "INSERT INTO testdb", "DROP TABLE"}
+
+    assert {t["name"] for t in spans} == expected_signatures
 
-    cursor.execute("""CREATE TABLE testdb (id integer, username text)""")
-    cursor.execute("""INSERT INTO testdb VALUES (1, "Ron")""")
-    cursor.execute("""DROP TABLE testdb""")
+    assert spans[0]["name"] == "CREATE TABLE"
+    assert spans[0]["type"] == "db.sqlite.sql"
 
+    assert spans[1]["name"] == "INSERT INTO testdb"
+    assert spans[1]["type"] == "db.sqlite.sql"
+
+    assert spans[2]["name"] == "INSERT INTO testdb"
+    assert spans[2]["type"] == "db.sqlite.sql"
+
+    assert spans[3]["name"] == "DROP TABLE"
+    assert spans[3]["type"] == "db.sqlite.sql"
+
+    assert len(spans) == 4
+
+
+def test_nonstandard_connection_execute(instrument, elasticapm_client):
+    conn = sqlite3.connect(":memory:")
+
+    elasticapm_client.begin_transaction("transaction.test")
+    conn.execute("CREATE TABLE testdb (id integer, username text)")
+    conn.execute('INSERT INTO testdb VALUES (1, "Ron")')
+    conn.executemany("INSERT INTO testdb VALUES (?, ?)", ((2, "Rasmus"), (3, "Shay")))
+    conn.execute("DROP TABLE testdb")
     elasticapm_client.end_transaction("MyView")
 
     transactions = elasticapm_client.transaction_store.get_all()
     spans = transactions[0]["spans"]
-
-    expected_signatures = {"sqlite3.connect :memory:", "CREATE TABLE", "INSERT INTO testdb", "DROP TABLE"}
+    expected_signatures = {"CREATE TABLE", "INSERT INTO testdb", "DROP TABLE"}
 
     assert {t["name"] for t in spans} == expected_signatures
 
-    assert spans[0]["name"] == "sqlite3.connect :memory:"
-    assert spans[0]["type"] == "db.sqlite.connect"
+    assert spans[0]["name"] == "CREATE TABLE"
+    assert spans[0]["type"] == "db.sqlite.sql"
 
-    assert spans[1]["name"] == "CREATE TABLE"
+    assert spans[1]["name"] == "INSERT INTO testdb"
     assert spans[1]["type"] == "db.sqlite.sql"
 
     assert spans[2]["name"] == "INSERT INTO testdb"