def test_set_and_get_context():
ctx = get_context()
assert "x" not in ctx
with set_context(x=1):
assert get_context()["x"] == 1
with set_context(x=2):
assert get_context()["x"] == 2
assert get_context()["x"] == 1
assert "x" not in ctx
async def resolve_set_flow_run_states(obj: Any, info: GraphQLResolveInfo,
input: dict) -> dict:
"""
Sets the flow run state, first deserializing a State from the provided input.
"""
server_context = context.get_context()
agent_id = server_context.get("headers", {}).get("x-prefect-agent-id")
async def check_size_and_set_state(state_input: dict) -> str:
state_size = sys.getsizeof(json.dumps(state_input["state"]))
if state_size > 1000000: # 1 mb max
raise ValueError("State payload is too large.")
state = state_schema.load(state_input["state"])
flow_run_id = state_input.get("flow_run_id")
await api.states.set_flow_run_state(
flow_run_id=flow_run_id,
version=state_input.get("version"),
state=state,
agent_id=agent_id,
)
return {"id": flow_run_id, "status": "SUCCESS", "message": None}
result = await asyncio.gather(*[
check_size_and_set_state(state_input)
for state_input in input["states"]
])
return {"states": result}
async def resolve_get_runs_in_queue(obj: Any, info: GraphQLResolveInfo,
input: dict) -> dict:
labels = input.get("labels", [])
labels.sort()
cloud_context = context.get_context()
agent_id = cloud_context.get("headers", {}).get("x-prefect-agent-id")
result = await api.runs.get_runs_in_queue(
tenant_id=input["tenant_id"],
before=input.get("before"),
labels=labels,
agent_id=agent_id,
)
return {"flow_run_ids": result}
def log_error(exc: Exception) -> None:
ctx = context.get_context()
ctx.pop("auth_token", None)
if config.env == "local":
logger.error(
textwrap.dedent(f"""
An application error occurred:
### --- Error ------------------------------
{textwrap.indent(traceback.format_exc(), " ")}
### --- Context ------------------------------
{textwrap.indent(str(ctx), " ")}
"""))
else:
logger.error({"traceback": traceback.format_exc(), "context": ctx})
async def f(x):
await asyncio.sleep(0.1)
with set_context(x=x):
assert x == get_context()["x"]
await asyncio.sleep(0.2)
assert x == get_context()["x"]
async def register_agent(
tenant_id: str,
labels: List[str],
agent_config_id: str = None,
name: str = None,
type: str = None,
) -> str:
"""
Register a new agent
Args:
- tenant_id (str): the id of a tenant
- labels (list): a list of strings representing the agent's labels
- agent_config_id (str): the id of an agent config to link this to
- name (str): the name of the agent
- type (str): the type of the agent
Returns:
- str: the agent id
"""
server_context = context.get_context()
core_version = server_context.get("headers",
{}).get("x-prefect-core-version")
if not tenant_id:
try:
tenant_id = await models.Tenant.where({
"id": {
"_eq": None
}
}).first().id
except:
raise ValueError("No tenant found.")
# Check for existing agents with these kwargs
agent = await models.Agent.where({
"_and": [
{
"tenant_id": {
"_eq": tenant_id
}
},
{
"name": {
"_eq": name
}
},
{
"type": {
"_eq": type
}
},
{
"core_version": {
"_eq": core_version
}
},
{
"labels": {
"_eq": sorted(labels or [])
}
},
]
}).first()
# Return existing agent ID
if agent:
return agent.id
# Insert new agent
return await models.Agent(
tenant_id=tenant_id,
agent_config_id=agent_config_id,
labels=sorted(labels or []),
name=name,
type=type,
core_version=core_version,
).insert()