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"]),