govoplan-dataflow/tests/test_padding_budget.py

from __future__ import annotations

import unittest

from govoplan_dataflow.backend.executor import (
    MAX_RESULT_BYTES,
    PipelineExecutionError,
    execute_preview,
)
from govoplan_dataflow.backend.expressions import ExpressionError, evaluate_expression
from govoplan_dataflow.backend.manifest import get_manifest
from govoplan_dataflow.backend.schemas import (
    GraphEdge,
    GraphNode,
    GraphPosition,
    PipelineGraph,
)


class _UnreadableFill:
    def __str__(self) -> str:
        raise AssertionError(
            "Oversized padding must be rejected before reading or multiplying its fill."
        )


def _expression_graph(expression: str, *, result_type: str = "string") -> PipelineGraph:
    return PipelineGraph(
        nodes=[
            GraphNode(
                id="source",
                type="source.inline",
                label="Source",
                position=GraphPosition(x=0, y=0),
                config={"source_name": "fixture", "rows": [{"value": "x"}]},
            ),
            GraphNode(
                id="padding",
                type="expression",
                label="Padding",
                position=GraphPosition(x=200, y=0),
                config={
                    "target_column": "padded",
                    "expression": expression,
                    "result_type": result_type,
                },
            ),
            GraphNode(
                id="output",
                type="output",
                label="Output",
                position=GraphPosition(x=400, y=0),
                config={},
            ),
        ],
        edges=[
            GraphEdge(id="source-padding", source="source", target="padding"),
            GraphEdge(id="padding-output", source="padding", target="output"),
        ],
    )


class PaddingBudgetTests(unittest.TestCase):
    def test_oversized_padding_is_rejected_before_fill_evaluation_or_allocation(
        self,
    ) -> None:
        for operation in ("lpad", "rpad"):
            for length in (MAX_RESULT_BYTES + 1, 10**100):
                with self.subTest(operation=operation, length=length):
                    with self.assertRaisesRegex(
                        ExpressionError, "Padding length.*1,000,000"
                    ):
                        evaluate_expression(
                            f"{operation}('x', {length}, fill)",
                            {"fill": _UnreadableFill()},
                        )

    def test_budget_cannot_be_bypassed_by_hiding_large_padding_in_a_small_scalar(
        self,
    ) -> None:
        for wrapper in ("length({})", "substring({}, 1, 1)"):
            with self.subTest(wrapper=wrapper):
                expression = wrapper.format(f"lpad('x', {MAX_RESULT_BYTES + 1}, fill)")
                with self.assertRaises(ExpressionError):
                    evaluate_expression(expression, {"fill": _UnreadableFill()})

    def test_existing_boundary_and_ordinary_padding_are_preserved(self) -> None:
        self.assertEqual(1_000_000, MAX_RESULT_BYTES)
        for operation in ("lpad", "rpad"):
            with self.subTest(operation=operation):
                value = evaluate_expression(
                    f"{operation}('x', {MAX_RESULT_BYTES}, '0')", {}
                )
                self.assertEqual(MAX_RESULT_BYTES, len(value))
                self.assertEqual(1, value.count("x"))
                self.assertEqual(
                    "abc", evaluate_expression(f"{operation}('abcdef', 3, '')", {})
                )
                self.assertEqual(
                    "", evaluate_expression(f"{operation}('abcdef', 0, '')", {})
                )
                self.assertEqual(
                    "abc", evaluate_expression(f"{operation}('abc', 3, '')", {})
                )

    def test_null_negative_and_empty_fill_semantics_are_unchanged(self) -> None:
        for operation in ("lpad", "rpad"):
            with self.subTest(operation=operation):
                self.assertIsNone(
                    evaluate_expression(f"{operation}(NULL, {10**100}, '')", {})
                )
                self.assertIsNone(evaluate_expression(f"{operation}(NULL, -1, '')", {}))
                with self.assertRaisesRegex(ValueError, "cannot be negative"):
                    evaluate_expression(f"{operation}('x', -1, '0')", {})
                with self.assertRaisesRegex(ValueError, "fill text cannot be empty"):
                    evaluate_expression(f"{operation}('x', 2, '')", {})

    def test_multibyte_fill_and_truncation_preserve_character_semantics(self) -> None:
        self.assertEqual("ö🙂öÄ", evaluate_expression("lpad('Ä', 4, 'ö🙂')", {}))
        self.assertEqual("Äö🙂ö", evaluate_expression("rpad('Ä', 4, 'ö🙂')", {}))
        self.assertEqual("🙂ä", evaluate_expression("lpad('🙂ä中', 2, '0')", {}))
        result = execute_preview(
            _expression_graph("rpad(value, 4, 'ö🙂')"), row_limit=10
        )
        self.assertEqual("xö🙂ö", result.rows[0]["padded"])

    def test_preview_reports_padding_guard_at_owning_node_and_retains_final_byte_limit(
        self,
    ) -> None:
        with self.assertRaisesRegex(PipelineExecutionError, "Padding length") as raised:
            execute_preview(
                _expression_graph(
                    f"length(lpad(value, {MAX_RESULT_BYTES + 1}, '0'))",
                    result_type="integer",
                ),
                row_limit=10,
            )
        self.assertEqual("padding", raised.exception.node_id)
        # Non-ASCII characters need several serialized bytes. The preallocation
        # character bound supplements, and never replaces, the node byte bound.
        with self.assertRaisesRegex(
            PipelineExecutionError, "one-megabyte result limit"
        ) as raised:
            execute_preview(_expression_graph("rpad(value, 200000, 'ö')"), row_limit=10)
        self.assertEqual("padding", raised.exception.node_id)

    def test_user_and_operator_documentation_explains_intermediate_padding_limit_in_both_languages(
        self,
    ) -> None:
        topic = next(
            topic
            for topic in get_manifest().documentation
            if topic.id == "dataflow.reference.nodes-and-expressions"
        )
        self.assertIn("user", topic.documentation_types)
        self.assertIn("admin", topic.documentation_types)
        for text in (topic.body, topic.translations["de"]["body"]):
            self.assertIn("LPAD", text)
            self.assertIn("RPAD", text)
            self.assertIn("LENGTH", text)
            self.assertIn("SUBSTRING", text)


if __name__ == "__main__":
    unittest.main()