from __future__ import annotations import hashlib import json import os from types import SimpleNamespace import unittest from unittest.mock import Mock, patch from sqlalchemy.orm import Session from govoplan_core.auth import ApiPrincipal from govoplan_core.core.access import PrincipalRef from govoplan_core.core.datasources import ( CAPABILITY_DATASOURCE_CATALOGUE, DatasourceAccessError, DatasourceCatalogueProvider, DatasourceDescriptor, DatasourceReadRequest, DatasourceReadResult, ) from govoplan_core.security.bounded_process import run_bounded_operation from govoplan_dataflow.backend.backends.reference import ReferenceExecutionBackend from govoplan_dataflow.backend.executor import EXECUTOR_VERSION, PipelineExecutionError from govoplan_dataflow.backend.graph import validate_graph from govoplan_dataflow.backend.schemas import GraphEdge, GraphNode, GraphPosition, PipelineGraph from govoplan_dataflow.backend.service import _execute_pipeline_preview OUTER_ROWS = ({"id": "outer", "amount": 15}, {"id": "outer-low", "amount": 5}) INNER_ROWS = ({"id": "inner", "amount": 25}, {"id": "inner-low", "amount": 2}) def node(node_id: str, node_type: str, config: dict) -> GraphNode: return GraphNode( id=node_id, type=node_type, label=node_id, position=GraphPosition(x=0, y=0), config=config, ) def nested_source_graph() -> PipelineGraph: # Both external sources deliberately share their node ID and logical name. # The pinned subflow still has exactly one distinct inline input binding. nested = PipelineGraph( nodes=[ node("input", "source.inline", { "source_name": "bound_input", "rows": [], "input_binding": True, }), node("shared", "source.reference", { "source_name": "records", "source_ref": {"$parameter": "source_ref"}, "expected_fingerprint": "inner-pinned", "consistency": "frozen", }), node("union", "combine.union", {"mode": "all"}), node("minimum", "filter.expression", {"expression": "amount >= ${minimum}"}), node("output", "output", {}), ], edges=[ GraphEdge(id="input-union", source="input", target="union"), GraphEdge(id="shared-union", source="shared", target="union"), GraphEdge(id="union-minimum", source="union", target="minimum"), GraphEdge(id="minimum-output", source="minimum", target="output"), ], ) return PipelineGraph( nodes=[ node("shared", "source.reference", { "source_name": "records", "source_ref": "datasource:outer", "expected_fingerprint": "outer-pinned", "consistency": "current", }), node("nested", "subflow", { "template_ref": "fixture-nested-source", "template_version": "1", "parameters": {"source_ref": "datasource:inner", "minimum": 10}, "graph": nested.model_dump(mode="python"), }), node("output", "output", {}), ], edges=[ GraphEdge(id="shared-nested", source="shared", target="nested"), GraphEdge(id="nested-output", source="nested", target="output"), ], ) class ReferenceSubflowProcessTests(unittest.TestCase): def setUp(self) -> None: self.graph = nested_source_graph() self.assertEqual([], [item.model_dump() for item in validate_graph(self.graph) if item.severity == "error"]) self.session = Mock(spec=Session) self.principal = ApiPrincipal( principal=PrincipalRef( account_id="fixture-account", membership_id="fixture-membership", tenant_id="fixture-tenant", scopes=frozenset(), ), account=object(), user=object(), ) self.provider = Mock(spec=DatasourceCatalogueProvider) self.provider.read_datasource.side_effect = self.read_source self.registry = SimpleNamespace( has_capability=lambda name: name == CAPABILITY_DATASOURCE_CATALOGUE, capability=lambda name: self.provider if name == CAPABILITY_DATASOURCE_CATALOGUE else None, ) self.reads: list[tuple[int, DatasourceReadRequest]] = [] self.denied = False def read_source(self, session, principal, *, request: DatasourceReadRequest) -> DatasourceReadResult: self.assertIs(session, self.session) self.assertIs(principal, self.principal) self.reads.append((os.getpid(), request)) if request.datasource_ref == "datasource:inner" and self.denied: raise DatasourceAccessError("Current principal cannot read datasource:inner.") rows, fingerprint = { "datasource:outer": (OUTER_ROWS, "outer-pinned"), "datasource:inner": (INNER_ROWS, "inner-pinned"), }[request.datasource_ref] self.assertEqual(fingerprint, request.expected_fingerprint) return DatasourceReadResult( datasource=DatasourceDescriptor( ref=request.datasource_ref, source_name="records", name="Fixture source", kind="custom", mode="static", shape="tabular", fingerprint=fingerprint, provider="fixture-catalogue", ), rows=rows, total_rows=len(rows), truncated=False, ) def preview(self, *, row_limit: int = 10): return _execute_pipeline_preview( self.graph, session=self.session, principal=self.principal, registry=self.registry, backend="reference", row_limit=row_limit, preview_node_id="nested", ) def test_nested_parameterized_sources_execute_in_real_child_without_id_aliasing(self) -> None: with ( patch.object(ReferenceExecutionBackend, "_execute_in_process", side_effect=AssertionError("parent evaluation")), patch("govoplan_dataflow.backend.backends.reference.run_bounded_operation", wraps=run_bounded_operation) as worker, ): result, version = self.preview() worker.assert_called_once() self.assertEqual(EXECUTOR_VERSION, version) self.assertEqual([OUTER_ROWS[0], INNER_ROWS[0]], result.rows) self.assertEqual(2, result.total_rows) self.assertFalse(result.truncated) self.assertEqual(2, result.input_row_count) # Root input count, not the nested binding again. self.assertIsNotNone(result.node_preview) self.assertEqual("nested", result.node_preview.node_id) self.assertEqual(result.rows, result.node_preview.rows) self.assertEqual(2, result.node_preview.total_rows) self.assertEqual( [("shared", "succeeded", 0, 2), ("nested", "succeeded", 2, 2), ("output", "succeeded", 2, 2)], [(item.node_id, item.status, item.input_rows, item.output_rows) for item in result.node_diagnostics], ) binding_hash = hashlib.sha256(json.dumps(list(OUTER_ROWS), sort_keys=True, separators=(",", ":")).encode()).hexdigest() self.assertEqual([ {"node_id": "shared", "source_ref": "datasource:outer", "source_name": "records", "kind": "datasource", "provider": "fixture-catalogue", "fingerprint": "outer-pinned", "row_count": 2, "preview_rows": 2, "truncated": False}, {"node_id": "input", "source_name": "bound_input", "kind": "inline", "fingerprint": binding_hash, "row_count": 2, "subflow_node_id": "nested"}, {"node_id": "shared", "source_ref": "datasource:inner", "source_name": "records", "kind": "datasource", "provider": "fixture-catalogue", "fingerprint": "inner-pinned", "row_count": 2, "preview_rows": 2, "truncated": False, "subflow_node_id": "nested"}, ], result.source_fingerprints) self.assertEqual( [("datasource:outer", "current", "outer-pinned"), ("datasource:inner", "frozen", "inner-pinned")], [(request.datasource_ref, request.consistency, request.expected_fingerprint) for _, request in self.reads], ) self.assertTrue(all(pid == os.getpid() for pid, _ in self.reads)) self.assertTrue(all(request.limit <= 500 and request.offset == 0 for _, request in self.reads)) self.assertEqual([], self.session.mock_calls) def test_nested_result_and_node_preview_keep_full_totals_when_output_is_bounded(self) -> None: result, _ = self.preview(row_limit=1) self.assertEqual([OUTER_ROWS[0]], result.rows) self.assertEqual(2, result.total_rows) self.assertTrue(result.truncated) self.assertEqual([OUTER_ROWS[0]], result.node_preview.rows) self.assertEqual(2, result.node_preview.total_rows) self.assertTrue(result.node_preview.truncated) self.assertEqual({"outer-pinned", "inner-pinned"}, { item["fingerprint"] for item in result.source_fingerprints if item["kind"] == "datasource" }) self.assertEqual([], self.session.mock_calls) def test_denied_nested_datasource_stops_before_worker_or_persistence(self) -> None: self.denied = True with ( patch("govoplan_dataflow.backend.service.execute_typed_graph") as execute, patch("govoplan_dataflow.backend.backends.reference.run_bounded_operation") as worker, self.assertRaises(PipelineExecutionError) as caught, ): self.preview() execute.assert_not_called() worker.assert_not_called() self.assertEqual("shared", caught.exception.node_id) self.assertEqual("Current principal cannot read datasource:inner.", str(caught.exception)) self.assertFalse(caught.exception.retryable) self.assertIsInstance(caught.exception.__cause__, DatasourceAccessError) self.assertEqual(["datasource:outer", "datasource:inner"], [request.datasource_ref for _, request in self.reads]) self.assertEqual([], self.session.mock_calls)