primitive 0.2.3__tar.gz → 0.2.4__tar.gz

{primitive-0.2.3 → primitive-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: primitive
-Version: 0.2.3
+Version: 0.2.4
 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.3 → primitive-0.2.4}/src/primitive/__about__.py

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

{primitive-0.2.3 → primitive-0.2.4}/src/primitive/agent/runner.py

@@ -1,6 +1,6 @@
 import os
-import typing
 import re
+import typing
 import shutil
 from typing import Dict, TypedDict, List
 from abc import abstractmethod
@@ -9,8 +9,8 @@ from pathlib import Path, PurePath
 from loguru import logger
 import yaml
 import asyncio
-from primitive.utils.shell import env_string_to_dict
 from ..utils.cache import get_artifacts_cache, get_sources_cache, get_logs_cache
+from ..utils.shell import env_to_dict
 
 try:
     from yaml import CLoader as Loader
@@ -20,6 +20,7 @@ except ImportError:
 if typing.TYPE_CHECKING:
     import primitive.client
 
+CHUNK_SIZE = 64 * 1024
 ENV_VAR_LOOKUP_START = "_ENV_VAR_LOOKUP_START"
 ENV_VAR_LOOKUP_END = "_ENV_VAR_LOOKUP_END"
 
@@ -49,6 +50,16 @@ 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,
@@ -158,10 +169,18 @@ class Runner:
 
     async def run_task(self, task: Task) -> bool:
         for cmd in task["cmd"].strip().split("\n"):
+            # 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
             args = [
                 "/bin/bash",
                 "-c",
-                f"{cmd} && echo '{ENV_VAR_LOOKUP_START}' && env && echo '{ENV_VAR_LOOKUP_END}'",
+                (
+                    f"{cmd} "
+                    f"&& stdbuf -oL echo && stdbuf -oL echo '{ENV_VAR_LOOKUP_START}' "
+                    "&& env "
+                    f"&& stdbuf -oL echo && echo '{ENV_VAR_LOOKUP_END}'"
+                ),
             ]
 
             process = await asyncio.create_subprocess_exec(
@@ -205,52 +224,80 @@ class Runner:
 
     async def log_cmd(self, process, stream, tags: Dict = {}) -> bool:
         failure_detected = False
-        while line := await stream.readline():
-            raw_data = line.decode()
-
-            # handle env vars
-            if ENV_VAR_LOOKUP_START in raw_data:
-                env_vars_string = ""
-                while env_line := await stream.readline():
-                    if ENV_VAR_LOOKUP_END in env_line.decode():
-                        break
-                    env_vars_string += env_line.decode()
+        parse_environment = False
+        env_lines = []
+        while chunk := await stream.read(CHUNK_SIZE):
+            processed_lines = await self.read_chunk(chunk)
+            single_line = len(processed_lines) == 1
+            empty_chunk = len(processed_lines) == 0
+
+            if empty_chunk:
+                # Ignore the empty lines
+                continue
+            elif single_line and ENV_VAR_LOOKUP_START in processed_lines[0]:
+                # Next chunk will contain the environent variables
+                parse_environment = True
+                continue
+            elif single_line and ENV_VAR_LOOKUP_END in processed_lines[0]:
+                # Done reading environment variables
+                parse_environment = False
+                self.modified_env = env_to_dict(env_lines)
+                continue
 
-                self.modified_env = env_string_to_dict(env_vars_string)
+            if parse_environment:
+                env_lines.extend(processed_lines)
                 continue
 
             # Handle logging
             parse_logs = self.job_settings["parseLogs"]
             parse_stderr = self.job_settings["parseStderr"]
 
-            level = LogLevel.INFO
-            tag = None
-            if (parse_logs and "error" in raw_data.lower()) or (
-                parse_stderr and stream is process.stderr
-            ):
-                level = LogLevel.ERROR
-            elif parse_logs and "warning" in raw_data.lower():
-                level = LogLevel.WARNING
-
-            failure_detected = (
-                level == LogLevel.ERROR
-                and self.job_settings["failureLevel"] >= FailureLevel.ERROR
-            ) or (
-                level == LogLevel.WARNING
-                and self.job_settings["failureLevel"] >= FailureLevel.WARNING
-            )
+            for line in processed_lines:
+                level = LogLevel.INFO
+                tag = None
+                if (parse_logs and "error" in line.lower()) or (
+                    parse_stderr and stream is process.stderr
+                ):
+                    level = LogLevel.ERROR
+                elif parse_logs and "warning" in line.lower():
+                    level = LogLevel.WARNING
+
+                failure_detected = (
+                    level == LogLevel.ERROR
+                    and self.job_settings["failureLevel"] >= FailureLevel.ERROR
+                ) or (
+                    level == LogLevel.WARNING
+                    and self.job_settings["failureLevel"] >= FailureLevel.WARNING
+                )
 
-            # Tag on the first matching regex in the list
-            for tag_key, regex in tags.items():
-                pattern = re.compile(regex)
-                if pattern.match(raw_data):
-                    tag = tag_key
-                    break
+                # Tag on the first matching regex in the list
+                for tag_key, regex in tags.items():
+                    pattern = re.compile(regex)
+                    if pattern.match(line):
+                        tag = tag_key
+                        break
 
-            logger.bind(tag=tag).log(level.value, raw_data.rstrip())
+                logger.bind(tag=tag).log(level.value, line)
 
         return failure_detected
 
+    async def read_chunk(self, chunk: bytes) -> List[str]:
+        """Converts a chunk of bytes into proper lines."""
+        lines = []
+        current_line = ""
+        for char in chunk.decode():
+            if char in ("\n", "\r"):
+                lines.append(current_line)
+                current_line = ""
+            else:
+                current_line += char
+
+        # Handle extra characters in the buffer
+        if len(current_line) > 0:
+            lines.append(current_line)
+
+        return [line for line in lines if len(line) > 0]
+
     async def monitor_cmd(self, process) -> bool:
         while process.returncode is None:
             status = await self.primitive.jobs.aget_job_status(self.job_run["id"])
@@ -304,6 +351,7 @@ class Runner:
         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> | "

{primitive-0.2.3 → primitive-0.2.4}/src/primitive/utils/shell.py

@@ -1,7 +1,7 @@
 import subprocess
 from pathlib import Path
 from shutil import which
-from typing import Dict
+from typing import Dict, Union, List
 
 
 def add_path_to_shell(path: Path):
@@ -35,8 +35,14 @@ def add_path_to_shell(path: Path):
     return True
 
 
-def env_string_to_dict(env_str: str) -> Dict:
-    lines = env_str.splitlines()
+def env_to_dict(env_vars: Union[str, List[str]]) -> Dict:
+    lines = None
+    if isinstance(env_vars, list):
+        lines = env_vars
+    elif isinstance(env_vars, str):
+        lines = env_vars.splitlines()
+    else:
+        raise ValueError("Unsupported type. Env_vars must be a list or a string")
 
     current_key = None
     current_value = []