#366 plus fixed tests

elasticapm/instrumentation/packages/dbapi2.py

@@ -7,6 +7,7 @@
 from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
 from elasticapm.traces import capture_span
 from elasticapm.utils import compat, wrapt
+from elasticapm.utils.encoding import force_text
 
 
 class Literal(object):
@@ -59,7 +60,7 @@ def _scan_for_table_with_tokens(tokens, keyword):
 
 def tokenize(sql):
     # split on anything that is not a word character, excluding dots
-    return [t for t in re.split("([^\w.])", sql) if t != ""]
+    return [t for t in re.split(r"([^\w.])", sql) if t != ""]
 
 
 def scan(tokens):
@@ -123,6 +124,7 @@ def extract_signature(sql):
     :param sql: the SQL statement
     :return: a string representing the signature
     """
+    sql = force_text(sql)
     sql = sql.strip()
     first_space = sql.find(" ")
     if first_space < 0:

tests/instrumentation/dbapi2_tests.py

@@ -1,4 +1,4 @@
-from elasticapm.instrumentation.packages.dbapi2 import Literal, scan, tokenize
+from elasticapm.instrumentation.packages.dbapi2 import Literal, extract_signature, scan, tokenize
 
 
 def test_scan_simple():
@@ -39,3 +39,17 @@ def test_scan_double_quotes_at_end():
     actual = [t[1] for t in scan(tokens)]
     expected = ["Hello", "Peter", "Pan", "at", "Disney", Literal("'", "World")]
     assert actual == expected
+
+
+def test_extract_signature_string():
+    sql = "Hello 'Peter Pan' at Disney World"
+    actual = extract_signature(sql)
+    expected = "HELLO"
+    assert actual == expected
+
+
+def test_extract_signature_bytes():
+    sql = b"Hello 'Peter Pan' at Disney World"
+    actual = extract_signature(sql)
+    expected = "HELLO"
+    assert actual == expected

tests/instrumentation/mysql_tests.py

@@ -45,10 +45,10 @@ def test_select_with_difficult_values():
 
 
 def test_select_with_difficult_table_name():
-    sql = "SELECT id FROM `myta\n-æøåble` WHERE id = 2323" ""
+    sql = u"""SELECT id FROM `myta\n-æøåble` WHERE id = 2323"""
     actual = extract_signature(sql)
 
-    assert "SELECT FROM myta\n-æøåble" == actual
+    assert u"SELECT FROM myta\n-æøåble" == actual
 
 
 def test_select_subselect():

tests/instrumentation/psycopg2_tests.py

@@ -121,10 +121,10 @@ def test_select_with_dollar_quotes_custom_token():
 
 
 def test_select_with_difficult_table_name():
-    sql_statement = 'SELECT id FROM "myta\n-æøåble" WHERE id = 2323' ""
+    sql_statement = u"""SELECT id FROM "myta\n-æøåble" WHERE id = 2323"""
     actual = extract_signature(sql_statement)
 
-    assert "SELECT FROM myta\n-æøåble" == actual
+    assert u"SELECT FROM myta\n-æøåble" == actual
 
 
 def test_select_subselect():