diff --git a/elasticapm/instrumentation/packages/dbapi2.py b/elasticapm/instrumentation/packages/dbapi2.py index dd75d4e11..059a41380 100644 --- a/elasticapm/instrumentation/packages/dbapi2.py +++ b/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: diff --git a/tests/instrumentation/dbapi2_tests.py b/tests/instrumentation/dbapi2_tests.py index 54d0b10c1..3814d612a 100644 --- a/tests/instrumentation/dbapi2_tests.py +++ b/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 diff --git a/tests/instrumentation/mysql_tests.py b/tests/instrumentation/mysql_tests.py index 945b2b321..a5a70d800 100644 --- a/tests/instrumentation/mysql_tests.py +++ b/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(): diff --git a/tests/instrumentation/psycopg2_tests.py b/tests/instrumentation/psycopg2_tests.py index a9270a7c4..ab7496e5e 100644 --- a/tests/instrumentation/psycopg2_tests.py +++ b/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():