Rap subset plan

pyproject.toml

@@ -27,6 +27,7 @@ rich = "^12.6.0"
 dataclasses-json = "^0.5.6"
 pandas = "^1.5.1"
 appdirs = "^1.4.4"
+deepdiff = "^6.3.1"
 
 [tool.poetry.dev-dependencies]
 pytest = "^6.2.5"

src/aerie_cli/aerie_client.py

@@ -20,6 +20,7 @@
 from .schemas.client import ExpansionSet
 from .schemas.client import ResourceType
 from .utils.serialization import postgres_duration_to_microseconds
+from .utils.serialization import timedelta_to_postgres_interval
 from .aerie_host import AerieHostSession
 
 
@@ -78,6 +79,64 @@ def get_activity_plan_by_id(self, plan_id: int, full_args: str = None) -> Activi
         plan = ActivityPlanRead.from_api_read(api_plan)
         return self.__expand_activity_arguments(plan, full_args)
 
+    def get_activity_plan_subset_by_id(
+        self, plan_id: int,
+        start_time: arrow.Arrow, 
+        end_time: arrow.Arrow,
+        full_args: str = None) -> ActivityPlanRead:
+        """Download a subset of an activity plan from Aerie
+
+        Args:
+            plan_id (int): ID of the plan in Aerie
+            full_args (str): comma separated list of activity types for which to
+            get full arguments, otherwise only modified arguments are returned.
+            Set to "true" to get full arguments for all activity types.
+            Disabled if missing, None, "false", or "".
+            start_time (arrow.Arrow)
+            end_time (arrow.Arrow)
+
+        Returns:
+            ActivityPlanRead: the activity plan subset
+        """
+
+        parent = self.get_activity_plan_by_id(plan_id)
+        parent_start = parent.start_time
+        start_offset = timedelta_to_postgres_interval(start_time - parent_start)
+        end_offset = timedelta_to_postgres_interval(end_time - parent_start)
+
+        query = """
+        query get_plans ($start_offset: interval!, $end_offset: interval!, $plan_id: Int!) {
+                plan_by_pk(id: $plan_id) {
+                    id
+                    model_id
+                    name
+                    start_time
+                    duration
+                    simulations{
+                        id
+                    }
+                    activity_directives(where: {start_offset: {_gte: $start_offset, _lt: $end_offset}}, order_by: { start_offset: asc }) {
+                        id
+                        name
+                        type
+                        start_offset
+                        arguments
+                        metadata
+                        tags
+                    }
+                }
+        }
+        """
+        resp = self.host_session.post_to_graphql(
+            query,
+            plan_id=plan_id,
+            start_offset=start_offset,
+            end_offset=end_offset
+        )
+        plan = ApiActivityPlanRead.from_dict(resp)
+        plan = ActivityPlanRead.from_api_read(plan)
+        return self.__expand_activity_arguments(plan, full_args)
+
     def list_all_activity_plans(self) -> List[ActivityPlanRead]:
         list_all_plans_query = """
         query list_all_plans {
@@ -278,6 +337,21 @@ def update_activity(
         )
         return resp["id"]
 
+    def delete_activity(self, activity_id: int, plan_id: int):
+        query = """
+        mutation DeleteActivityDirective($id: Int!, $plan_id: Int!) {
+            delete_activity_directive_by_pk(id: $id, plan_id: $plan_id) {
+                name
+            }
+        }
+        """
+        resp = self.host_session.post_to_graphql(
+            query,
+            id=activity_id,
+            plan_id=plan_id
+        )
+        return resp
+
     def simulate_plan(self, plan_id: int, poll_period: int = 5) -> int:
 
         simulate_query = """
@@ -1658,7 +1732,7 @@ def add_directive_metadata_schemas(self, schemas: list) -> list:
         )
         return resp
 
-    def delete_directive_metadata_schema(self, key) -> list:
+    def delete_directive_metadata_schema(self, key: str) -> list:
         """Delete metadata schemas
 
         Returns:
@@ -1676,4 +1750,108 @@ def delete_directive_metadata_schema(self, key) -> list:
             delete_schema_query,
             key=key
         )
-        return resp["key"]
+        return resp["key"]
+
+    def get_activity_directive_preset(self, activity_id: int, plan_id: int) -> dict:
+        """Get the preset applied on a activity directive
+
+        Returns:
+            dict: an object representing the activity directive.
+                dict["applied_preset"] == None if there is no preset applied to the activity
+        """
+        query = """
+        query MyQuery($id: Int!, $plan_id: Int!) {
+            activity_directive_by_pk(id: $id, plan_id: $plan_id) {
+                applied_preset {
+                presets_applied {
+                    name
+                }
+                preset_id
+                }
+            }
+        }
+        """
+        resp = self.host_session.post_to_graphql(
+            query,
+            id=activity_id,
+            plan_id=plan_id
+        )
+        return resp
+
+    def apply_activity_directive_preset(self, activity_id, plan_id, preset_id) -> int:
+        """Apply the a preset to an activity directive
+
+        Returns:
+            int: the ID of the applied preset
+        """
+        mutation = """
+        mutation MyMutation($_activity_id: Int!, $_plan_id: Int!, $_preset_id: Int!) {
+            apply_preset_to_activity(args: {_activity_id: $_activity_id, _plan_id: $_plan_id, _preset_id: $_preset_id}) {
+                id
+            }
+        }
+        """
+        resp = self.host_session.post_to_graphql(
+            mutation,
+            _activity_id=activity_id,
+            _plan_id=plan_id,
+            _preset_id=preset_id
+        )
+        return resp["id"]
+
+    def delete_activity_directive_preset(self, activity_id, plan_id) -> dict:
+        """Remove a preset from an activity directive
+
+        Returns:
+            dict: an object representing the preset that was deleted
+        """
+        mutation = """
+        mutation MyMutation($_activity_id: Int!, $_plan_id: Int!) {
+            delete_preset_to_directive(where: {activity_id: {_eq: $_activity_id}, _and: {plan_id: {_eq: $_plan_id}}}) {
+                returning {
+                    preset_id
+                }
+            }
+        }
+        """
+        resp = self.host_session.post_to_graphql(
+            mutation,
+            _activity_id=activity_id,
+            _plan_id=plan_id
+        )
+        return resp
+
+    # TODO: the two functions below are very specific to the plans merge command
+    # do we still want them here or is there another way to integrate them into the cli?
+
+    def get_plan_created_date(self, plan_id: int):
+        # for some reason, plan doesnt have a `created_at` field on local, but its there 
+        # in the dev venue. i need to look at the activity `created_at` for now.....
+        plan_query = """
+        query MyQuery($plan_id: Int!) {
+            activity_directive(where: {plan_id: {_eq: $plan_id}}) {
+                created_at
+            }
+        }
+        """
+        resp = self.host_session.post_to_graphql(
+            plan_query,
+            plan_id=plan_id,
+        )
+        return arrow.get(resp[0]["created_at"]) if len(resp) > 0 else None
+
+    def get_plan_recently_updated_activities(self, plan_id: int, time: arrow.Arrow):
+        activity_query = """
+        query GetRecentlyUpdatedActivities($plan_id: Int!, $time: timestamptz!) {
+            activity_directive(where: {_and: {plan_id: {_eq: $plan_id}}, last_modified_at: {_gt: $time}}) {
+                id
+                name
+            }
+        }
+        """
+        resp = self.host_session.post_to_graphql(
+            activity_query,
+            plan_id=plan_id,
+            time=str(time)
+        )
+        return resp