optovia/prefect/flows/sync_offer.py

"""
Sync offer node to Optovia Maps graph.
- Upsert offer node
- Create edge offer -> location
- Signal Temporal workflow on completion
"""
import asyncio
import os
from dataclasses import dataclass
from datetime import datetime

from arango import ArangoClient
from prefect import flow, task, get_run_logger
from temporalio.client import Client as TemporalClient

ARANGODB_URL = os.environ.get("ARANGODB_URL", "http://217.216.32.39:32844")
ARANGODB_DATABASE = os.environ.get("ARANGODB_DATABASE", "optovia_maps")
ARANGODB_PASSWORD = os.environ.get("ARANGODB_PASSWORD", "")
TEMPORAL_HOST = os.environ.get("TEMPORAL_HOST", "temporal:7233")
TEMPORAL_NAMESPACE = os.environ.get("TEMPORAL_NAMESPACE", "default")


@dataclass
class PrefectSyncCompletedSignal:
    """Signal data when Prefect flow completes."""
    success: bool
    flow_run_id: str = ""
    message: str = ""


def get_db():
    client = ArangoClient(hosts=ARANGODB_URL)
    return client.db(ARANGODB_DATABASE, username="root", password=ARANGODB_PASSWORD)


@task
def upsert_offer_node(offer_data: dict) -> str:
    logger = get_run_logger()
    db = get_db()

    if not db.has_collection("nodes"):
        db.create_collection("nodes")

    nodes = db.collection("nodes")

    offer_uuid = offer_data["uuid"]
    doc = {
        "_key": offer_uuid,
        "name": offer_data.get("name") or f"Offer {offer_uuid}",
        "latitude": offer_data.get("location_latitude"),
        "longitude": offer_data.get("location_longitude"),
        "country": offer_data.get("location_country"),
        "country_code": offer_data.get("location_country_code"),
        "transport_types": [],
        "node_type": "offer",
        "supplier_uuid": offer_data.get("supplier_uuid"),
        "team_uuid": offer_data.get("team_uuid"),
        "product_uuid": offer_data.get("product_uuid"),
        "price_per_unit": offer_data.get("price_per_unit"),
        "currency": offer_data.get("currency"),
        "quantity": offer_data.get("quantity"),
        "unit": offer_data.get("unit"),
        "synced_at": datetime.utcnow().isoformat(),
    }

    if nodes.has(offer_uuid):
        nodes.update(doc)
        logger.info("Updated offer node %s", offer_uuid)
    else:
        nodes.insert(doc)
        logger.info("Inserted offer node %s", offer_uuid)

    return offer_uuid


@task(retries=3, retry_delay_seconds=5)
def signal_temporal_completion(
    workflow_id: str,
    success: bool,
    flow_run_id: str,
    message: str = "",
) -> bool:
    """Signal Temporal workflow that Prefect flow completed."""
    logger = get_run_logger()

    async def _signal():
        client = await TemporalClient.connect(TEMPORAL_HOST, namespace=TEMPORAL_NAMESPACE)
        handle = client.get_workflow_handle(workflow_id)
        await handle.signal(
            "prefect_sync_completed",
            PrefectSyncCompletedSignal(
                success=success,
                flow_run_id=flow_run_id,
                message=message,
            ),
        )

    try:
        asyncio.run(_signal())
        logger.info("Signaled Temporal workflow %s: success=%s", workflow_id, success)
        return True
    except Exception as e:
        logger.error("Failed to signal Temporal workflow %s: %s", workflow_id, e)
        raise


@task
def upsert_offer_edge(offer_uuid: str, location_uuid: str) -> bool:
    logger = get_run_logger()
    db = get_db()

    if not db.has_collection("edges"):
        db.create_collection("edges", edge=True)

    edges = db.collection("edges")

    edge_key = f"{offer_uuid}_{location_uuid}_offer"
    doc = {
        "_key": edge_key,
        "_from": f"nodes/{offer_uuid}",
        "_to": f"nodes/{location_uuid}",
        "type": "offer",
        "transport_type": "offer",
        "distance_km": 0,
        "travel_time_seconds": 0,
        "source": "offer",
        "priority": 1,
        "calculated_at": datetime.utcnow().isoformat(),
    }

    if edges.has(edge_key):
        edges.update(doc)
        logger.info("Updated offer edge %s -> %s", offer_uuid, location_uuid)
    else:
        edges.insert(doc)
        logger.info("Created offer edge %s -> %s", offer_uuid, location_uuid)

    return True


@flow(name="sync-offer", log_prints=True)
def sync_offer(offer_data: dict):
    """Sync offer to graph database and signal Temporal workflow."""
    from prefect.context import get_run_context

    logger = get_run_logger()

    offer_uuid = offer_data.get("uuid")
    location_uuid = offer_data.get("location_uuid")
    workflow_id = offer_data.get("workflow_id")

    if not offer_uuid or not location_uuid:
        raise ValueError("offer_data must include uuid and location_uuid")

    logger.info("Syncing offer %s to graph", offer_uuid)

    success = True
    message = ""

    try:
        upsert_offer_node(offer_data)
        upsert_offer_edge(offer_uuid, location_uuid)
        logger.info("Offer %s synced", offer_uuid)
    except Exception as e:
        success = False
        message = str(e)
        logger.error("Failed to sync offer %s: %s", offer_uuid, e)

    # Signal Temporal workflow if workflow_id provided
    if workflow_id:
        ctx = get_run_context()
        flow_run_id = str(ctx.flow_run.id) if ctx.flow_run else ""
        signal_temporal_completion(workflow_id, success, flow_run_id, message)

    if not success:
        raise RuntimeError(f"Offer sync failed: {message}")

    return {"offer_uuid": offer_uuid, "status": "ok"}