|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | import re |
4 | | -from typing import Generator |
| 4 | +from typing import Any, Generator |
5 | 5 |
|
6 | 6 | import sqlglot |
7 | 7 | import sqlparse |
@@ -77,7 +77,7 @@ def is_subselect(parsed: TokenList) -> bool: |
77 | 77 | return False |
78 | 78 |
|
79 | 79 |
|
80 | | -def extract_from_part(parsed: TokenList, stop_at_punctuation: bool = True) -> Generator[str, None, None]: |
| 80 | +def extract_from_part(parsed: TokenList, stop_at_punctuation: bool = True) -> Generator[Any, None, None]: |
81 | 81 | tbl_prefix_seen = False |
82 | 82 | for item in parsed.tokens: |
83 | 83 | if tbl_prefix_seen: |
@@ -123,7 +123,7 @@ def extract_from_part(parsed: TokenList, stop_at_punctuation: bool = True) -> Ge |
123 | 123 | break |
124 | 124 |
|
125 | 125 |
|
126 | | -def extract_table_identifiers(token_stream: TokenList) -> Generator[tuple[str | None, str, str], None, None]: |
| 126 | +def extract_table_identifiers(token_stream: Generator[Any, None, None]) -> Generator[tuple[str | None, str, str], None, None]: |
127 | 127 | """yields tuples of (schema_name, table_name, table_alias)""" |
128 | 128 |
|
129 | 129 | for item in token_stream: |
@@ -187,15 +187,15 @@ def extract_tables_from_complete_statements(sql: str) -> list[tuple[str | None, |
187 | 187 | return [] |
188 | 188 |
|
189 | 189 | finely_parsed = [] |
190 | | - for statement in roughly_parsed: |
| 190 | + for rough_statement in roughly_parsed: |
191 | 191 | try: |
192 | | - finely_parsed.append(sqlglot.parse_one(str(statement), read='mysql')) |
| 192 | + finely_parsed.append(sqlglot.parse_one(str(rough_statement), read='mysql')) |
193 | 193 | except sqlglot.errors.ParseError: |
194 | 194 | pass |
195 | 195 |
|
196 | 196 | tables = [] |
197 | | - for statement in finely_parsed: |
198 | | - for identifier in statement.find_all(sqlglot.exp.Table): |
| 197 | + for fine_statement in finely_parsed: |
| 198 | + for identifier in fine_statement.find_all(sqlglot.exp.Table): |
199 | 199 | if identifier.parent_select and identifier.parent_select.sql().startswith('WITH'): |
200 | 200 | continue |
201 | 201 | tables.append(( |
|
0 commit comments