Skip to content

rs_workflows/utils/prefect.md

<< Back to index

Utilities for working with Prefect variables.

update_prefect_variable(variable_name, updates) async

Merge updates into a Prefect variable and verify that they were persisted.

Source code in docs/rs-client-libraries/rs_workflows/utils/prefect.py
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
async def update_prefect_variable(variable_name: str, updates: dict[str, Any]) -> dict[str, Any]:
    """Merge updates into a Prefect variable and verify that they were persisted."""
    logger = get_run_logger()
    logger.info("Reading current Prefect variable %s", variable_name)

    raw_value = await cast(Awaitable[Any], Variable.get(variable_name, default={}))
    logger.info(
        "Read Prefect variable %s: type=%s, keys=%s",
        variable_name,
        type(raw_value).__name__,
        sorted(raw_value) if isinstance(raw_value, dict) else [],
    )

    value = _deep_merge(raw_value if isinstance(raw_value, dict) else {}, updates)
    logger.info("Updating Prefect variable %s with keys=%s", variable_name, sorted(updates))
    await cast(Awaitable[Any], Variable.set(variable_name, value, overwrite=True))

    saved_value = await cast(Awaitable[Any], Variable.get(variable_name, default={}))
    if not isinstance(saved_value, dict):
        raise RuntimeError(
            f"Prefect variable {variable_name!r} was not updated: expected a dictionary, "
            f"got {type(saved_value).__name__}",
        )

    if not _contains_updates(saved_value, updates):
        raise RuntimeError(f"Prefect variable {variable_name!r} was not updated: expected nested updates {updates!r}")

    logger.info("Verified Prefect variable %s update for keys=%s", variable_name, sorted(updates))
    return saved_value