def test_flow_run_task_submit_args(self, client, cloud_api,
idempotency_key, task_run_id):
run_config = UniversalRun()
# verify that create_flow_run was called
task = StartFlowRun(
project_name="Test Project",
flow_name="Test Flow",
parameters={"test": "ing"},
run_config=run_config,
run_name="test-run",
)
# verify that run returns the new flow run ID
with prefect.context(task_run_id=task_run_id):
assert task.run(idempotency_key=idempotency_key) == "xyz890"
# verify the GraphQL query was called with the correct arguments
query_args = list(
client.graphql.call_args_list[0][0][0]["query"].keys())[0]
assert "Test Project" in query_args
assert "Test Flow" in query_args
# verify create_flow_run was called with the correct arguments
assert client.create_flow_run.call_args[1] == dict(
flow_id="abc123",
parameters={"test": "ing"},
run_config=run_config,
idempotency_key=idempotency_key or task_run_id,
context=None,
run_name="test-run",
scheduled_start_time=None,
)
def test_flow_run_task_with_wait(self, client, server_api):
# verify that create_flow_run was called
task = StartFlowRun(
flow_name="Test Flow",
project_name="Demo",
parameters={"test": "ing"},
run_name="test-run",
wait=True,
poll_interval=timedelta(seconds=3),
)
assert task.poll_interval == timedelta(seconds=3)
# Run flow, and assert that signals a success
with pytest.raises(signals.SUCCESS) as exc_info:
task.run()
flow_state_signal = exc_info.value
assert isinstance(flow_state_signal.state, state.Success)
# Check flow ID
assert str(flow_state_signal).split(" ")[0] == "xyz890"
# verify the GraphQL query was called with the correct arguments
query_args = list(
client.graphql.call_args_list[0][0][0]["query"].keys())[0]
assert "Test Flow" in query_args
# verify create_flow_run was called with the correct arguments
client.create_flow_run.assert_called_once_with(
flow_id="abc123",
parameters={"test": "ing"},
idempotency_key=None,
context=None,
run_name="test-run",
scheduled_start_time=None,
run_config=None,
)
def test_flow_run_task_with_no_matching_flow(self, client, server_api):
# verify a ValueError is raised if the client returns no flows
task = StartFlowRun(flow_name="Test Flow", project_name="Demo")
client.graphql = MagicMock(return_value=MagicMock(data=MagicMock(
flow=[])))
with pytest.raises(ValueError, match="Flow 'Test Flow' not found."):
task.run()
def test_flow_run_task(self, client, cloud_api):
# verify that create_flow_run was called
task = StartFlowRun(
project_name="Test Project",
flow_name="Test Flow",
parameters={"test": "ing"},
run_name="test-run",
)
# verify that run returns the new flow run ID
assert task.run() == "xyz890"
# verify the GraphQL query was called with the correct arguments
query_args = list(
client.graphql.call_args_list[0][0][0]["query"].keys())[0]
assert "Test Project" in query_args
assert "Test Flow" in query_args
# verify create_flow_run was called with the correct arguments
client.create_flow_run.assert_called_once_with(
flow_id="abc123",
parameters={"test": "ing"},
idempotency_key=None,
context=None,
run_name="test-run",
scheduled_start_time=None,
)
def test_flow_run_link_artifact(self, client, cloud_api):
task = StartFlowRun(
project_name="Test Project",
flow_name="Test Flow",
parameters={"test": "ing"},
run_name="test-run",
)
with prefect.context(running_with_backend=True, task_run_id="trid"):
task.run()
client.create_task_run_artifact.assert_called_once_with(
data={"link": "/flow/run/url"}, kind="link", task_run_id="trid"
)
def test_initialization(self, cloud_api):
now = pendulum.now()
run_config = UniversalRun()
# verify that the task is initialized as expected
task = StartFlowRun(
name="My Flow Run Task",
checkpoint=False,
project_name="Test Project",
flow_name="Test Flow",
new_flow_context={"foo": "bar"},
parameters={"test": "ing"},
run_config=run_config,
run_name="test-run",
scheduled_start_time=now,
)
assert task.name == "My Flow Run Task"
assert task.checkpoint is False
assert task.project_name == "Test Project"
assert task.flow_name == "Test Flow"
assert task.new_flow_context == {"foo": "bar"}
assert task.parameters == {"test": "ing"}
assert task.run_config == run_config
assert task.run_name == "test-run"
assert task.scheduled_start_time == now
def test_idempotency_key_uses_map_index_if_present(self, client,
cloud_api):
# verify that create_flow_run was called
task = StartFlowRun(project_name="Test Project", flow_name="Test Flow")
# verify that run returns the new flow run ID
with prefect.context(flow_run_id="test-id", map_index=4):
assert task.run() == "xyz890"
# verify create_flow_run was called with the correct arguments
client.create_flow_run.assert_called_once_with(
flow_id="abc123",
idempotency_key="test-id-4",
parameters=None,
context=None,
run_name=None,
)
def test_init_errors_if_tasks_passed_to_parameters(self, cloud_api):
with pytest.raises(TypeError,
match="An instance of `Task` was passed"):
StartFlowRun(name="testing",
parameters={
"a": 1,
"b": prefect.Parameter("b")
})
def test_flow_run_task_poll_interval_too_short(self):
with pytest.raises(ValueError):
StartFlowRun(
flow_name="Test Flow",
project_name="Demo",
parameters={"test": "ing"},
run_name="test-run",
wait=True,
poll_interval=timedelta(seconds=2),
)
def create_main_flow(flows: List[Flow], project_name):
schedule = CronSchedule("0 */3 * * *")
with Flow("MainFlow", schedule) as main_flow:
tasks = []
for flow in flows:
task = StartFlowRun(flow_name=flow.name,
project_name=project_name,
wait=True)
tasks.append(task)
parquet_flow = StartFlowRun(flow_name="UpdateParquetFiles",
project_name=project_name,
wait=True)
for task in tasks:
task.set_downstream(parquet_flow)
return main_flow
def test_flow_run_task_uses_scheduled_start_time(self, client, cloud_api):
in_one_hour = pendulum.now().add(hours=1)
# verify that create_flow_run was called
task = StartFlowRun(
project_name="Test Project",
flow_name="Test Flow",
scheduled_start_time=in_one_hour,
)
# verify that run returns the new flow run ID
assert task.run() == "xyz890"
# verify create_flow_run was called with the correct arguments
client.create_flow_run.assert_called_once_with(
flow_id="abc123",
parameters=None,
idempotency_key=None,
context=None,
run_name=None,
scheduled_start_time=in_one_hour,
)
def test_flow_run_task_with_flow_run_id(self, client, cloud_api):
# verify that create_flow_run was called
task = StartFlowRun(
project_name="Test Project",
flow_name="Test Flow",
parameters={"test": "ing"},
)
# verify that run returns the new flow run ID
with prefect.context(flow_run_id="test-id"):
assert task.run() == "xyz890"
# verify create_flow_run was called with the correct arguments
client.create_flow_run.assert_called_once_with(
flow_id="abc123",
parameters={"test": "ing"},
idempotency_key="test-id",
context=None,
run_name=None,
)
def create_cdc_all_states_flow():
"""Creates a flow that runs the CDC data update on all states."""
sched = CronSchedule("17 */4 * * *")
flow = Flow("CDCAllStatesDataUpdate", sched)
for state in ALL_STATES_PLUS_DC:
task = StartFlowRun(
flow_name=CDCCovidDataTracker.__name__,
project_name="can-scrape",
wait=True,
parameters={"state": state.abbr},
)
flow.add_task(task)
return flow
def create_prefect_flow_run(flow_name: str, project_name: str, task_refs: List,
params: Mapping) -> str:
"""Creates new prefect flow run for given flow id, parameters, task references
and API server URL to send GraphQL requests to.
Returns results value and state from a Prefect flow run.
"""
try:
flow_run = StartFlowRun(flow_name=flow_name,
project_name=project_name,
parameters=params)
flow_run_id = flow_run.run()
client = Client()
while True:
time.sleep(10)
flow_run_info = client.get_flow_run_info(flow_run_id)
flow_state = flow_run_info.state
task_runs_info = flow_run_info.task_runs
if flow_state.is_finished():
task_res_locs = {}
for task_run in task_runs_info:
# Return ref if ref string is a substring of any task slug
ref = next((ref_str for ref_str in task_refs
if ref_str in task_run.task_slug), None)
if ref:
task_id = task_run.id
task_state = client.get_task_run_state(task_id)
task_res_locs[ref] = task_state._result.location
task_results = {}
for ref, loc in task_res_locs.items():
local_res = LocalResult()
result = local_res.read(loc)
task_results[ref] = result.value
return task_results, flow_state, task_res_locs
except ValueError as err:
raise err
def test_initialization(self, server_api):
# verify that the task is initialized as expected
task = StartFlowRun(
name="My Flow Run Task",
project_name="Demo",
checkpoint=False,
flow_name="Test Flow",
new_flow_context={"foo": "bar"},
parameters={"test": "ing"},
run_name="test-run",
)
assert task.name == "My Flow Run Task"
assert task.checkpoint is False
assert task.flow_name == "Test Flow"
assert task.new_flow_context == {"foo": "bar"}
assert task.parameters == {"test": "ing"}
assert task.run_name == "test-run"
def test_flow_run_task_without_flow_name(self, server_api):
# verify that a ValueError is raised without a flow name
task = StartFlowRun()
with pytest.raises(ValueError, match="Must provide a flow name."):
task.run()
def test_flow_run_task_without_project_name(self, cloud_api):
# verify that a ValueError is raised without a project name
task = StartFlowRun(flow_name="Test Flow")
with pytest.raises(ValueError, match="Must provide a project name."):
task.run()
from prefect import task, Flow, Parameter
from prefect.storage import GitHub
# from prefect.engine.signals import FAIL
import time
# from prefect.run_configs import ECSRun
# from prefect.run_configs import LocalRun
from prefect.run_configs import KubernetesRun
import prefect
from prefect.tasks.prefect.flow_run import StartFlowRun
new_flow = StartFlowRun(flow_name="Skip Flow", project_name="Jenny")
@task(name="")
def sleep_for_x(x):
time.sleep(x)
prefect.artifacts.create_link("ftp://ftp-server/my-file.csv")
with Flow(name="Start Flow") as flow:
x = Parameter('x', default=22, required=True)
sleep_for_x(x)
new_flow()
flow.run_config = KubernetesRun(cpu_request=2, memory_request="2Gi")
# flow.run_config = LocalRun(
# labels=['runConfig']
# )
flow.storage = GitHub(
@task
def run_locally_or_in_cloud(input_string, manual_switch):
if manual_switch:
return manual_switch
length = len(input_string)
if length > 10:
return "local"
return "cloud"
cloud = StartFlowRun(
flow_name="ETL - Docker",
project_name="PGR Examples",
)
local = StartFlowRun(
flow_name="ETL - Local",
project_name="PGR Examples",
)
with Flow(
"Orchestrator Flow",
storage=GitHub(
repo="dylanbhughes/pgr_examples_3",
path="orchestrator.py",
secrets=["GITHUB_ACCESS_TOKEN"],
),
run_config=DockerRun(image="prefecthq/prefect:latest",
labels=["pgr docker"]),
