diff --git a/sql_metadata/generalizator.py b/sql_metadata/generalizator.py index f0639517..cdb8283a 100644 --- a/sql_metadata/generalizator.py +++ b/sql_metadata/generalizator.py @@ -9,6 +9,9 @@ import re +from sqlglot.errors import TokenError +from sqlglot.tokens import Tokenizer, TokenType + from sql_metadata.comments import strip_comments @@ -36,12 +39,33 @@ def __init__(self, sql: str = ""): self._raw_query = sql # SQL queries normalization (#16) + @staticmethod + def _normalize_strings(sql: str) -> str: + """Replace whole string tokens, including doubled quote escapes.""" + if "'" not in sql: + return sql + try: + tokens = Tokenizer().tokenize(sql) + except TokenError: + # Preserve best-effort normalization of lexically malformed SQL. + return re.sub(r"'[^\']*'", "X", sql) + + parts: list[str] = [] + previous_end = 0 + for token in tokens: + if token.token_type in (TokenType.STRING, TokenType.NATIONAL_STRING): + # Keep the N/n prefix, as the previous normalization did. + start = token.start + (token.token_type == TokenType.NATIONAL_STRING) + parts.extend((sql[previous_end:start], "X")) + previous_end = token.end + 1 + return "".join(parts) + sql[previous_end:] + @staticmethod def _normalize_likes(sql: str) -> str: """Normalise and collapse repeated ``LIKE`` clauses. - Strips ``%`` wildcards, replaces ``LIKE '...'`` with ``LIKE X``, - and collapses consecutive ``or/and ... LIKE X`` clauses into a + Strips ``%`` wildcards and collapses consecutive + ``or/and ... LIKE X`` clauses into a single instance with ``...`` suffix. :param sql: SQL string with LIKE clauses. @@ -51,9 +75,6 @@ def _normalize_likes(sql: str) -> str: """ sql = sql.replace("%", "") - # LIKE '%bot' - sql = re.sub(r"LIKE '[^\']+'", "LIKE X", sql) - # or all_groups LIKE X or all_groups LIKE X found = re.finditer(r"(or|and) [^\s]+ LIKE X", sql, flags=re.IGNORECASE) like_matches = [m.group(0) for m in found] @@ -86,9 +107,9 @@ def generalize(self) -> str: 1. Strip comments. 2. Remove double-quotes. 3. Collapse multiple spaces. - 4. Normalise ``LIKE`` clauses. - 5. Replace escaped characters. - 6. Replace string literals with ``X``. + 4. Replace escaped characters. + 5. Replace string literals with ``X``. + 6. Normalise ``LIKE`` clauses. 7. Collapse whitespace to single spaces. 8. Replace numbers with ``N``. 9. Collapse ``IN (...)`` / ``VALUES (...)`` lists to ``(XYZ)``. @@ -107,14 +128,11 @@ def generalize(self) -> str: # multiple spaces sql = re.sub(r"\s{2,}", " ", sql) - # handle LIKE statements - sql = self._normalize_likes(sql) - sql = re.sub(r"\\\\", "", sql) sql = re.sub(r"\\'", "", sql) sql = re.sub(r'\\"', "", sql) - sql = re.sub(r"'[^\']*'", "X", sql) - sql = re.sub(r'"[^\"]*"', "X", sql) + sql = self._normalize_strings(sql) + sql = self._normalize_likes(sql) # All newlines, tabs, etc replaced by single space sql = re.sub(r"\s+", " ", sql) diff --git a/test/test_generalize_escaped_strings.py b/test/test_generalize_escaped_strings.py new file mode 100644 index 00000000..09487de0 --- /dev/null +++ b/test/test_generalize_escaped_strings.py @@ -0,0 +1,20 @@ +import pytest + +from sql_metadata import Parser + + +@pytest.mark.parametrize("literal", ["'O''Reilly'", "'a''b''c'", "''''"]) +@pytest.mark.parametrize( + "prefix", + ["SELECT ", "SELECT N", "SELECT n", "SELECT name FROM authors WHERE name LIKE "], +) +def test_generalize_doubled_quotes_as_one_literal(literal: str, prefix: str) -> None: + assert Parser(prefix + literal).generalize == prefix + "X" + + +def test_generalize_separate_literals() -> None: + assert Parser("SELECT 'O''Reilly', '', 'ordinary'").generalize == "SELECT X, X, X" + + +def test_generalize_unterminated_literal_remains_best_effort() -> None: + assert Parser("SELECT 'unterminated").generalize == "SELECT 'unterminated"