posthog/common/hogvm/python/utils.py

from typing import Any

import re2

from common.hogvm.python.objects import is_hog_date, is_hog_datetime
from common.hogvm.python.stl.date import date_string_to_seconds, to_hog_datetime

_CASE_INSENSITIVE_OPTS = re2.Options()
_CASE_INSENSITIVE_OPTS.case_sensitive = False

COST_PER_UNIT = 8
MAX_MEMORY = 64 * 1024 * 1024  # 64 MB
MAX_REGEX_PATTERN_LENGTH = 16 * 1024


def _temporal_seconds(value: Any) -> float | None:
    """Epoch seconds for a Hog datetime/date value, else None. A bare HogDate is UTC midnight."""
    if is_hog_datetime(value):
        return value["dt"]
    if is_hog_date(value):
        return to_hog_datetime(value)["dt"]
    return None


class HogVMException(Exception):
    pass


class UncaughtHogVMException(HogVMException):
    type: str
    message: str
    payload: Any

    def __init__(self, type, message, payload):
        super().__init__(message)
        self.type = type
        self.message = message
        self.payload = payload

    def __str__(self):
        msg = self.message.replace("'", "\\'")
        return f"{self.type}('{msg}')"


class HogVMRuntimeExceededException(HogVMException):
    """Exception thrown when HogVM code exceeds its runtime limit"""

    def __init__(self, timeout_seconds: float, ops_performed: int):
        self.timeout_seconds = timeout_seconds
        self.ops_performed = ops_performed
        super().__init__(f"Runtime exceeded {timeout_seconds} seconds after {ops_performed} operations")


class HogVMMemoryExceededException(HogVMException):
    """Exception thrown when HogVM code exceeds its memory limit"""

    def __init__(self, memory_limit: int, attempted_memory: int):
        self.memory_limit = memory_limit
        self.attempted_memory = attempted_memory
        super().__init__(f"Memory limit of {memory_limit} bytes exceeded. Attempted to use {attempted_memory} bytes")


def _require_string(value: Any, name: str, function_name: str) -> str:
    if not isinstance(value, str):
        raise HogVMException(f"Function {function_name} requires {name} to be a string, got {type(value).__name__}")
    return value


def _format_regex_error(error: Exception) -> str:
    if error.args and isinstance(error.args[0], bytes):
        return error.args[0].decode("utf-8", errors="replace")
    return str(error)


def _validate_regex_pattern(pattern: str) -> None:
    # Pattern compilation runs inside one opcode, before the VM can check its time budget again.
    if len(pattern) > MAX_REGEX_PATTERN_LENGTH:
        raise HogVMException(f"Pattern exceeds {MAX_REGEX_PATTERN_LENGTH} characters. Use a shorter pattern.")


def _compile_regex(pattern: str, case_insensitive: bool = False) -> Any:
    # re2 matches in linear time, unlike Python's backtracking re engine. It also makes the character
    # classes ASCII-only, so `\w+` extracts "caf" from "café" and `^\w+$` does not match "Müller".
    # The =~ operator, like(), the Node VM and ClickHouse all use re2, so the Hog surfaces agree.
    try:
        return re2.compile(pattern, options=_CASE_INSENSITIVE_OPTS) if case_insensitive else re2.compile(pattern)
    except re2.error as e:
        raise HogVMException(f"Invalid regex pattern: {_format_regex_error(e)}") from e


def regex_match(string: Any, pattern: Any, case_insensitive: bool = False) -> bool:
    if not string or not pattern:
        return False

    string = _require_string(string, "input", "match")
    pattern = _require_string(pattern, "pattern", "match")
    _validate_regex_pattern(pattern)
    return _compile_regex(pattern, case_insensitive).search(string) is not None


def regex_extract(string: Any, pattern: Any) -> str:
    # Matches ClickHouse extract(): first capture group if the pattern has groups, else the whole
    # match, else empty.
    if string is None or pattern is None:
        return ""
    haystack = str(string)
    pattern = str(pattern)
    _validate_regex_pattern(pattern)
    try:
        compiled = _compile_regex(pattern)
    except HogVMException:
        return ""
    found = compiled.search(haystack)
    if not found:
        return ""
    # Select on the pattern's static group count, not on which groups participated. A group that
    # captured nothing still wins over the whole match, so `(a)?b` on "b" gives "". The Node and
    # Rust VMs branch on the same static count.
    if compiled.groups > 0:
        return found.group(1) or ""
    return found.group(0) or ""


def like(string: Any, pattern: Any, case_insensitive: bool = False) -> bool:
    _validate_regex_pattern(pattern)
    # re2.escape backslash-escapes % (and not _), so undo that before the wildcards turn into regex.
    pattern = re2.escape(pattern).replace("\\%", "%").replace("%", ".*").replace("_", ".")
    re_pattern = re2.compile(pattern, options=_CASE_INSENSITIVE_OPTS) if case_insensitive else re2.compile(pattern)
    return re_pattern.search(string) is not None


def get_nested_value(obj, chain, nullish=False) -> Any:
    if obj is None:
        return None
    for key in chain:
        if nullish and obj is None:
            return None
        if isinstance(key, int):
            if key == 0:
                raise HogVMException(f"Hog arrays start from index 1")
            elif key > 0:
                if key > len(obj):
                    return None
                obj = obj[key - 1]
            elif key < 0:
                if -key > len(obj):
                    return None
                obj = obj[key]
        else:
            obj = obj.get(key, None)
    return obj


def set_nested_value(obj, chain, value) -> Any:
    if obj is None:
        return None
    for key in chain[:-1]:
        if isinstance(key, int):
            obj = obj[key]
        else:
            obj = obj.get(key, None)

    if isinstance(obj, dict):
        obj[chain[-1]] = value
    elif isinstance(obj, list):
        if not isinstance(chain[-1], int):
            raise HogVMException(f"Invalid index: {chain[-1]}")
        if chain[-1] <= 0:
            raise HogVMException(f"Hog arrays start from index 1")
        if chain[-1] > len(obj):
            raise HogVMException(f"Index {chain[-1]} out of range for array of length {len(obj)}")
        obj[chain[-1] - 1] = value
    else:
        raise HogVMException(f'Can not set property "{chain[-1]}" on object of type "{type(obj).__name__}"')

    return obj


def calculate_cost(object, marked: set | None = None) -> int:
    if marked is None:
        marked = set()
    if isinstance(object, dict) or isinstance(object, list) or isinstance(object, tuple):
        if id(object) in marked:
            return COST_PER_UNIT
        marked.add(id(object))
        try:
            if isinstance(object, dict):
                return COST_PER_UNIT + sum(
                    [calculate_cost(key, marked) + calculate_cost(value, marked) for key, value in object.items()]
                )
            elif isinstance(object, list) or isinstance(object, tuple):
                return COST_PER_UNIT + sum([calculate_cost(val, marked) for val in object])
        finally:
            marked.remove(id(object))
    elif isinstance(object, str):
        return COST_PER_UNIT + len(object)
    return COST_PER_UNIT


def unify_comparison_types(left, right):
    # Two temporal values order by epoch seconds (matching ClickHouse and the TS/Rust VMs). Without
    # this a HogDateTime/HogDate dict falls through unchanged and ordering operators end up comparing
    # dicts, which Python can't order.
    left_seconds = _temporal_seconds(left)
    right_seconds = _temporal_seconds(right)
    if left_seconds is not None and right_seconds is not None:
        return left_seconds, right_seconds
    # A bare-field SQL comparison like `timestamp > toDateTime(...)` puts a plain date-like string
    # against a HogDateTime/HogDate object. Parse the string the same way `toDateTime` would rather
    # than falling through to the string/number branches below, which leave it unordered.
    if left_seconds is not None and isinstance(right, str):
        right_seconds_from_string = date_string_to_seconds(right)
        if right_seconds_from_string is not None:
            return left_seconds, right_seconds_from_string
    if right_seconds is not None and isinstance(left, str):
        left_seconds_from_string = date_string_to_seconds(left)
        if left_seconds_from_string is not None:
            return left_seconds_from_string, right_seconds

    # Handle boolean cases FIRST since bool is a subclass of int in Python
    if isinstance(left, bool) and isinstance(right, str):
        # Convert string to boolean: 'true'/'false' strings, or truthy/falsy
        if right.lower() == "true":
            return left, True
        elif right.lower() == "false":
            return left, False
        else:
            return left, bool(right)
    if isinstance(left, str) and isinstance(right, bool):
        # Convert string to boolean: 'true'/'false' strings, or truthy/falsy
        if left.lower() == "true":
            return True, right
        elif left.lower() == "false":
            return False, right
        else:
            return bool(left), right
    if isinstance(left, bool) and isinstance(right, int | float):
        return int(left), right
    if isinstance(left, int | float) and isinstance(right, bool):
        return left, int(right)

    # Handle numeric conversions (after boolean checks)
    if isinstance(left, int | float) and isinstance(right, str):
        try:
            return left, float(right)
        except ValueError:
            return left, right
    if isinstance(left, str) and isinstance(right, int | float):
        try:
            return float(left), right
        except ValueError:
            return left, right

    return left, right