Use cursor.mogrify when available (getsentry#203)

mlennon-lumere · untitaker · commit 807c31932307 · 2019-03-21T16:18:13.000+01:00
* Use cursor.mogrify when available instead of trying to reimplement sql formatting * fix: Port changes from getsentry#202 * fix: Formatting * adds test for psycopg2 methods * use old style django postgres backend * have travis install test database as well for older versions of django * fix: Avoid bytestrings, refactor postgres setup * fix: Formatting
diff --git a/.travis.yml b/.travis.yml
@@ -7,6 +7,10 @@ python:
   - "3.5"
   - "3.6"
 
+env:
+  - SENTRY_PYTHON_TEST_POSTGRES_USER=postgres
+  - SENTRY_PYTHON_TEST_POSTGRES_NAME=travis_ci_test
+
 cache:
   pip: true
   cargo: true
@@ -48,6 +52,13 @@ matrix:
         - source $HOME/osx-py/bin/activate
         - export TRAVIS_PYTHON_VERSION=2.7
 
+before_script:
+ - psql -c 'create database travis_ci_test;' -U postgres
+ - psql -c 'create database test_travis_ci_test;' -U postgres
+
+services:
+  - postgresql
+
 install:
   - pip install tox
   - pip install codecov
diff --git a/sentry_sdk/integrations/django/__init__.py b/sentry_sdk/integrations/django/__init__.py
@@ -312,7 +312,6 @@ def format_sql(sql, params):
         # convert sql with named parameters to sql with unnamed parameters
         conv = _FormatConverter(params)
         if params:
-            sql = sql_to_string(sql)
             sql = sql % conv
             params = conv.params
         else:
@@ -327,20 +326,27 @@ def format_sql(sql, params):
     return sql, rv
 
 
-def record_sql(sql, params):
-    # type: (Any, Any) -> None
+def record_sql(sql, params, cursor=None):
+    # type: (Any, Any, Any) -> None
     hub = Hub.current
     if hub.get_integration(DjangoIntegration) is None:
         return
-    real_sql, real_params = format_sql(sql, params)
 
-    if real_params:
-        try:
-            real_sql = format_and_strip(real_sql, real_params)
-        except Exception:
-            pass
+    with capture_internal_exceptions():
+        if cursor and hasattr(cursor, "mogrify"):  # psycopg2
+            real_sql = cursor.mogrify(sql, params)
+            with capture_internal_exceptions():
+                if isinstance(real_sql, bytes):
+                    real_sql = real_sql.decode(cursor.connection.encoding)
+        else:
+            real_sql, real_params = format_sql(sql, params)
 
-    hub.add_breadcrumb(message=real_sql, category="query")
+            if real_params:
+                try:
+                    real_sql = format_and_strip(real_sql, real_params)
+                except Exception:
+                    pass
+        hub.add_breadcrumb(message=real_sql, category="query")
 
 
 def install_sql_hook():
@@ -358,21 +364,21 @@ def install_sql_hook():
         # This won't work on Django versions < 1.6
         return
 
-    def record_many_sql(sql, param_list):
+    def record_many_sql(sql, param_list, cursor):
         for params in param_list:
-            record_sql(sql, params)
+            record_sql(sql, params, cursor)
 
     def execute(self, sql, params=None):
         try:
             return real_execute(self, sql, params)
         finally:
-            record_sql(sql, params)
+            record_sql(sql, params, self.cursor)
 
     def executemany(self, sql, param_list):
         try:
             return real_executemany(self, sql, param_list)
         finally:
-            record_many_sql(sql, param_list)
+            record_many_sql(sql, param_list, self.cursor)
 
     CursorWrapper.execute = execute
     CursorWrapper.executemany = executemany
diff --git a/tests/integrations/django/myapp/settings.py b/tests/integrations/django/myapp/settings.py
@@ -94,6 +94,15 @@ def process_response(self, request, response):
 
 DATABASES = {"default": {"ENGINE": "django.db.backends.sqlite3", "NAME": ":memory:"}}
 
+try:
+    DATABASES["postgres"] = {
+        "ENGINE": "django.db.backends.postgresql_psycopg2",
+        "NAME": os.environ["SENTRY_PYTHON_TEST_POSTGRES_NAME"],
+        "USER": os.environ["SENTRY_PYTHON_TEST_POSTGRES_USER"],
+    }
+except KeyError:
+    pass
+
 
 # Password validation
 # https://docs.djangoproject.com/en/2.0/ref/settings/#auth-password-validators
diff --git a/tests/integrations/django/test_basic.py b/tests/integrations/django/test_basic.py
@@ -3,7 +3,7 @@
 from werkzeug.test import Client
 from django.contrib.auth.models import User
 from django.core.management import execute_from_command_line
-from django.db.utils import OperationalError
+from django.db.utils import OperationalError, ProgrammingError, DataError
 
 
 try:
@@ -186,49 +186,95 @@ def test_sql_queries(sentry_init, capture_events):
 @pytest.mark.django_db
 def test_sql_dict_query_params(sentry_init, capture_events):
     sentry_init(integrations=[DjangoIntegration()], send_default_pii=True)
-    from django.db import connection
 
-    sql = connection.cursor()
+    from django.db import connections
+
+    if "postgres" not in connections:
+        pytest.skip("postgres tests disabled")
+
+    sql = connections["postgres"].cursor()
 
     events = capture_events()
-    with pytest.raises(OperationalError):
-        # This really only works with postgres. sqlite will crash with syntax error
+    with pytest.raises(ProgrammingError):
         sql.execute(
             """SELECT count(*) FROM people_person WHERE foo = %(my_foo)s""",
             {"my_foo": 10},
         )
 
     capture_message("HI")
-
     event, = events
 
     crumb, = event["breadcrumbs"]
     assert crumb["message"] == ("SELECT count(*) FROM people_person WHERE foo = 10")
 
 
+@pytest.mark.parametrize(
+    "query",
+    [
+        lambda sql: sql.SQL("SELECT %(my_param)s FROM {mytable}").format(
+            mytable=sql.Identifier("foobar")
+        ),
+        lambda sql: sql.SQL('SELECT %(my_param)s FROM "foobar"'),
+    ],
+)
 @pytest.mark.django_db
-def test_sql_psycopg2_string_composition(sentry_init, capture_events):
+def test_sql_psycopg2_string_composition(sentry_init, capture_events, query):
     sentry_init(integrations=[DjangoIntegration()], send_default_pii=True)
-    from django.db import connection
+    from django.db import connections
 
-    psycopg2 = pytest.importorskip("psycopg2")
-    psycopg2_sql = psycopg2.sql
+    if "postgres" not in connections:
+        pytest.skip("postgres tests disabled")
 
-    sql = connection.cursor()
+    import psycopg2
+
+    sql = connections["postgres"].cursor()
 
     events = capture_events()
-    with pytest.raises(TypeError):
-        # crashes because we use sqlite
-        sql.execute(
-            psycopg2_sql.SQL("SELECT %(my_param)s FROM people_person"), {"my_param": 10}
-        )
+    with pytest.raises(ProgrammingError):
+        sql.execute(query(psycopg2.sql), {"my_param": 10})
 
     capture_message("HI")
 
     event, = events
-
     crumb, = event["breadcrumbs"]
-    assert crumb["message"] == ("SELECT 10 FROM people_person")
+    assert crumb["message"] == ('SELECT 10 FROM "foobar"')
+
+
+@pytest.mark.django_db
+def test_sql_psycopg2_placeholders(sentry_init, capture_events):
+    sentry_init(integrations=[DjangoIntegration()], send_default_pii=True)
+    from django.db import connections
+
+    if "postgres" not in connections:
+        pytest.skip("postgres tests disabled")
+
+    import psycopg2
+
+    sql = connections["postgres"].cursor()
+
+    events = capture_events()
+    with pytest.raises(DataError):
+        names = ["foo", "bar"]
+        identifiers = [psycopg2.sql.Identifier(name) for name in names]
+        placeholders = [
+            psycopg2.sql.Placeholder(var) for var in ["first_var", "second_var"]
+        ]
+        sql.execute("create table my_test_table (foo text, bar date)")
+
+        query = psycopg2.sql.SQL("insert into my_test_table ({}) values ({})").format(
+            psycopg2.sql.SQL(", ").join(identifiers),
+            psycopg2.sql.SQL(", ").join(placeholders),
+        )
+        sql.execute(query, {"first_var": "fizz", "second_var": "not a date"})
+
+    capture_message("HI")
+
+    event, = events
+    crumb1, crumb2 = event["breadcrumbs"]
+    assert crumb1["message"] == ("create table my_test_table (foo text, bar date)")
+    assert crumb2["message"] == (
+        """insert into my_test_table ("foo", "bar") values ('fizz', 'not a date')"""
+    )
 
 
 @pytest.mark.django_db
diff --git a/tox.ini b/tox.ini
@@ -126,6 +126,8 @@ passenv =
     SENTRY_PYTHON_TEST_AWS_ACCESS_KEY_ID
     SENTRY_PYTHON_TEST_AWS_SECRET_ACCESS_KEY
     SENTRY_PYTHON_TEST_AWS_IAM_ROLE
+    SENTRY_PYTHON_TEST_POSTGRES_USER
+    SENTRY_PYTHON_TEST_POSTGRES_NAME
 usedevelop = True
 extras =
     flask: flask
