Merge pull request #108 from NASA-AMMOS/download-upload-plan-tags

Download/Upload Plan Tags

src/aerie_cli/aerie_client.py

@@ -62,6 +62,12 @@ def get_activity_plan_by_id(self, plan_id: int, full_args: str = None) -> Activi
                 simulations{
                     id
                 }
+                tags {
+                    tag {
+                        id
+                        name
+                    }
+                }
                 activity_directives(order_by: { start_offset: asc }) {
                     id
                     name
@@ -92,6 +98,12 @@ def list_all_activity_plans(self) -> List[ActivityPlanRead]:
                 simulations{
                     id
                 }
+                tags {
+                    tag {
+                        id
+                        name
+                    }
+                }
             }
         }
         """
@@ -147,6 +159,60 @@ def get_plan_id_by_sim_id(self, simulation_dataset_id: int) -> int:
             simulation_dataset_id=simulation_dataset_id
         )
         return resp['simulation']['plan']['id']
+
+    def get_tag_id_by_name(self, tag_name: str):
+        get_tags_by_name_query = """
+        query GetTagByName($name: String) {
+            tags(where: {name: {_eq: $name}}) {
+                id
+            }
+        }
+        """
+
+        #make default color of tag white
+        create_new_tag = """
+        mutation CreateNewTag($name: String, $color: String = "#FFFFFF") {
+            insert_tags_one(object: {name: $name, color: $color}) {
+                id
+            }
+        }
+        """
+
+        resp = self.aerie_host.post_to_graphql(
+            get_tags_by_name_query, 
+            name=tag_name
+        )
+
+        #if a tag with the specified name exists then returns the ID, else creates a new tag with this name
+        if len(resp) > 0: 
+            return resp[0]["id"]
+        else: 
+            new_tag_resp = self.aerie_host.post_to_graphql(
+                create_new_tag, 
+                name=tag_name
+            )
+
+            return new_tag_resp["id"]
+
+    def add_plan_tag(self, plan_id: int, tag_name: str):
+        add_tag_to_plan = """
+        mutation AddTagToPlan($plan_id: Int, $tag_id: Int) {
+            insert_plan_tags(objects: {plan_id: $plan_id, tag_id: $tag_id}) {
+                returning {
+                    tag_id
+                }
+            }
+        }
+        """
+
+        #add tag to plan
+        resp = self.aerie_host.post_to_graphql(
+            add_tag_to_plan, 
+            plan_id=plan_id, 
+            tag_id=self.get_tag_id_by_name(tag_name)
+        )
+
+        return resp['returning'][0]
 
     def create_activity_plan(
         self, model_id: int, plan_to_create: ActivityPlanCreate
@@ -167,6 +233,11 @@ def create_activity_plan(
         )
         plan_id = plan_resp["id"]
         plan_revision = plan_resp["revision"]
+
+        #add plan tags if exists from plan_to_create
+        for tag in plan_to_create.tags:
+            self.add_plan_tag(plan_id, tag["tag"]["name"])
+
         # This loop exists to make sure all anchor IDs are updated as necessary
 
         # Deep copy activities so we can augment and pop from the list

src/aerie_cli/schemas/api.py

@@ -138,6 +138,12 @@ class ApiActivityPlanRead(ApiActivityPlanBase):
         converter=converters.optional(
             lambda listOfDicts: [ApiActivityRead.from_dict(d) if isinstance(d, dict) else d for d in listOfDicts])
     )
+    tags: Optional[List[Dict]] = field(
+        default = [], 
+        converter=converters.optional(
+            lambda listOfDicts: [d for d in listOfDicts]
+        )
+    )
 
 
 @define

src/aerie_cli/schemas/client.py

@@ -113,6 +113,7 @@ class EmptyActivityPlan(ClientSerialize):
     )
     end_time: Arrow = field(
         converter = arrow.get)
+    tags: Optional[List[Dict]] 
 
     def duration(self) -> timedelta:
         return self.end_time - self.start_time
@@ -145,7 +146,12 @@ class ActivityPlanCreate(EmptyActivityPlan):
     sim_id: Optional[int] = field(
         default=None
     )
-
+    tags: Optional[List[Dict]] = field(
+        default=[], 
+        converter=converters.optional(
+            lambda listOfDicts: [d for d in listOfDicts]
+        )
+    )
 
     @classmethod
     def from_plan_read(cls, plan_read: "ActivityPlanRead") -> "ActivityPlanCreate":
@@ -154,6 +160,7 @@ def from_plan_read(cls, plan_read: "ActivityPlanRead") -> "ActivityPlanCreate":
             start_time=plan_read.start_time,
             end_time=plan_read.end_time,
             activities=plan_read.activities,
+            tags=plan_read.tags
         )
 
     def to_api_create(self, model_id: int) -> "ApiActivityPlanCreate":
@@ -170,6 +177,12 @@ class ActivityPlanRead(EmptyActivityPlan):
     id: int
     model_id: int
     sim_id: int
+    tags: Optional[List[Dict]] = field(
+        default = [], 
+        converter=converters.optional(
+            lambda listOfDicts: [d for d in listOfDicts]
+        )
+    )
     activities: Optional[List[Activity]] = field(
         default = None,
         converter=converters.optional(
@@ -220,6 +233,7 @@ def from_api_read(cls, api_plan_read: ApiActivityPlanRead) -> "ActivityPlanRead"
             sim_id=api_plan_read.simulations[0]["id"],
             start_time=plan_start,
             end_time=plan_start + api_plan_read.duration,
+            tags=api_plan_read.tags,
             activities= None if api_plan_read.activity_directives is None else [
                 Activity.from_api_read(api_activity)
                 for api_activity in api_plan_read.activity_directives

tests/unit_tests/files/expected_results/get_activity_plan_by_id.json

@@ -5,6 +5,7 @@
   "id": 1,
   "model_id": 1,
   "sim_id": 1,
+  "tags": [],
   "activities": [
     {
       "type": "ACT_One",

tests/unit_tests/files/mock_responses/get_activity_plan_by_id.json

@@ -1,7 +1,7 @@
 [
     {
         "request": {
-            "query": "query get_plans ($plan_id: Int!) { plan_by_pk(id: $plan_id) { id model_id name start_time duration simulations{ id } activity_directives(order_by: { start_offset: asc }) { id name type start_offset arguments metadata anchor_id anchored_to_start } } }",
+            "query": "query get_plans ($plan_id: Int!) { plan_by_pk(id: $plan_id) { id model_id name start_time duration simulations{ id } tags { tag { id name } } activity_directives(order_by: { start_offset: asc }) { id name type start_offset arguments metadata anchor_id anchored_to_start } } }",
             "variables": {
                 "plan_id": 1
             }

tests/unit_tests/files/mock_responses/list_all_activity_plans.json

@@ -1,7 +1,7 @@
 [
     {
         "request": {
-            "query": "\n query list_all_plans {\n plan(order_by: { id: asc }) {\n id\n model_id\n name\n start_time\n duration\n simulations{\n id\n }\n }\n }",
+            "query": "\n query list_all_plans {\n plan(order_by: { id: asc }) {\n id\n model_id\n name\n start_time\n duration\n simulations{\n id\n }\n tags {\n tag {\n id\n name\n }\n }\n }\n }",
             "variables": {}
         },
         "response": [
@@ -15,7 +15,8 @@
                     {
                         "id": 1
                     }
-                ]
+                ], 
+                "tags": []
             },
             {
                 "id": 2,
@@ -27,6 +28,14 @@
                     {
                         "id": 2
                     }
+                ], 
+                "tags": [
+                    {
+                        "tag": {
+                            "id": 1, 
+                            "name": "Test"
+                        }
+                    }
                 ]
             }
         ]

tests/unit_tests/test_aerie_client.py

@@ -88,7 +88,8 @@ def test_list_all_activity_plans():
                 {
                     "id": 1
                 }
-            ]
+            ], 
+            "tags": []
         },
         {
             "id": 2,
@@ -100,6 +101,14 @@ def test_list_all_activity_plans():
                 {
                     "id": 2
                 }
+            ],
+            "tags": [
+                {
+                    "tag": {
+                        "id": 1, 
+                        "name": "Test"
+                    }
+                }
             ]
         }
     ]""")