ast_classes.py

from collections.abc import Generator, Sequence
from datetime import datetime

import attrs
from typing_extensions import Self

from data_diff.abcs.compiler import Compilable
from data_diff.abcs.database_types import DbPath
from data_diff.queries.base import SKIP, SqeletonError, args_as_tuple
from data_diff.schema import Schema
from data_diff.utils import ArithString, safezip


class QueryBuilderError(SqeletonError):
    pass


class QB_TypeError(QueryBuilderError):
    pass


@attrs.define(frozen=True)
class Root:
    "Nodes inheriting from Root can be used as root statements in SQL (e.g. SELECT yes, RANDOM() no)"


@attrs.define(frozen=False, eq=False)
class ExprNode(Compilable):
    "Base class for query expression nodes"

    @property
    def type(self) -> type | None:
        return None

    def _dfs_values(self):
        yield self
        for k, vs in attrs.asdict(self, recurse=False).items():
            if k == "source_table":
                # Skip data-sources, we're only interested in data-parameters
                continue
            if not isinstance(vs, list | tuple):
                vs = [vs]
            for v in vs:
                if isinstance(v, ExprNode):
                    yield from v._dfs_values()

    def cast_to(self, to) -> "Cast":
        return Cast(self, to)


# Query expressions can only interact with objects that are an instance of 'Expr'
Expr = ExprNode | str | bool | int | float | datetime | ArithString | None


@attrs.define(frozen=True, eq=False)
class Code(ExprNode, Root):
    code: str
    args: dict[str, Expr] | None = None


def _expr_type(e: Expr) -> type:
    if isinstance(e, ExprNode):
        return e.type
    return type(e)


@attrs.define(frozen=True, eq=False)
class Alias(ExprNode):
    expr: Expr
    name: str

    @property
    def type(self):
        return _expr_type(self.expr)


def _drop_skips(exprs):
    return [e for e in exprs if e is not SKIP]


def _drop_skips_dict(exprs_dict):
    return {k: v for k, v in exprs_dict.items() if v is not SKIP}


@attrs.define(frozen=True)
class ITable:
    @property
    def source_table(self) -> "ITable":  # not always Self, it can be a substitute
        return self

    @property
    def schema(self) -> Schema | None:
        return None

    def select(self, *exprs, distinct=SKIP, optimizer_hints=SKIP, **named_exprs) -> "ITable":
        """Choose new columns, based on the old ones. (aka Projection)

        Parameters:
            exprs: List of expressions to constitute the columns of the new table.
                    If not provided, returns all columns in source table (i.e. ``select *``)
            distinct: 'select' or 'select distinct'
            named_exprs: More expressions to constitute the columns of the new table, aliased to keyword name.

        """
        exprs = args_as_tuple(exprs)
        exprs = _drop_skips(exprs)
        named_exprs = _drop_skips_dict(named_exprs)
        exprs += _named_exprs_as_aliases(named_exprs)
        resolve_names(self.source_table, exprs)
        return Select.make(self, columns=exprs, distinct=distinct, optimizer_hints=optimizer_hints)

    def where(self, *exprs) -> "Select":
        """Filter the rows, based on the given predicates. (aka Selection)"""
        exprs = args_as_tuple(exprs)
        exprs = _drop_skips(exprs)
        if not exprs:
            return self

        resolve_names(self.source_table, exprs)
        return Select.make(self, where_exprs=exprs)

    def order_by(self, *exprs) -> "Select":
        """Order the rows lexicographically, according to the given expressions."""
        exprs = _drop_skips(exprs)
        if not exprs:
            return self

        resolve_names(self.source_table, exprs)
        return Select.make(self, order_by_exprs=exprs)

    def limit(self, limit: int) -> "Select":
        """Stop yielding rows after the given limit. i.e. take the first 'n=limit' rows"""
        if limit is SKIP:
            return self

        return Select.make(self, limit_expr=limit)

    def join(self, target: "ITable") -> "Join":
        """Join the current table with the target table, returning a new table containing both side-by-side.

        When joining, it's recommended to use explicit tables names, instead of `this`, in order to avoid potential name collisions.

        Example:
            ::

                person = table('person')
                city = table('city')

                name_and_city = (
                    person
                    .join(city)
                    .on(person['city_id'] == city['id'])
                    .select(person['id'], city['name'])
                )
        """
        return Join([self, target])

    def group_by(self, *keys) -> "GroupBy":
        """Behaves like in SQL, except for a small change in syntax:

        A call to `.agg()` must follow every call to `.group_by()`.

        Example:
            ::

                # SELECT a, sum(b) FROM tmp GROUP BY 1
                table('tmp').group_by(this.a).agg(this.b.sum())

                # SELECT a, sum(b) FROM a GROUP BY 1 HAVING (b > 10)
                (table('tmp')
                    .group_by(this.a)
                    .agg(this.b.sum())
                    .having(this.b > 10)
                )

        """
        keys = _drop_skips(keys)
        resolve_names(self.source_table, keys)

        return GroupBy(self, keys)

    def _get_column(self, name: str) -> "Column":
        if self.schema:
            name = self.schema.get_key(name)  # Get the actual name. Might be case-insensitive.
        return Column(self, name)

    # def __getattr__(self, column):
    #     return self._get_column(column)

    def __getitem__(self, column) -> "Column":
        if not isinstance(column, str):
            raise TypeError()
        return self._get_column(column)

    def count(self) -> "Select":
        """SELECT count() FROM self"""
        return Select(self, [Count()])

    def union(self, other: "ITable") -> "TableOp":
        """SELECT * FROM self UNION other"""
        return TableOp("UNION", self, other)

    def union_all(self, other: "ITable") -> "TableOp":
        """SELECT * FROM self UNION ALL other"""
        return TableOp("UNION ALL", self, other)

    def minus(self, other: "ITable") -> "TableOp":
        """SELECT * FROM self EXCEPT other"""
        # aka
        return TableOp("EXCEPT", self, other)

    def intersect(self, other: "ITable") -> "TableOp":
        """SELECT * FROM self INTERSECT other"""
        return TableOp("INTERSECT", self, other)


@attrs.define(frozen=True, eq=False)
class Concat(ExprNode):
    exprs: list
    sep: str | None = None


@attrs.define(frozen=True, eq=False)
class Count(ExprNode):
    expr: Expr = None
    distinct: bool = False

    @property
    def type(self) -> type | None:
        return int


@attrs.define(frozen=False, eq=False)
class LazyOps:
    def __add__(self, other) -> "BinOp":
        return BinOp("+", [self, other])

    def __sub__(self, other) -> "BinOp":
        return BinOp("-", [self, other])

    def __neg__(self) -> "UnaryOp":
        return UnaryOp("-", self)

    def __gt__(self, other) -> "BinBoolOp":
        return BinBoolOp(">", [self, other])

    def __ge__(self, other) -> "BinBoolOp":
        return BinBoolOp(">=", [self, other])

    def __eq__(self, other) -> "BinBoolOp":
        if other is None:
            return BinBoolOp("IS", [self, None])
        return BinBoolOp("=", [self, other])

    def __lt__(self, other) -> "BinBoolOp":
        return BinBoolOp("<", [self, other])

    def __le__(self, other) -> "BinBoolOp":
        return BinBoolOp("<=", [self, other])

    def __or__(self, other) -> "BinBoolOp":
        return BinBoolOp("OR", [self, other])

    def __and__(self, other) -> "BinBoolOp":
        return BinBoolOp("AND", [self, other])

    def is_distinct_from(self, other) -> "IsDistinctFrom":
        return IsDistinctFrom(self, other)

    def like(self, other) -> "BinBoolOp":
        return BinBoolOp("LIKE", [self, other])

    def sum(self) -> "Func":
        return Func("SUM", [self])

    def max(self) -> "Func":
        return Func("MAX", [self])

    def min(self) -> "Func":
        return Func("MIN", [self])


@attrs.define(frozen=True, eq=False)
class Func(LazyOps, ExprNode):
    name: str
    args: Sequence[Expr]


@attrs.define(frozen=True, eq=False)
class WhenThen(ExprNode):
    when: Expr
    then: Expr


@attrs.define(frozen=True, eq=False)
class CaseWhen(ExprNode):
    cases: Sequence[WhenThen]
    else_expr: Expr | None = None

    @property
    def type(self):
        then_types = {_expr_type(case.then) for case in self.cases}
        if self.else_expr:
            then_types |= {_expr_type(self.else_expr)}
        if len(then_types) > 1:
            raise QB_TypeError(f"Non-matching types in when: {then_types}")
        (t,) = then_types
        return t

    def when(self, *whens: Expr) -> "QB_When":
        """Add a new 'when' clause to the case expression

        Must be followed by a call to `.then()`
        """
        whens = args_as_tuple(whens)
        whens = _drop_skips(whens)
        if not whens:
            raise QueryBuilderError("Expected valid whens")

        # XXX reimplementing api.and_()
        if len(whens) == 1:
            return QB_When(self, whens[0])
        return QB_When(self, BinBoolOp("AND", whens))

    def else_(self, then: Expr) -> Self:
        """Add an 'else' clause to the case expression.

        Can only be called once!
        """
        if self.else_expr is not None:
            raise QueryBuilderError(f"Else clause already specified in {self}")

        return attrs.evolve(self, else_expr=then)


@attrs.define(frozen=True, eq=False)
class QB_When:
    "Partial case-when, used for query-building"

    casewhen: CaseWhen
    when: Expr

    def then(self, then: Expr) -> CaseWhen:
        """Add a 'then' clause after a 'when' was added."""
        case = WhenThen(self.when, then)
        return attrs.evolve(self.casewhen, cases=self.casewhen.cases + [case])


@attrs.define(frozen=True, eq=False)
class IsDistinctFrom(LazyOps, ExprNode):
    a: Expr
    b: Expr

    @property
    def type(self) -> type | None:
        return bool


@attrs.define(frozen=True, eq=False)
class BinOp(LazyOps, ExprNode):
    op: str
    args: Sequence[Expr]

    @property
    def type(self):
        types = {_expr_type(i) for i in self.args}
        if len(types) > 1:
            raise TypeError(f"Expected all args to have the same type, got {types}")
        (t,) = types
        return t


@attrs.define(frozen=True, eq=False)
class UnaryOp(LazyOps, ExprNode):
    op: str
    expr: Expr


@attrs.define(frozen=True)
class BinBoolOp(BinOp):
    @property
    def type(self) -> type | None:
        return bool


@attrs.define(frozen=True, eq=False)
class Column(LazyOps, ExprNode):
    source_table: ITable
    name: str

    @property
    def type(self):
        if self.source_table.schema is None:
            raise QueryBuilderError(f"Schema required for table {self.source_table}")
        return self.source_table.schema[self.name]


@attrs.define(frozen=False, eq=False)
class TablePath(ExprNode, ITable):
    path: DbPath
    schema: Schema | None = None  # overrides the inherited property

    # Statement shorthands
    def create(self, source_table: ITable = None, *, if_not_exists: bool = False, primary_keys: list[str] = None):
        """Returns a query expression to create a new table.

        Parameters:
            source_table: a table expression to use for initializing the table.
                          If not provided, the table must have a schema specified.
            if_not_exists: Add a 'if not exists' clause or not. (note: not all dbs support it!)
            primary_keys: List of column names which define the primary key
        """

        if source_table is None and not self.schema:
            raise ValueError("Either schema or source table needed to create table")
        if isinstance(source_table, TablePath):
            source_table = source_table.select()
        return CreateTable(self, source_table, if_not_exists=if_not_exists, primary_keys=primary_keys)

    def drop(self, if_exists=False):
        """Returns a query expression to delete the table.

        Parameters:
            if_not_exists: Add a 'if not exists' clause or not. (note: not all dbs support it!)
        """
        return DropTable(self, if_exists=if_exists)

    def truncate(self):
        """Returns a query expression to truncate the table. (remove all rows)"""
        return TruncateTable(self)

    def insert_rows(self, rows: Sequence, *, columns: list[str] = None):
        """Returns a query expression to insert rows to the table, given as Python values.

        Parameters:
            rows: A list of tuples. Must all have the same width.
            columns: Names of columns being populated. If specified, must have the same length as the tuples.
        """
        rows = list(rows)
        return InsertToTable(self, ConstantTable(rows), columns=columns)

    def insert_row(self, *values, columns: list[str] = None):
        """Returns a query expression to insert a single row to the table, given as Python values.

        Parameters:
            columns: Names of columns being populated. If specified, must have the same length as 'values'
        """
        return InsertToTable(self, ConstantTable([values]), columns=columns)

    def insert_expr(self, expr: Expr):
        """Returns a query expression to insert rows to the table, given as a query expression.

        Parameters:
            expr: query expression to from which to read the rows
        """
        if isinstance(expr, TablePath):
            expr = expr.select()
        return InsertToTable(self, expr)


@attrs.define(frozen=True, eq=False)
class TableAlias(ExprNode, ITable):
    table: ITable
    name: str

    @property
    def source_table(self) -> ITable:
        return self.table

    @property
    def schema(self) -> Schema:
        return self.table.schema


@attrs.define(frozen=True, eq=False)
class Join(ExprNode, ITable, Root):
    source_tables: Sequence[ITable]
    op: str | None = None
    on_exprs: Sequence[Expr] | None = None
    columns: Sequence[Expr] | None = None

    @property
    def schema(self) -> Schema:
        if not self.columns:
            raise QueryBuilderError("Join must specify columns explicitly (SELECT * not yet implemented).")
        # No cross-table type validation needed: join combines columns from both tables rather than unioning rows
        s = self.source_tables[0].schema
        if s is None:
            raise QueryBuilderError("Cannot resolve Join schema: source table has no schema defined")
        return s.new({c.name: c.type for c in self.columns})

    def on(self, *exprs) -> Self:
        """Add an ON clause, for filtering the result of the cartesian product (i.e. the JOIN)"""
        if len(exprs) == 1:
            (e,) = exprs
            if isinstance(e, Generator):
                exprs = tuple(e)

        exprs = _drop_skips(exprs)
        if not exprs:
            return self

        return attrs.evolve(self, on_exprs=(self.on_exprs or []) + exprs)

    def select(self, *exprs, **named_exprs) -> Self | ITable:
        """Select fields to return from the JOIN operation

        See Also: ``ITable.select()``
        """
        if self.columns is not None:
            # join-select already applied
            return super().select(*exprs, **named_exprs)

        exprs = _drop_skips(exprs)
        named_exprs = _drop_skips_dict(named_exprs)
        exprs += _named_exprs_as_aliases(named_exprs)
        resolve_names(self.source_table, exprs)
        # TODO Ensure exprs <= self.columns ?
        return attrs.evolve(self, columns=exprs)


@attrs.define(frozen=True, eq=False)
class GroupBy(ExprNode, ITable, Root):
    table: ITable
    keys: Sequence[Expr] | None = None  # IKey?
    values: Sequence[Expr] | None = None
    having_exprs: Sequence[Expr] | None = None

    def __attrs_post_init__(self) -> None:
        if not self.keys and not self.values:
            raise ValueError("GroupBy must have at least one key or value expression.")

    def having(self, *exprs) -> Self:
        """Add a 'HAVING' clause to the group-by"""
        exprs = args_as_tuple(exprs)
        exprs = _drop_skips(exprs)
        if not exprs:
            return self

        resolve_names(self.table, exprs)
        return attrs.evolve(self, having_exprs=(self.having_exprs or []) + exprs)

    def agg(self, *exprs) -> Self:
        """Select aggregated fields for the group-by."""
        exprs = args_as_tuple(exprs)
        exprs = _drop_skips(exprs)
        resolve_names(self.table, exprs)
        return attrs.evolve(self, values=(self.values or []) + exprs)


@attrs.define(frozen=True, eq=False)
class TableOp(ExprNode, ITable, Root):
    op: str
    table1: ITable
    table2: ITable

    @property
    def type(self):
        t1 = self.table1.type
        t2 = self.table2.type
        if t1 is None or t2 is None:
            return None
        if type(t1) is not type(t2):
            raise QueryBuilderError(f"Type mismatch in {self.op}: got {type(t1).__name__} and {type(t2).__name__}")
        return t1

    @property
    def schema(self) -> Schema:
        s1 = self.table1.schema
        s2 = self.table2.schema
        if s1 is None or s2 is None:
            raise QueryBuilderError(f"Cannot validate {self.op}: one or both tables have no schema defined")
        if len(s1) != len(s2):
            raise QueryBuilderError(f"Schema length mismatch in {self.op}: got {len(s1)} and {len(s2)} columns")
        for (name1, type1), (name2, type2) in zip(s1.items(), s2.items()):
            if type(type1) is not type(type2):
                raise QueryBuilderError(
                    f"Type mismatch in {self.op}: column {name1!r} is {type(type1).__name__} "
                    f"but column {name2!r} is {type(type2).__name__}"
                )
        return s1


@attrs.define(frozen=True, eq=False)
class Select(ExprNode, ITable, Root):
    table: Expr | None = None
    columns: Sequence[Expr] | None = None
    where_exprs: Sequence[Expr] | None = None
    order_by_exprs: Sequence[Expr] | None = None
    group_by_exprs: Sequence[Expr] | None = None
    having_exprs: Sequence[Expr] | None = None
    limit_expr: int | None = None
    distinct: bool = False
    optimizer_hints: Sequence[Expr] | None = None

    @property
    def schema(self) -> Schema:
        s = self.table.schema
        if s is None or self.columns is None:
            return s
        return s.new({c.name: c.type for c in self.columns})

    @classmethod
    def make(cls, table: ITable, distinct: bool = SKIP, optimizer_hints: str = SKIP, **kwargs):
        if "table" in kwargs:
            raise ValueError(
                "Do not pass 'table' as a keyword argument to Select.make(); pass it as the first positional argument."
            )

        if not isinstance(table, cls):  # If not Select
            if distinct is not SKIP:
                kwargs["distinct"] = distinct
            if optimizer_hints is not SKIP:
                kwargs["optimizer_hints"] = optimizer_hints
            return cls(table, **kwargs)

        # We can safely assume isinstance(table, Select)
        if optimizer_hints is not SKIP:
            kwargs["optimizer_hints"] = optimizer_hints

        if distinct is not SKIP:
            if distinct == False and table.distinct:
                return cls(table, **kwargs)
            kwargs["distinct"] = distinct

        if table.limit_expr or table.group_by_exprs:
            return cls(table, **kwargs)

        # Fill in missing attributes, instead of nesting instances
        for k, v in kwargs.items():
            if getattr(table, k) is not None:
                if k == "where_exprs":  # Additive attribute
                    kwargs[k] = getattr(table, k) + v
                elif k in ["distinct", "optimizer_hints"]:
                    pass
                else:
                    raise ValueError(k)

        return attrs.evolve(table, **kwargs)


@attrs.define(frozen=True, eq=False)
class Cte(ExprNode, ITable):
    table: Expr
    name: str | None = None
    params: Sequence[str] | None = attrs.field(default=None)

    @params.validator
    def _validate_params(self, attribute, value):
        if value is not None:
            for i, p in enumerate(value):
                if not isinstance(p, str):
                    raise QB_TypeError(f"CTE params[{i}] must be str, got {type(p).__name__}")

    @property
    def source_table(self) -> "ITable":
        return self.table

    @property
    def schema(self) -> Schema | None:
        try:
            s = self.table.schema
        except (QueryBuilderError, ValueError) as exc:
            # ValueError caught because some ITable.schema implementations (e.g. TableOp)
            # still raise ValueError for validation errors pre-dating QueryBuilderError.
            name_hint = f" '{self.name}'" if self.name else ""
            raise QueryBuilderError(f"Failed to resolve schema for CTE{name_hint}: {exc}") from exc
        if not self.params:
            return s
        if s is None:
            raise QueryBuilderError(f"CTE params were provided ({self.params!r}) but the source table has no schema")
        try:
            pairs = dict(safezip(self.params, s.values()))
        except ValueError as e:
            raise QueryBuilderError(
                f"CTE params length ({len(self.params)}) does not match source schema length ({len(s)})"
            ) from e
        result = s.new(pairs)
        if len(result) != len(s):
            raise QueryBuilderError(f"CTE params contain duplicate column names: {self.params!r}")
        return result


def _named_exprs_as_aliases(named_exprs):
    return [Alias(expr, name) for name, expr in named_exprs.items()]


def resolve_names(source_table, exprs):
    i = 0
    for expr in exprs:
        # Iterate recursively and update _ResolveColumn instances with the right expression
        if isinstance(expr, ExprNode):
            for v in expr._dfs_values():
                if isinstance(v, _ResolveColumn):
                    v.resolve(source_table._get_column(v.resolve_name))
                    i += 1


@attrs.define(frozen=False, eq=False)
class _ResolveColumn(LazyOps, ExprNode):
    resolve_name: str
    resolved: Expr | None = None

    def resolve(self, expr: Expr):
        if self.resolved is not None:
            raise QueryBuilderError("Already resolved!")
        self.resolved = expr

    def _get_resolved(self) -> Expr:
        if self.resolved is None:
            raise QueryBuilderError(f"Column not resolved: {self.resolve_name}")
        return self.resolved

    @property
    def type(self):
        return self._get_resolved().type

    @property
    def name(self):
        return self._get_resolved().name


@attrs.define(frozen=True)
class This:
    """Builder object for accessing table attributes.

    Automatically evaluates to the the 'top-most' table during compilation.
    """

    def __getattr__(self, name):
        return _ResolveColumn(name)

    def __getitem__(self, name):
        if isinstance(name, list | tuple):
            return [_ResolveColumn(n) for n in name]
        return _ResolveColumn(name)


@attrs.define(frozen=True, eq=False)
class In(ExprNode):
    expr: Expr
    list: Sequence[Expr]

    @property
    def type(self) -> type | None:
        return bool


@attrs.define(frozen=True, eq=False)
class Cast(ExprNode):
    expr: Expr
    target_type: Expr


@attrs.define(frozen=True, eq=False)
class Random(LazyOps, ExprNode):
    @property
    def type(self) -> type | None:
        return float


@attrs.define(frozen=True, eq=False)
class ConstantTable(ExprNode):
    rows: Sequence[Sequence]


@attrs.define(frozen=True, eq=False)
class Explain(ExprNode, Root):
    select: Select

    @property
    def type(self) -> type | None:
        return str


@attrs.define(frozen=True)
class CurrentTimestamp(ExprNode):
    @property
    def type(self) -> type | None:
        return datetime


# DDL


@attrs.define(frozen=True)
class Statement(Compilable, Root):
    @property
    def type(self) -> type | None:
        return None


@attrs.define(frozen=True, eq=False)
class CreateTable(Statement):
    path: TablePath
    source_table: Expr | None = None
    if_not_exists: bool = False
    primary_keys: list[str] | None = None


@attrs.define(frozen=True, eq=False)
class DropTable(Statement):
    path: TablePath
    if_exists: bool = False


@attrs.define(frozen=True, eq=False)
class TruncateTable(Statement):
    path: TablePath


@attrs.define(frozen=True, eq=False)
class InsertToTable(Statement):
    path: TablePath
    expr: Expr
    columns: list[str] | None = None
    returning_exprs: list[str] | None = None

    def returning(self, *exprs) -> Self:
        """Add a 'RETURNING' clause to the insert expression.

        Note: Not all databases support this feature!
        """
        if self.returning_exprs:
            raise ValueError("A returning clause is already specified")

        exprs = args_as_tuple(exprs)
        exprs = _drop_skips(exprs)
        if not exprs:
            return self

        resolve_names(self.path, exprs)
        return attrs.evolve(self, returning_exprs=exprs)


@attrs.define(frozen=True, eq=False)
class Commit(Statement):
    """Generate a COMMIT statement, if we're in the middle of a transaction, or in auto-commit. Otherwise SKIP."""