Repository: elastic/apm-agent-python

Issue:
Capture `rows_affected` in DB-API2 integrations
See https://github.com/elastic/apm/issues/112


Discussion:

Resolving patch:
diff --git a/elasticapm/instrumentation/packages/dbapi2.py b/elasticapm/instrumentation/packages/dbapi2.py
--- a/elasticapm/instrumentation/packages/dbapi2.py
+++ b/elasticapm/instrumentation/packages/dbapi2.py
@@ -197,6 +197,7 @@ def extract_signature(sql):
 
 class CursorProxy(wrapt.ObjectProxy):
     provider_name = None
+    DML_QUERIES = ("INSERT", "DELETE", "UPDATE")
 
     def __init__(self, wrapped, destination_info=None):
         super(CursorProxy, self).__init__(wrapped)
@@ -231,11 +232,15 @@ def _trace_sql(self, method, sql, params, action=QUERY_ACTION):
             span_action=action,
             extra={"db": {"type": "sql", "statement": sql_string}, "destination": self._self_destination_info},
             skip_frames=1,
-        ):
+        ) as span:
             if params is None:
-                return method(sql)
+                result = method(sql)
             else:
-                return method(sql, params)
+                result = method(sql, params)
+            # store "rows affected", but only for DML queries like insert/update/delete
+            if span and self.rowcount not in (-1, None) and signature.startswith(self.DML_QUERIES):
+                span.update_context("db", {"rows_affected": self.rowcount})
+            return result
 
     def extract_signature(self, sql):
         raise NotImplementedError()
diff --git a/elasticapm/traces.py b/elasticapm/traces.py
--- a/elasticapm/traces.py
+++ b/elasticapm/traces.py
@@ -458,6 +458,17 @@ def end(self, skip_frames=0, duration=None):
                 self.type, self.subtype, self.duration - self._child_durations.duration
             )
 
+    def update_context(self, key, data):
+        """
+        Update the context data for given key
+        :param key: the key, e.g. "db"
+        :param data: a dictionary
+        :return: None
+        """
+        current = self.context.get(key, {})
+        current.update(data)
+        self.context[key] = current
+
     def __str__(self):
         return u"{}/{}/{}".format(self.name, self.type, self.subtype)
 
@@ -480,6 +491,9 @@ def child_started(self, timestamp):
     def child_ended(self, timestamp):
         pass
 
+    def update_context(self, key, data):
+        pass
+
     @property
     def type(self):
         return None


Test patch:
diff --git a/tests/instrumentation/psycopg2_tests.py b/tests/instrumentation/psycopg2_tests.py
--- a/tests/instrumentation/psycopg2_tests.py
+++ b/tests/instrumentation/psycopg2_tests.py
@@ -456,3 +456,31 @@ def test_psycopg_context_manager(instrument, elasticapm_client):
 
     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_psycopg2_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
diff --git a/tests/instrumentation/pymssql_tests.py b/tests/instrumentation/pymssql_tests.py
--- a/tests/instrumentation/pymssql_tests.py
+++ b/tests/instrumentation/pymssql_tests.py
@@ -106,3 +106,30 @@ def test_host_port_parsing(args, kwargs, expected):
     host, port = get_host_port(args, kwargs)
     assert host == expected["host"]
     assert port == expected["port"]
+
+
+@pytest.mark.integrationtest
+def test_pymssql_rows_affected(instrument, pymssql_connection, elasticapm_client):
+    cursor = pymssql_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
diff --git a/tests/instrumentation/pyodbc_tests.py b/tests/instrumentation/pyodbc_tests.py
--- a/tests/instrumentation/pyodbc_tests.py
+++ b/tests/instrumentation/pyodbc_tests.py
@@ -86,3 +86,30 @@ def test_pyodbc_select(instrument, pyodbc_postgres_connection, elasticapm_client
         assert "db" in span["context"]
         assert span["context"]["db"]["type"] == "sql"
         assert span["context"]["db"]["statement"] == query
+
+
+@pytest.mark.integrationtest
+def test_psycopg2_rows_affected(instrument, pyodbc_postgres_connection, elasticapm_client):
+    cursor = pyodbc_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