DataOps-platform/tests/integration/test_kestra_runner_execution.py

"""Opt-in Kestra -> signed token -> Runner -> governed pool acceptance."""

import os
import time

import pytest
import requests
from neo4j import GraphDatabase
from sqlalchemy import create_engine, text

from app.core.common.identifiers import new_governance_uid
from app.core.data_source.credentials import (
    CredentialCodec,
    DataSourceCredentialRepository,
)
from app.core.data_source.definitions import DataSourceDefinitionRepository
from app.core.data_source.models import (
    DataSourceCredential,
    DataSourceDefinition,
)
from app.core.orchestration.compilers import compile_kestra_flow
from app.core.orchestration.engines import KestraAdapter
from app.runner.auth import TaskTokenIssuer


pytestmark = pytest.mark.integration


def test_kestra_passes_one_signed_task_to_the_runner():
    if os.environ.get("RUN_RUNNER_INTEGRATION") != "1":
        pytest.skip("set RUN_RUNNER_INTEGRATION=1 for Kestra Runner acceptance")

    platform_engine = create_engine(
        "postgresql://dataops:dataops-test-password@127.0.0.1:15432/dataops"
    )
    graph_driver = GraphDatabase.driver(
        "bolt://127.0.0.1:17687",
        auth=("neo4j", "Passw0rd"),
    )
    source_uid = new_governance_uid()
    dataflow_uid = new_governance_uid()
    node = {
        "id": "read_customers",
        "type": "sql.query",
        "data_source_uid": source_uid,
        "purpose": "read",
        "config": {
            "statement": (
                "SELECT customer_name FROM acceptance_customers "
                "WHERE id >= :minimum_id ORDER BY id"
            ),
            "parameters": {"minimum_id": "${parameters.minimum_id}"},
        },
    }
    task_uid = new_governance_uid()
    token = TaskTokenIssuer(
        "dataops-local-runner-task-token-secret-change-me",
        ttl_seconds=120,
    ).issue(
        task_uid=task_uid,
        dataflow_uid=dataflow_uid,
        workflow_version=1,
        correlation_id=new_governance_uid(),
        node=node,
    )
    compiled = compile_kestra_flow(
        {
            "schema_version": "1.0",
            "dataflow_uid": dataflow_uid,
            "name": "V52 Kestra Runner acceptance",
            "nodes": [node],
            "edges": [],
            "parameters": {
                "minimum_id": {"type": "integer", "required": True}
            },
        },
        {
            "schema_version": "1.0",
            "timezone": "Asia/Shanghai",
            "triggers": [{"type": "manual"}],
            "max_concurrency": 1,
            "conflict_policy": "skip",
            "timeout_seconds": 60,
            "retry": {"max_attempts": 1, "delay_seconds": 0},
            "backfill": {"max_days": 1, "max_runs": 1},
        },
        "test",
        1,
    )
    kestra = KestraAdapter(
        base_url="http://127.0.0.1:18080/api/v1",
        username="admin@dataops.local",
        password="DataOpsKestra1!",
    )
    flow_url = (
        "http://127.0.0.1:18080/api/v1/main/flows/"
        f"{compiled.namespace}/{compiled.flow_id}"
    )
    codec = CredentialCodec.from_base64(
        "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=",
        "v1",
    )
    try:
        with platform_engine.begin() as connection:
            sealed = DataSourceCredentialRepository(codec).create_version(
                connection,
                data_source_uid=source_uid,
                credential=DataSourceCredential(
                    "source_reader",
                    "source-test-password",
                ),
                actor_uid=None,
            )
        with graph_driver.session() as session:
            DataSourceDefinitionRepository(session).save(
                DataSourceDefinition(
                    uid=source_uid,
                    name_en="v52-kestra-runner-postgres",
                    database_type="postgresql",
                    host="source-postgres",
                    port=5432,
                    database="acceptance",
                    credential_ref=source_uid,
                    credential_version=sealed.credential_version,
                    pool_size=1,
                    max_overflow=1,
                )
            )

        kestra.deploy_disabled(compiled.yaml)
        kestra.activate(compiled.namespace, compiled.flow_id)
        execution = kestra.execute(
            compiled.namespace,
            compiled.flow_id,
            {
                "minimum_id": 1,
                "dataops_task_tokens": {"read_customers": token},
            },
        )
        execution_id = execution["id"]
        current = ""
        for _attempt in range(40):
            current = kestra.get_execution(execution_id)["state"]["current"]
            if current in {"SUCCESS", "FAILED", "KILLED", "CANCELLED"}:
                break
            time.sleep(0.25)
        assert current == "SUCCESS", kestra.get_logs(execution_id)

        with platform_engine.connect() as connection:
            ledger = connection.execute(
                text(
                    """
                    SELECT status, commit_outcome
                    FROM public.runner_task_executions
                    WHERE task_uid = CAST(:task_uid AS uuid)
                    """
                ),
                {"task_uid": task_uid},
            ).mappings().one()
        assert dict(ledger) == {
            "status": "success",
            "commit_outcome": "not_applicable",
        }
    finally:
        requests.delete(
            flow_url,
            auth=("admin@dataops.local", "DataOpsKestra1!"),
            timeout=20,
        )
        with graph_driver.session() as session:
            DataSourceDefinitionRepository(session).delete(source_uid)
        with platform_engine.begin() as connection:
            connection.execute(
                text(
                    "DELETE FROM public.runner_task_executions "
                    "WHERE task_uid = CAST(:task_uid AS uuid)"
                ),
                {"task_uid": task_uid},
            )
            connection.execute(
                text(
                    "DELETE FROM public.datasource_credential_audit_events "
                    "WHERE data_source_uid = CAST(:uid AS uuid)"
                ),
                {"uid": source_uid},
            )
            connection.execute(
                text(
                    "DELETE FROM public.datasource_credentials "
                    "WHERE data_source_uid = CAST(:uid AS uuid)"
                ),
                {"uid": source_uid},
            )
        graph_driver.close()
        platform_engine.dispose()