primitive 0.2.34__py3-none-any.whl → 0.2.36__py3-none-any.whl

primitive/__about__.py

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

primitive/agent/actions.py

@@ -102,4 +102,4 @@ class Agent(BaseAction):
 
                 sleep(5)
         except KeyboardInterrupt:
-            logger.info("[agent] Stopping primitive agent...")
+            logger.info("Stopping primitive agent...")

primitive/agent/runner.py

@@ -3,11 +3,12 @@ import os
 import re
 import shutil
 import typing
-from abc import abstractmethod
 from enum import Enum, IntEnum
 from pathlib import Path, PurePath
 from typing import Dict, List, TypedDict
 
+from ..utils.logging import log_context, fmt
+
 import yaml
 from loguru import logger
 
@@ -57,16 +58,6 @@ class LogLevel(Enum):
     WARNING = "WARNING"
 
 
-# Log Counter
-class LogCounter:
-    count = 0
-
-    @classmethod
-    def next(cls) -> int:
-        cls.count += 1
-        return cls.count
-
-
 class Runner:
     def __init__(
         self,
@@ -79,7 +70,7 @@ class Runner:
         self.job_run = job_run
         self.job_settings = job_run["jobSettings"]
         self.config = None
-        self.source_dir: Path = None
+        self.source_dir: Path | None = None
         self.initial_env = {}
         self.modified_env = {}
         self.file_logger = None
@@ -91,11 +82,11 @@ class Runner:
             self.file_logger = logger.add(
                 Path(get_logs_cache(self.job_run["id"]) / log_name),
                 rotation=max_log_size,
-                format=Runner.fmt,
+                format=fmt,
                 backtrace=True,
-                diagnose=True,
             )
 
+    @log_context(label="setup")
     def setup(self) -> None:
         # Attempt to download the job source code
         git_repo_full_name = self.job_run["gitCommit"]["repoFullName"]
@@ -155,66 +146,62 @@ class Runner:
             self.job_run["gitCommit"]["repoFullName"]
         )
 
+    @log_context(label="execute")
     def execute(self) -> None:
         logger.info(f"Executing {self.job['slug']} job")
         self.primitive.jobs.job_run_update(
             self.job_run["id"], status="request_in_progress"
         )
-        self.modified_env = {**self.initial_env}
 
+        self.modified_env = {**self.initial_env}
         task_failed = False
         cancelled = False
 
         for task in self.config["executes"]:
-            # the get status check here is to ensure that if cancel is called
-            # while one task is running, we do not run any OTHER laebeled tasks
-            # THIS is required for MULTI STEP JOBS
+            # Everything inside this loop should be contextualized with the task label
+            # this way we aren't jumping back and forth between the task label and "execute"
+            with logger.contextualize(label=task["label"]):
+                # the get status check here is to ensure that if cancel is called
+                # while one task is running, we do not run any OTHER labeled tasks
+                # THIS is required for MULTI STEP JOBS
+                status = self.primitive.jobs.get_job_status(self.job_run["id"])
+                status_value = status.data["jobRun"]["status"]
+                conclusion_value = status.data["jobRun"]["conclusion"]
+
+                if status_value == "completed" and conclusion_value == "cancelled":
+                    cancelled = True
+                    break
+
+                # Everything within this block should be contextualized as user logs
+                with logger.contextualize(type="user"):
+                    with asyncio.Runner() as async_runner:
+                        if task_failed := async_runner.run(self.run_task(task)):
+                            break
+
+        # FOR NONE MULTI STEP JOBS
+        # we still have to check that the job was cancelled here as well
+        with logger.contextualize(label="conclusion"):
             status = self.primitive.jobs.get_job_status(self.job_run["id"])
             status_value = status.data["jobRun"]["status"]
             conclusion_value = status.data["jobRun"]["conclusion"]
-
             if status_value == "completed" and conclusion_value == "cancelled":
                 cancelled = True
-                break
 
-            with logger.contextualize(label=task["label"]):
-                with asyncio.Runner() as async_runner:
-                    if task_failed := async_runner.run(self.run_task(task)):
-                        break
+            if cancelled:
+                logger.warning("Job cancelled by user")
+                return
 
-        number_of_files_produced = self.get_number_of_files_produced()
-        logger.info(
-            f"Produced {number_of_files_produced} files for {self.job['slug']} job"
-        )
+            conclusion = "success"
+            if task_failed:
+                conclusion = "failure"
+            else:
+                logger.success(f"Completed {self.job['slug']} job")
 
-        # FOR NONE MULTI STEP JOBS
-        # we still have to check that the job was cancelled here as well
-        status = self.primitive.jobs.get_job_status(self.job_run["id"])
-        status_value = status.data["jobRun"]["status"]
-        conclusion_value = status.data["jobRun"]["conclusion"]
-        if status_value == "completed" and conclusion_value == "cancelled":
-            cancelled = True
-
-        if cancelled:
-            logger.warning("Job cancelled by user")
             self.primitive.jobs.job_run_update(
                 self.job_run["id"],
-                number_of_files_produced=number_of_files_produced,
+                status="request_completed",
+                conclusion=conclusion,
             )
-            return
-
-        conclusion = "success"
-        if task_failed:
-            conclusion = "failure"
-        else:
-            logger.success(f"Completed {self.job['slug']} job")
-
-        self.primitive.jobs.job_run_update(
-            self.job_run["id"],
-            status="request_completed",
-            conclusion=conclusion,
-            number_of_files_produced=number_of_files_produced,
-        )
 
     def get_number_of_files_produced(self) -> int:
         """Returns the number of files produced by the job."""
@@ -261,7 +248,9 @@ class Runner:
         return number_of_files_produced
 
     async def run_task(self, task: Task) -> bool:
-        for cmd in task["cmd"].strip().split("\n"):
+        logger.info(f"Running step '{task['label']}'")
+        commands = task["cmd"].strip().split("\n")
+        for i, cmd in enumerate(commands):
             # Adding an additional echo and utilizing stdbuf to force line buffering
             # This ensures that the environment variables and starting delimiter are
             # always in a new chunk, vastly simplifying our parsing logic
@@ -271,6 +260,8 @@ class Runner:
                 f"{cmd} && echo -n '{ENV_VAR_LOOKUP_START}' && env && echo -n '{ENV_VAR_LOOKUP_END}'",
             ]
 
+            logger.info(f"Executing command {i + 1}/{len(commands)}: {cmd}")
+
             process = await asyncio.create_subprocess_exec(
                 *args,
                 env=self.modified_env,
@@ -294,6 +285,10 @@ class Runner:
 
             returncode = await process.wait()
 
+            logger.info(
+                f"Finished executing command {i + 1}/{len(commands)}: {cmd} with return code {returncode}"
+            )
+
             JobRun.objects.filter_by(job_run_id=self.job_run["id"]).update(
                 {"pid": None}
             )
@@ -307,6 +302,7 @@ class Runner:
                 logger.error(f"Task {task['label']} failed on '{cmd}'")
                 return True
 
+        logger.success(f"Completed {task['label']} task")
         return False
 
     async def log_cmd(self, process, stream, tags: Dict = {}) -> bool:
@@ -401,9 +397,8 @@ class Runner:
 
         return [line for line in lines if len(line) > 0]
 
+    @log_context(label="cleanup")
     def cleanup(self) -> None:
-        logger.remove(self.file_logger)
-
         if "stores" not in self.config:
             return
 
@@ -419,23 +414,13 @@ class Runner:
 
         shutil.rmtree(path=self.source_dir)
 
-    @abstractmethod
-    def fmt(record) -> str:
-        extra = record["extra"]
-        # Delimiters with empty space MUST exist for LogQL pattern matching
-        label = extra.get("label", None)
-        tag = extra.get("tag", None)
-        context = f"{label} | " if label else " | "
-        context += f"{tag} | " if tag else " | "
-
-        log = (
-            f"{LogCounter.next()} | "
-            "<green>{time:YYYY-MM-DD HH:mm:ss.SSS!UTC}</green> | "
-            "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
-            "<level>{level}</level> | "
-            f"{context}"
-            "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - "
-            "<level>{message}</level>\n"
+        number_of_files_produced = self.get_number_of_files_produced()
+        logger.info(
+            f"Produced {number_of_files_produced} files for {self.job['slug']} job"
+        )
+        self.primitive.jobs.job_run_update(
+            self.job_run["id"],
+            number_of_files_produced=number_of_files_produced,
         )
 
-        return log
+        logger.remove(self.file_logger)

primitive/utils/logging.py

@@ -0,0 +1,44 @@
+from functools import wraps
+from loguru import logger
+from datetime import timezone
+import json
+
+
+def log_context(**context):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with logger.contextualize(**context):
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def fmt(record) -> str:
+    extra = record["extra"]
+    label = extra["label"]
+    tag = extra.get("tag", None)
+    type = extra.get("type", "system")
+
+    context_object = {
+        "label": label,
+        "type": type,
+        "utc": record["time"]
+        .astimezone(timezone.utc)
+        .strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
+        "level": record["level"].name,
+        "name": record["name"],
+        "function": record["function"],
+        "line": record["line"],
+        "message": record["message"],
+    }
+
+    if tag:
+        context_object["tag"] = tag
+
+    # Loguru will fail if you return a string that doesn't select
+    # something within its record
+    record["extra"]["serialized"] = json.dumps(context_object)
+    return "{extra[serialized]}\n"

{primitive-0.2.34.dist-info → primitive-0.2.36.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: primitive
-Version: 0.2.34
+Version: 0.2.36
 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

{primitive-0.2.34.dist-info → primitive-0.2.36.dist-info}/RECORD

@@ -1,11 +1,11 @@
-primitive/__about__.py,sha256=s58wIznYy_MiaxrfZEt9pFQ41lrUgI4Bc_OF1gkNEwg,130
+primitive/__about__.py,sha256=5-Cyrivu3AqsJWwi4il_NEU7wvPqzCfGtM5hHx7QXoM,130
 primitive/__init__.py,sha256=bwKdgggKNVssJFVPfKSxqFMz4IxSr54WWbmiZqTMPNI,106
 primitive/cli.py,sha256=g7EtHI9MATAB0qQu5w-WzbXtxz_8zu8z5E7sETmMkKU,2509
 primitive/client.py,sha256=h8WZVnQylVe0vbpuyC8YZHl2JyITSPC-1HbUcmrE5pc,3623
 primitive/agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-primitive/agent/actions.py,sha256=JLnHCtu0AgeHRZccCc2I39mziLixmj9EKtKeYmvBE7A,3700
+primitive/agent/actions.py,sha256=VHlNTw2M-T1MRajpIBu2weT8AsdDdYtlgbBnRBYYeco,3692
 primitive/agent/commands.py,sha256=cK7d3OcN5Z65gQWVZFQ-Y9ddw9Pes4f9OVBpeMsj5sE,255
-primitive/agent/runner.py,sha256=Vuz0bJjfq3IytLI6HxC9Gfd1TdARjvNsG273GEWxjKI,15462
+primitive/agent/runner.py,sha256=UMLCF0BhyBRJLGntB1C5dDCmIVYjE7AiI2RqpupJRd0,15464
 primitive/agent/uploader.py,sha256=ZzrzsajNBogwEC7mT6Ejy0h2Jd9axMYGzt9pbCvVMlk,3171
 primitive/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 primitive/auth/actions.py,sha256=9NIEXJ1BNJutJs6AMMSjMN_ziONUAUhY_xHwojYJCLA,942
@@ -92,12 +92,13 @@ primitive/utils/chunk_size.py,sha256=PAuVuirUTA9oRXyjo1c6MWxo31WVBRkWMuWw-AS58Bw
 primitive/utils/config.py,sha256=DlFM5Nglo22WPtbpZSVtH7NX-PTMaKYlcrUE7GPRG4c,1058
 primitive/utils/daemons.py,sha256=mSoSHitiGfS4KYAEK9sKsiv_YcACHKgY3qISnDpUUIE,1086
 primitive/utils/exceptions.py,sha256=DrYHTcCAJGC7cCUwOx_FmdlVLWRdpzvDvpLb82heppE,311
+primitive/utils/logging.py,sha256=W-MY6B---nG7A5lg17t5ZWTGytO0Y8_KPUvFCs5MBcs,1121
 primitive/utils/memory_size.py,sha256=4xfha21kW82nFvOTtDFx9Jk2ZQoEhkfXii-PGNTpIUk,3058
 primitive/utils/printer.py,sha256=f1XUpqi5dkTL3GWvYRUGlSwtj2IxU1q745T4Fxo7Tn4,370
 primitive/utils/shell.py,sha256=Z4zxmOaSyGCrS0D6I436iQci-ewHLt4UxVg1CD9Serc,2171
 primitive/utils/text.py,sha256=XiESMnlhjQ534xE2hMNf08WehE1SKaYFRNih0MmnK0k,829
-primitive-0.2.34.dist-info/METADATA,sha256=tXBOqx2I7GoxBe13HXcWi0qSby1oIwlX4_hdwNicsYQ,3569
-primitive-0.2.34.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-primitive-0.2.34.dist-info/entry_points.txt,sha256=p1K8DMCWka5FqLlqP1sPek5Uovy9jq8u51gUsP-z334,48
-primitive-0.2.34.dist-info/licenses/LICENSE.txt,sha256=B8kmQMJ2sxYygjCLBk770uacaMci4mPSoJJ8WoDBY_c,1098
-primitive-0.2.34.dist-info/RECORD,,
+primitive-0.2.36.dist-info/METADATA,sha256=fSyVyuB8t19XnCtCHqh5gcP1QPhBdl_MSWvRC0ybcFo,3569
+primitive-0.2.36.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+primitive-0.2.36.dist-info/entry_points.txt,sha256=p1K8DMCWka5FqLlqP1sPek5Uovy9jq8u51gUsP-z334,48
+primitive-0.2.36.dist-info/licenses/LICENSE.txt,sha256=B8kmQMJ2sxYygjCLBk770uacaMci4mPSoJJ8WoDBY_c,1098
+primitive-0.2.36.dist-info/RECORD,,