primitive 0.2.10__py3-none-any.whl → 0.2.11__py3-none-any.whl

primitive/daemons/launch_service.py

@@ -2,178 +2,238 @@ import os
 import configparser
 import subprocess
 from pathlib import Path
+from loguru import logger
+from ..utils.daemons import Daemon
 
 HOME_DIRECTORY = Path.home()
-
 PRIMITIVE_BINARY_PATH = Path(HOME_DIRECTORY / ".pyenv" / "shims" / "primitive")
 
-PRIMITIVE_AGENT_SERVICE = "tech.primitive.agent.service"
-PRIMITIVE_AGENT_LOGS = "tech.primitive.agent.log"
-PRIMITIVE_AGENT_SERVICE_FILEPATH = Path(
-    HOME_DIRECTORY / ".config" / "systemd" / "user" / PRIMITIVE_AGENT_SERVICE
-)
-PRIMITIVE_AGENT_LOGS_FILEPATH = Path(
-    HOME_DIRECTORY / ".cache" / "primitive" / PRIMITIVE_AGENT_LOGS
-)
-
-
-def stop_service():
-    try:
-        if is_service_active():
-            stop_existing_service = f"systemctl --user stop {PRIMITIVE_AGENT_SERVICE}"
-            subprocess.check_output(stop_existing_service.split(" "))
-        return True
-    except subprocess.CalledProcessError as exception:
-        print("stop_service: ", exception)
-        return False
-
-
-def is_service_active():
-    try:
-        is_service_active = (
-            f"systemctl --user show {PRIMITIVE_AGENT_SERVICE} -p ActiveState --value"
+
+class LaunchService(Daemon):
+    def __init__(self, label: str):
+        self.label = label
+        self.name = label.split(".")[-1]
+
+    @property
+    def service_name(self) -> str:
+        return f"{self.label}.service"
+
+    @property
+    def file_path(self) -> Path:
+        return Path(HOME_DIRECTORY / ".config" / "systemd" / "user" / self.service_name)
+
+    @property
+    def logs(self) -> Path:
+        return Path(HOME_DIRECTORY / ".cache" / "primitive" / f"{self.label}.log")
+
+    def stop(self) -> bool:
+        try:
+            if self.is_active():
+                stop_existing_service = f"systemctl --user stop {self.service_name}"
+                subprocess.check_output(
+                    stop_existing_service.split(" "), stderr=subprocess.DEVNULL
+                )
+                logger.info(f":white_check_mark: {self.label} stopped successfully!")
+            return True
+        except subprocess.CalledProcessError as exception:
+            if exception.returncode == 4:
+                logger.debug(f"{self.label} is not running or does not exist.")
+                return True
+            else:
+                logger.error(f"Unable to stop {self.label}, {exception.returncode}")
+                logger.error(exception)
+                return False
+
+    def start(self) -> bool:
+        try:
+            start_new_service = f"systemctl --user start {self.service_name}"
+            subprocess.check_output(start_new_service.split(" "))
+            logger.info(f":white_check_mark: {self.label} started successfully!")
+            return True
+        except subprocess.CalledProcessError as exception:
+            logger.error(f"Unable to start {self.label}")
+            logger.error(exception)
+            return False
+
+    def disable(self) -> bool:
+        try:
+            if self.is_installed():
+                disable_existing_service = (
+                    f"systemctl --user disable {self.service_name}"
+                )
+                subprocess.check_output(
+                    disable_existing_service.split(" "), stderr=subprocess.DEVNULL
+                )
+            return True
+        except subprocess.CalledProcessError as exception:
+            logger.error(f"Unable to disable {self.label}")
+            logger.error(exception)
+            return False
+
+    def enable(self) -> bool:
+        try:
+            enable_service = f"systemctl --user enable {self.service_name}"
+            subprocess.check_output(
+                enable_service.split(" "), stderr=subprocess.DEVNULL
+            )
+            return True
+        except subprocess.CalledProcessError as exception:
+            logger.error(f"Unable to enable {self.label}")
+            logger.error(exception)
+            return False
+
+    def verify(self) -> bool:
+        systemctl_check = (
+            f"systemctl --user show {self.service_name} -p CanStart --value"
         )
-        output = subprocess.check_output(is_service_active.split(" ")).decode().strip()
-        return output == "active"
-    except subprocess.CalledProcessError as exception:
-        print("is_service_active: ", exception)
-        return False
-
-
-def disable_service():
-    try:
-        if is_service_enabled():
-            disable_existing_service = (
-                f"systemctl --user disable {PRIMITIVE_AGENT_SERVICE}"
+        try:
+            output = (
+                subprocess.check_output(systemctl_check.split(" ")).decode().strip()
+            )
+            if output == "no":
+                raise Exception(f"{systemctl_check} yielded {output}")
+            return True
+        except subprocess.CalledProcessError as exception:
+            logger.error(f"Unable to verify {self.label}")
+            logger.error(exception)
+            return False
+
+    def view_logs(self) -> None:
+        follow_logs = f"tail -f -n +1 {self.logs}"
+        os.system(follow_logs)
+
+    def populate(self) -> bool:
+        self.file_path.parent.mkdir(parents=True, exist_ok=True)
+        self.file_path.touch()
+
+        if self.file_path.exists():
+            self.file_path.unlink()
+        self.file_path.parent.mkdir(parents=True, exist_ok=True)
+        self.file_path.touch()
+
+        config = configparser.ConfigParser()
+        config.optionxform = str  # type: ignore
+
+        config["Unit"] = {
+            "Description": "Primitive Agent",
+            "After": "network.target",
+        }
+
+        found_primitive_binary_path = PRIMITIVE_BINARY_PATH
+        if not PRIMITIVE_BINARY_PATH.exists():
+            result = subprocess.run(["which", "primitive"], capture_output=True)
+            if result.returncode == 0:
+                found_primitive_binary_path = result.stdout.decode().rstrip("\n")
+            else:
+                print("primitive binary not found")
+                return False
+
+        config["Service"] = {
+            "ExecStart": f'/bin/sh -lc "{found_primitive_binary_path} agent"',
+            "Restart": "always",
+            "StandardError": f"append:{self.logs}",
+            "StandardOutput": f"append:{self.logs}",
+        }
+
+        config["Install"] = {
+            "WantedBy": "default.target",
+        }
+
+        try:
+            with open(self.file_path, "w") as service_file:
+                config.write(service_file)
+        except IOError as exception:
+            print(f"populate_service_file: {exception}")
+
+        self.file_path.chmod(0o644)
+        return self.verify()
+
+    def create_stdout_file(self) -> bool:
+        try:
+            if not self.logs.exists():
+                self.logs.parent.mkdir(parents=True, exist_ok=True)
+                self.logs.touch()
+
+            return True
+        except Exception as e:
+            logger.error(
+                f"Unable to create log file at {self.logs} for daemon {self.label}"
+            )
+            logger.error(e)
+            return False
+
+    def delete_stdout_file(self) -> bool:
+        try:
+            if self.logs.exists():
+                self.logs.unlink()
+
+            return True
+        except Exception as e:
+            logger.error(
+                f"Unable to delete log file at {self.logs} for daemon {self.label}"
             )
-            subprocess.check_output(disable_existing_service.split(" "))
-        return True
-    except subprocess.CalledProcessError as exception:
-        print("disable_service: ", exception)
-        return False
+            logger.error(e)
+            return False
+
+    def delete_service_file(self) -> bool:
+        try:
+            if self.file_path.exists():
+                self.file_path.unlink()
+
+            return True
+        except Exception as e:
+            logger.error(
+                f"Unable to delete service file at {self.file_path} for daemon {self.label}"
+            )
+            logger.error(e)
+            return False
 
+    def install(self) -> bool:
+        return all(
+            [
+                self.stop(),
+                self.disable(),
+                self.populate(),
+                self.create_stdout_file(),
+                self.enable(),
+                self.start(),
+            ]
+        )
 
-def is_service_enabled():
-    try:
-        is_service_active = (
-            f"systemctl --user show {PRIMITIVE_AGENT_SERVICE} -p UnitFileState --value"  # noqa
+    def uninstall(self) -> bool:
+        return all(
+            [
+                self.stop(),
+                self.disable(),
+                self.delete_service_file(),
+                self.delete_stdout_file(),
+            ]
         )
-        output = subprocess.check_output(is_service_active.split(" ")).decode().strip()
-        return output == "enabled"
-    except subprocess.CalledProcessError as exception:
-        print("is_service_enabled: ", exception)
-        return False
-
-
-def populate_service_file():
-    PRIMITIVE_AGENT_LOGS_FILEPATH.parent.mkdir(parents=True, exist_ok=True)
-    PRIMITIVE_AGENT_LOGS_FILEPATH.touch()
-
-    if PRIMITIVE_AGENT_SERVICE_FILEPATH.exists():
-        PRIMITIVE_AGENT_SERVICE_FILEPATH.unlink()
-    PRIMITIVE_AGENT_SERVICE_FILEPATH.parent.mkdir(parents=True, exist_ok=True)
-    PRIMITIVE_AGENT_SERVICE_FILEPATH.touch()
-
-    config = configparser.ConfigParser()
-    config.optionxform = str  # type: ignore
-
-    config["Unit"] = {
-        "Description": "Primitive Agent",
-        "After": "network.target",
-    }
-
-    found_primitive_binary_path = PRIMITIVE_BINARY_PATH
-    if not PRIMITIVE_BINARY_PATH.exists():
-        result = subprocess.run(["which", "primitive"], capture_output=True)
-        if result.returncode == 0:
-            found_primitive_binary_path = result.stdout.decode().rstrip("\n")
-        else:
-            print("primitive binary not found")
+
+    def is_active(self) -> bool:
+        try:
+            is_service_active = (
+                f"systemctl --user show {self.service_name} -p ActiveState --value"
+            )
+            output = (
+                subprocess.check_output(is_service_active.split(" ")).decode().strip()
+            )
+            return output == "active"
+        except subprocess.CalledProcessError as exception:
+            logger.error(f"Unable to check if {self.label} is active")
+            logger.error(exception)
             return False
 
-    config["Service"] = {
-        "ExecStart": f'/bin/sh -lc "{found_primitive_binary_path} agent"',
-        "Restart": "always",
-        "StandardError": f"append:{PRIMITIVE_AGENT_LOGS_FILEPATH}",
-        "StandardOutput": f"append:{PRIMITIVE_AGENT_LOGS_FILEPATH}",
-    }
-
-    config["Install"] = {
-        "WantedBy": "default.target",
-    }
-
-    try:
-        with open(PRIMITIVE_AGENT_SERVICE_FILEPATH, "w") as service_file:
-            config.write(service_file)
-    except IOError as exception:
-        print(f"populate_service_file: {exception}")
-
-    PRIMITIVE_AGENT_SERVICE_FILEPATH.chmod(0o644)
-    verify_service_file()
-
-
-def verify_service_file():
-    systemctl_check = (
-        f"systemctl --user show {PRIMITIVE_AGENT_SERVICE} -p CanStart --value"
-    )
-    try:
-        output = subprocess.check_output(systemctl_check.split(" ")).decode().strip()
-        if output == "no":
-            raise Exception(f"{systemctl_check} yielded {output}")
-        return True
-    except subprocess.CalledProcessError as exception:
-        print("verify_service_file: ", exception)
-        return False
-
-
-def create_stdout_file():
-    if not PRIMITIVE_AGENT_LOGS_FILEPATH.exists():
-        PRIMITIVE_AGENT_LOGS_FILEPATH.parent.mkdir(parents=True, exist_ok=True)
-        PRIMITIVE_AGENT_LOGS_FILEPATH.touch()
-
-
-def delete_stdout_file():
-    if PRIMITIVE_AGENT_LOGS_FILEPATH.exists():
-        PRIMITIVE_AGENT_LOGS_FILEPATH.unlink()
-
-
-def enable_service():
-    try:
-        enable_service = f"systemctl --user enable {PRIMITIVE_AGENT_SERVICE}"
-        subprocess.check_output(enable_service.split(" "))
-        return True
-    except subprocess.CalledProcessError as exception:
-        print("enable_service: ", exception)
-        return False
-
-
-def start_service():
-    try:
-        start_new_service = f"systemctl --user start {PRIMITIVE_AGENT_SERVICE}"
-        subprocess.check_output(start_new_service.split(" "))
-        return True
-    except subprocess.CalledProcessError as exception:
-        print("start_service: ", exception)
-        return False
-
-
-def view_service_logs():
-    follow_logs = f"tail -f -n +1 {PRIMITIVE_AGENT_LOGS_FILEPATH}"
-    os.system(follow_logs)
-
-
-def full_service_install():
-    stop_service()
-    disable_service()
-    populate_service_file()
-    create_stdout_file()
-    enable_service()
-    start_service()
-
-
-def full_service_uninstall():
-    stop_service()
-    disable_service()
-    if PRIMITIVE_AGENT_SERVICE_FILEPATH.exists():
-        PRIMITIVE_AGENT_SERVICE_FILEPATH.unlink()
-    delete_stdout_file()
+    def is_installed(self) -> bool:
+        try:
+            is_service_active = (
+                f"systemctl --user show {self.service_name} -p UnitFileState --value"  # noqa
+            )
+            output = (
+                subprocess.check_output(is_service_active.split(" ")).decode().strip()
+            )
+            return output == "enabled"
+        except subprocess.CalledProcessError as exception:
+            logger.error(f"Unable to check if {self.label} is enabled")
+            logger.error(exception)
+            return False

primitive/daemons/ui.py

@@ -0,0 +1,41 @@
+from rich.console import Console
+from rich.table import Table
+from typing import List
+
+from ..utils.daemons import Daemon
+
+
+def render_daemon_list(daemons: List[Daemon]) -> None:
+    console = Console()
+
+    table = Table(show_header=True, header_style="bold #FFA800")
+    table.add_column("Name")
+    table.add_column("Label")
+    table.add_column("Installed")
+    table.add_column("Active")
+    table.add_column("File Path")
+    table.add_column("Log Path")
+
+    for daemon in daemons:
+        child_table = Table(show_header=False, header_style="bold #FFA800")
+        child_table.add_column("Name")
+        child_table.add_column("Label")
+        child_table.add_column("Installed")
+        child_table.add_column("Active")
+        child_table.add_column("File Path")
+        child_table.add_column("Log Path")
+
+        table.add_row(
+            daemon.name,
+            daemon.label,
+            "[bold green]Yes[/bold green]"
+            if daemon.is_installed()
+            else "[bold red]No[/bold red]",
+            "[bold green]Yes[/bold green]"
+            if daemon.is_active()
+            else "[bold red]No[/bold red]",
+            str(daemon.file_path),
+            str(daemon.logs),
+        )
+
+    console.print(table)

primitive/db/base.py

@@ -0,0 +1,5 @@
+from sqlalchemy.orm import DeclarativeBase
+
+
+class Base(DeclarativeBase):
+    pass

primitive/db/models.py

@@ -0,0 +1,78 @@
+from sqlalchemy import Column, Integer, String
+from sqlalchemy.orm import Mapped, mapped_column, Query
+from typing import Optional, Generic, TypeVar, Callable, Type, List, Union, Dict, Any
+from .sqlite import Session
+from .base import Base
+
+T = TypeVar("T", bound="Base")
+
+
+class Manager(Generic[T]):
+    def __init__(self, model_cls_lambda: Callable[[], Type[T]]) -> None:
+        self.model_cls_lambda = model_cls_lambda
+        self.filters: Dict[str, Any] = {}
+
+    def create(self, **kwargs) -> T:
+        with Session() as session:
+            model = self.model_cls_lambda()
+            obj = model(**kwargs)
+            session.add(obj)
+            session.commit()
+            session.refresh(obj)
+            return obj
+
+    def filter_by(self, **kwargs) -> "Manager[T]":
+        self.filters = kwargs
+        return self
+
+    def all(self) -> List[T]:
+        with Session() as session:
+            model = self.model_cls_lambda()
+            query = session.query(model)
+            query.filter_by(**self.filters)
+            self.filters.clear()
+            return query.all()
+
+    def first(self) -> Union[T, None]:
+        with Session() as session:
+            model = self.model_cls_lambda()
+            query = session.query(model)
+            query.filter_by(**self.filters)
+            self.filters.clear()
+            return query.first()
+
+    def update(self, update: Dict[Any, Any]) -> Query[T]:
+        with Session() as session:
+            model = self.model_cls_lambda()
+            query = session.query(model).filter_by(**self.filters)
+
+            if query.count() > 0:
+                query.update(update)
+                session.commit()
+                return query
+            else:
+                raise ValueError(f"{model.__name__} not found")
+
+    def delete(self) -> None:
+        with Session() as session:
+            model = self.model_cls_lambda()
+            query = session.query(model).filter_by(**self.filters)
+
+            if query.count() > 0:
+                query.delete()
+                session.commit()
+            else:
+                raise ValueError(f"{model.__name__} not found")
+
+
+class JobRun(Base):
+    __tablename__ = "JobRun"
+
+    id = Column(Integer, primary_key=True)
+    job_run_id: Mapped[str] = mapped_column(String, nullable=False)
+    pid: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
+
+    objects: Manager["JobRun"] = Manager(lambda: JobRun)
+
+    def __repr__(self):
+        return f"<JobRun(id={self.id} job_run_id={self.job_run_id}, pid={self.pid})>"

primitive/db/sqlite.py

@@ -0,0 +1,27 @@
+from pathlib import Path
+from sqlalchemy import create_engine, Engine
+from sqlalchemy.orm import Session as SQLAlchemySession
+from ..utils.cache import get_cache_dir
+from .base import Base
+
+
+def init() -> None:
+    db_path: Path = get_cache_dir() / "primitive.sqlite3"
+
+    # Drop DB existing database if it exists
+    if db_path.exists():
+        db_path.unlink()
+
+    engine = create_engine(f"sqlite:///{db_path}", echo=False)
+    Base.metadata.create_all(engine)
+
+
+def engine() -> Engine:
+    db_path: Path = get_cache_dir() / "primitive.sqlite3"
+    return create_engine(f"sqlite:///{db_path}", echo=False)
+
+
+def Session() -> SQLAlchemySession:
+    from sqlalchemy.orm import sessionmaker
+
+    return sessionmaker(bind=engine())()

primitive/exec/actions.py

@@ -58,7 +58,6 @@ class Exec(BaseAction):
 
         reservation = reservation_result.data["reservation"]
         if reservation.get("status") != "in_progress":
-            logger.enable("primitive")
             logger.info(
                 f"Reservation {reservation.get('id')} is in status {reservation.get('status')}, cannot execute command at this time."
             )

primitive/files/actions.py

@@ -200,7 +200,6 @@ class Files(BaseAction):
         key_prefix: str = "",
         file_id: Optional[str] = None,
     ):
-        logger.enable("primitive")
         if path.exists() is False:
             raise Exception(f"File {path} does not exist.")
 

primitive/hardware/actions.py

@@ -293,7 +293,6 @@ class Hardware(BaseAction):
         )
         if messages := result.data.get("registerHardware").get("messages"):
             for message in messages:
-                logger.enable("primitive")
                 if message.get("kind") == "ERROR":
                     logger.error(message.get("message"))
                 else:
@@ -327,7 +326,6 @@ class Hardware(BaseAction):
 
         if messages := result.data.get("unregisterHardware").get("messages"):
             for message in messages:
-                logger.enable("primitive")
                 if message.get("kind") == "ERROR":
                     logger.error(message.get("message"))
                 else:
@@ -368,11 +366,11 @@ class Hardware(BaseAction):
                 mutation, variable_values=variables, get_execution_result=True
             )
         except client_exceptions.ClientConnectorError as exception:
-            message = " [*] Failed to update hardware system info! "
+            message = "[*] Failed to update hardware system info! "
             logger.exception(message)
             raise exception
 
-        message = " [*] Updated hardware system info successfully! "
+        message = "[*] Updated hardware system info successfully! "
         logger.info(message)
 
         return result
@@ -416,7 +414,6 @@ class Hardware(BaseAction):
             checkin_success = result.data.get("checkIn").get("lastCheckIn")
             if messages := result.data.get("checkIn").get("messages"):
                 for message in messages:
-                    logger.enable("primitive")
                     if message.get("kind") == "ERROR":
                         logger.error(message.get("message"))
                     else:
@@ -435,7 +432,7 @@ class Hardware(BaseAction):
                     }
                 self.status_cache[fingerprint] = new_state.copy()
 
-                message = f" [*] Checked in successfully for {fingerprint}: "
+                message = f"[*] Checked in successfully for {fingerprint}: "
                 is_new_status = False
                 for key, value in new_state.items():
                     if value != previous_status.get(key, None):
@@ -463,7 +460,7 @@ class Hardware(BaseAction):
             return result
         except client_exceptions.ClientConnectorError as exception:
             if not stopping_agent:
-                message = " [*] Failed to check in! "
+                message = "[*] Failed to check in! "
                 logger.error(message)
                 raise exception
             else:
@@ -578,7 +575,6 @@ class Hardware(BaseAction):
 
         if messages := result.data.get("registerChildHardware").get("messages"):
             for message in messages:
-                logger.enable("primitive")
                 if message.get("kind") == "ERROR":
                     logger.error(message.get("message"))
                 else: