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

primitive/hardware/actions.py

@@ -2,18 +2,27 @@ import csv
 import io
 import json
 import platform
-from shutil import which
 import subprocess
+import typing
+from shutil import which
 from typing import Dict, List, Optional
+
 import click
+from aiohttp import client_exceptions
+from gql import gql
 from loguru import logger
+
+from primitive.graphql.relay import from_base64
 from primitive.utils.memory_size import MemorySize
-from gql import gql
-from aiohttp import client_exceptions
-from ..utils.config import update_config_file
-from ..utils.auth import guard
 
-import typing
+from ..utils.auth import guard
+from ..utils.config import update_config_file
+from .graphql.mutations import (
+    hardware_checkin_mutation,
+    hardware_update_mutation,
+    register_hardware_mutation,
+)
+from .graphql.queries import hardware_list
 
 if typing.TYPE_CHECKING:
     pass
@@ -269,29 +278,13 @@ class Hardware(BaseAction):
     @guard
     def register(self):
         system_info = self.get_system_info()
-        mutation = gql(
-            """
-            mutation registerHardware($input: RegisterHardwareInput!) {
-                registerHardware(input: $input) {
-                    ... on Hardware {
-                        fingerprint
-                    }
-                    ... on OperationInfo {
-                          messages {
-                            kind
-                            message
-                            field
-                            code
-                        }
-                    }
-                }
-            }
-        """
-        )
+        mutation = gql(register_hardware_mutation)
         input = {"systemInfo": system_info}
         variables = {"input": input}
-        result = self.primitive.session.execute(mutation, variable_values=variables)
-        if messages := result.get("registerHardware").get("messages"):
+        result = self.primitive.session.execute(
+            mutation, variable_values=variables, get_execution_result=True
+        )
+        if messages := result.data.get("registerHardware").get("messages"):
             for message in messages:
                 logger.enable("primitive")
                 if message.get("kind") == "ERROR":
@@ -300,7 +293,7 @@ class Hardware(BaseAction):
                     logger.debug(message.get("message"))
             return False
 
-        fingerprint = result.get("registerHardware").get("fingerprint")
+        fingerprint = result.data.get("registerHardware").get("fingerprint")
 
         self.primitive.host_config["fingerprint"] = fingerprint
         self.primitive.full_config[self.primitive.host] = self.primitive.host_config
@@ -335,30 +328,14 @@ class Hardware(BaseAction):
             "systemInfo": system_info,
         }
 
-        mutation = gql(
-            """
-            mutation hardwareUpdate($input: HardwareUpdateInput!) {
-                hardwareUpdate(input: $input) {
-                    ... on Hardware {
-                        systemInfo
-                    }
-                    ... on OperationInfo {
-                          messages {
-                            kind
-                            message
-                            field
-                            code
-                        }
-                    }
-                }
-            }
-            """
-        )
+        mutation = gql(hardware_update_mutation)
 
         input = new_state
         variables = {"input": input}
         try:
-            result = self.primitive.session.execute(mutation, variable_values=variables)
+            result = self.primitive.session.execute(
+                mutation, variable_values=variables, get_execution_result=True
+            )
         except client_exceptions.ClientConnectorError as exception:
             message = " [*] Failed to update hardware system info! "
             logger.error(message)
@@ -391,27 +368,7 @@ class Hardware(BaseAction):
             "isOnline": is_online,
         }
 
-        mutation = gql(
-            """
-            mutation checkIn($input: CheckInInput!) {
-                checkIn(input: $input) {
-                    ... on Hardware {
-                        createdAt
-                        updatedAt
-                        lastCheckIn
-                    }
-                    ... on OperationInfo {
-                          messages {
-                            kind
-                            message
-                            field
-                            code
-                        }
-                    }
-                }
-            }
-        """  # noqa
-        )
+        mutation = gql(hardware_checkin_mutation)
         input = new_state
         variables = {"input": input}
         try:
@@ -459,70 +416,66 @@ class Hardware(BaseAction):
             raise exception
 
     @guard
-    def get_hardware_list(self, fingerprint: Optional[str] = None):
-        query = gql(
-            """
-fragment HardwareFragment on Hardware {
-  id
-  pk
-  name
-  slug
-  createdAt
-  updatedAt
-  isAvailable
-  isOnline
-  isQuarantined
-  isHealthy
-  capabilities {
-    id
-    pk
-  }
-  activeReservation {
-    id
-    pk
-  }
-}
-
-query hardwareList(
-  $before: String
-  $after: String
-  $first: Int
-  $last: Int
-  $filters: HardwareFilters
-) {
-  hardwareList(
-    before: $before
-    after: $after
-    first: $first
-    last: $last
-    filters: $filters
-  ) {
-    totalCount
-    edges {
-      cursor
-      node {
-        ...HardwareFragment
-      }
-    }
-  }
-}
-"""
-        )
+    def get_hardware_list(
+        self, fingerprint: Optional[str] = None, slug: Optional[str] = None
+    ):
+        query = gql(hardware_list)
 
         filters = {}
         if fingerprint:
             filters["fingerprint"] = {"exact": fingerprint}
+        if slug:
+            filters["slug"] = {"exact": slug}
 
         variables = {
             "first": 1,
             "filters": filters,
         }
-        result = self.primitive.session.execute(query, variable_values=variables)
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
         return result
 
     def get_own_hardware_details(self):
-        hardware_list_data = self.get_hardware_list(
+        hardware_list_result = self.get_hardware_list(
             fingerprint=self.primitive.host_config.get("fingerprint")
         )
-        hardware = hardware_list_data.get("hardwareList").get("edges")[0].get("node")
+        hardware = (
+            hardware_list_result.data.get("hardwareList").get("edges")[0].get("node")
+        )
+        return hardware
+
+    def get_hardware_from_slug_or_id(self, hardware_identifier: str):
+        is_id = False
+        is_slug = False
+        # first check if the hardware_identifier is a slug or ID
+        try:
+            type_name, id = from_base64(hardware_identifier)
+
+            is_id = True
+            if type_name == "Hardware":
+                is_hardware_type = True
+            else:
+                raise Exception(
+                    f"ID was not for Hardware, you supplied an ID for a {type_name}"
+                )
+
+        except ValueError:
+            is_slug = True
+
+        hardware = None
+
+        if is_id and is_hardware_type:
+            hardware_list_result = self.get_hardware_list(slug=hardware_identifier)
+            if edges := hardware_list_result.data["hardwareList"]["edges"]:
+                hardware = edges[0]["node"]
+            else:
+                raise Exception(f"No hardware found with slug {hardware_identifier}")
+        elif is_slug:
+            hardware_list_result = self.get_hardware_list(slug=hardware_identifier)
+            if edges := hardware_list_result.data["hardwareList"]["edges"]:
+                hardware = edges[0]["node"]
+            else:
+                raise Exception(f"No hardware found with slug {hardware_identifier}")
+
         return hardware

primitive/hardware/commands.py

@@ -1,9 +1,9 @@
+import typing
+
 import click
 
 from ..utils.printer import print_result
 
-import typing
-
 if typing.TYPE_CHECKING:
     from ..client import Primitive
 
@@ -31,7 +31,7 @@ def register_command(context):
     primitive: Primitive = context.obj.get("PRIMITIVE")
     result = primitive.hardware.register()
     color = "green" if result else "red"
-    if result:
+    if result.data.get("registerHardware"):
         message = "Hardware registered successfully"
     else:
         message = (
@@ -45,9 +45,19 @@ def register_command(context):
 def checkin_command(context):
     """Checkin Hardware with Primitive"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    result = primitive.hardware.check_in_http()
-    if messages := result.get("checkIn").get("messages"):
+    check_in_http_result = primitive.hardware.check_in_http()
+    if messages := check_in_http_result.data.get("checkIn").get("messages"):
         print_result(message=messages, context=context, fg="yellow")
     else:
         message = "Hardware checked in successfully"
         print_result(message=message, context=context, fg="green")
+
+
+@cli.command("list")
+@click.pass_context
+def list_command(context):
+    """List Hardware"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    get_hardware_list_result = primitive.hardware.get_hardware_list()
+    message = get_hardware_list_result.data
+    print_result(message=message, context=context)

primitive/hardware/graphql/fragments.py

@@ -0,0 +1,22 @@
+hardware_fragment = """
+fragment HardwareFragment on Hardware {
+  id
+  pk
+  name
+  slug
+  createdAt
+  updatedAt
+  isAvailable
+  isOnline
+  isQuarantined
+  isHealthy
+  capabilities {
+    id
+    pk
+  }
+  activeReservation {
+    id
+    pk
+  }
+}
+"""

primitive/hardware/graphql/mutations.py

@@ -0,0 +1,45 @@
+from primitive.graphql.utility_fragments import operation_info_fragment
+
+register_hardware_mutation = (
+    operation_info_fragment
+    + """
+mutation registerHardware($input: RegisterHardwareInput!) {
+    registerHardware(input: $input) {
+        ... on Hardware {
+            fingerprint
+        }
+        ...OperationInfoFragment
+    }
+}
+"""
+)
+
+hardware_update_mutation = (
+    operation_info_fragment
+    + """
+mutation hardwareUpdate($input: HardwareUpdateInput!) {
+    hardwareUpdate(input: $input) {
+        ... on Hardware {
+            systemInfo
+        }
+        ...OperationInfoFragment
+    }
+}
+"""
+)
+
+hardware_checkin_mutation = (
+    operation_info_fragment
+    + """
+mutation checkIn($input: CheckInInput!) {
+    checkIn(input: $input) {
+        ... on Hardware {
+            createdAt
+            updatedAt
+            lastCheckIn
+        }
+        ...OperationInfoFragment
+    }
+}
+"""
+)

primitive/hardware/graphql/queries.py

@@ -0,0 +1,31 @@
+from .fragments import hardware_fragment
+
+hardware_list = (
+    hardware_fragment
+    + """
+
+query hardwareList(
+  $before: String
+  $after: String
+  $first: Int
+  $last: Int
+  $filters: HardwareFilters
+) {
+  hardwareList(
+    before: $before
+    after: $after
+    first: $first
+    last: $last
+    filters: $filters
+  ) {
+    totalCount
+    edges {
+      cursor
+      node {
+        ...HardwareFragment
+      }
+    }
+  }
+}
+"""
+)

primitive/jobs/actions.py

@@ -1,9 +1,18 @@
 from typing import List, Optional
-from gql import gql
 
+from gql import gql
 
 from primitive.utils.actions import BaseAction
+
 from ..utils.auth import guard
+from .graphql.mutations import job_run_update_mutation
+from .graphql.queries import (
+    github_app_token_for_job_run_query,
+    job_run_query,
+    job_run_status_query,
+    job_runs_query,
+    jobs_query,
+)
 
 
 class Jobs(BaseAction):
@@ -17,52 +26,7 @@ class Jobs(BaseAction):
         first: Optional[int] = 1,
         last: Optional[int] = None,
     ):
-        query = gql(
-            """ 
-            fragment PageInfoFragment on PageInfo {
-                hasNextPage
-                hasPreviousPage
-                startCursor
-                endCursor
-            }
-
-            fragment JobFragment on Job {
-                id
-                pk
-                slug
-                name
-                createdAt
-                updatedAt
-            }
-
-            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
-                        }
-                    }
-                }
-            }
-            """
-        )
+        query = gql(jobs_query)
 
         filters = {}
         if organization_id:
@@ -101,73 +65,7 @@ class Jobs(BaseAction):
         first: Optional[int] = 1,
         last: Optional[int] = None,
     ):
-        query = gql(
-            """
-fragment PageInfoFragment on PageInfo {
-  hasNextPage
-  hasPreviousPage
-  startCursor
-  endCursor
-}
-
-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
-  }
-}
-
-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
-      }
-    }
-  }
-}
-"""
-        )
+        query = gql(job_runs_query)
 
         filters = {}
         if organization_id:
@@ -194,50 +92,18 @@ query jobRuns(
             },
         }
 
-        result = self.primitive.session.execute(query, variable_values=variables)
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
         return result
 
     @guard
     def get_job_run(self, id: str):
-        query = gql(
-            """
-            fragment JobRunFragment on JobRun {
-                id
-                pk
-                createdAt
-                updatedAt
-                completedAt
-                startedAt
-                status
-                conclusion
-                job {
-                    id
-                    pk
-                    slug
-                    name
-                    createdAt
-                    updatedAt
-                }
-                gitCommit {
-                    sha
-                    branch
-                    repoFullName
-                }
-                jobSettings {
-                    containerArgs
-                    rootDirectory
-                }
-            }
-
-            query jobRun($id: GlobalID!) {
-                jobRun(id: $id) {
-                    ...JobRunFragment
-                }
-            }
-            """
-        )
+        query = gql(job_run_query)
         variables = {"id": id}
-        result = self.primitive.session.execute(query, variable_values=variables)
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
         return result
 
     @guard
@@ -248,19 +114,7 @@ query jobRuns(
         conclusion: str = None,
         file_ids: Optional[List[str]] = [],
     ):
-        mutation = gql(
-            """
-            mutation jobRunUpdate($input: JobRunUpdateInput!) {
-                jobRunUpdate(input: $input) {
-                    ... on JobRun {
-                        id
-                        status
-                        conclusion
-                    }
-                }
-            }
-        """
-        )
+        mutation = gql(job_run_update_mutation)
         input = {"id": id}
         if status:
             input["status"] = status
@@ -269,21 +123,19 @@ query jobRuns(
         if file_ids and len(file_ids) > 0:
             input["files"] = file_ids
         variables = {"input": input}
-        result = self.primitive.session.execute(mutation, variable_values=variables)
+        result = self.primitive.session.execute(
+            mutation, variable_values=variables, get_execution_result=True
+        )
         return result
 
     @guard
     def github_access_token_for_job_run(self, job_run_id: str):
-        query = gql(
-            """
-query ghAppTokenForJobRun($jobRunId: GlobalID!) {
-    ghAppTokenForJobRun(jobRunId: $jobRunId)
-}
-"""
-        )
+        query = gql(github_app_token_for_job_run_query)
         variables = {"jobRunId": job_run_id}
-        result = self.primitive.session.execute(query, variable_values=variables)
-        return result["ghAppTokenForJobRun"]
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
+        return result.data["ghAppTokenForJobRun"]
 
     def get_latest_job_run_for_job(
         self, job_slug: Optional[str] = None, job_id: Optional[str] = None
@@ -300,32 +152,11 @@ query ghAppTokenForJobRun($jobRunId: GlobalID!) {
         ]
         return job_run
 
-    def job_run_start(
-        self,
-        job_slug: str,
-        job_id: str,
-    ):
-        if not job_slug and not job_id:
-            raise ValueError("job_slug or job_id is required")
-
-        self.get_jobs(slug=job_slug)
-
     @guard
     def get_job_status(self, id: str):
-        query = gql(
-            """
-            fragment JobRunFragment on JobRun {
-                id
-                status
-            }
-
-            query jobRun($id: GlobalID!) {
-                jobRun(id: $id) {
-                    ...JobRunFragment
-                }
-            }
-            """
-        )
+        query = gql(job_run_status_query)
         variables = {"id": id}
-        result = self.primitive.session.execute(query, variable_values=variables)
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
         return result