get-objects-lib 0.1.0__py3-none-any.whl

get_objects_lib/dialect/parser/non_semicolon.py

@@ -0,0 +1,182 @@
+import typing as t
+from collections import defaultdict
+from typing import Generator
+
+from sqlglot import Parser, Token, exp, TokenType
+
+from get_objects_lib.dialect.parser.base import BaseParser
+
+
+class NonSemiColonParser(BaseParser):
+    @staticmethod
+    def iter_around(items: t.Sequence[Token]) -> Generator[tuple[int, Token | None, Token, Token | None]]:
+        for i, item in enumerate(items):
+            _prev = items[i - 1] if i > 0 else None
+            _next = items[i + 1] if i + 1 < len(items) else None
+            yield i, _prev, item, _next
+
+    def _parse(
+            self,
+            parse_method: t.Callable[[Parser], exp.Expr | None],
+            raw_tokens: list[Token],
+            sql: str | None = None,
+    ) -> list[exp.Expr | None]:
+        self.reset()
+        self.sql = sql or ""
+
+        chunks: list[list[Token]] = [[]]
+        raw_tokens = list(filter(lambda x: x.token_type != TokenType.SEMICOLON, raw_tokens))
+        levels = defaultdict(lambda: False)
+        begins = []
+
+        block_without_begin = False
+        # Profundidad de parentesis: un ")" interno no cierra el exterior
+        paren_depth = 0
+        updating = False
+        inserting = False
+        added = False
+        skip = 0
+        case_depth = 0
+        prev_closed_case = False
+        # Un IF/ELSE/WHILE espera su cuerpo: un BEGIN aqui es ese cuerpo
+        pending_body = False
+        # Tras WITH cte AS (...) la sentencia principal no abre otra
+        cte_pending = False
+
+        for i, _prev, token, _next in self.iter_around(raw_tokens):
+            if skip:
+                skip -= 1
+                continue
+
+            # GO cierra el lote: todo lo abierto termina aqui y GO [n] queda
+            # como su propia sentencia.
+            separator_length = self._batch_separator_length(raw_tokens, i)
+            if separator_length:
+                chunks.append(raw_tokens[i:i + separator_length])
+                chunks.append([])
+                skip = separator_length - 1
+                levels.clear()
+                begins.clear()
+                block_without_begin = updating = inserting = added = False
+                case_depth = paren_depth = 0
+                pending_body = False
+                continue
+
+            # Dentro de un CASE, ELSE y END son parte de la expresion, no
+            # cierran ni abren sentencias.
+            in_case = bool(case_depth)
+            closes_case = bool(case_depth) and self._is_end(token)
+            if token.token_type == TokenType.CASE:
+                case_depth += 1
+            elif closes_case:
+                case_depth -= 1
+
+            # FOR UPDATE de un cursor, el evento UPDATE de un trigger y la funcion
+            # UPDATE(columna) no son un UPDATE
+            if (
+                token.token_type == TokenType.UPDATE
+                and not (_prev and _prev.token_type == TokenType.FOR)
+                and self._is_open_token(token, _prev, _next)
+            ):
+                updating = True
+
+            # INSERT ... SELECT / INSERT ... EXEC: el primer SELECT o EXEC es el
+            # origen del INSERT, salvo que antes aparezca VALUES.
+            if token.token_type == TokenType.INSERT:
+                inserting = True
+            if token.token_type == TokenType.VALUES:
+                inserting = False
+
+            if token.token_type == TokenType.L_PAREN:
+                paren_depth += 1
+
+            if token.token_type == TokenType.R_PAREN and paren_depth:
+                paren_depth -= 1
+
+            if self._is_end(token) and not in_case:
+                levels[len(begins)] = False
+                chunks.append([])
+                if begins: begins.pop()
+
+            # Lo que sigue a "AS" es el cuerpo de un VIEW/PROCEDURE/FUNCTION,
+            # no una sentencia nueva.
+            # Lo mismo despues de FOR (DECLARE c CURSOR FOR SELECT / FOR UPDATE) y
+            # de THEN (MERGE ... WHEN MATCHED THEN UPDATE / INSERT / DELETE)
+            after_as = _prev is not None and _prev.token_type in (TokenType.ALIAS, TokenType.FOR, TokenType.THEN)
+
+            # BEGIN que nadie esperaba (ni IF/ELSE/WHILE ni AS): bloque suelto,
+            # abre una sentencia nueva
+            is_block_begin = self._is_begin(token, _next) and not in_case and not paren_depth
+            bare_begin = is_block_begin and not pending_body and not after_as
+            starts_cte = self._starts_cte(raw_tokens, i) and not paren_depth and not in_case
+            is_open = self._is_open_token(token, _prev, _next) or bare_begin or starts_cte
+
+            # SELECT/INSERT/UPDATE/DELETE/MERGE que usa el CTE: sigue en su sentencia
+            if (
+                cte_pending and is_open and not paren_depth and not in_case
+                and token.token_type in (TokenType.SELECT, TokenType.INSERT, TokenType.UPDATE, TokenType.DELETE, TokenType.MERGE)
+            ):
+                cte_pending = False
+                is_open = False
+                if token.token_type == TokenType.UPDATE:
+                    updating = True
+                if token.token_type == TokenType.INSERT:
+                    inserting = True
+
+            if is_open and not paren_depth and not after_as and not in_case:
+                if block_without_begin:
+                    block_without_begin = False
+                    begins.append(token)
+                    added = True
+
+                if updating and token.token_type == TokenType.SET:
+                    updating = False
+                elif inserting and token.token_type in (TokenType.SELECT, TokenType.EXECUTE):
+                    inserting = False
+                else:
+                    if token.token_type != TokenType.INSERT:
+                        inserting = False
+
+                    level = len(begins)
+                    if levels[level]:
+                        # END ELSE de un bloque se queda junto; el END de un CASE no
+                        block_end = _prev.token_type == TokenType.END and not prev_closed_case
+                        if not block_end or token.token_type != TokenType.ELSE:
+                            chunks.append([])
+
+                    levels[level] = True
+
+            if starts_cte:
+                cte_pending = True
+
+            starts_body = (
+                self._is_word(token, 'IF', 'ELSE', 'WHILE')
+                and not in_case
+                and self._is_open_token(token, _prev, _next)  # no el IF de DROP TABLE IF EXISTS
+            )
+            if starts_body:
+                block_without_begin = self._is_block_without_begin(raw_tokens, i)
+
+            if starts_body:
+                pending_body = True
+            elif is_block_begin or (is_open and not paren_depth and not in_case):
+                pending_body = False
+
+            chunks[-1].append(token)
+
+            if self._is_begin(token, _next):
+                begins.append(token)
+
+            prev_closed_case = closes_case
+
+            _next_next = raw_tokens[i + 2] if i + 2 < len(raw_tokens) else None
+            # Si seguimos dentro de un CASE, el siguiente token es parte de el
+            if added and _next and not case_depth and (self._is_end(_next) or self._is_open_token(_next, token, _next_next)):
+                levels[len(begins)] = False
+                added = False
+                begins.pop()
+
+        chunks = list(filter(lambda x: x, chunks))
+        self._chunks = chunks
+
+        return self._parse_batch_statements(parse_method=parse_method, sep_first_statement=False, top_level=True)