|
| 1 | +"""Dialect compatibility normalization for Polyglot analysis input.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import re |
| 6 | + |
| 7 | +from sqlbuild.compiler.sql_analysis._helpers.scanning import ( |
| 8 | + skip_block_comment_impl, |
| 9 | + skip_line_comment_impl, |
| 10 | + skip_quoted_text_impl, |
| 11 | +) |
| 12 | +from sqlbuild.compiler.sql_analysis.constants import ( |
| 13 | + SNOWFLAKE_DIALECT_NAME, |
| 14 | + SQL_CLOSE_BRACKET, |
| 15 | + SQL_DOLLAR_QUOTE_DELIMITER, |
| 16 | + SQL_ESCAPE_CHARACTER, |
| 17 | + SQL_OPEN_BRACKET, |
| 18 | + SQL_QUOTE_CHARACTERS, |
| 19 | + SQL_STRING_QUOTE_CHARACTER, |
| 20 | +) |
| 21 | + |
| 22 | +_IDENTIFIER: str = r'(?:[A-Za-z_][A-Za-z0-9_$]*|"(?:""|[^"])+")' |
| 23 | +_VARIANT_PATH_SUFFIX: re.Pattern[str] = re.compile( |
| 24 | + rf"(?P<base>{_IDENTIFIER}(?:\.{_IDENTIFIER})*:{_IDENTIFIER}" |
| 25 | + rf"(?:(?::|\.){_IDENTIFIER})*)$" |
| 26 | +) |
| 27 | +_SIMPLE_BRACKET_KEY: re.Pattern[str] = re.compile(rf"\s*(?:{_IDENTIFIER}|[0-9]+|'(?:''|[^'])*')\s*") |
| 28 | +_SCAN_SPECIAL: re.Pattern[str] = re.compile(r"[-/$'\"`\[]") |
| 29 | + |
| 30 | + |
| 31 | +def normalize_sql_for_polyglot_impl(*, sql: str, dialect: str | None) -> str: |
| 32 | + """Return analysis-equivalent SQL for Polyglot's supported dialect surface.""" |
| 33 | + |
| 34 | + if dialect is None or dialect.casefold() != SNOWFLAKE_DIALECT_NAME: |
| 35 | + return sql |
| 36 | + replacements: list[tuple[int, int, str]] = [] |
| 37 | + index: int = 0 |
| 38 | + while index < len(sql): |
| 39 | + special: re.Match[str] | None = _SCAN_SPECIAL.search(sql, index) |
| 40 | + if special is None: |
| 41 | + break |
| 42 | + index = special.start() |
| 43 | + if sql.startswith("--", index): |
| 44 | + index = skip_line_comment_impl(sql=sql, start=index) |
| 45 | + continue |
| 46 | + if sql.startswith("/*", index): |
| 47 | + index = skip_block_comment_impl( |
| 48 | + sql=sql, |
| 49 | + start=index, |
| 50 | + context="Snowflake SQL analysis", |
| 51 | + ) |
| 52 | + continue |
| 53 | + if sql.startswith(SQL_DOLLAR_QUOTE_DELIMITER, index): |
| 54 | + index = _skip_dollar_quoted_text(sql=sql, start=index) |
| 55 | + continue |
| 56 | + if sql[index] in SQL_QUOTE_CHARACTERS: |
| 57 | + index = _skip_snowflake_quoted_text(sql=sql, start=index) |
| 58 | + continue |
| 59 | + close_bracket: int | None = _matching_bracket(sql=sql, open_bracket=index) |
| 60 | + if close_bracket is None: |
| 61 | + index += 1 |
| 62 | + continue |
| 63 | + key_sql: str = sql[index + 1 : close_bracket] |
| 64 | + normalized_key: str = normalize_sql_for_polyglot_impl(sql=key_sql, dialect=dialect) |
| 65 | + base_end: int = _trailing_trivia_start(sql=sql, end=index) |
| 66 | + base_match: re.Match[str] | None = _VARIANT_PATH_SUFFIX.search(sql[:base_end]) |
| 67 | + if base_match is not None and _SIMPLE_BRACKET_KEY.fullmatch(key_sql) is None: |
| 68 | + base_sql: str = base_match.group("base") |
| 69 | + trivia: str = sql[base_end:index] |
| 70 | + replacements.append( |
| 71 | + ( |
| 72 | + base_match.start("base"), |
| 73 | + close_bracket + 1, |
| 74 | + f"GET({base_sql}{trivia}, {normalized_key})", |
| 75 | + ) |
| 76 | + ) |
| 77 | + elif normalized_key != key_sql: |
| 78 | + replacements.append((index + 1, close_bracket, normalized_key)) |
| 79 | + index = close_bracket + 1 |
| 80 | + normalized: str = sql |
| 81 | + for start, end, replacement in reversed(replacements): |
| 82 | + normalized = f"{normalized[:start]}{replacement}{normalized[end:]}" |
| 83 | + return normalized |
| 84 | + |
| 85 | + |
| 86 | +def _matching_bracket(*, sql: str, open_bracket: int) -> int | None: |
| 87 | + depth: int = 1 |
| 88 | + index: int = open_bracket + 1 |
| 89 | + while index < len(sql): |
| 90 | + if sql.startswith("--", index): |
| 91 | + index = skip_line_comment_impl(sql=sql, start=index) |
| 92 | + continue |
| 93 | + if sql.startswith("/*", index): |
| 94 | + index = skip_block_comment_impl( |
| 95 | + sql=sql, |
| 96 | + start=index, |
| 97 | + context="Snowflake SQL analysis", |
| 98 | + ) |
| 99 | + continue |
| 100 | + if sql.startswith(SQL_DOLLAR_QUOTE_DELIMITER, index): |
| 101 | + index = _skip_dollar_quoted_text(sql=sql, start=index) |
| 102 | + continue |
| 103 | + if sql[index] in SQL_QUOTE_CHARACTERS: |
| 104 | + index = _skip_snowflake_quoted_text(sql=sql, start=index) |
| 105 | + continue |
| 106 | + if sql[index] == SQL_OPEN_BRACKET: |
| 107 | + depth += 1 |
| 108 | + elif sql[index] == SQL_CLOSE_BRACKET: |
| 109 | + depth -= 1 |
| 110 | + if depth == 0: |
| 111 | + return index |
| 112 | + index += 1 |
| 113 | + return None |
| 114 | + |
| 115 | + |
| 116 | +def _skip_snowflake_quoted_text(*, sql: str, start: int) -> int: |
| 117 | + if sql[start] != SQL_STRING_QUOTE_CHARACTER: |
| 118 | + return skip_quoted_text_impl(sql=sql, start=start, context="Snowflake SQL analysis") |
| 119 | + index: int = start + 1 |
| 120 | + while index < len(sql): |
| 121 | + if sql[index] == SQL_ESCAPE_CHARACTER: |
| 122 | + index += 2 |
| 123 | + continue |
| 124 | + if sql[index] == SQL_STRING_QUOTE_CHARACTER: |
| 125 | + if index + 1 < len(sql) and sql[index + 1] == SQL_STRING_QUOTE_CHARACTER: |
| 126 | + index += 2 |
| 127 | + continue |
| 128 | + return index + 1 |
| 129 | + index += 1 |
| 130 | + return len(sql) |
| 131 | + |
| 132 | + |
| 133 | +def _skip_dollar_quoted_text(*, sql: str, start: int) -> int: |
| 134 | + closing: int = sql.find(SQL_DOLLAR_QUOTE_DELIMITER, start + len(SQL_DOLLAR_QUOTE_DELIMITER)) |
| 135 | + return len(sql) if closing == -1 else closing + len(SQL_DOLLAR_QUOTE_DELIMITER) |
| 136 | + |
| 137 | + |
| 138 | +def _trailing_trivia_start(*, sql: str, end: int) -> int: |
| 139 | + current: int = end |
| 140 | + while True: |
| 141 | + trimmed: int = current |
| 142 | + while trimmed > 0 and sql[trimmed - 1].isspace(): |
| 143 | + trimmed -= 1 |
| 144 | + if not sql[:trimmed].endswith("*/"): |
| 145 | + return trimmed |
| 146 | + opening: int = sql.rfind("/*", 0, trimmed - 2) |
| 147 | + if opening == -1: |
| 148 | + return trimmed |
| 149 | + current = opening |
0 commit comments