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