insights/products/logs/backend/column_expressions.py

import hashlib

from insights.insightsql import ast
from insights.insightsql.functions.mapping import find_insightsql_aggregation
from insights.insightsql.parser import parse_expr
from insights.insightsql.visitor import TraversingVisitor

# A recognized source prefix selects the lowering strategy (source-driven):
# flat map lookup for attribute maps, JSON dig for the body string.
FLAT_MAP_SOURCES = ("attributes", "resource_attributes")
JSON_BODY_SOURCE = "body"


def path_to_expr(source: str, path: str) -> ast.Expr:
    """Lower a custom-column descriptor `(source, path)` into a InsightsQL AST expression.

    - `attributes` / `resource_attributes`: the whole `path` is a single map key
      (dots are part of the OTel key, never split) -> `attributes['http.url']`.
    - `body`: `path` is a dot-separated JSON path -> `JSONExtractString(body, 'user', 'id')`.

    `path` is always carried as a bound chain member / Constant, so it cannot escape
    the map-index or JSONExtract boundary as interpolated SQL.
    """
    if source in FLAT_MAP_SOURCES:
        return ast.Field(chain=[source, path])

    if source == JSON_BODY_SOURCE:
        segments = path.split(".")
        if any(segment == "" for segment in segments):
            raise ValueError(f"Body path {path!r} contains an empty segment")
        keys = [ast.Constant(value=segment) for segment in segments]
        return ast.Call(name="JSONExtractString", args=[ast.Field(chain=[JSON_BODY_SOURCE]), *keys])

    raise ValueError(f"Unknown custom-column source: {source!r}")


def parse_shorthand(text: str) -> ast.Expr | None:
    """Lower `<source>.<path>` shorthand, or return `None` if `text` has no recognized prefix.

    A bare source name with no path (`attributes`) is not shorthand — it falls through
    to full-expression parsing like any other field.
    """
    source, dot, path = text.partition(".")
    if not dot or not path:
        return None
    if source not in FLAT_MAP_SOURCES and source != JSON_BODY_SOURCE:
        return None
    return path_to_expr(source, path)


# Datastore functions that emit more than one row per input row. They have no place in a per-row
# custom column: a reader could submit e.g. `arrayJoin(range(1000000000))` and force the logs query
# to expand every matching row into billions, a low-friction availability hit. Matched case-insensitively.
# Mirrors insights.insightsql.database.schema.util.where_clause_extractor._ROW_MULTIPLYING_FUNCTIONS.
_ROW_MULTIPLYING_FUNCTIONS = {"arrayjoin"}

# Datastore functions whose output size is driven by a (typically constant) size argument rather than by
# the row's own data — a reader could submit e.g. `range(1000000000)` or `repeat(body, 1000000000)` and make
# Datastore build a huge array/string for every matching row without multiplying rows, a low-friction
# availability hit that the row-multiplier check above does not cover. They have no legitimate use in a
# per-row log column, so we reject them outright rather than trying to bound their argument. Matched
# case-insensitively.
_VALUE_GENERATING_FUNCTIONS = {
    "range",
    "arraywithconstant",
    "arrayresize",
    "repeat",
    "space",
    "randomstring",
    "randomstringutf8",
    "randomprintableascii",
    "randomfixedstring",
    "leftpad",
    "rightpad",
    "leftpadutf8",
    "rightpadutf8",
}


class _ScalarValidator(TraversingVisitor):
    def visit_select_query(self, node: ast.SelectQuery) -> None:
        raise ValueError("Custom columns cannot contain subqueries")

    def visit_select_set_query(self, node: ast.SelectSetQuery) -> None:
        raise ValueError("Custom columns cannot contain subqueries")

    def visit_placeholder(self, node: ast.Placeholder) -> None:
        raise ValueError("Custom columns cannot contain placeholders")

    def visit_window_function(self, node: ast.WindowFunction) -> None:
        # Window functions (e.g. `row_number() over (order by timestamp)`) evaluate over a frame of
        # rows, not per-row, forcing Datastore to process every matching log before the page limit
        # applies — the same availability risk as aggregations, which we already reject.
        raise ValueError("Custom columns cannot contain window functions")

    def visit_call(self, node: ast.Call) -> None:
        name = node.name.lower()
        if find_insightsql_aggregation(node.name) is not None:
            raise ValueError(f"Custom columns must be per-row: aggregation {node.name!r} is not allowed")
        if name in _ROW_MULTIPLYING_FUNCTIONS:
            raise ValueError(f"Custom columns must be per-row: {node.name!r} can multiply rows and is not allowed")
        if name in _VALUE_GENERATING_FUNCTIONS:
            raise ValueError(
                f"Custom columns must be bounded: {node.name!r} can generate arbitrarily large values and is not allowed"
            )
        super().visit_call(node)


def _validate_scalar(expr: ast.Expr) -> None:
    """Reject subqueries, aggregations, row-multiplying and value-generating functions, and unresolved
    placeholders anywhere in `expr`.

    Scalar function calls, field access, arithmetic, conditionals, and constants all
    pass. This is an AST check — the expression is never inspected as a string.
    """
    _ScalarValidator().visit(expr)


def column_to_expr(text: str) -> ast.Expr:
    """Resolve a custom-column string into a validated per-row InsightsQL AST expression.

    Recognized source prefix -> shorthand lowering; otherwise the whole string is
    parsed as InsightsQL (`parse_expr`, so input is never interpolated as SQL). Raises
    `ValueError` for non-scalar expressions and InsightsQL syntax errors for unparsable input.
    """
    text = text.strip()
    expr = parse_shorthand(text)
    if expr is None:
        expr = parse_expr(text)
    _validate_scalar(expr)
    return expr


def canonical_key(text: str) -> str:
    """Derive a stable `col_<hash>` alias from the expression string.

    Independent of any client-generated uuid so identical column setups hash to the
    same query (and therefore the same query cache key).
    """
    digest = hashlib.sha256(text.strip().encode()).hexdigest()
    return f"col_{digest[:12]}"