Artifex/graph/bootstrap.py

from __future__ import annotations

from django.utils import timezone

from graph.models import ExecutionGraphDefinition, ExecutionGraphVersion, ExecutionGraphVersionStatus
from graph.steward import steward_run_graph_v1
from graph.task_execution import task_execution_graph_v1


def champion_task_execution_graph_v1() -> ExecutionGraphVersion:
    spec = task_execution_graph_v1()
    definition, _ = ExecutionGraphDefinition.objects.get_or_create(
        name=spec.name,
        defaults={"graph_type": spec.graph_type, "description": str(spec.metadata.get("description", ""))},
    )
    version, created = ExecutionGraphVersion.objects.get_or_create(
        graph=definition,
        version=spec.version,
        defaults={
            "status": ExecutionGraphVersionStatus.CHAMPION,
            "graph_spec": spec.to_dict(),
            "metadata": {"immutable_after_use": True},
            "promoted_at": timezone.now(),
        },
    )
    if not created and version.status != ExecutionGraphVersionStatus.CHAMPION:
        version.status = ExecutionGraphVersionStatus.CHAMPION
        version.promoted_at = timezone.now()
        version.save(update_fields=["status", "promoted_at"])
    return version


def champion_steward_run_graph_v1() -> ExecutionGraphVersion:
    spec = steward_run_graph_v1()
    definition, _ = ExecutionGraphDefinition.objects.get_or_create(
        name=spec.name,
        defaults={"graph_type": spec.graph_type, "description": str(spec.metadata.get("description", ""))},
    )
    version, created = ExecutionGraphVersion.objects.get_or_create(
        graph=definition,
        version=spec.version,
        defaults={
            "status": ExecutionGraphVersionStatus.CHAMPION,
            "graph_spec": spec.to_dict(),
            "metadata": {"immutable_after_use": True},
            "promoted_at": timezone.now(),
        },
    )
    if not created and version.status != ExecutionGraphVersionStatus.CHAMPION:
        version.status = ExecutionGraphVersionStatus.CHAMPION
        version.promoted_at = timezone.now()
        version.save(update_fields=["status", "promoted_at"])
    return version