insights/insights/insightsql/base.py

import re
import copy
from dataclasses import dataclass, field, fields
from enum import Enum
from typing import TYPE_CHECKING, Any, Literal, Optional, TypeVar

from insights.insightsql.constants import ConstantDataType
from insights.insightsql.errors import NotImplementedError

if TYPE_CHECKING:
    from insights.insightsql.context import InsightsQLContext

# Given a string like "CorrectHorseBS", match the "H" and "B", so that we can convert this to "correct_horse_bs"
camel_case_pattern = re.compile(r"(?<!^)(?<![A-Z])(?=[A-Z])")


# Class-name → snake-case overrides that don't fall out of the default regex.
# NOTE: Sync with ./test/test_visitor.py#test_insightsql_visitor_naming_exceptions
#
# These keys are the REGEX OUTPUT, not the class name, and that is the whole
# trap. `InsightsQLXTag` reduces to `insights_qlxtag`, because the pattern only
# breaks before a capital that follows a lower-case letter — so the run `QLXT`
# stays welded to the `tag`. The keys said `script_qlxtag`, which is what the
# class reduced to under an earlier name; the class was renamed and the key was
# not, so the replacement stopped matching anything.
#
# Nothing failed at import. accept() computes the name, finds no such method,
# and raises only when a query actually contains the tag — so the whole
# InsightsQLX surface was unreachable while every test that avoided it passed.
_VISIT_NAME_REPLACEMENTS = {
    "insights_qlxtag": "insightsqlx_tag",
    "insights_qlxattribute": "insightsqlx_attribute",
    "uuidtype": "uuid_type",
    "string_jsontype": "string_json_type",
}


def _compute_visit_method_name(cls_name: str) -> str:
    """Camel → snake conversion of a class name, with the replacements above
    applied. Computed once per AST subclass at class-creation time (see
    ``AST.__init_subclass__``) and cached on the class as
    ``_visit_method_name``."""
    name = camel_case_pattern.sub("_", cls_name).lower()
    for old, new in _VISIT_NAME_REPLACEMENTS.items():
        name = name.replace(old, new)
    return f"visit_{name}"


@dataclass(kw_only=True, slots=True)
class AST:
    start: Optional[int] = field(default=None)
    end: Optional[int] = field(default=None)

    # Populated by __init_subclass__ — every concrete AST class gets its own
    # method name computed once at class-creation time, so accept() is just
    # an attribute read and a getattr per visit instead of a regex + four
    # str.replace + f-string on every node walk.
    _visit_method_name = "visit_ast"

    def __init_subclass__(cls, **kwargs):
        # Deliberately not calling super().__init_subclass__: AST's parent is
        # `object` (whose hook is a no-op anyway), and `super()` here trips
        # the same class-identity quirk that affects AST.__str__ below — the
        # method's __class__ closure cell points at the pre-
        # @dataclass(slots=True) AST, which isn't in the new MRO.
        #
        # This body runs twice per slotted subclass (once before the
        # @dataclass(slots=True) rebuild, once after) — keep it idempotent.
        # Extension point for overriding the cached name is
        # `_VISIT_NAME_REPLACEMENTS` above, not a class-level assignment
        # of `_visit_method_name` (which would be clobbered here).
        cls._visit_method_name = _compute_visit_method_name(cls.__name__)

    # This is part of the visitor pattern from visitor.py.
    def accept(self, visitor):
        method_name = self._visit_method_name
        visit = getattr(visitor, method_name, None)
        if visit is not None:
            return visit(self)
        visit_unknown = getattr(visitor, "visit_unknown", None)
        if visit_unknown is not None:
            return visit_unknown(self)
        raise NotImplementedError(f"{visitor.__class__.__name__} has no method {method_name}")

    def __deepcopy__(self, memo: dict[int, Any]) -> "AST":
        # Faster than stdlib pickle; memo-before-recurse (here + _clone_value) preserves cycles and shared subtrees.
        cls = self.__class__
        plan = _clone_plan_cache.get(cls)
        if plan is None:
            plan = tuple((f.name, f.name not in ("start", "end")) for f in fields(cls))
            _clone_plan_cache[cls] = plan
        new = cls.__new__(cls)  # bare instance, no __init__/__post_init__
        memo[id(self)] = new  # register BEFORE recursing so cycles terminate
        for name, needs_clone in plan:
            # Assumes every slot is set, true for any node built via the dataclass __init__ (the only path in practice).
            value = getattr(self, name)
            setattr(new, name, _clone_value(value, memo) if needs_clone else value)
        return new

    def to_insightsql(self):
        from insights.insightsql.context import InsightsQLContext
        from insights.insightsql.printer import print_prepared_ast

        return print_prepared_ast(
            node=self,
            context=InsightsQLContext(enable_select_queries=True, limit_top_select=False),
            dialect="insightsql",
        )

    def __str__(self):
        if isinstance(self, Type):
            # `repr(self)` instead of `super().__str__()` because
            # @dataclass(slots=True) rebuilds the class object, and the
            # method's super() closure cell points at the pre-slots class
            # which is no longer in self's MRO. object.__str__ falls back
            # to __repr__ anyway, and __repr__ is generated by @dataclass.
            return repr(self)
        return f"sql({self.to_insightsql()})"


# --- deep-copy fast path (see AST.__deepcopy__) ---------------------------------
# Immutable scalars deepcopy returns as-is; everything else is rebuilt.
_CLONE_ATOMIC_TYPES = frozenset({str, bytes, int, float, bool, complex, type(None)})
# Per-class clone plan: tuple of (field_name, needs_clone), built once from dataclasses.fields.
_clone_plan_cache: dict[type, tuple[tuple[str, bool], ...]] = {}
_NOT_CLONED = object()  # memo-miss sentinel; a real clone may itself be None/falsy


def _clone_value(value: Any, memo: dict[int, Any]) -> Any:
    cls = type(value)
    if cls in _CLONE_ATOMIC_TYPES:
        return value
    # Dedup via memo so shared subobjects clone once and cycles terminate; containers register before recursing.
    existing = memo.get(id(value), _NOT_CLONED)
    if existing is not _NOT_CLONED:
        return existing
    if cls is list:
        new_list: list = []
        memo[id(value)] = new_list
        for item in value:
            new_list.append(_clone_value(item, memo))
        return new_list
    if isinstance(value, AST):
        return value.__deepcopy__(memo)  # registers itself in memo before recursing
    if cls is dict:
        new_dict: dict = {}
        memo[id(value)] = new_dict
        for key, item in value.items():
            new_dict[_clone_value(key, memo)] = _clone_value(item, memo)  # keys too, like stdlib
        return new_dict
    if isinstance(value, Enum):
        return value
    # Tuples/sets/embedded non-AST -> stdlib (shares memo) for exact cycle identity; AST children stay fast.
    return copy.deepcopy(value, memo)


_T_AST = TypeVar("_T_AST", bound=AST)


@dataclass(kw_only=True, slots=True)
class Type(AST):
    def get_child(self, name: str, context: "InsightsQLContext") -> "Type":
        raise NotImplementedError("Type.get_child not overridden")

    def has_child(self, name: str, context: "InsightsQLContext") -> bool:
        return self.get_child(name, context) is not None

    def resolve_constant_type(self, context: "InsightsQLContext") -> "ConstantType":
        raise NotImplementedError(f"{self.__class__.__name__}.resolve_constant_type not overridden")

    def resolve_column_constant_type(self, name: str, context: "InsightsQLContext") -> "ConstantType":
        raise NotImplementedError(f"{self.__class__.__name__}.resolve_column_constant_type not overridden")


@dataclass(kw_only=True, slots=True)
class Expr(AST):
    type: Optional[Type] = field(default=None)


@dataclass(kw_only=True, slots=True)
class CTE(Expr):
    """A common table expression."""

    name: str
    expr: Expr
    # Whether the CTE is an inlined column "WITH 1 AS a" or a subquery "WITH a AS (SELECT 1)"
    cte_type: Literal["column", "subquery"]
    columns: Optional[list[str]] = field(default=None)  # e.g. ["a", "b"] for `WITH x(a, b) AS (SELECT 1, 2)`
    recursive: bool = field(default=False)
    materialized: Optional[bool] = field(default=None)
    using_key: Optional[list[str]] = field(default=None)  # e.g. ["a"] for USING KEY (a)


@dataclass(kw_only=True, slots=True)
class ConstantType(Type):
    data_type: ConstantDataType
    nullable: bool = field(default=True)

    def resolve_constant_type(self, context: "InsightsQLContext") -> "ConstantType":
        return self

    def print_type(self) -> str:
        raise NotImplementedError("ConstantType.print_type not implemented")


@dataclass(kw_only=True, slots=True)
class UnknownType(ConstantType):
    data_type: ConstantDataType = field(default="unknown", init=False)
    # True when the type genuinely could not be determined (e.g. an unmapped function), so the
    # value could be anything. Such a type poisons supertype unification rather than being absorbed.
    # The default (False) marks a vacuous unknown — a null literal or empty container — that imposes
    # no constraint and is absorbed, letting a known sibling type win.
    unanalyzable: bool = field(default=False)

    def print_type(self) -> str:
        return "Unknown"