Skip to content

Commit

Permalink
Add a --replace option to "schedule create" command (#12123)
Browse files Browse the repository at this point in the history
  • Loading branch information
@abrookins
abrookins authored Feb 29, 2024
1 parent 6dcd3f4 commit e297260
Show file tree
Hide file tree
Showing 2 changed files with 208 additions and 6 deletions.
43 changes: 37 additions & 6 deletions src/prefect/cli/deployment.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -340,6 +340,17 @@ async def create_schedule(
"--active",
help="Whether the schedule is active. Defaults to True.",
),
replace: Optional[bool] = typer.Option(
False,
"--replace",
help="Replace the deployment's current schedule(s) with this new schedule.",
),
assume_yes: Optional[bool] = typer.Option(
False,
"--accept-yes",
"-y",
help="Accept the confirmation prompt without prompting",
),
):
"""
Create a schedule for a given deployment.
Expand Down Expand Up @@ -409,8 +420,29 @@ async def create_schedule(
except ObjectNotFound:
return exit_with_error(f"Deployment {name!r} not found!")

num_schedules = len(deployment.schedules)
noun = "schedule" if num_schedules == 1 else "schedules"

if replace and num_schedules > 0:
if not assume_yes and not typer.confirm(
f"Are you sure you want to replace {num_schedules} {noun} for {name}?"
):
return exit_with_error("Schedule replacement cancelled.")

for existing_schedule in deployment.schedules:
try:
await client.delete_deployment_schedule(
deployment.id, existing_schedule.id
)
except ObjectNotFound:
pass

await client.create_deployment_schedules(deployment.id, [(schedule, active)])
exit_with_success("Created deployment schedule!")

if replace and num_schedules > 0:
exit_with_success(f"Replaced existing deployment {noun} with new schedule!")
else:
exit_with_success("Created deployment schedule!")


@schedule_app.command("delete")
Expand Down Expand Up @@ -440,8 +472,8 @@ async def delete_schedule(
except IndexError:
return exit_with_error("Deployment schedule not found!")

if not assume_yes and not typer.prompt(
f"Are you sure you want to delete this schedule: {schedule.schedule} (y/n)",
if not assume_yes and not typer.confirm(
f"Are you sure you want to delete this schedule: {schedule.schedule}",
):
return exit_with_error("Deletion cancelled.")

Expand Down Expand Up @@ -578,9 +610,8 @@ async def clear_schedules(
await client.read_flow(deployment.flow_id)

# Get input from user: confirm removal of all schedules
if not assume_yes and not typer.prompt(
"Are you sure you want to clear all schedules for this deployment? (y/n)",
type=bool,
if not assume_yes and not typer.confirm(
"Are you sure you want to clear all schedules for this deployment?",
):
exit_with_error("Clearing schedules cancelled.")

Expand Down
171 changes: 171 additions & 0 deletions tests/cli/deployment/test_deployment_cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -943,6 +943,177 @@ def test_set_schedule_invalid_interval_anchor_raises(self, flojo, commands, erro
expected_output_contains=error,
)

def test_create_schedule_replace_replaces_existing_schedules_many(self, flojo):
create_args = [
"deployment",
"schedule",
"create",
"rence-griffith/test-deployment",
"--interval",
]

invoke_and_assert(
[
*create_args,
"90",
],
expected_code=0,
expected_output_contains="Created deployment schedule!",
)

invoke_and_assert(
[
*create_args,
"1800",
],
expected_code=0,
expected_output_contains="Created deployment schedule!",
)

invoke_and_assert(
[
"deployment",
"inspect",
"rence-griffith/test-deployment",
],
expected_output_contains=["'interval': 90.0,", "'interval': 1800.0,"],
expected_code=0,
)

invoke_and_assert(
[*create_args, "10", "--replace", "-y"],
expected_code=0,
expected_output_contains="Replaced existing deployment schedules with new schedule!",
)

invoke_and_assert(
[
"deployment",
"inspect",
"rence-griffith/test-deployment",
],
expected_output_contains=["'interval': 10.0,"],
expected_output_does_not_contain=[
"'interval': 90.0,",
"'interval': 1800.0,",
],
expected_code=0,
)

def test_create_schedule_replace_replaces_existing_schedule(self, flojo):
invoke_and_assert(
[
"deployment",
"schedule",
"clear",
"rence-griffith/test-deployment",
"-y",
],
expected_code=0,
expected_output_contains="Cleared all schedules",
)

create_args = [
"deployment",
"schedule",
"create",
"rence-griffith/test-deployment",
"--interval",
]

invoke_and_assert(
[
*create_args,
"90",
],
expected_code=0,
expected_output_contains="Created deployment schedule!",
)

invoke_and_assert(
[
"deployment",
"inspect",
"rence-griffith/test-deployment",
],
expected_output_contains=["'interval': 90.0,"],
expected_code=0,
)

invoke_and_assert(
[*create_args, "10", "--replace", "-y"],
expected_code=0,
expected_output_contains="Replaced existing deployment schedule with new schedule!",
)

invoke_and_assert(
[
"deployment",
"inspect",
"rence-griffith/test-deployment",
],
expected_output_contains=["'interval': 10.0,"],
expected_output_does_not_contain=["'interval': 1800.0,"],
expected_code=0,
)

def test_create_schedule_replace_seeks_confirmation(self, flojo):
deployment_name = "rence-griffith/test-deployment"
invoke_and_assert(
[
"deployment",
"schedule",
"create",
deployment_name,
"--interval",
"60",
"--replace",
],
user_input="N",
expected_code=1,
expected_output_contains=f"Are you sure you want to replace 1 schedule for {deployment_name}?",
)

invoke_and_assert(
[
"deployment",
"inspect",
"rence-griffith/test-deployment",
],
expected_output_contains=["'interval': 10.76,"], # original schedule
expected_code=0,
)

def test_create_schedule_replace_accepts_confirmation(self, flojo):
deployment_name = "rence-griffith/test-deployment"
invoke_and_assert(
[
"deployment",
"schedule",
"create",
deployment_name,
"--interval",
"60",
"--replace",
],
user_input="y",
expected_code=0,
expected_output_contains=f"Are you sure you want to replace 1 schedule for {deployment_name}?",
)

invoke_and_assert(
[
"deployment",
"inspect",
"rence-griffith/test-deployment",
],
expected_output_contains=["'interval': 60.0,"], # new schedule
expected_output_does_not_contain=[
"'interval': 10.76,"
], # original schedule
expected_code=0,
)

def test_clear_schedule_deletes(self, flojo):
invoke_and_assert(
[
Expand Down

0 comments on commit e297260

Please sign in to comment.