primitive 0.1.60__tar.gz → 0.1.62__tar.gz

{primitive-0.1.60 → primitive-0.1.62}/.vscode/settings.json

@@ -3,6 +3,7 @@
     "Checkin",
     "loguru",
     "machdep",
+    "pkey",
     "testsuites",
     "untar",
     "Untaring",

{primitive-0.1.60 → primitive-0.1.62}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: primitive
-Version: 0.1.60
+Version: 0.1.62
 Project-URL: Documentation, https://github.com//primitivecorp/primitive-cli#readme
 Project-URL: Issues, https://github.com//primitivecorp/primitive-cli/issues
 Project-URL: Source, https://github.com//primitivecorp/primitive-cli
@@ -21,6 +21,7 @@ Requires-Python: >=3.11
 Requires-Dist: click
 Requires-Dist: gql[all]
 Requires-Dist: loguru
+Requires-Dist: paramiko[all]
 Requires-Dist: primitive-pal==0.1.4
 Requires-Dist: pyyaml
 Description-Content-Type: text/markdown

{primitive-0.1.60 → primitive-0.1.62}/pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
   "loguru",
   "pyyaml",
   "primitive-pal == 0.1.4",
+  "paramiko[all]"
 ]
 
 [tool.uv]

{primitive-0.1.60 → primitive-0.1.62}/src/primitive/__about__.py

@@ -1,4 +1,4 @@
 # SPDX-FileCopyrightText: 2024-present Dylan Stein <dylan@primitive.tech>
 #
 # SPDX-License-Identifier: MIT
-__version__ = "0.1.60"
+__version__ = "0.1.62"

{primitive-0.1.60 → primitive-0.1.62}/src/primitive/agent/runner.py

@@ -191,7 +191,8 @@ class AgentRunner:
             logger.remove(self.logger_handle)
 
         self.logger_handle = logger.add(
-            Path(self.logs_dir / f"{label}_{{time}}.primitive.log"), rotation=self.max_log_size
+            Path(self.logs_dir / f"{label}_{{time}}.primitive.log"),
+            rotation=self.max_log_size,
         )
 
     def provision(self) -> Optional[Dict]:

{primitive-0.1.60 → primitive-0.1.62}/src/primitive/agent/uploader.py

@@ -2,8 +2,10 @@ import typing
 from typing import Dict
 import shutil
 import os
-from loguru import logger
 from pathlib import Path, PurePath
+
+from loguru import logger
+
 from ..utils.cache import get_artifacts_cache, get_logs_cache
 
 if typing.TYPE_CHECKING:
@@ -17,12 +19,6 @@ class Uploader:
     ):
         self.primitive = primitive
 
-    def upload_file(self, path: Path, prefix: str, job_run_id: str) -> str:
-        file_upload_response = self.primitive.files.file_upload(
-            path, key_prefix=prefix, job_run_id=job_run_id
-        )
-        return file_upload_response.json()["data"]["fileUpload"]["id"]
-    
     def upload_dir(self, cache: Path) -> Dict:
         file_ids = []
         job_run_id = cache.name
@@ -49,11 +45,12 @@ class Uploader:
             )
 
         for file in files:
-            upload_id = self.upload_file(
-                file,
-                prefix=str(PurePath(file).relative_to(cache.parent).parent),
+            response = self.primitive.files.upload_file_via_api(
+                path=file,
+                key_prefix=str(PurePath(file).relative_to(cache.parent).parent),
                 job_run_id=job_run_id,
             )
+            upload_id = response.json()["data"]["fileUpload"]["id"]
 
             if upload_id:
                 file_ids.append(upload_id)
@@ -66,23 +63,24 @@ class Uploader:
 
         return {job_run_id: file_ids}
 
-
     def scan(self) -> None:
         # Scan artifacts directory
         artifacts_dir = get_artifacts_cache()
         logs_dir = get_logs_cache()
 
-        artifacts = sorted([
-            artifacts_cache
-            for artifacts_cache in artifacts_dir.iterdir()
-            if artifacts_cache.is_dir()
-        ], key=lambda p: p.stat().st_ctime)
+        artifacts = sorted(
+            [
+                artifacts_cache
+                for artifacts_cache in artifacts_dir.iterdir()
+                if artifacts_cache.is_dir()
+            ],
+            key=lambda p: p.stat().st_ctime,
+        )
 
-        logs = sorted([
-            logs_cache 
-            for logs_cache in logs_dir.iterdir() 
-            if logs_cache.is_dir()
-        ], key=lambda p: p.stat().st_ctime)
+        logs = sorted(
+            [logs_cache for logs_cache in logs_dir.iterdir() if logs_cache.is_dir()],
+            key=lambda p: p.stat().st_ctime,
+        )
 
         log_files = {
             job_id: files
@@ -97,7 +95,7 @@ class Uploader:
         }
 
         files_by_id = {
-            job_id: log_files.get(job_id, []) + artifact_files.get(job_id, []) 
+            job_id: log_files.get(job_id, []) + artifact_files.get(job_id, [])
             for job_id in log_files.keys() | artifact_files.keys()
         }
 

{primitive-0.1.60 → primitive-0.1.62}/src/primitive/exec/actions.py

@@ -1,9 +1,14 @@
+import tempfile
 import typing
 
+from primitive.exec.interactive import interactive_shell
+
 if typing.TYPE_CHECKING:
     pass
 
 
+from paramiko import SSHClient
+
 from primitive.utils.actions import BaseAction
 
 
@@ -41,7 +46,48 @@ class Exec(BaseAction):
             reservation_id=reservation["id"], desired_status="in_progress"
         )
 
-        print(f"Executing command: {command} on {hardware['name']}")
+        ssh_credentials = self.primitive.hardware.get_hardware_ssh_credentials(
+            hardware_id=hardware["id"]
+        )
+
+        ssh_hostname = ssh_credentials["ssh_hostname"]
+        ssh_username = ssh_credentials["ssh_username"]
+        ssh_password = ssh_credentials["ssh_password"]
+        ssh_private_key = ssh_credentials["ssh_private_key"]
+
+        ssh_client = SSHClient()
+        ssh_client.load_system_host_keys()
+        keyfile = None
+        if ssh_private_key:
+            keyfile = tempfile.NamedTemporaryFile()
+            keyfile.write(ssh_private_key.encode())
+            keyfile.flush()
+            ssh_client.connect(
+                hostname=ssh_hostname,
+                username=ssh_username,
+                key_filename=keyfile.name,
+            )
+        else:
+            ssh_client.connect(
+                hostname=ssh_hostname,
+                username=ssh_username,
+                pkey=ssh_private_key,
+                password=ssh_password,
+            )
+
+        if command:
+            formatted_command = " ".join(command)
+            stdin, stdout, stderr = ssh_client.exec_command(formatted_command)
+            print(stdout.read())
+            ssh_client.close()
+        else:
+            channel = ssh_client.get_transport().open_session()
+            channel.get_pty()
+            channel.invoke_shell()
+            interactive_shell(channel)
+            ssh_client.close()
+            if keyfile:
+                keyfile.close()
 
         if created_reservation_on_behalf_of_user:
             print("Cleaning up reservation.")

{primitive-0.1.60 → primitive-0.1.62}/src/primitive/exec/commands.py

@@ -13,7 +13,7 @@ if typing.TYPE_CHECKING:
     type=str,
     required=True,
 )
-@click.argument("command", nargs=-1, required=True)
+@click.argument("command", nargs=-1, required=False)
 def cli(context, hardware_identifier: str, command: str) -> None:
     """Exec"""
     primitive: Primitive = context.obj.get("PRIMITIVE")

primitive-0.1.62/src/primitive/exec/interactive.py

@@ -0,0 +1,99 @@
+# https://github.com/paramiko/paramiko/blob/main/demos/interactive.py
+
+import socket
+import sys
+
+
+def _to_unicode(s):
+    """
+    decode a string as ascii or utf8 if possible (as required by the sftp
+    protocol).  if neither works, just return a byte string because the server
+    probably doesn't know the filename's encoding.
+    """
+    try:
+        return s.encode("ascii")
+    except (UnicodeError, AttributeError):
+        try:
+            return s.decode("utf-8")
+        except UnicodeError:
+            return s
+
+
+# windows does not have termios...
+try:
+    import termios
+    import tty
+
+    has_termios = True
+except ImportError:
+    has_termios = False
+
+
+def interactive_shell(chan):
+    if has_termios:
+        posix_shell(chan)
+    else:
+        windows_shell(chan)
+
+
+def posix_shell(chan):
+    import select
+
+    oldtty = termios.tcgetattr(sys.stdin)
+    try:
+        tty.setraw(sys.stdin.fileno())
+        tty.setcbreak(sys.stdin.fileno())
+        chan.settimeout(0.0)
+
+        while True:
+            r, w, e = select.select([chan, sys.stdin], [], [])
+            if chan in r:
+                try:
+                    x = _to_unicode(chan.recv(1024))
+                    if len(x) == 0:
+                        sys.stdout.write("\r\n*** EOF\r\n")
+                        break
+                    sys.stdout.write(x)
+                    sys.stdout.flush()
+                except socket.timeout:
+                    pass
+            if sys.stdin in r:
+                x = sys.stdin.read(1)
+                if len(x) == 0:
+                    break
+                chan.send(x)
+
+    finally:
+        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)
+
+
+# thanks to Mike Looijmans for this code
+def windows_shell(chan):
+    import threading
+
+    sys.stdout.write(
+        "Line-buffered terminal emulation. Press F6 or ^Z to send EOF.\r\n\r\n"
+    )
+
+    def writeall(sock):
+        while True:
+            data = sock.recv(256)
+            if not data:
+                sys.stdout.write("\r\n*** EOF ***\r\n\r\n")
+                sys.stdout.flush()
+                break
+            sys.stdout.write(data)
+            sys.stdout.flush()
+
+    writer = threading.Thread(target=writeall, args=(chan,))
+    writer.start()
+
+    try:
+        while True:
+            d = sys.stdin.read(1)
+            if not d:
+                break
+            chan.send(d)
+    except EOFError:
+        # user hit ^Z or F6
+        pass

primitive-0.1.62/src/primitive/files/actions.py

@@ -0,0 +1,189 @@
+import hashlib
+import sys
+import threading
+from pathlib import Path
+from typing import Dict, Optional
+
+import requests
+from gql import gql
+from loguru import logger
+
+from primitive.graphql.sdk import create_requests_session
+from primitive.utils.actions import BaseAction
+
+from ..utils.auth import guard
+from .graphql.mutations import (
+    file_update_mutation,
+    pending_file_create_mutation,
+)
+
+
+# this class can be used in multithreaded S3 client uploader
+# this requires getting an S3 access token to this machine however
+# we are using presigned urls instead at this time Oct 29th, 2024
+class ProgressPercentage(object):
+    def __init__(self, filepath: Path) -> None:
+        self._filename = filepath.name
+        self._size = float(filepath.stat().st_size)
+        self._seen_so_far = 0
+        self._lock = threading.Lock()
+
+    def __call__(self, bytes_amount):
+        # To simplify, assume this is hooked up to a single filename
+        with self._lock:
+            self._seen_so_far += bytes_amount
+            percentage = (self._seen_so_far / self._size) * 100
+            sys.stdout.write(
+                "\r%s  %s / %s  (%.2f%%)"
+                % (self._filename, self._seen_so_far, self._size, percentage)
+            )
+            sys.stdout.flush()
+
+
+class Files(BaseAction):
+    def _pending_file_create(
+        self,
+        file_name: str,
+        file_size: int,
+        file_checksum: str,
+        file_path: str,
+        key_prefix: str,
+        is_public: bool = False,
+    ):
+        mutation = gql(pending_file_create_mutation)
+        input = {
+            "filePath": file_path,
+            "fileName": file_name,
+            "fileSize": file_size,
+            "fileChecksum": file_checksum,
+            "keyPrefix": key_prefix,
+            "isPublic": is_public,
+        }
+        variables = {"input": input}
+        result = self.primitive.session.execute(
+            mutation, variable_values=variables, get_execution_result=True
+        )
+        return result.data.get("pendingFileCreate")
+
+    def _update_file_status(
+        self,
+        file_id: str,
+        is_uploading: Optional[bool] = None,
+        is_complete: Optional[bool] = None,
+    ):
+        mutation = gql(file_update_mutation)
+        input: Dict[str, str | bool] = {
+            "id": file_id,
+        }
+        if is_uploading is not None:
+            input["isUploading"] = is_uploading
+        if is_complete is not None:
+            input["isComplete"] = is_complete
+
+        variables = {"input": input}
+        result = self.primitive.session.execute(
+            mutation, variable_values=variables, get_execution_result=True
+        )
+        return result
+
+    @guard
+    def upload_file_direct(
+        self,
+        path: Path,
+        is_public: False,
+        key_prefix: str = "",
+        file_id: Optional[str] = None,
+    ):
+        logger.enable("primitive")
+        if path.exists() is False:
+            raise Exception(f"File {path} does not exist.")
+
+        file_size = path.stat().st_size
+        if file_size == 0:
+            raise Exception(f"{path} is empty.")
+
+        file_checksum = hashlib.md5(path.read_bytes()).hexdigest()
+
+        if not file_id:
+            pending_file_create = self._pending_file_create(
+                file_name=path.name,
+                file_size=path.stat().st_size,
+                file_checksum=file_checksum,
+                file_path=str(path),
+                key_prefix=key_prefix,
+                is_public=is_public,
+            )
+        file_id = pending_file_create.get("id")
+        presigned_url = pending_file_create.get("presignedUrlForUpload")
+
+        if not file_id:
+            raise Exception("No file_id found or provided.")
+        if not presigned_url:
+            raise Exception("No presigned_url returned.")
+
+        self._update_file_status(file_id, is_uploading=True)
+        with open(path, "rb") as object_file:
+            object_text = object_file.read()
+            response = requests.put(presigned_url, data=object_text)
+            if response.ok:
+                logger.info(f"File {path} uploaded successfully.")
+                update_file_status_result = self._update_file_status(
+                    file_id, is_uploading=False, is_complete=True
+                )
+            else:
+                message = f"Failed to upload file {path}. {response.status_code}: {response.text}"
+                logger.error(message)
+                raise Exception(message)
+        file_pk = update_file_status_result.data.get("fileUpdate").get("pk")
+        file_access_url = f"{self.primitive.host_config.get("transport")}://{self.primitive.host}/files/{file_pk}/presigned-url/"
+        logger.info(f"Available at: {file_access_url}")
+        return update_file_status_result
+
+    @guard
+    def upload_file_via_api(
+        self,
+        path: Path,
+        is_public: bool = False,
+        key_prefix: str = "",
+        job_run_id: str = "",
+    ):
+        """
+        This method uploads a file via the Primitive API.
+        This does NOT upload the file straight to S3
+        """
+        file_path = str(path.resolve())
+        if path.exists() is False:
+            raise FileNotFoundError(f"File not found at {file_path}")
+
+        if is_public:
+            operations = (
+                """{ "query": "mutation fileUpload($input: FileUploadInput!) { fileUpload(input: $input) { ... on File { id } } }", "variables": { "input": { "fileObject": null, "isPublic": true, "filePath": \""""
+                + file_path
+                + """\", "keyPrefix": \""""
+                + key_prefix
+                + """\", "jobRunId": \""""
+                + job_run_id
+                + """\" } } }"""
+            )  # noqa
+
+        else:
+            operations = (
+                """{ "query": "mutation fileUpload($input: FileUploadInput!) { fileUpload(input: $input) { ... on File { id } } }", "variables": { "input": { "fileObject": null, "isPublic": false, "filePath": \""""
+                + file_path
+                + """\", "keyPrefix": \""""
+                + key_prefix
+                + """\", "jobRunId": \""""
+                + job_run_id
+                + """\" } } }"""
+            )  # noqa
+        body = {
+            "operations": ("", operations),
+            "map": ("", '{"fileObject": ["variables.input.fileObject"]}'),
+            "fileObject": (path.name, open(path, "rb")),
+        }
+
+        session = create_requests_session(self.primitive.host_config)
+        transport = self.primitive.host_config.get("transport")
+        url = f"{transport}://{self.primitive.host}/"
+        response = session.post(url, files=body)
+        return response

{primitive-0.1.60 → primitive-0.1.62}/src/primitive/files/commands.py

@@ -1,7 +1,10 @@
+import json
+import typing
+from pathlib import Path
+
 import click
+
 from ..utils.printer import print_result
-from pathlib import Path
-import typing
 
 if typing.TYPE_CHECKING:
     from ..client import Primitive
@@ -19,10 +22,19 @@ def cli(context):
 @click.argument("path", type=click.Path(exists=True))
 @click.option("--public", "-p", help="Is this a Public file", is_flag=True)
 @click.option("--key-prefix", "-k", help="Key Prefix", default="")
-def file_upload_command(context, path, public, key_prefix):
+@click.option("--direct", "-k", help="direct", is_flag=True)
+def file_upload_command(context, path, public, key_prefix, direct):
     """File Upload"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
     path = Path(path)
-    result = primitive.files.file_upload(path, is_public=public, key_prefix=key_prefix)
-    message = result.json()
+    if direct:
+        result = primitive.files.upload_file_direct(
+            path, is_public=public, key_prefix=key_prefix
+        )
+    else:
+        result = primitive.files.upload_file_via_api(
+            path, is_public=public, key_prefix=key_prefix
+        )
+
+    message = json.dumps(result.data)
     print_result(message=message, context=context)

primitive-0.1.62/src/primitive/files/graphql/mutations.py

@@ -0,0 +1,32 @@
+from primitive.graphql.utility_fragments import operation_info_fragment
+
+file_update_mutation = (
+    operation_info_fragment
+    + """
+mutation fileUpdate($input: FileInputPartial!) {
+    fileUpdate(input: $input) {
+        ... on File {
+            id
+            pk
+        }
+        ...OperationInfoFragment
+    }
+}
+"""
+)
+
+pending_file_create_mutation = (
+    operation_info_fragment
+    + """
+mutation pendingFileCreate($input: PendingFileCreateInput!) {
+    pendingFileCreate(input: $input) {
+        ... on File {
+            id
+            pk
+            presignedUrlForUpload
+        }
+        ...OperationInfoFragment
+    }
+}
+"""
+)

{primitive-0.1.60 → primitive-0.1.62}/src/primitive/hardware/actions.py

@@ -22,7 +22,7 @@ from .graphql.mutations import (
     hardware_update_mutation,
     register_hardware_mutation,
 )
-from .graphql.queries import hardware_list
+from .graphql.queries import hardware_list, hardware_ssh_credentials
 
 if typing.TYPE_CHECKING:
     pass
@@ -479,3 +479,16 @@ class Hardware(BaseAction):
                 raise Exception(f"No hardware found with slug {hardware_identifier}")
 
         return hardware
+
+    @guard
+    def get_hardware_ssh_credentials(self, hardware_id: str):
+        query = gql(hardware_ssh_credentials)
+
+        variables = {
+            "id": hardware_id,
+        }
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
+        ssh_credentials = result.data.get("hardware").get("sshCredentials")
+        return ssh_credentials

{primitive-0.1.60 → primitive-0.1.62}/src/primitive/hardware/graphql/queries.py

@@ -29,3 +29,13 @@ query hardwareList(
 }
 """
 )
+
+hardware_ssh_credentials = """
+query hardwareSSHCredentials($id: GlobalID!) {
+  hardware(id: $id) {
+    id
+    pk
+    sshCredentials
+  }
+}
+"""

{primitive-0.1.60 → primitive-0.1.62}/src/primitive/sim/actions.py

@@ -46,10 +46,6 @@ class Sim(BaseAction):
 
         return True, message
 
-    def upload_file(self, path: Path, prefix: str) -> str:
-        file_upload_response = self.primitive.files.file_upload(path, key_prefix=prefix)
-        return file_upload_response.json()["data"]["fileUpload"]["id"]
-
     def collect_artifacts(self, source: Path, job_run_id: str) -> None:
         # Parse VCD artifacts using rust binding
         # TODO: eventually make this smarter, only parsing VCDs for failed tests
@@ -72,12 +68,12 @@ class Sim(BaseAction):
         )
         for file_path in files:
             try:
-                file_ids.append(
-                    self.upload_file(
-                        file_path,
-                        prefix=f"{job_run_id}/{str(PurePath(file_path).relative_to(Path(source)).parent)}",
-                    )
+                key_prefix = f"{job_run_id}/{str(PurePath(file_path).relative_to(Path(source)).parent)}"
+                file_upload_response = self.primitive.files.upload_file_via_api(
+                    file_path, key_prefix=key_prefix
                 )
+                file_id = file_upload_response.json()["data"]["fileUpload"]["id"]
+                file_ids.append(file_id)
             except FileNotFoundError:
                 logger.warning(f"{file_path} not found...")