sqlglot.anonymize

  1from __future__ import annotations
  2
  3import string
  4
  5from sqlglot.dialects.dialect import Dialect, DialectType
  6from sqlglot.errors import TokenError
  7from sqlglot.tokens import Token, TokenType
  8
  9
 10ALPHABET = string.ascii_lowercase
 11ALPHABET_SIZE = len(ALPHABET)
 12ANONYMIZED_TYPES = {
 13    TokenType.BIT_STRING,
 14    TokenType.BYTE_STRING,
 15    TokenType.HEX_STRING,
 16    TokenType.HEREDOC_STRING,
 17    TokenType.IDENTIFIER,
 18    TokenType.NATIONAL_RAW_STRING,
 19    TokenType.NATIONAL_STRING,
 20    TokenType.NUMBER,
 21    TokenType.RAW_STRING,
 22    TokenType.STRING,
 23    TokenType.UNICODE_STRING,
 24    TokenType.VAR,
 25}
 26QUOTED_TYPES = ANONYMIZED_TYPES - {TokenType.NUMBER, TokenType.VAR}
 27REWRITTEN_TYPES = {TokenType.HINT, TokenType.UNKNOWN}
 28
 29
 30def anonymize(
 31    sql_or_tokens: list[Token] | str,
 32    dialect: DialectType = None,
 33) -> list[Token]:
 34    """Replaces sensitive tokens (identifiers, strings, numbers) with fixed-width,
 35    length-preserving, consistent aliases, and blanks out comments and hint bodies. When a
 36    SQL string is given, it is tokenized with `dialect` first; any un-tokenized remainder
 37    (e.g. an unterminated literal) is appended as a blanked UNKNOWN token. Mutates and
 38    returns `sql_or_tokens`.
 39
 40    Args:
 41        sql_or_tokens: The SQL string to anonymize, or its token list.
 42        dialect: The dialect used to tokenize a SQL string.
 43    """
 44    dialect = Dialect.get_or_raise(dialect)
 45    tokenizer_class = dialect.tokenizer_class
 46    parser_class = dialect.parser_class
 47
 48    errored = False
 49    if isinstance(sql_or_tokens, str):
 50        sql = sql_or_tokens
 51        tokenizer = dialect.tokenizer()
 52        try:
 53            tokens = tokenizer.tokenize(sql)
 54        except TokenError:
 55            tokens = tokenizer.tokens
 56            errored = True
 57    else:
 58        tokens = sql_or_tokens
 59        sql = None
 60
 61    hint_start = tokenizer_class.HINT_START
 62    hint_end = tokenizer_class._COMMENTS.get(hint_start)
 63    nested = tokenizer_class.NESTED_COMMENTS
 64
 65    seen: dict[tuple[bool, str], str] = {}
 66    counter = 0
 67
 68    for i, token in enumerate(tokens):
 69        token.comments = [_blank(comment) for comment in token.comments]
 70
 71        if token.token_type == TokenType.HINT:
 72            # A hint's text is the whole /*+ ... */ comment, so its body is blanked as well
 73            text, stop = _blank_comment(token.text, 0, hint_start, hint_end, nested)
 74            token.text = text + _blank(token.text[stop:])
 75            continue
 76
 77        if (
 78            token.token_type == TokenType.VAR
 79            and i + 1 < len(tokens)
 80            and tokens[i + 1].token_type == TokenType.L_PAREN
 81        ):
 82            # A function name can live in either registry, e.g. JSON_OBJECT is only in
 83            # FUNCTION_PARSERS. They're consulted separately to avoid building their union
 84            name = token.text.upper()
 85            if name in parser_class.FUNCTIONS or name in parser_class.FUNCTION_PARSERS:
 86                continue
 87        if token.token_type not in ANONYMIZED_TYPES:
 88            continue
 89        if not token.text:
 90            continue
 91
 92        is_number = token.token_type == TokenType.NUMBER
 93        key = (is_number, token.text)
 94        alias = seen.get(key)
 95        if alias is None:
 96            seen[key] = (
 97                _number_alias(counter, token.text) if is_number else _alias(counter, token.text)
 98            )
 99            counter += 1
100        token.text = seen[key]
101
102    if sql is not None and errored:
103        start = tokens[-1].end + 1 if tokens else 0
104        length = len(sql)
105        while start < length and sql[start].isspace():
106            start += 1
107        if start < length:
108            # The first two characters are kept so that the delimiter the tokenizer choked
109            # on is still visible, e.g. 'u, /*, $$, `u
110            tokens.append(
111                Token(
112                    TokenType.UNKNOWN,
113                    sql[start : start + 2] + "." * (length - start - 2),
114                    start=start,
115                    end=length - 1,
116                    line=sql.count("\n", 0, start) + 1,
117                    col=start - sql.rfind("\n", 0, start),
118                )
119            )
120
121    return tokens
122
123
124def render(sql: str, tokens: list[Token], dialect: DialectType = None) -> str:
125    """Recreates the (anonymized) SQL string from the original `sql` and token positions.
126
127    Every token is rendered over its own source span, so the result is always as long as
128    `sql`: a token that wasn't anonymized is emitted verbatim, keeping the spelling the
129    tokenizer normalized away, and an anonymized one has its alias fitted between the
130    quotes of its span. The gaps between tokens hold only whitespace and comments, so
131    they're redacted rather than reconstructed: comment markers and whitespace survive,
132    everything else is blanked. That covers anything the tokenizer didn't reach as well,
133    which is a trailing gap whenever `anonymize` wasn't the one to tokenize `sql`.
134
135    Args:
136        sql: The original SQL string.
137        tokens: The anonymized tokens to render.
138        dialect: The dialect used to identify comments and quotes in `sql`.
139    """
140    tokenizer_class = Dialect.get_or_raise(dialect).tokenizer_class
141    comments = sorted(tokenizer_class._COMMENTS.items(), key=lambda c: len(c[0]), reverse=True)
142    nested = tokenizer_class.NESTED_COMMENTS
143    quotes = sorted(
144        {
145            **tokenizer_class._QUOTES,
146            **{start: end for start, (end, _) in tokenizer_class._FORMAT_STRINGS.items()},
147            **tokenizer_class._IDENTIFIERS,
148        }.items(),
149        key=lambda quote: (len(quote[0]), len(quote[1])),
150        reverse=True,
151    )
152
153    result = []
154    prev = 0
155
156    for token in tokens:
157        result.append(_redact(sql[prev : token.start], comments, nested))
158        prev = token.end + 1
159        span = sql[token.start : prev]
160        if token.token_type in REWRITTEN_TYPES:
161            # Blanked in place by `anonymize`, or synthesized by it, so already span-shaped
162            result.append(token.text)
163        elif token.token_type in ANONYMIZED_TYPES:
164            result.append(_fit(span, token.text, quotes, token.token_type))
165        else:
166            result.append(span)
167
168    result.append(_redact(sql[prev:], comments, nested))
169
170    return "".join(result)
171
172
173def _alias(counter: int, text: str) -> str:
174    digits = []
175    while counter:
176        counter, digit = divmod(counter, ALPHABET_SIZE)
177        digits.append(ALPHABET[digit])
178
179    letters = "".join(reversed(digits)).rjust(sum(not char.isspace() for char in text), "a")
180
181    alias = []
182    i = 0
183    for char in text:
184        if char.isspace():
185            alias.append(char)
186        else:
187            alias.append(letters[i])
188            i += 1
189
190    return "".join(alias)
191
192
193def _number_alias(counter: int, text: str) -> str:
194    if len(text) > 4000:
195        digits_seen = False
196        blanked = []
197        for char in text:
198            if char.isdigit():
199                blanked.append("0" if digits_seen else "1")
200                digits_seen = True
201            else:
202                blanked.append(char)
203        return "".join(blanked)
204
205    sep = "e" if "e" in text else ("E" if "E" in text else "")
206    mantissa, _, exponent = text.partition(sep) if sep else (text, "", "")
207    sign = ""
208    if exponent.startswith(("-", "+")):
209        sign, exponent = exponent[0], exponent[1:]
210    exponent_length = len(exponent)
211
212    integer, dot, fraction = mantissa.partition(".")
213    integer_length = len(integer)
214    digits = integer_length + len(fraction)
215    mantissa_value = 10 ** (digits - 1) + counter % (9 * 10 ** (digits - 1))
216    result = str(mantissa_value)
217    if dot:
218        result = result[:integer_length] + "." + result[integer_length:]
219    if sep:
220        # The exponent marker is kept even when there are no digits after it, e.g. 1e
221        result += sep + sign
222        if exponent_length:
223            exponent_value = 10 ** (exponent_length - 1) + (counter // (9 * 10 ** (digits - 1))) % (
224                9 * 10 ** (exponent_length - 1)
225            )
226            result += str(exponent_value)
227
228    return result
229
230
231def _fit(span: str, alias: str, quotes: list[tuple[str, str]], token_type: TokenType) -> str:
232    """Fits `alias` into the quoted region of `span`, so the two are always the same length."""
233    start = end = 0
234
235    if token_type in QUOTED_TYPES:
236        for open_quote, close_quote in quotes:
237            if (
238                len(span) >= len(open_quote) + len(close_quote)
239                and span.startswith(open_quote)
240                and span.endswith(close_quote)
241            ):
242                start, end = len(open_quote), len(close_quote)
243
244                if token_type == TokenType.HEREDOC_STRING:
245                    # A heredoc's tag is part of its delimiter, e.g. $tag$body$tag$
246                    open_end = span.find(close_quote, start)
247                    close_start = span.rfind(open_quote, 0, len(span) - end)
248                    if start <= open_end < close_start:
249                        start, end = open_end + len(close_quote), len(span) - close_start
250
251                break
252
253    width = len(span) - start - end
254    pad = "0" if token_type == TokenType.NUMBER else "a"
255
256    return span[:start] + alias[:width].rjust(width, pad) + span[len(span) - end :]
257
258
259def _blank(sql: str) -> str:
260    return "".join(char if char.isspace() else "." for char in sql)
261
262
263def _blank_comment(sql: str, i: int, start: str, end: str | None, nested: bool) -> tuple[str, int]:
264    """Blanks the body of the comment at `i`, returning its text and the index past it."""
265    body = i + len(start)
266
267    if not end:
268        stop = sql.find("\n", body)
269        stop = len(sql) if stop == -1 else stop
270        return start + _blank(sql[body:stop]), stop
271
272    depth = 1
273    j = body
274    while j < len(sql):
275        if nested and sql.startswith(start, j):
276            depth += 1
277            j += len(start)
278        elif sql.startswith(end, j):
279            j += len(end)
280            depth -= 1
281            if not depth:
282                return start + _blank(sql[body : j - len(end)]) + end, j
283        else:
284            j += 1
285
286    return start + _blank(sql[body:]), len(sql)
287
288
289def _redact(sql: str, comments: list[tuple[str, str | None]], nested: bool) -> str:
290    if not sql or sql.isspace():
291        return sql
292
293    result = []
294    i = 0
295    length = len(sql)
296
297    while i < length:
298        for start, end in comments:
299            if sql.startswith(start, i):
300                text, i = _blank_comment(sql, i, start, end, nested)
301                result.append(text)
302                break
303        else:
304            char = sql[i]
305            result.append(char if char.isspace() else ".")
306            i += 1
307
308    return "".join(result)