| 7 | |
| 8 | |
| 9 | def statements_from_filehandle(file_h: IO) -> Generator[tuple[str, int], None, None]: |
| 10 | statements = '' |
| 11 | line_counter = 0 |
| 12 | batch_counter = 0 |
| 13 | for batch_text in file_h: |
| 14 | line_counter += 1 |
| 15 | if line_counter > MAX_MULTILINE_BATCH_STATEMENT: |
| 16 | raise ValueError(f'Saw single input statement greater than {MAX_MULTILINE_BATCH_STATEMENT} lines; assuming a parsing error.') |
| 17 | statements += batch_text |
| 18 | try: |
| 19 | tokens = sqlglot.tokenize(statements, read='mysql') |
| 20 | if not tokens: |
| 21 | continue |
| 22 | # we don't yet handle changing the delimiter within the batch input |
| 23 | if tokens[-1].text == ';': |
| 24 | # The advantage of sqlparse for splitting is that it preserves the input. |
| 25 | # https://github.com/tobymao/sqlglot/issues/2587#issuecomment-1823109501 |
| 26 | for statement in sqlparse.split(statements): |
| 27 | yield (statement, batch_counter) |
| 28 | batch_counter += 1 |
| 29 | statements = '' |
| 30 | line_counter = 0 |
| 31 | except sqlglot.errors.TokenError: |
| 32 | continue |
| 33 | if statements: |
| 34 | for statement in sqlparse.split(statements): |
| 35 | yield (statement, batch_counter) |
| 36 | batch_counter += 1 |