Skip to content

Commit

Permalink
Updated constraints queries to match v2.6.0 schema
Browse files Browse the repository at this point in the history
  • Loading branch information
@joshhaug
joshhaug committed Apr 3, 2024
1 parent 07df905 commit 4ba94af
Show file tree
Hide file tree
Showing 4 changed files with 198 additions and 62 deletions.
155 changes: 115 additions & 40 deletions src/aerie_cli/aerie_client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1272,7 +1272,7 @@ def list_sequences(self, simulation_dataset_id: int) -> List[str]:
"""

list_sequences_query = """
query MyQuery($simulation_dataset_id: Int!) {
query ListSequences($simulation_dataset_id: Int!) {
sequence(where: { simulation_dataset_id: { _eq: $simulation_dataset_id } }) {
seq_id
}
Expand Down Expand Up @@ -1427,7 +1427,7 @@ def get_typescript_dictionary(self, command_dictionary_id: int) -> str:
"""

get_command_dictionary_metadata_query = """
query MyQuery($command_dictionary_id: Int!) {
query GetCommandDictionaryMetadata($command_dictionary_id: Int!) {
command_dictionary(where: { id: { _eq: $command_dictionary_id } }) {
mission
version
Expand Down Expand Up @@ -1514,24 +1514,22 @@ def upload_scheduling_goal(self, model_id, name, definition):
def upload_scheduling_goals(self, upload_object):
"""
Bulk upload operation for uploading scheduling goals.
@param upload_object should be JSON-like with keys name, model_id, definition
@param upload_object should be JSON-like with a definition key and metadata containing the goal name and model id
[
{
name: str,
model_id: int,
definition: str
metadata: {
name: str,
models_using: {
data: {
model_id: int
}
}
}
},
...
]
"""

# Note that as of Aerie v2.3.0, the metadata (incl. model_id and goal name) are stored in a separate table
formatted_upload_objects = []
for entry in upload_object:
new_entry = {}
new_entry["definition"] = entry["definition"]
new_entry["metadata"] = {"data": {"name": entry["name"], "models_using": {"data": {"model_id": entry["model_id"]}}}}
formatted_upload_objects.append(new_entry)

upload_scheduling_goals_query = """
mutation InsertGoal($input:[scheduling_goal_definition_insert_input]!){
Expand All @@ -1542,13 +1540,13 @@ def upload_scheduling_goals(self, upload_object):

resp = self.aerie_host.post_to_graphql(
upload_scheduling_goals_query,
input=formatted_upload_objects
input=upload_object
)

return resp["returning"]

def get_specification_for_plan(self, plan_id):
get_specification_for_plan_query = """
def get_scheduling_specification_for_plan(self, plan_id):
get_scheduling_specification_for_plan_query = """
query GetSpecificationForPlan($plan_id: Int!) {
scheduling_specification(where: {plan_id: {_eq: $plan_id}}) {
id
Expand All @@ -1557,7 +1555,7 @@ def get_specification_for_plan(self, plan_id):
"""

resp = self.aerie_host.post_to_graphql(
get_specification_for_plan_query,
get_scheduling_specification_for_plan_query,
plan_id=plan_id
)
return resp[0]["id"]
Expand All @@ -1576,7 +1574,7 @@ def add_goals_to_specifications(self, upload_object):
"""

add_goal_to_specification_query = """
mutation MyMutation($object: [scheduling_specification_goals_insert_input!]!) {
mutation AddGoalToSpec($object: [scheduling_specification_goals_insert_input!]!) {
insert_scheduling_specification_goals(objects: $object) {
returning {
enabled
Expand Down Expand Up @@ -1681,20 +1679,62 @@ def __expand_activity_arguments(self, plan: ActivityPlanRead, full_args: str = N

def upload_constraint(self, constraint):
upload_constraint_query = """
mutation CreateConstraint($constraint: constraint_insert_input!) {
createConstraint: insert_constraint_one(object: $constraint) {
id
mutation CreateConstraint($constraint: constraint_definition_insert_input!) {
createConstraint: insert_constraint_definition_one(object: $constraint) {
constraint_id
}
}
"""

resp = self.aerie_host.post_to_graphql(upload_constraint_query, constraint=constraint)
return resp["id"]
return resp["constraint_id"]


def add_constraint_to_plan(self, constraint_id, plan_id):
"""
Add a constraint to a plan's constraint specification.
@param constraint_id The constraint ID to add to the given plan
@param plan_id The plan ID to add this constraint to
"""

add_constraint_to_specification_query = """
mutation InsertConstraintSpec($constraint_id: Int!, $plan_id: Int!) {
insert_constraint_specification_one(object: {constraint_id: $constraint_id, plan_id: $plan_id}) {
constraint_id
}
}
"""
resp = self.aerie_host.post_to_graphql(
add_constraint_to_specification_query,
constraint_id = constraint_id,
plan_id = plan_id
)

return resp['constraint_id']

def delete_constraint(self, id):

# We must remove the constraint from any specifications before deleting it
delete_constraint_from_all_specs_query = """
mutation DeleteConstraintsFromAllSpecs($id: Int!) {
delete_constraint_specification(where: {constraint_id: {_eq: $id}}) {
returning {
constraint_id
plan_id
}
}
}
"""

resp_for_deleting_from_specs = self.aerie_host.post_to_graphql(
delete_constraint_from_all_specs_query,
id=id
)

delete_constraint_query = """
mutation DeleteConstraint($id: Int!) {
deleteConstraint: delete_constraint_by_pk(id: $id) {
deleteConstraint: delete_constraint_metadata_by_pk(id: $id) {
id
}
}
Expand All @@ -1703,43 +1743,78 @@ def delete_constraint(self, id):
resp = self.aerie_host.post_to_graphql(delete_constraint_query, id=id)
return resp["id"]

def update_constraint(self, id, constraint):
update_constraint_query = """
mutation UpdateConstraint($id: Int!, $constraint: constraint_set_input!) {
updateConstraint: update_constraint_by_pk(
pk_columns: { id: $id }, _set: $constraint
def update_constraint(self, id, definition):
old_update_constraint_query = """
mutation UpdateConstraint($constarint_id: Int!, $constraint: constraint_definition_set_input!) {
update_constraint_definition_by_pk(
pk_columns: { constarint_id: $constraint_id }, _set: $constraint
) {
id
constarint_id
definition
author
created_at
}
}
"""

resp = self.aerie_host.post_to_graphql(update_constraint_query, id=id, constraint=constraint)
return resp["id"]
update_constraint_query = """
mutation UpdateConstraint($constarint_id: Int!, $definition: String!) {
update_constraint_definition_many(
updates: {_set: {definition: $definition}, where: {constraint_id: {_eq: $constarint_id}}}) {
returning {
constraint_id
}
}
}
"""

resp = self.aerie_host.post_to_graphql(update_constraint_query, constarint_id=id, definition=definition)
return resp

def get_constraint_by_id(self, id):
get_constraint_by_id_query = """
query get_constraint($id: Int!) {
constraint_by_pk(id: $id) {
model_id
plan_id
name
query get_constraint($constraint_id: Int!) {
constraint_definition(where: {constraint_id: {_eq: $constraint_id}}) {
author
definition
description
metadata {
description
name
plans_using {
plan_id
}
models_using {
model_id
}
}
}
}
"""

resp = self.aerie_host.post_to_graphql(get_constraint_by_id_query, id=id)
resp = self.aerie_host.post_to_graphql(get_constraint_by_id_query, constraint_id=id)
return resp

def get_constraint_specification_for_plan(self, plan_id):
get_constraint_specification_for_plan_query = """
query GetConstraintSpecificationForPlan($plan_id: Int!) {
constraint_specification(where: {plan_id: {_eq: $plan_id}}) {
constraint_id
}
}
"""

resp = self.aerie_host.post_to_graphql(
get_constraint_specification_for_plan_query,
plan_id=plan_id
)
return resp[0]["id"]

def get_constraint_violations(self, plan_id):
get_violations_query = """
query ($plan_id: Int!) {
constraintResponses: constraintViolations(planId: $plan_id) {
constraintId
constraintName
type
success
results {
resourceIds
Expand Down Expand Up @@ -1879,7 +1954,7 @@ def delete_directive_metadata_schema(self, key) -> list:
list: a list of the metadata keys that were deleted
"""
delete_schema_query = """
mutation MyMutation($key: String!) {
mutation DeleteDirectiveMetadataSchema($key: String!) {
delete_activity_directive_metadata_schema_by_pk(key: $key) {
key
}
Expand Down
72 changes: 56 additions & 16 deletions src/aerie_cli/commands/constraints.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,33 @@

app = typer.Typer()

def construct_constraint_metadata(name, description, model_id, plan_id):
metadata = {
"data": {
"name": name,
"description": description
}
}
if model_id != None:
if model_id is int:
metadata["data"]["models_using"] = {
"data": {
"model_id": model_id
}
}
elif model_id is list:
metadata["data"]["models_using"] = model_id
if plan_id != None:
if plan_id is int:
metadata["data"]["plans_using"] = {
"data": {
"plan_id": plan_id
}
}
elif plan_id is list:
metadata["data"]["plans_using"] = plan_id
return metadata

@app.command()
def upload(
model_id: int = typer.Option(None, help="The model id associated with the constraint (do not input plan id)"),
Expand All @@ -25,17 +52,29 @@ def upload(
client = CommandContext.get_client()
with open(constraint_file) as in_file:
contents = in_file.readlines()
str_contents = " ".join(contents)
str_contents = " ".join(contents)
constraint = {
"model_id": model_id,
"plan_id": plan_id,
"name": name,
"description": description,
"definition": str_contents
"definition": str_contents,
"metadata": construct_constraint_metadata(name, description, model_id, plan_id)
}
constraint_id = client.upload_constraint(constraint)
typer.echo(f"Created constraint: {constraint_id}")


@app.command()
def add_to_plan(
plan_id: int = typer.Option(None, help="The plan id associated with the constraint (do not input model id)", prompt=True),
constraint_id: int = typer.Option(..., help="The id of the constraint", prompt=True)
):
"""Associate a constraint with a plan's constraint specification."""

client = CommandContext.get_client()
#specification = client.get_constraint_specification_for_plan(plan_id)
#upload_to_spec = [{"constraint_id": constraint_id, "specification_id": specification}]
client.add_constraint_to_plan(constraint_id=constraint_id, plan_id=plan_id)
typer.echo(f"Added constraint: {constraint_id} to plan: {plan_id} constraint specification")


@app.command()
def delete(
id: int = typer.Option(..., help="The constraint id to be deleted", prompt=True)
Expand All @@ -48,7 +87,7 @@ def delete(

@app.command()
def update(
id: int = typer.Option(..., help="The constraint id to be modifyed", prompt=True),
id: int = typer.Option(..., help="The constraint id to be modified", prompt=True),
constraint_file: str = typer.Option(..., help="The new constraint for the id", prompt=True)
):
"""Update a constraint"""
Expand All @@ -58,15 +97,16 @@ def update(
with open(constraint_file) as in_file:
contents = in_file.readlines()
str_contents = " ".join(contents)
constraint = {
"model_id": constraint["model_id"],
"plan_id": constraint["plan_id"],
"name": constraint["name"],
"description": constraint["description"],
"definition": str_contents
}
constraint_id = client.update_constraint(id, constraint)
typer.echo(f"Updated constraint: {constraint_id}")
# desc = constraint["metadata"]["description"]
# models = constraint["metadata"]["models_using"]
# plans = constraint["metadata"]["plans_using"]
# constraint_def = {
# "definition": str_contents,
# "metadata": construct_constraint_metadata(constraint["name"], constraint["metadata"]["description"], constraint["metadata"]["models_using"], constraint["metadata"]["plans_using"])
# }
response = client.update_constraint(id, str_contents)
for r in response:
typer.echo(f"Updated constraint: {r["returning"]}")

@app.command()
def violations(
Expand Down
Loading

0 comments on commit 4ba94af

Please sign in to comment.