diff --git a/src/aerie_cli/aerie_client.py b/src/aerie_cli/aerie_client.py index f010555b..2443a42e 100644 --- a/src/aerie_cli/aerie_client.py +++ b/src/aerie_cli/aerie_client.py @@ -1654,8 +1654,28 @@ def __expand_activity_arguments(self, plan: ActivityPlanRead, full_args: str = N activity.arguments = ApiEffectiveActivityArguments.from_dict( resp).arguments return plan + + def add_constraint_tag(self, constraint_id: int, tag_name: str): + add_constraint_tag_query = """ + mutation AddTagToConstraint($constraint_id: Int, $tag_id: Int) { + insert_constraint_tags(objects: {constraint_id: $constraint_id, tag_id: $tag_id}) { + returning { + tag_id + } + } + } + """ + + #add tag to constraint + resp = self.aerie_host.post_to_graphql( + add_constraint_tag_query, + constraint_id=constraint_id, + tag_id=self.get_tag_id_by_name(tag_name) + ) - def upload_constraint(self, constraint): + return resp['returning'][0]["tag_id"] + + def upload_constraint(self, constraint, tags=None): upload_constraint_query = """ mutation CreateConstraint($constraint: constraint_insert_input!) { createConstraint: insert_constraint_one(object: $constraint) { @@ -1665,6 +1685,12 @@ def upload_constraint(self, constraint): """ resp = self.aerie_host.post_to_graphql(upload_constraint_query, constraint=constraint) + + #add each tag to constraint + if tags is not None: + for tag in tags: + new_tag_id = self.add_constraint_tag(resp["id"], tag) + return resp["id"] def delete_constraint(self, id): diff --git a/src/aerie_cli/commands/constraints.py b/src/aerie_cli/commands/constraints.py index 60e2ceda..cd45aed4 100644 --- a/src/aerie_cli/commands/constraints.py +++ b/src/aerie_cli/commands/constraints.py @@ -4,6 +4,7 @@ import typer from rich.console import Console from rich.table import Table +from typing import List from aerie_cli.commands.command_context import CommandContext @@ -15,7 +16,8 @@ def upload( plan_id: int = typer.Option(None, help="The plan id associated with the constraint (do not input model id)"), name: str = typer.Option(..., help="The name of the constraint", prompt=True), description: str = typer.Option("", help="The description of the constraint"), # optional - constraint_file: str = typer.Option(..., help="The file that holds the constraint", prompt=True) + constraint_file: str = typer.Option(..., help="The file that holds the constraint", prompt=True), + tags: List[str] = typer.Option(None, help="A list of tag names to upload with the constraint (existing or new tags)") ): """Upload a constraint""" @@ -33,7 +35,7 @@ def upload( "description": description, "definition": str_contents } - constraint_id = client.upload_constraint(constraint) + constraint_id = client.upload_constraint(constraint, tags) typer.echo(f"Created constraint: {constraint_id}") @app.command()