Skip to content

Commit

Permalink
Update scheduling queries and unit tests to conform to Aerie v2.6 (#129)
Browse files Browse the repository at this point in the history 
* Updating scheduling module to work with Aerie 2.6

* Update to include a 2.6 banananation jar, and updated .env to use 2.6

* Refactor to replace deprecated datetime.utcnow() calls

* Removed banananation 2.2.0 jar

* Updated to use 2.6.0 instead of 2.2.0

* Updated aerie_client based on PR feedback

* Updating conftest to use v2.6.0

* Updated constraints queries to match v2.6.0 schema

* Updated aerie_host.is_auth_enabled()

* Updated constraints.py to remove unnecessary comments/add docstrings

* Removed upload_scheduling_goal wrapper function, as it is not being used and the same functionality is available with upload_scheduling_goals
  • Loading branch information
@joshhaug
joshhaug authored Apr 8, 2024
1 parent 353ffbf commit a9de995
Show file tree
Hide file tree
Showing 11 changed files with 248 additions and 91 deletions.
2 changes: 1 addition & 1 deletion .env
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
DOCKER_TAG=v2.2.0
DOCKER_TAG=v2.6.0
REPOSITORY_DOCKER_URL=ghcr.io/nasa-ammos

AERIE_USERNAME=aerie
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/test.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ jobs:
strategy:
matrix:
python-version: ["3.6.15", "3.11"]
aerie-version: ["2.2.0"]
aerie-version: ["2.6.0"]
steps:
- uses: actions/checkout@v3
- name: Set up Python ${{ matrix.python-version }}
Expand Down
196 changes: 145 additions & 51 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 @@ -1488,49 +1488,47 @@ def get_scheduling_goals_by_specification(self, spec_id):
scheduling_specification_goals(where: {
specification_id:{_eq:$spec}
}){
goal{
goal_metadata {
id
model_id
name
description
author
last_modified_by
created_date
modified_date
public
owner
updated_by
created_at
updated_at
}
}
}
"""
resp = self.aerie_host.post_to_graphql(list_all_goals_by_spec_query, spec=spec_id)

return resp

def upload_scheduling_goal(self, model_id, name, definition):
obj = dict()
obj["name"] = name
obj["model_id"] = model_id
obj["definition"] = definition

return self.upload_scheduling_goals([obj])


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
}
}
}
},
...
]
"""

upload_scheduling_goals_query = """
mutation InsertGoals($input:[scheduling_goal_insert_input!]!){
insert_scheduling_goal(objects: $input){
returning {id}
mutation InsertGoal($input:[scheduling_goal_definition_insert_input]!){
insert_scheduling_goal_definition(objects: $input){
returning {goal_id}
}
}"""

Expand All @@ -1541,8 +1539,8 @@ def upload_scheduling_goals(self, 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 @@ -1551,7 +1549,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 @@ -1570,7 +1568,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 All @@ -1593,9 +1591,28 @@ def delete_scheduling_goal(self, goal_id):
return self.delete_scheduling_goals(list([goal_id]))

def delete_scheduling_goals(self, goal_id_list):
# We must remove the goal(s) from any specifications before deleting them
delete_scheduling_goals_from_all_specs_query = """
mutation DeleteSchedulingGoalsFromAllSpecs($id_list: [Int!]!) {
delete_scheduling_model_specification_goals (where: {goal_id: {_in:$id_list}}){
returning {goal_id}
}
delete_scheduling_specification_goals (where: {goal_id: {_in:$id_list}}){
returning {goal_id}
}
}
"""

resp_for_deleting_from_specs = self.aerie_host.post_to_graphql(
delete_scheduling_goals_from_all_specs_query,
id_list=goal_id_list
)

# Note that deleting the scheduling goal metadata entry will take care of the
# scheduling goal definition entry too
delete_scheduling_goals_query = """
mutation DeleteSchedulingGoals($id_list: [Int!]!) {
delete_scheduling_goal (where: {id: {_in:$id_list}}){
delete_scheduling_goal_metadata (where: {id: {_in:$id_list}}){
returning {id}
}
}
Expand Down Expand Up @@ -1656,20 +1673,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 @@ -1678,43 +1737,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 @@ -1854,7 +1948,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
8 changes: 6 additions & 2 deletions src/aerie_cli/aerie_host.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -242,15 +242,19 @@ def is_auth_enabled(self) -> bool:
Returns:
bool: False if authentication is disabled, otherwise True
"""
resp = self.session.get(self.gateway_url + "/auth/session")
# Try to login using blank credentials. If "Authentication is disabled" is returned, we can safely skip auth
resp = self.session.post(self.gateway_url + "/auth/login", json={"username": "", "password": ""},)

if resp.ok:
try:
resp_json = resp.json()
resp_json = process_gateway_response(resp)
if (
"message" in resp_json.keys()
and resp_json["message"] == "Authentication is disabled"
):
return False
except RuntimeError:
pass
except requests.exceptions.JSONDecodeError:
pass

Expand Down
Loading

0 comments on commit a9de995

Please sign in to comment.