primitive 0.1.58__py3-none-any.whl → 0.1.59__py3-none-any.whl

primitive/jobs/graphql/fragments.py

@@ -0,0 +1,47 @@
+job_fragment = """
+fragment JobFragment on Job {
+    id
+    pk
+    slug
+    name
+    createdAt
+    updatedAt
+}
+"""
+
+job_run_fragment = """
+fragment JobRunFragment on JobRun {
+  id
+  pk
+  createdAt
+  updatedAt
+  completedAt
+  startedAt
+  status
+  conclusion
+  job {
+    id
+    pk
+    slug
+    name
+    createdAt
+    updatedAt
+  }
+  jobSettings {
+    containerArgs
+    rootDirectory
+  }
+  gitCommit {
+    sha
+    branch
+    repoFullName
+  }
+}
+"""
+
+job_run_status_fragment = """
+fragment JobRunStatusFragment on JobRun {
+    id
+    status
+}
+"""

primitive/jobs/graphql/mutations.py

@@ -0,0 +1,11 @@
+job_run_update_mutation = """
+mutation jobRunUpdate($input: JobRunUpdateInput!) {
+    jobRunUpdate(input: $input) {
+        ... on JobRun {
+            id
+            status
+            conclusion
+        }
+    }
+}
+"""

primitive/jobs/graphql/queries.py

@@ -0,0 +1,100 @@
+from primitive.graphql.utility_fragments import page_info_fragment
+
+from .fragments import job_fragment, job_run_fragment, job_run_status_fragment
+
+jobs_query = (
+    page_info_fragment
+    + job_fragment
+    + """
+query jobs(
+    $before: String
+    $after: String
+    $first: Int
+    $last: Int
+    $filters: JobFilters
+) {
+    jobs(
+        before: $before
+        after: $after
+        first: $first
+        last: $last
+        filters: $filters
+    ) {
+        totalCount
+        pageInfo {
+            ...PageInfoFragment
+        }
+        edges {
+            cursor
+            node {
+                ...JobFragment
+            }
+        }
+    }
+}
+"""
+)
+
+
+job_runs_query = (
+    page_info_fragment
+    + job_run_fragment
+    + """
+query jobRuns(
+  $before: String
+  $after: String
+  $first: Int
+  $last: Int
+  $filters: JobRunFilters
+  $order: JobRunOrder
+) {
+  jobRuns(
+    before: $before
+    after: $after
+    first: $first
+    last: $last
+    filters: $filters
+    order: $order
+  ) {
+    totalCount
+    pageInfo {
+      ...PageInfoFragment
+    }
+    edges {
+      cursor
+      node {
+        ...JobRunFragment
+      }
+    }
+  }
+}
+"""
+)
+
+job_run_query = (
+    job_run_fragment
+    + """
+query jobRun($id: GlobalID!) {
+    jobRun(id: $id) {
+        ...JobRunFragment
+    }
+}
+"""
+)
+
+github_app_token_for_job_run_query = """
+query ghAppTokenForJobRun($jobRunId: GlobalID!) {
+    ghAppTokenForJobRun(jobRunId: $jobRunId)
+}
+"""
+
+job_run_status_query = (
+    job_run_status_fragment
+    + """
+query jobRun($id: GlobalID!) {
+    jobRun(id: $id) {
+        ...JobRunStatusFragment
+    }
+}
+"""
+)

primitive/organizations/actions.py

@@ -1,9 +1,11 @@
 from typing import Optional
-from gql import gql
 
+from gql import gql
 
 from primitive.utils.actions import BaseAction
+
 from ..utils.auth import guard
+from .graphql.queries import organizations_query
 
 
 class Organizations(BaseAction):
@@ -15,54 +17,7 @@ class Organizations(BaseAction):
         first: Optional[int] = 1,
         last: Optional[int] = None,
     ):
-        query = gql(
-            """ 
-            fragment PageInfoFragment on PageInfo {
-                hasNextPage
-                hasPreviousPage
-                startCursor
-                endCursor
-            }
-
-            fragment OrganizationFragment on Organization {
-                id
-                pk
-                slug
-                name
-                createdAt
-                updatedAt
-            }
-
-            query organizations(
-                $before: String
-                $after: String
-                $first: Int
-                $last: Int
-                $filters: OrganizationFilters
-                $order: OrganizationOrder
-            ) {
-                organizations(
-                    before: $before
-                    after: $after
-                    first: $first
-                    last: $last
-                    filters: $filters
-                    order: $order
-                ) {
-                    totalCount
-                    pageInfo {
-                        ...PageInfoFragment
-                    }
-                    edges {
-                        cursor
-                        node {
-                            ...OrganizationFragment
-                        }
-                    }
-                }
-            }
-            """
-        )
+        query = gql(organizations_query)
 
         filters = {}
         if organization_id:

primitive/organizations/graphql/fragments.py

@@ -0,0 +1,10 @@
+organization_fragment = """
+fragment OrganizationFragment on Organization {
+    id
+    pk
+    slug
+    name
+    createdAt
+    updatedAt
+}
+"""

primitive/organizations/graphql/queries.py

@@ -0,0 +1,38 @@
+from primitive.graphql.utility_fragments import page_info_fragment
+
+from .fragments import organization_fragment
+
+organizations_query = (
+    page_info_fragment
+    + organization_fragment
+    + """
+query organizations(
+    $before: String
+    $after: String
+    $first: Int
+    $last: Int
+    $filters: OrganizationFilters
+    $order: OrganizationOrder
+) {
+    organizations(
+        before: $before
+        after: $after
+        first: $first
+        last: $last
+        filters: $filters
+        order: $order
+    ) {
+        totalCount
+        pageInfo {
+            ...PageInfoFragment
+        }
+        edges {
+            cursor
+            node {
+                ...OrganizationFragment
+            }
+        }
+    }
+}
+"""
+)

primitive/projects/actions.py

@@ -1,8 +1,11 @@
+from typing import Optional
+
 from gql import gql
 
-from typing import Optional
 from primitive.utils.actions import BaseAction
+
 from ..utils.auth import guard
+from .graphql.queries import projects_query
 
 
 class Projects(BaseAction):
@@ -14,52 +17,7 @@ class Projects(BaseAction):
         first: Optional[int] = 1,
         last: Optional[int] = None,
     ):
-        query = gql(
-            """ 
-            fragment PageInfoFragment on PageInfo {
-                hasNextPage
-                hasPreviousPage
-                startCursor
-                endCursor
-            }
-
-            fragment ProjectFragment on Project {
-                id
-                pk
-                slug
-                name
-                createdAt
-                updatedAt
-            }
-
-            query projects(
-                $before: String
-                $after: String
-                $first: Int
-                $last: Int
-                $filters: ProjectFilters
-            ) {
-                projects(
-                    before: $before
-                    after: $after
-                    first: $first
-                    last: $last
-                    filters: $filters
-                ) {
-                    totalCount
-                    pageInfo {
-                        ...PageInfoFragment
-                    }
-                    edges {
-                        cursor
-                        node {
-                            ...ProjectFragment
-                        }
-                    }
-                }
-            }
-            """
-        )
+        query = gql(projects_query)
 
         filters = {}
         if organization_id:

primitive/projects/graphql/fragments.py

@@ -0,0 +1,10 @@
+project_fragment = """
+fragment ProjectFragment on Project {
+    id
+    pk
+    slug
+    name
+    createdAt
+    updatedAt
+}
+"""

primitive/projects/graphql/queries.py

@@ -0,0 +1,36 @@
+from primitive.graphql.utility_fragments import page_info_fragment
+
+from .fragments import project_fragment
+
+projects_query = (
+    page_info_fragment
+    + project_fragment
+    + """
+query projects(
+    $before: String
+    $after: String
+    $first: Int
+    $last: Int
+    $filters: ProjectFilters
+) {
+    projects(
+        before: $before
+        after: $after
+        first: $first
+        last: $last
+        filters: $filters
+    ) {
+        totalCount
+        pageInfo {
+            ...PageInfoFragment
+        }
+        edges {
+            cursor
+            node {
+                ...ProjectFragment
+            }
+        }
+    }
+}
+"""
+)

primitive/reservations/actions.py

@@ -0,0 +1,134 @@
+import typing
+from time import sleep
+
+from primitive.graphql.relay import from_base64
+
+if typing.TYPE_CHECKING:
+    pass
+
+from typing import List, Optional
+
+from gql import gql
+
+from primitive.utils.actions import BaseAction
+
+from ..utils.auth import guard
+from .graphql.mutations import reservation_create_mutation, reservation_release_mutation
+from .graphql.queries import reservation_query, reservations_query
+
+
+class Reservations(BaseAction):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    @guard
+    def get_reservations(
+        self,
+        status: str = "in_progress",
+    ):
+        query = gql(reservations_query)
+
+        filters = {}
+        if status:
+            filters["status"] = {"exact": status}
+
+        variables = {
+            "filters": filters,
+        }
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
+        return result
+
+    @guard
+    def get_reservation(self, reservation_id: str):
+        query = gql(reservation_query)
+
+        variables = {
+            "id": reservation_id,
+        }
+
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
+        return result
+
+    @guard
+    def create_reservation(
+        self,
+        reason: str,
+        requested_hardware_ids: Optional[List[str]] = None,
+        organization_id: Optional[str] = None,
+        hardware_identifier: Optional[str] = None,
+    ):
+        mutation = gql(reservation_create_mutation)
+
+        if hardware_identifier and not requested_hardware_ids:
+            hardware = self.primitive.hardware.get_hardware_from_slug_or_id(
+                hardware_identifier=hardware_identifier
+            )
+            requested_hardware_ids = [hardware["id"]]
+
+        if not organization_id:
+            whoami_result = self.primitive.auth.whoami()
+            default_organization = whoami_result.data["whoami"]["defaultOrganization"]
+            organization_id = default_organization["id"]
+
+        input = {
+            "requestedHardwareIds": requested_hardware_ids,
+            "reason": reason,
+            "organizationId": organization_id,
+        }
+
+        variables = {"input": input}
+        result = self.primitive.session.execute(
+            mutation, variable_values=variables, get_execution_result=True
+        )
+        return result
+
+    @guard
+    def release_reservation(self, reservation_or_hardware_identifier: str):
+        mutation = gql(reservation_release_mutation)
+        try:
+            # check if it is a base64 encoded id
+            type_name, _id = from_base64(reservation_or_hardware_identifier)
+            if type_name == "Reservation":
+                reservation_id = reservation_or_hardware_identifier
+            elif type_name == "Hardware":
+                hardware = self.primitive.hardware.get_hardware_from_slug_or_id(
+                    hardware_identifier=reservation_or_hardware_identifier
+                )
+                reservation_id = hardware["activeReservation"]["id"]
+        except ValueError:
+            # if not, its a string and check for it here
+            hardware = self.primitive.hardware.get_hardware_from_slug_or_id(
+                hardware_identifier=reservation_or_hardware_identifier
+            )
+            reservation_id = hardware["activeReservation"]["id"]
+
+        input = {
+            "reservationId": reservation_id,
+        }
+        variables = {"input": input}
+        result = self.primitive.session.execute(
+            mutation, variable_values=variables, get_execution_result=True
+        )
+        return result
+
+    @guard
+    def wait_for_reservation_status(self, reservation_id: str, desired_status: str):
+        reservation_result = self.get_reservation(reservation_id=reservation_id)
+        reservation = reservation_result.data["reservation"]
+        current_status = reservation["status"]
+
+        sleep_amount = 1
+        while current_status != desired_status:
+            reservation_result = self.get_reservation(reservation_id=reservation_id)
+            reservation = reservation_result.data["reservation"]
+            current_status = reservation["status"]
+            if current_status == desired_status:
+                break
+            sleep(sleep_amount)
+            sleep_amount += 1
+
+        return reservation

primitive/reservations/commands.py

@@ -0,0 +1,67 @@
+import typing
+
+import click
+
+if typing.TYPE_CHECKING:
+    from ..client import Primitive
+from typing import Optional
+
+from ..utils.printer import print_result
+
+
+@click.group()
+@click.pass_context
+def cli(context):
+    """Reservations"""
+    pass
+
+
+@cli.command("list")
+@click.pass_context
+@click.option("--status", default="in_progress", type=str, help="Filter by status")
+def list(context, status: Optional[str] = "in_progress"):
+    """List reservations"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    get_reservations_result = primitive.reservations.get_reservations(status=status)
+    message = get_reservations_result.data
+    print_result(message=message, context=context)
+
+
+@cli.command("get")
+@click.pass_context
+@click.argument("reservation_id", type=str)
+def get(context, reservation_id: str):
+    """Get a reservation"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    get_reservation_result = primitive.reservations.get_reservation(
+        reservation_id=reservation_id
+    )
+    message = get_reservation_result.data
+    print_result(message=message, context=context)
+
+
+@cli.command("create")
+@click.pass_context
+@click.argument("hardware_identifier", type=str)
+@click.argument("reason", type=str)
+def create_reservation(context, hardware_identifier: str, reason: str):
+    """Crate a reservation by a Hardware's ID or Slug"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    create_reservation_result = primitive.reservations.create_reservation(
+        hardware_identifier=hardware_identifier, reason=reason
+    )
+    message = create_reservation_result.data
+    print_result(message=message, context=context)
+
+
+@cli.command("release")
+@click.pass_context
+@click.argument("reservation_or_hardware_identifier", type=str)
+def release_reservation(context, reservation_or_hardware_identifier: str):
+    """Release a reservation by Reservation ID, Hardware ID or Hardware Slug"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    release_reservation_result = primitive.reservations.release_reservation(
+        reservation_or_hardware_identifier=reservation_or_hardware_identifier
+    )
+    message = release_reservation_result.data
+    print_result(message=message, context=context)

primitive/reservations/graphql/fragments.py

@@ -0,0 +1,40 @@
+reservation_fragment = """
+fragment ReservationFragment on Reservation {
+  id
+  pk
+  createdAt
+  updatedAt
+  createdBy {
+    id
+    pk
+    email
+    username
+    displayName
+  }
+  hardware {
+    id
+    pk
+    name
+    slug
+    createdAt
+    updatedAt
+    isAvailable
+    isOnline
+    isQuarantined
+    isHealthy
+    capabilities {
+      id
+      pk
+      slug
+      name
+    }
+  }
+  reason
+  startedAt
+  endedAt
+  elapsedTime
+  status
+  conclusion
+  conclusionMessage
+}
+"""

primitive/reservations/graphql/mutations.py

@@ -0,0 +1,29 @@
+from primitive.graphql.utility_fragments import operation_info_fragment
+
+from .fragments import reservation_fragment
+
+reservation_create_mutation = (
+    operation_info_fragment
+    + reservation_fragment
+    + """
+mutation reservationCreate($input: ReservationCreateInput!) {
+  reservationCreate(input: $input) {
+    ...ReservationFragment
+    ...OperationInfoFragment
+  }
+}
+"""
+)
+
+reservation_release_mutation = (
+    operation_info_fragment
+    + reservation_fragment
+    + """
+mutation reservationRelease($input: ReservationReleaseInput!) {
+  reservationRelease(input: $input) {
+    ...ReservationFragment
+    ...OperationInfoFragment
+  }
+}
+"""
+)