DataOps-platform/tests/runner/test_nodes.py

import json

import pytest

from app.runner.nodes import (
    GovernedHttpExecutor,
    NodeExecutionError,
    NodeRegistry,
    RestrictedPythonExecutor,
    SqlExecuteExecutor,
    SqlQueryExecutor,
)


class Rows:
    def __init__(self, rows=(), rowcount=0):
        self._rows = rows
        self.rowcount = rowcount

    def mappings(self):
        return self

    def all(self):
        return list(self._rows)


class Connection:
    def __init__(self, result):
        self.result = result
        self.calls = []

    def execute(self, statement, parameters):
        self.calls.append((str(statement), parameters))
        return self.result


class Manager:
    def __init__(self, result):
        self.connection = Connection(result)
        self.purposes = []

    def connect(self, uid, purpose):
        manager = self

        class Context:
            def __enter__(self):
                manager.purposes.append((uid, purpose))
                return manager.connection

            def __exit__(self, *_args):
                return False

        return Context()


def test_sql_query_is_parameterized_read_only_and_bounded():
    manager = Manager(Rows([{"id": 1}, {"id": 2}]))
    executor = SqlQueryExecutor(manager, max_rows=1)
    node = {
        "id": "read",
        "type": "sql.query",
        "data_source_uid": "source-1",
        "purpose": "read",
        "config": {
            "statement": "SELECT id FROM orders WHERE day = :day",
            "parameters": {"day": "2026-07-19"},
        },
    }

    result = executor.execute(node, {})

    assert result == {"rows": [{"id": 1}], "row_count": 1, "truncated": True}
    assert manager.purposes == [("source-1", "dataflow_read")]
    assert manager.connection.calls == [
        ("SELECT id FROM orders WHERE day = :day", {"day": "2026-07-19"})
    ]


def test_sql_parameters_resolve_only_declared_runtime_placeholders():
    manager = Manager(Rows([{"id": 1}]))
    executor = SqlQueryExecutor(manager)
    node = {
        "id": "read",
        "type": "sql.query",
        "data_source_uid": "source-1",
        "purpose": "read",
        "config": {
            "statement": "SELECT id FROM orders WHERE day = :day",
            "parameters": {"day": "${parameters.day}"},
        },
    }

    executor.execute(node, {"day": "2026-07-19"})

    assert manager.connection.calls[-1][1] == {"day": "2026-07-19"}
    with pytest.raises(NodeExecutionError, match="runtime parameter"):
        executor.execute(node, {})


@pytest.mark.parametrize(
    "statement",
    [
        "DELETE FROM orders",
        "SELECT 1; DELETE FROM orders",
        "COPY orders TO PROGRAM 'id'",
        "WITH deleted AS (DELETE FROM orders RETURNING *) SELECT * FROM deleted",
    ],
)
def test_sql_query_rejects_non_read_statements(statement):
    executor = SqlQueryExecutor(Manager(Rows()))
    node = {
        "id": "read",
        "type": "sql.query",
        "data_source_uid": "source-1",
        "purpose": "read",
        "config": {"statement": statement, "parameters": {}},
    }

    with pytest.raises(NodeExecutionError, match="read-only"):
        executor.execute(node, {})


def test_sql_execute_requires_governed_write_and_idempotency():
    manager = Manager(Rows(rowcount=3))
    executor = SqlExecuteExecutor(manager)
    node = {
        "id": "write",
        "type": "sql.execute",
        "data_source_uid": "source-1",
        "purpose": "write",
        "config": {
            "statement": "UPDATE orders SET ready = :ready WHERE day = :day",
            "parameters": {"ready": True, "day": "2026-07-19"},
        },
        "idempotency": {
            "strategy": "deduplication_key",
            "key": "orders:2026-07-19",
        },
    }

    result = executor.execute(node, {}, write_authorized=True)

    assert result == {"affected_rows": 3, "commit_outcome": "committed"}
    assert manager.purposes == [("source-1", "dataflow_write")]

    with pytest.raises(NodeExecutionError, match="authorization"):
        executor.execute(node, {}, write_authorized=False)

    with pytest.raises(NodeExecutionError, match="idempotency"):
        executor.execute({**node, "idempotency": {}}, {}, write_authorized=True)


def test_sql_execute_never_converts_unknown_commit_into_a_retryable_failure():
    class UnknownManager:
        def connect(self, _uid, purpose):
            assert purpose == "dataflow_write"

            class Context:
                def __enter__(self):
                    return Connection(Rows(rowcount=1))

                def __exit__(self, *_args):
                    error = RuntimeError("connection lost during commit")
                    error.commit_outcome = "unknown"
                    raise error

            return Context()

    node = {
        "id": "write",
        "type": "sql.execute",
        "data_source_uid": "source-1",
        "purpose": "write",
        "config": {
            "statement": "UPDATE orders SET ready = :ready",
            "parameters": {"ready": True},
        },
        "idempotency": {
            "strategy": "deduplication_key",
            "key": "orders:2026-07-19",
        },
    }

    with pytest.raises(NodeExecutionError) as captured:
        SqlExecuteExecutor(UnknownManager()).execute(
            node,
            {},
            write_authorized=True,
        )

    assert captured.value.commit_outcome == "unknown"


@pytest.mark.parametrize(
    "statement",
    [
        "DROP TABLE orders",
        "GRANT ALL ON orders TO public",
        "UPDATE orders SET ready = true; DELETE FROM orders",
        "SELECT * FROM orders",
    ],
)
def test_sql_execute_accepts_dml_only(statement):
    node = {
        "id": "write",
        "type": "sql.execute",
        "data_source_uid": "source-1",
        "purpose": "write",
        "config": {"statement": statement, "parameters": {}},
        "idempotency": {
            "strategy": "deduplication_key",
            "key": "orders:2026-07-19",
        },
    }

    with pytest.raises(NodeExecutionError, match="DML"):
        SqlExecuteExecutor(Manager(Rows(rowcount=1))).execute(
            node,
            {},
            write_authorized=True,
        )


def test_restricted_python_executes_only_registered_handlers():
    executor = RestrictedPythonExecutor(
        {"normalize_orders": lambda values: {"count": len(values["orders"])}}
    )
    node = {
        "id": "normalize",
        "type": "python",
        "config": {
            "handler": "normalize_orders",
            "arguments": {"orders": [1, 2]},
        },
    }

    assert executor.execute(node, {}) == {"count": 2}

    for config in (
        {"source": "import os"},
        {"handler": "unknown", "arguments": {}},
    ):
        with pytest.raises(NodeExecutionError):
            executor.execute({**node, "config": config}, {})


def test_restricted_python_enforces_timeout_and_no_network():
    import socket
    import time

    def use_network(_arguments):
        socket.socket()

    def run_forever(_arguments):
        time.sleep(1)

    node = {
        "id": "restricted",
        "type": "python",
        "config": {"handler": "blocked", "arguments": {}},
    }

    with pytest.raises(NodeExecutionError, match="network"):
        RestrictedPythonExecutor(
            {"blocked": use_network},
            timeout_seconds=1,
        ).execute(node, {})

    with pytest.raises(NodeExecutionError, match="time limit"):
        RestrictedPythonExecutor(
            {"blocked": run_forever},
            timeout_seconds=0.05,
        ).execute(node, {})


def test_http_node_uses_allowlist_and_never_accepts_inline_authorization():
    calls = []

    class Session:
        def request(self, method, url, **kwargs):
            calls.append((method, url, kwargs))

            class Response:
                status_code = 200
                content = b'{"ok":true}'
                headers = {"content-type": "application/json"}

                def json(self):
                    return json.loads(self.content)

            return Response()

    executor = GovernedHttpExecutor(
        allowed_hosts={"api.internal.example"},
        session=Session(),
    )
    node = {
        "id": "notify",
        "type": "http",
        "config": {
            "method": "POST",
            "url": "https://api.internal.example/events",
            "json": {"state": "ready"},
        },
    }
    assert executor.execute(node, {})["body"] == {"ok": True}

    for config in (
        {**node["config"], "url": "http://127.0.0.1/admin"},
        {**node["config"], "headers": {"Authorization": "Bearer secret"}},
        {**node["config"], "url": "https://user:pass@api.internal.example/events"},
    ):
        with pytest.raises(NodeExecutionError):
            executor.execute({**node, "config": config}, {})


def test_node_registry_fails_closed_for_unregistered_types():
    registry = NodeRegistry({"sql.query": object()})
    with pytest.raises(NodeExecutionError, match="not registered"):
        registry.execute({"type": "notify"}, {})