sqlleaf 0.0.1__py3-none-any.whl

sqlleaf/__init__.py

@@ -0,0 +1,4 @@
+from sqlleaf.lineage import Lineage as Lineage
+from sqlleaf.processors.collector.collector import collect_queries as collect_queries
+from sqlleaf.processors.generator.generator import generate_lineage_for_query as generate_lineage_for_query
+from sqlleaf.processors.transformer.transformer import transform_query as transform_query

sqlleaf/dialects/plpgsql/p_classes.py

@@ -0,0 +1,275 @@
+from __future__ import annotations
+
+from sqlglot import exp
+
+
+class PGBlock(exp.Block):
+    """A BEGIN ... END block."""
+    arg_types = {
+        "expressions": False,
+        "begin": False,
+        "declare": False,
+        "exception": False,
+        "label": False,
+    }
+
+
+class PGPerform(exp.Select):
+    """A PERFORM statement."""
+
+
+class PGDeclareItem(exp.DeclareItem):
+    """A variable declaration item inside a DECLARE section."""
+    arg_types = {
+        **exp.DeclareItem.arg_types,
+        "expression": False,
+        "value": False,
+        "assign": False,
+        "constant": False,
+        "collate": False,
+        "not_null": False,
+        "alias_for": False,
+        "cursor": False,
+        "scroll": False,
+        "cursor_args": False,
+        "query": False,
+    }
+
+
+class PGTypeRef(exp.Expression):
+    """A `<reference>%TYPE` or `<reference>%ROWTYPE` type expression used in DECLARE items."""
+    arg_types = {"this": True, "array": False, "rowtype": False}
+
+
+class PGCursorArg(exp.Expression):
+    """A cursor argument declaration."""
+    arg_types = {"this": True, "kind": True}
+
+
+class PGException(exp.Expression):
+    """An EXCEPTION section."""
+    arg_types = {"whens": True}
+
+
+class PGWhen(exp.Expression):
+    """A WHEN ... THEN ... entry inside an EXCEPTION section."""
+    arg_types = {"condition": True, "then": True}
+
+
+class PGCase(exp.CaseStatement):
+    """A CASE statement."""
+
+
+class PGIfBranch(exp.Expression):
+    """An IF/ELSIF branch."""
+    arg_types = {"condition": True, "then": True}
+
+
+class PGIf(exp.CaseStatement):
+    """An IF statement."""
+
+
+class PGLoop(exp.LoopBlock):
+    """A LOOP statement."""
+
+    arg_types = {
+        "body": True,
+        "label": False,
+        "end_label": False,
+    }
+
+
+class PGForIn(exp.Expression):
+    """A FOR statement."""
+    arg_types = {
+        "this": True,
+        "query": False,
+        "reverse": False,
+        "start": False,
+        "end": False,
+        "step": False,
+        "expressions": False,
+        "label": False,
+    }
+
+
+class PGCursorCall(exp.Expression):
+    """A bound cursor invocation."""
+    arg_types = {
+        "this": True,
+        "expressions": False,
+    }
+
+
+class PGForEach(exp.Expression):
+    """A FOREACH statement."""
+    arg_types = {
+        "this": True,
+        "slice": False,
+        "expression": True,
+        "expressions": False,
+        "label": False,
+    }
+
+
+class PGExit(exp.Leave):
+    """An EXIT statement."""
+    arg_types = {"this": False, "when": False}
+
+
+class PGContinue(exp.Iterate):
+    """A CONTINUE statement.
+    """
+    arg_types = {"this": False, "when": False}
+
+
+class PGRaise(exp.Expression):
+    """A RAISE statement."""
+
+    arg_types = {
+        "level": False,
+        "message": False,
+        "expressions": False,
+        "condition": False,
+        "sqlstate": False,
+        "using": False,
+    }
+
+
+class PGAssert(exp.Expression):
+    """An ASSERT statement."""
+    arg_types = {"condition": True, "message": False}
+
+
+class PGGetDiagnostics(exp.Expression):
+    """A GET DIAGNOSTICS statement."""
+    arg_types = {"current": False, "stacked": False, "expressions": True}
+
+
+class PGOpenCursor(exp.Expression):
+    """An OPEN cursor statement."""
+    arg_types = {"this": True, "expression": False, "scroll": False, "expressions": False}
+
+
+class PGFetch(exp.Fetch):
+    """A FETCH statement."""
+    arg_types = {
+        **exp.Fetch.arg_types,
+        "this": True,
+        "expressions": True,
+        "preposition": False,
+    }
+
+
+class PGExecute(exp.Execute):
+    """A dynamic EXECUTE statement."""
+    arg_types = {
+        **exp.Execute.arg_types,
+        "strict": False,
+        "using": False,
+    }
+
+
+class PGInto(exp.Into):
+    """An INTO clause."""
+    arg_types = {
+        **exp.Into.arg_types,
+        "strict": False,
+    }
+
+
+class PGReturning(exp.Returning):
+    """A RETURNING clause that may carry an INTO [STRICT] target list.
+
+    The existing `into` argument will hold a `PGInto` instance when present.
+    """
+    # Inherit arg_types from exp.Returning
+
+
+class PGMove(exp.Expression):
+    """A MOVE statement."""
+    arg_types = {
+        "this": True,
+        "direction": False,
+        "preposition": False,
+    }
+
+
+class PGClose(exp.Expression):
+    """A CLOSE cursor statement."""
+    arg_types = {"this": True}
+
+
+class PGNext(exp.Expression):
+    """The NEXT direction."""
+    arg_types = {"this": False}
+
+
+class PGPrior(exp.Expression):
+    """The PRIOR direction."""
+    arg_types = {"this": False}
+
+
+class PGFirst(exp.Expression):
+    """The FIRST direction."""
+    arg_types = {"this": False}
+
+
+class PGLast(exp.Expression):
+    """The LAST direction."""
+    arg_types = {"this": False}
+
+
+class PGForward(exp.Expression):
+    """The FORWARD direction."""
+    arg_types = {"this": False, "all": False}
+
+
+class PGBackward(exp.Expression):
+    """The BACKWARD direction."""
+    arg_types = {"this": False, "all": False}
+
+
+class PGAbsolute(exp.Expression):
+    """The ABSOLUTE <count> direction."""
+    arg_types = {"this": True}
+
+
+class PGRelative(exp.Expression):
+    """The RELATIVE <count> direction."""
+    arg_types = {"this": True}
+
+
+class PGAll(exp.Expression):
+    """The ALL direction."""
+    arg_types = {"this": False}
+
+
+class PGSqlState(exp.Expression):
+    """The SQLSTATE statement."""
+    arg_types = {"this": False}
+
+
+class PGFound(exp.Expression):
+    """The special PL/pgSQL boolean status variable FOUND."""
+    arg_types = {"this": False}
+
+
+class PGUpdate(exp.Update):
+    """An UPDATE statement with optional WHERE CURRENT OF cursor."""
+    arg_types = {**exp.Update.arg_types, "current_of": False}
+
+
+class PGDelete(exp.Delete):
+    """A DELETE statement with optional WHERE CURRENT OF cursor."""
+    arg_types = {**exp.Delete.arg_types, "current_of": False}
+
+
+class PGReturn(exp.Return):
+    """A RETURN statement."""
+    arg_types = {"this": False, "query": False, "next": False}
+
+
+class PGOthers(exp.Expression):
+    """The OTHERS keyword in EXCEPTION WHEN OTHERS THEN ..."""
+    arg_types = {"this": False}

sqlleaf/dialects/plpgsql/p_dialect.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+from sqlglot.dialects.postgres import Postgres
+from sqlglot.tokens import Token
+
+from sqlleaf.dialects.plpgsql.p_generator import PlPgSQLGenerator
+from sqlleaf.dialects.plpgsql.p_parser import PlPgSQLParser
+from sqlleaf.dialects.plpgsql.p_tokenizer import PLPGSQL_KEYWORD_TOKEN_NAMES, TokenType
+
+
+class PlPgSQL(Postgres):
+    """
+    A custom dialect that extends Postgres to parse PL/pgSQL.
+    """
+    class Tokenizer(Postgres.Tokenizer):
+        KEYWORDS = {
+            **Postgres.Tokenizer.KEYWORDS,
+            **{keyword: getattr(TokenType, keyword) for keyword in PLPGSQL_KEYWORD_TOKEN_NAMES},
+        }
+
+        # In base sqlglot, FETCH and EXECUTE are treated as COMMANDs, which swallow the rest
+        # of the statement as a single STRING token, so we exclude them.
+        COMMANDS = Postgres.Tokenizer.COMMANDS - {TokenType.FETCH, TokenType.EXECUTE}
+
+        def tokenize(self, sql: str) -> list[Token]:
+            tokens = super().tokenize(sql)
+            return self._merge_range_tokens(tokens)
+
+        def _merge_range_tokens(self, tokens: list[Token]) -> list[Token]:
+            def adjacent(a: Token, b: Token) -> bool:
+                return a.end + 1 == b.start
+
+            result: list[Token] = []
+            i = 0
+            n = len(tokens)
+            while i < n:
+                tok = tokens[i]
+                nxt = tokens[i + 1] if i + 1 < n else None
+
+                # Pattern 0: LABEL full pattern << id >>
+                # Merge only when we see LT LT <identifier-like> GT GT with strict adjacency
+                if (
+                    i + 4 < n
+                    and tok.token_type == TokenType.LT
+                    and nxt is not None and nxt.token_type == TokenType.LT and adjacent(tok, nxt)
+                ):
+                    ident_tok = tokens[i + 2]
+                    gt1 = tokens[i + 3]
+                    gt2 = tokens[i + 4]
+
+                    if (
+                        ident_tok is not None
+                        and ident_tok.text
+                        and (ident_tok.text[0].isalpha() or ident_tok.text[0] == "_")
+                        and all(ch.isalnum() or ch == "_" for ch in ident_tok.text)
+                        and gt1.token_type == TokenType.GT
+                        and gt2.token_type == TokenType.GT
+                        and adjacent(nxt, ident_tok)
+                        and adjacent(ident_tok, gt1)
+                        and adjacent(gt1, gt2)
+                    ):
+                        # Create LABEL_BEGIN spanning the two '<'
+                        lbegin = Token(
+                            TokenType.LABEL_BEGIN,
+                            "<<",
+                            line=tok.line,
+                            col=tok.col,
+                            start=tok.start,
+                            end=nxt.end,
+                            comments=tok.comments,
+                        )
+                        # Keep the identifier token as-is
+                        # Create LABEL_END spanning the two '>'
+                        lend = Token(
+                            TokenType.LABEL_END,
+                            ">>",
+                            line=gt1.line,
+                            col=gt1.col,
+                            start=gt1.start,
+                            end=gt2.end,
+                            comments=gt1.comments,
+                        )
+                        result.append(lbegin)
+                        result.append(ident_tok)
+                        result.append(lend)
+                        i += 5
+                        continue
+
+                # Pattern 1: NUMBER('...' endswith '.') immediately followed by DOT -> DDOT
+                if (
+                    nxt is not None
+                    and tok.token_type == TokenType.NUMBER
+                    and tok.text.endswith(".")
+                    and nxt.token_type == TokenType.DOT
+                    and adjacent(tok, nxt)
+                ):
+                    number = Token(
+                        TokenType.NUMBER,
+                        tok.text[:-1],
+                        line=tok.line,
+                        col=tok.col,
+                        start=tok.start,
+                        end=tok.end - 1,
+                        comments=tok.comments,
+                    )
+                    ddot = Token(
+                        TokenType.DDOT,
+                        "..",
+                        line=nxt.line,
+                        col=nxt.col,
+                        start=tok.end,
+                        end=nxt.end,
+                        comments=nxt.comments,
+                    )
+                    result.append(number)
+                    result.append(ddot)
+                    i += 2
+                    continue
+
+                # Pattern 2: DOT immediately followed by DOT -> DDOT
+                if (
+                    nxt is not None
+                    and tok.token_type == TokenType.DOT
+                    and nxt.token_type == TokenType.DOT
+                    and adjacent(tok, nxt)
+                ):
+                    ddot = Token(
+                        TokenType.DDOT,
+                        "..",
+                        line=tok.line,
+                        col=tok.col,
+                        start=tok.start,
+                        end=nxt.end,
+                        comments=tok.comments,
+                    )
+                    result.append(ddot)
+                    i += 2
+                    continue
+
+                result.append(tok)
+                i += 1
+
+            return result
+
+    Generator = PlPgSQLGenerator
+    Parser = PlPgSQLParser
+
+plpgsql = PlPgSQL