Repository: elastic/apm-agent-python

Issue:
APM agent crashes if the SQL query is too large
**Describe the bug**: ...
Same as https://github.com/elastic/apm-agent-python/issues/827, but with `psycopg2` library instead of Django ORM.

**To Reproduce**

1. use any of `psycopg2` extra utilities, for example `execute_values`, with a query that is sufficiently long.

**Environment (please complete the following information)**
- OS: mac
- Python version: 3.7
- APM Server version: ?
- Agent version: 5.10.0


**Additional context**

The issue is that `psycopg2` is passing the sql as bytes object (bytes array). The agent sql instrumentation [code](https://github.com/elastic/apm-agent-python/blob/master/elasticapm/instrumentation/packages/dbapi2.py#L222-L231) is calling the [`shorten()`](https://github.com/elastic/apm-agent-python/blob/f8977bf5c5d9921f83eed174b94cc48922bc72b2/elasticapm/utils/encoding.py#L182-L217) function which does not support bytes object.

I'm not sure if same behavior applies to `psycopg2` core `execute()` function, or only in the "extras" the query is passed as bytes object.
APM agent crashes if the SQL query is too large
**Describe the bug**: ...
Same as https://github.com/elastic/apm-agent-python/issues/827, but with `psycopg2` library instead of Django ORM.

**To Reproduce**

1. use any of `psycopg2` extra utilities, for example `execute_values`, with a query that is sufficiently long.

**Environment (please complete the following information)**
- OS: mac
- Python version: 3.7
- APM Server version: ?
- Agent version: 5.10.0


**Additional context**

The issue is that `psycopg2` is passing the sql as bytes object (bytes array). The agent sql instrumentation [code](https://github.com/elastic/apm-agent-python/blob/master/elasticapm/instrumentation/packages/dbapi2.py#L222-L231) is calling the [`shorten()`](https://github.com/elastic/apm-agent-python/blob/f8977bf5c5d9921f83eed174b94cc48922bc72b2/elasticapm/utils/encoding.py#L182-L217) function which does not support bytes object.

I'm not sure if same behavior applies to `psycopg2` core `execute()` function, or only in the "extras" the query is passed as bytes object.


Discussion:
@basepi I see you fixed it last time. it's an easy fix here as well. I can send PR if want...
@basepi I see you fixed it last time. it's an easy fix here as well. I can send PR if want...

Resolving patch:
diff --git a/elasticapm/instrumentation/packages/psycopg2.py b/elasticapm/instrumentation/packages/psycopg2.py
--- a/elasticapm/instrumentation/packages/psycopg2.py
+++ b/elasticapm/instrumentation/packages/psycopg2.py
@@ -27,6 +27,7 @@
 #  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 #  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+from __future__ import absolute_import
 
 from elasticapm.instrumentation.packages.dbapi2 import (
     ConnectionProxy,
@@ -42,10 +43,15 @@ class PGCursorProxy(CursorProxy):
     provider_name = "postgresql"
 
     def _bake_sql(self, sql):
+        from psycopg2 import extensions as psycopg2_extensions
+
         # 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__)
+            sql = sql.as_string(self.__wrapped__)
+        # if the sql string is already a byte string, we need to decode it using the connection encoding
+        if isinstance(sql, compat.binary_type):
+            sql = sql.decode(psycopg2_extensions.encodings[self.__wrapped__.connection.encoding])
         return sql
 
     def extract_signature(self, sql):


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
@@ -42,12 +42,11 @@
 
 
 try:
-    from psycopg2 import sql
+    from psycopg2 import compat
 
-    has_sql_module = True
+    is_cffi = True
 except ImportError:
-    # as of Jan 2018, psycopg2cffi doesn't have this module
-    has_sql_module = False
+    is_cffi = False
 
 try:
     import psycopg2.extensions
@@ -386,11 +385,13 @@ def test_psycopg2_select_LIKE(instrument, postgres_connection, elasticapm_client
 
 @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")
+@pytest.mark.skipif(is_cffi, reason="psycopg2cffi does not have the sql module")
 def test_psycopg2_composable_query_works(instrument, postgres_connection, elasticapm_client):
     """
     Check that we parse queries that are psycopg2.sql.Composable correctly
     """
+    from psycopg2 import sql
+
     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")
@@ -484,3 +485,25 @@ def test_psycopg2_rows_affected(instrument, postgres_connection, elasticapm_clie
 
         assert spans[3]["name"] == "DELETE FROM test"
         assert spans[3]["context"]["db"]["rows_affected"] == 1
+
+
+@pytest.mark.integrationtest
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
+@pytest.mark.skipif(is_cffi, reason="psycopg2cffi doesn't have execute_values")
+def test_psycopg2_execute_values(instrument, postgres_connection, elasticapm_client):
+    from psycopg2.extras import execute_values
+
+    cursor = postgres_connection.cursor()
+    try:
+        elasticapm_client.begin_transaction("web.django")
+        query = "INSERT INTO test VALUES %s"
+        data = tuple((i, "xxxxx") for i in range(999))
+        # this creates a long (~14000 characters) query, encoded as byte string.
+        # This tests that we shorten already-encoded strings.
+        execute_values(cursor, query, data, page_size=1000)
+        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 len(spans[0]["context"]["db"]["statement"]) == 10000, spans[0]["context"]["db"]["statement"]