bqsqlparse 0.1.0__py3-none-any.whl

bqsqlparse/nodes.py

@@ -0,0 +1,728 @@
+"""AST node definitions for bqsqlparse.
+
+Every node is a dataclass inheriting from :class:`Node`, which provides
+generic child iteration (``children``), tree walking (``walk``) and SQL
+re-generation (``sql``).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field, fields
+from typing import Iterator, List, Optional, Tuple, Union
+
+
+class Node:
+    """Base class for all AST nodes."""
+
+    def children(self) -> Iterator["Node"]:
+        for f in fields(self):  # type: ignore[arg-type]
+            v = getattr(self, f.name)
+            if isinstance(v, Node):
+                yield v
+            elif isinstance(v, (list, tuple)):
+                for x in v:
+                    if isinstance(x, Node):
+                        yield x
+                    elif isinstance(x, (list, tuple)):
+                        for y in x:
+                            if isinstance(y, Node):
+                                yield y
+
+    def walk(self) -> Iterator["Node"]:
+        yield self
+        for c in self.children():
+            yield from c.walk()
+
+    def find_all(self, *types) -> Iterator["Node"]:
+        for n in self.walk():
+            if isinstance(n, types):
+                yield n
+
+    def sql(self) -> str:
+        from .unparse import to_sql
+
+        return to_sql(self)
+
+    def __repr__(self) -> str:  # pragma: no cover - debugging aid
+        parts = []
+        for f in fields(self):  # type: ignore[arg-type]
+            v = getattr(self, f.name)
+            if v in (None, [], (), False):
+                continue
+            parts.append(f"{f.name}={v!r}")
+        return f"{type(self).__name__}({', '.join(parts)})"
+
+
+Expr = Node  # alias for readability
+
+
+# ---------------------------------------------------------------------------
+# Types
+# ---------------------------------------------------------------------------
+
+@dataclass(repr=False)
+class StructFieldType(Node):
+    name: Optional[str]
+    type: "TypeNode"
+
+
+@dataclass(repr=False)
+class TypeNode(Node):
+    name: str
+    params: List[str] = field(default_factory=list)
+    element: Optional["TypeNode"] = None  # ARRAY<T> / RANGE<T>
+    fields_: List[StructFieldType] = field(default_factory=list)  # STRUCT<...>
+
+
+# ---------------------------------------------------------------------------
+# Expressions
+# ---------------------------------------------------------------------------
+
+@dataclass(repr=False)
+class Literal(Node):
+    value: object
+    kind: str  # 'string' | 'bytes' | 'number' | 'bool' | 'null' | 'default'
+
+
+@dataclass(repr=False)
+class TypedLiteral(Node):
+    """DATE '2024-01-01', TIMESTAMP '...', NUMERIC '1.5', JSON '{"a":1}' ..."""
+
+    type_name: str
+    value: str
+
+
+@dataclass(repr=False)
+class ColumnRef(Node):
+    """A dotted identifier path: column, table.column, alias.struct.field ..."""
+
+    path: List[str]
+
+
+@dataclass(repr=False)
+class Param(Node):
+    name: str  # '@name', '@@sysvar' or '?'
+
+
+@dataclass(repr=False)
+class ReplaceItem(Node):
+    expr: Expr
+    name: str
+
+
+@dataclass(repr=False)
+class Star(Node):
+    prefix: List[str] = field(default_factory=list)  # for t.* / t.struct.*
+    except_: List[str] = field(default_factory=list)
+    replace: List[ReplaceItem] = field(default_factory=list)
+
+
+@dataclass(repr=False)
+class BinaryOp(Node):
+    op: str
+    left: Expr
+    right: Expr
+
+
+@dataclass(repr=False)
+class UnaryOp(Node):
+    op: str  # '-', '+', '~', 'NOT'
+    operand: Expr
+
+
+@dataclass(repr=False)
+class Between(Node):
+    expr: Expr
+    low: Expr
+    high: Expr
+    negated: bool = False
+
+
+@dataclass(repr=False)
+class InExpr(Node):
+    expr: Expr
+    values: Optional[List[Expr]] = None
+    query: Optional["Query"] = None
+    unnest: Optional[Expr] = None
+    negated: bool = False
+
+
+@dataclass(repr=False)
+class LikeExpr(Node):
+    expr: Expr
+    patterns: List[Expr]
+    quantifier: Optional[str] = None  # ANY | ALL | SOME
+    negated: bool = False
+
+
+@dataclass(repr=False)
+class IsExpr(Node):
+    expr: Expr
+    value: str  # NULL | TRUE | FALSE | UNKNOWN
+    negated: bool = False
+
+
+@dataclass(repr=False)
+class WhenClause(Node):
+    condition: Expr
+    result: Expr
+
+
+@dataclass(repr=False)
+class Case(Node):
+    operand: Optional[Expr]
+    whens: List[WhenClause]
+    else_: Optional[Expr] = None
+
+
+@dataclass(repr=False)
+class Cast(Node):
+    expr: Expr
+    to_type: TypeNode
+    safe: bool = False
+    format: Optional[Expr] = None
+
+
+@dataclass(repr=False)
+class Extract(Node):
+    part: str
+    expr: Expr
+    at_time_zone: Optional[Expr] = None
+
+
+@dataclass(repr=False)
+class IntervalExpr(Node):
+    value: Expr
+    unit: str
+    to_unit: Optional[str] = None
+
+
+@dataclass(repr=False)
+class ArrayExpr(Node):
+    elements: List[Expr]
+    elem_type: Optional[TypeNode] = None
+
+
+@dataclass(repr=False)
+class StructFieldValue(Node):
+    expr: Expr
+    name: Optional[str] = None
+
+
+@dataclass(repr=False)
+class StructExpr(Node):
+    fields_: List[StructFieldValue]
+    type_fields: Optional[List[StructFieldType]] = None
+
+
+@dataclass(repr=False)
+class Subscript(Node):
+    base: Expr
+    index: Expr
+    mode: Optional[str] = None  # OFFSET | ORDINAL | SAFE_OFFSET | SAFE_ORDINAL | None
+
+
+@dataclass(repr=False)
+class FieldAccess(Node):
+    base: Expr
+    field_name: str
+
+
+@dataclass(repr=False)
+class OrderItem(Node):
+    expr: Expr
+    desc: Optional[bool] = None
+    nulls: Optional[str] = None  # FIRST | LAST
+
+
+@dataclass(repr=False)
+class FrameBound(Node):
+    kind: str  # 'UNBOUNDED PRECEDING' | 'CURRENT ROW' | 'PRECEDING' | 'FOLLOWING'
+    value: Optional[Expr] = None
+
+
+@dataclass(repr=False)
+class WindowFrame(Node):
+    unit: str  # ROWS | RANGE | GROUPS
+    start: FrameBound
+    end: Optional[FrameBound] = None
+
+
+@dataclass(repr=False)
+class WindowSpec(Node):
+    name: Optional[str] = None  # base named window
+    partition_by: List[Expr] = field(default_factory=list)
+    order_by: List[OrderItem] = field(default_factory=list)
+    frame: Optional[WindowFrame] = None
+    is_ref: bool = False  # OVER window_name (no parens)
+
+
+@dataclass(repr=False)
+class NamedArg(Node):
+    name: str
+    value: Expr
+
+
+@dataclass(repr=False)
+class HavingModifier(Node):
+    kind: str  # MAX | MIN
+    expr: Expr
+
+
+@dataclass(repr=False)
+class FuncCall(Node):
+    name: List[str]  # e.g. ['SAFE', 'SUBSTR'] or ['my_ds', 'my_udf']
+    args: List[Expr] = field(default_factory=list)
+    distinct: bool = False
+    nulls: Optional[str] = None  # IGNORE | RESPECT
+    order_by: List[OrderItem] = field(default_factory=list)
+    limit: Optional[Expr] = None
+    having: Optional[HavingModifier] = None
+    over: Optional[WindowSpec] = None
+
+    @property
+    def name_str(self) -> str:
+        return ".".join(self.name).upper()
+
+
+@dataclass(repr=False)
+class ScalarSubquery(Node):
+    query: "Query"
+
+
+@dataclass(repr=False)
+class ArraySubquery(Node):
+    query: "Query"
+
+
+@dataclass(repr=False)
+class ExistsSubquery(Node):
+    query: "Query"
+
+
+# ---------------------------------------------------------------------------
+# Query structure
+# ---------------------------------------------------------------------------
+
+@dataclass(repr=False)
+class SelectItem(Node):
+    expr: Expr
+    alias: Optional[str] = None
+
+
+@dataclass(repr=False)
+class GroupBy(Node):
+    kind: str  # 'exprs' | 'rollup' | 'cube' | 'sets' | 'all'
+    exprs: List[Expr] = field(default_factory=list)
+    sets: List[List[Expr]] = field(default_factory=list)
+
+
+@dataclass(repr=False)
+class NamedWindow(Node):
+    name: str
+    spec: WindowSpec
+
+
+@dataclass(repr=False)
+class Select(Node):
+    items: List[Node]  # SelectItem | Star
+    distinct: bool = False
+    as_mode: Optional[str] = None  # STRUCT | VALUE
+    from_: Optional[Node] = None
+    where: Optional[Expr] = None
+    group_by: Optional[GroupBy] = None
+    having: Optional[Expr] = None
+    qualify: Optional[Expr] = None
+    windows: List[NamedWindow] = field(default_factory=list)
+
+
+@dataclass(repr=False)
+class SetOp(Node):
+    op: str  # UNION | INTERSECT | EXCEPT
+    all: bool
+    left: Node
+    right: Node
+
+
+@dataclass(repr=False)
+class CTE(Node):
+    name: str
+    query: "Query"
+
+
+@dataclass(repr=False)
+class Query(Node):
+    body: Node  # Select | SetOp | Query
+    ctes: List[CTE] = field(default_factory=list)
+    recursive: bool = False
+    order_by: List[OrderItem] = field(default_factory=list)
+    limit: Optional[Expr] = None
+    offset: Optional[Expr] = None
+
+
+# ---------------------------------------------------------------------------
+# FROM clause items
+# ---------------------------------------------------------------------------
+
+@dataclass(repr=False)
+class TableSample(Node):
+    method: str
+    value: Expr
+    unit: str = "PERCENT"
+
+
+@dataclass(repr=False)
+class TableRef(Node):
+    path: List[str]
+    alias: Optional[str] = None
+    system_time: Optional[Expr] = None
+    sample: Optional[TableSample] = None
+
+    @property
+    def full_name(self) -> str:
+        return ".".join(self.path)
+
+
+@dataclass(repr=False)
+class SubqueryRef(Node):
+    query: Query
+    alias: Optional[str] = None
+
+
+@dataclass(repr=False)
+class UnnestRef(Node):
+    expr: Expr
+    alias: Optional[str] = None
+    with_offset: bool = False
+    offset_alias: Optional[str] = None
+
+
+@dataclass(repr=False)
+class TableFuncRef(Node):
+    name: List[str]
+    args: List[Expr] = field(default_factory=list)
+    alias: Optional[str] = None
+
+
+@dataclass(repr=False)
+class Join(Node):
+    kind: str  # INNER | LEFT | RIGHT | FULL | CROSS
+    left: Node
+    right: Node
+    on: Optional[Expr] = None
+    using: List[str] = field(default_factory=list)
+
+
+@dataclass(repr=False)
+class PivotAgg(Node):
+    func: FuncCall
+    alias: Optional[str] = None
+
+
+@dataclass(repr=False)
+class PivotValue(Node):
+    value: Expr
+    alias: Optional[str] = None
+
+
+@dataclass(repr=False)
+class PivotRef(Node):
+    input: Node
+    aggregates: List[PivotAgg]
+    for_col: Expr
+    in_values: List[PivotValue]
+    alias: Optional[str] = None
+
+
+@dataclass(repr=False)
+class UnpivotGroup(Node):
+    columns: List[List[str]]  # one or more column paths per group
+    label: Optional[Expr] = None
+
+
+@dataclass(repr=False)
+class UnpivotRef(Node):
+    input: Node
+    value_columns: List[str]
+    name_column: str
+    groups: List[UnpivotGroup]
+    include_nulls: Optional[bool] = None
+    alias: Optional[str] = None
+
+
+# ---------------------------------------------------------------------------
+# Statements
+# ---------------------------------------------------------------------------
+
+@dataclass(repr=False)
+class ColumnDef(Node):
+    name: str
+    type: Optional[TypeNode] = None
+    options_sql: Optional[str] = None
+
+
+@dataclass(repr=False)
+class CreateTableAsSelect(Node):
+    kind: str  # TABLE | VIEW | MATERIALIZED VIEW
+    name: List[str]
+    query: Optional[Query] = None
+    replace: bool = False
+    temp: bool = False
+    if_not_exists: bool = False
+    columns: List[ColumnDef] = field(default_factory=list)
+    partition_by: Optional[Expr] = None
+    cluster_by: List[Expr] = field(default_factory=list)
+    options_sql: Optional[str] = None
+
+    @property
+    def full_name(self) -> str:
+        return ".".join(self.name)
+
+
+@dataclass(repr=False)
+class InsertStmt(Node):
+    table: List[str]
+    columns: List[str] = field(default_factory=list)
+    query: Optional[Query] = None
+    values: List[List[Expr]] = field(default_factory=list)
+
+    @property
+    def full_name(self) -> str:
+        return ".".join(self.table)
+
+
+@dataclass(repr=False)
+class Assignment(Node):
+    target: List[str]
+    value: Expr
+
+
+@dataclass(repr=False)
+class UpdateStmt(Node):
+    table: List[str]
+    assignments: List[Assignment]
+    alias: Optional[str] = None
+    from_: Optional[Node] = None
+    where: Optional[Expr] = None
+
+    @property
+    def full_name(self) -> str:
+        return ".".join(self.table)
+
+
+@dataclass(repr=False)
+class DeleteStmt(Node):
+    table: List[str]
+    alias: Optional[str] = None
+    where: Optional[Expr] = None
+
+    @property
+    def full_name(self) -> str:
+        return ".".join(self.table)
+
+
+@dataclass(repr=False)
+class MergeWhen(Node):
+    match_kind: str  # MATCHED | NOT_MATCHED | NOT_MATCHED_BY_SOURCE
+    action_kind: str  # UPDATE | DELETE | INSERT | INSERT_ROW
+    condition: Optional[Expr] = None
+    assignments: List[Assignment] = field(default_factory=list)
+    columns: List[str] = field(default_factory=list)
+    values: List[Expr] = field(default_factory=list)
+
+
+@dataclass(repr=False)
+class MergeStmt(Node):
+    target: List[str]
+    source: Node  # TableRef | SubqueryRef
+    on: Expr
+    whens: List[MergeWhen]
+    alias: Optional[str] = None
+
+    @property
+    def full_name(self) -> str:
+        return ".".join(self.target)
+
+
+# ---------------------------------------------------------------------------
+# Scripting (procedural language)
+# ---------------------------------------------------------------------------
+
+@dataclass(repr=False)
+class DeclareStmt(Node):
+    names: List[str]
+    type: Optional[TypeNode] = None
+    default: Optional[Expr] = None
+
+
+@dataclass(repr=False)
+class SetStmt(Node):
+    targets: List[str]  # variable names (or @@system vars)
+    value: Expr  # a StructExpr models the (a, b) = (1, 2) form
+
+
+@dataclass(repr=False)
+class ScriptBlock(Node):
+    """BEGIN ... [EXCEPTION WHEN ERROR THEN ...] END"""
+
+    statements: List[Node]
+    exception_statements: List[Node] = field(default_factory=list)
+    has_exception_handler: bool = False
+    label: Optional[str] = None
+
+
+@dataclass(repr=False)
+class IfBranch(Node):
+    condition: Expr
+    statements: List[Node]
+
+
+@dataclass(repr=False)
+class IfStmt(Node):
+    branches: List[IfBranch]  # IF + ELSEIFs, in order
+    else_statements: List[Node] = field(default_factory=list)
+
+
+@dataclass(repr=False)
+class LoopStmt(Node):
+    statements: List[Node]
+    label: Optional[str] = None
+
+
+@dataclass(repr=False)
+class WhileStmt(Node):
+    condition: Expr
+    statements: List[Node]
+    label: Optional[str] = None
+
+
+@dataclass(repr=False)
+class RepeatStmt(Node):
+    statements: List[Node]
+    until: Expr = None
+    label: Optional[str] = None
+
+
+@dataclass(repr=False)
+class ForInStmt(Node):
+    var: str
+    query: "Query"
+    statements: List[Node]
+    label: Optional[str] = None
+
+
+@dataclass(repr=False)
+class BreakContinueStmt(Node):
+    kind: str  # BREAK | LEAVE | CONTINUE | ITERATE
+    label: Optional[str] = None
+
+
+@dataclass(repr=False)
+class CallStmt(Node):
+    name: List[str]
+    args: List[Expr] = field(default_factory=list)
+
+    @property
+    def full_name(self) -> str:
+        return ".".join(self.name)
+
+
+@dataclass(repr=False)
+class ReturnStmt(Node):
+    pass
+
+
+@dataclass(repr=False)
+class RaiseStmt(Node):
+    message: Optional[Expr] = None
+
+
+@dataclass(repr=False)
+class ExecuteImmediate(Node):
+    sql_expr: Expr
+    into: List[str] = field(default_factory=list)
+    using: List[Tuple[Expr, Optional[str]]] = field(default_factory=list)
+
+
+@dataclass(repr=False)
+class AssertStmt(Node):
+    condition: Expr
+    message: Optional[Expr] = None
+
+
+@dataclass(repr=False)
+class TransactionStmt(Node):
+    kind: str  # BEGIN | COMMIT | ROLLBACK
+
+
+@dataclass(repr=False)
+class TruncateStmt(Node):
+    table: List[str]
+
+    @property
+    def full_name(self) -> str:
+        return ".".join(self.table)
+
+
+@dataclass(repr=False)
+class DropStmt(Node):
+    kind: str  # TABLE | VIEW | MATERIALIZED VIEW | FUNCTION | PROCEDURE | ...
+    name: List[str]
+    if_exists: bool = False
+
+    @property
+    def full_name(self) -> str:
+        return ".".join(self.name)
+
+
+@dataclass(repr=False)
+class RoutineParam(Node):
+    name: str
+    type: Optional[TypeNode] = None
+    mode: Optional[str] = None  # IN | OUT | INOUT (procedures only)
+
+
+@dataclass(repr=False)
+class CreateProcedure(Node):
+    name: List[str]
+    params: List[RoutineParam]
+    body: ScriptBlock
+    replace: bool = False
+    if_not_exists: bool = False
+    options_sql: Optional[str] = None
+
+    @property
+    def full_name(self) -> str:
+        return ".".join(self.name)
+
+
+@dataclass(repr=False)
+class CreateFunction(Node):
+    name: List[str]
+    params: List[RoutineParam] = field(default_factory=list)
+    returns: Optional[TypeNode] = None
+    returns_table: List[StructFieldType] = field(default_factory=list)
+    language: Optional[str] = None
+    body_expr: Optional[Expr] = None      # SQL UDF: AS (expr)
+    body_query: Optional["Query"] = None  # TVF: AS (SELECT ...)
+    body_string: Optional[str] = None     # JS UDF: AS "..."
+    replace: bool = False
+    temp: bool = False
+    if_not_exists: bool = False
+    table_function: bool = False
+    aggregate: bool = False
+    deterministic: Optional[bool] = None
+    options_sql: Optional[str] = None
+
+    @property
+    def full_name(self) -> str:
+        return ".".join(self.name)
+
+
+Statement = Union[
+    Query, CreateTableAsSelect, InsertStmt, UpdateStmt, DeleteStmt, MergeStmt,
+    DeclareStmt, SetStmt, ScriptBlock, IfStmt, LoopStmt, WhileStmt, RepeatStmt,
+    ForInStmt, BreakContinueStmt, CallStmt, ReturnStmt, RaiseStmt,
+    ExecuteImmediate, AssertStmt, TransactionStmt, TruncateStmt, DropStmt,
+    CreateProcedure, CreateFunction,
+]