OpenOrchestrator 1.3.1__py3-none-any.whl → 2.0.0rc2__py3-none-any.whl

OpenOrchestrator/__main__.py

@@ -1,20 +1,74 @@
-"""This module is used to run Orchestrator or Scheduler from the command line.
-Usage: python -m OpenOrchestrator [-o|-s]"""
+"""This module is used to run Orchestrator or Scheduler from the command line."""
+
+import argparse
+import subprocess
 
-import sys
 from OpenOrchestrator.scheduler.application import Application as s_app
 from OpenOrchestrator.orchestrator.application import Application as o_app
 
 
-def _print_usage():
-    print("Usage: -o to start Orchestrator. -s to start Scheduler")
+def main():
+    """The main entry point of the CLI."""
+    parser = argparse.ArgumentParser(
+        prog="OpenOrchestrator",
+        description="OpenOrchestrator is used to orchestrate, monitor and run Python-based automation scripts in Windows."
+    )
+
+    subparsers = parser.add_subparsers(title="Subcommands", required=True)
+
+    o_parser = subparsers.add_parser("orchestrator", aliases=["o"], help="Start the Orchestrator application.")
+    o_parser.add_argument("-p", "--port", type=int, help="Set the desired port for Orchestrator.")
+    o_parser.add_argument("-d", "--dont_show", action="store_false", help="Set if you don't want Orchestrator to open in the browser automatically.")
+    o_parser.set_defaults(func=orchestrator_command)
+
+    s_parser = subparsers.add_parser("scheduler", aliases=["s"], help="Start the Scheduler application.")
+    s_parser.set_defaults(func=scheduler_command)
+
+    u_parser = subparsers.add_parser("upgrade", aliases=["u"], help="Upgrade the database to the newest revision or create a new database from scratch.")
+    u_parser.add_argument("connection_string", type=str, help="The connection string to the database")
+    u_parser.set_defaults(func=upgrade_command)
+
+    args = parser.parse_args()
+    args.func(args)
+
+
+def orchestrator_command(args: argparse.Namespace):
+    """Start the Orchestrator app.
 
+    Args:
+        args: The arguments Namespace object.
+    """
+    o_app(port=args.port, show=args.dont_show)
 
-if len(sys.argv) != 2:
-    _print_usage()
-elif sys.argv[1] == '-s':
+
+def scheduler_command(args: argparse.Namespace):  # pylint: disable=unused-argument
+    """Start the Scheduler app."""
     s_app()
-elif sys.argv[1] == '-o':
-    o_app()
-else:
-    _print_usage()
+
+
+def upgrade_command(args: argparse.Namespace):
+    """Upgrade the database to the newest revision using alembic.
+
+    Args:
+        args: The arguments Namespace object.
+    """
+    confirmation = input("Are you sure you want to upgrade the database to the newest revision? This cannot be undone. (y/n)").strip()
+    if confirmation == "y":
+        # Check if database is already stamped
+        result = subprocess.run(["alembic", "-x", args.connection_string, "current"], check=True, capture_output=True)
+        if result.stdout == b'':
+            response = input("It seems like you are either creating a new database or trying to upgrade a legacy database. Please choose an option: (new/legacy/cancel)").strip()
+            if response == "legacy":
+                subprocess.run(["alembic", "-x", args.connection_string, "stamp", "b67b7649b282"], check=True)
+            elif response != "new":
+                print("Upgrade canceled")
+                return
+
+        subprocess.run(["alembic", "-x", args.connection_string, "upgrade", "head"], check=True)
+        print("Database upgraded to the newest revision!")
+    else:
+        print("Upgrade canceled")
+
+
+if __name__ == '__main__':
+    main()

OpenOrchestrator/common/connection_frame.py

@@ -6,6 +6,8 @@ from nicegui import ui
 
 from OpenOrchestrator.common import crypto_util
 from OpenOrchestrator.database import db_util
+from OpenOrchestrator.orchestrator import test_helper
+from OpenOrchestrator.orchestrator.popups import generic_popups
 
 
 # pylint: disable-next=too-few-public-methods
@@ -21,12 +23,13 @@ class ConnectionFrame():
             self.disconn_button.disable()
 
         self._initial_connect()
+        test_helper.set_automation_ids(self, "connection_frame")
 
     def _define_validation(self):
-        self.conn_input.validation = {"Please enter a connection string": bool}
-        self.key_input.validation = {"Invalid AES key": crypto_util.validate_key}
+        self.conn_input._validation = {"Please enter a connection string": bool}  # pylint: disable=protected-access
+        self.key_input._validation = {"Invalid AES key": crypto_util.validate_key}  # pylint: disable=protected-access
 
-    def _connect(self) -> None:
+    async def _connect(self) -> None:
         """Validate the connection string and encryption key
         and connect to the database.
         """
@@ -40,7 +43,10 @@ class ConnectionFrame():
         if db_util.connect(conn_string):
             crypto_util.set_key(crypto_key)
             self._set_state(True)
-            ui.notify("Connected!", type='positive')
+            ui.notify("Connected!", type='positive', timeout=1000)
+
+            if not os.getenv("ORCHESTRATOR_TEST") and not db_util.check_database_revision():
+                await generic_popups.info_popup("This version of Orchestrator doesn't match the version of the connected database. Unexpected errors might occur.")
 
     def _disconnect(self) -> None:
         db_util.disconnect()

OpenOrchestrator/database/base.py

@@ -0,0 +1,8 @@
+"""Contains the Base class for all ORM classes"""
+
+from sqlalchemy.orm import DeclarativeBase
+
+
+# pylint: disable=too-few-public-methods
+class Base(DeclarativeBase):
+    """SqlAlchemy base class for all ORM classes in this project."""

OpenOrchestrator/database/constants.py

@@ -2,19 +2,16 @@
 
 from datetime import datetime
 
-from sqlalchemy import String, Engine
-from sqlalchemy.orm import Mapped, DeclarativeBase, mapped_column
+from sqlalchemy import String
+from sqlalchemy.orm import Mapped, mapped_column
 
 from OpenOrchestrator.common import datetime_util
+from OpenOrchestrator.database.base import Base
 
 # All classes in this module are effectively dataclasses without methods.
 # pylint: disable=too-few-public-methods
 
 
-class Base(DeclarativeBase):
-    """SqlAlchemy base class for all ORM classes in this module."""
-
-
 class Constant(Base):
     """A class representing a constant object in the ORM."""
     __tablename__ = "Constants"
@@ -56,12 +53,3 @@ class Credential(Base):
         lower_length = int(((length-100)/20)*16)
         upper_length = lower_length + 15
         return f"{length} encrypted bytes. {lower_length}-{upper_length} decrypted bytes."
-
-
-def create_tables(engine: Engine):
-    """Create all SQL tables related to ORM classes in this module.
-
-    Args:
-        engine: The SqlAlchemy connection engine used to create the tables.
-    """
-    Base.metadata.create_all(engine)

OpenOrchestrator/database/db_util.py

@@ -3,19 +3,21 @@
 from datetime import datetime
 from typing import TypeVar, ParamSpec
 from uuid import UUID
+import json
 
 from cronsim import CronSim
-from sqlalchemy import Engine, create_engine, select, insert, desc
+from sqlalchemy import Engine, create_engine, select, insert, desc, text
+
 from sqlalchemy import exc as alc_exc
 from sqlalchemy import func as alc_func
 from sqlalchemy.orm import Session, selectin_polymorphic
 
 from OpenOrchestrator.common import crypto_util
-from OpenOrchestrator.database import logs, triggers, constants, queues
 from OpenOrchestrator.database.logs import Log, LogLevel
 from OpenOrchestrator.database.constants import Constant, Credential
 from OpenOrchestrator.database.triggers import Trigger, SingleTrigger, ScheduledTrigger, QueueTrigger, TriggerStatus
 from OpenOrchestrator.database.queues import QueueElement, QueueStatus
+from OpenOrchestrator.database.schedulers import Scheduler
 from OpenOrchestrator.database.truncated_string import truncate_message
 
 # Type hint helpers for decorators
@@ -55,6 +57,19 @@ def disconnect() -> None:
     _connection_engine = None
 
 
+def check_database_revision() -> bool:
+    """Check if the revision number of the connected database matches the expected revision."""
+    try:
+        with _get_session() as session:
+            version = session.execute(text(
+                """SELECT version_num FROM alembic_version"""
+            )).scalar()
+    except alc_exc.ProgrammingError:
+        return False
+
+    return version == "90d46abd44a3"
+
+
 def _get_session() -> Session:
     """Check if theres a database connection and return a
     session to it.
@@ -83,17 +98,6 @@ def get_conn_string() -> str:
     return str(_connection_engine.url)
 
 
-def initialize_database() -> None:
-    """Initializes the database with all the needed tables."""
-    if not _connection_engine:
-        raise RuntimeError("Not connected to database.")
-
-    logs.create_tables(_connection_engine)
-    triggers.create_tables(_connection_engine)
-    constants.create_tables(_connection_engine)
-    queues.create_tables(_connection_engine)
-
-
 def get_trigger(trigger_id: UUID | str) -> Trigger:
     """Get the trigger with the given id.
 
@@ -276,8 +280,10 @@ def get_unique_log_process_names() -> tuple[str, ...]:
         return tuple(result)
 
 
+# pylint: disable=too-many-positional-arguments
 def create_single_trigger(trigger_name: str, process_name: str, next_run: datetime,
-                          process_path: str, process_args: str, is_git_repo: bool, is_blocking: bool) -> None:
+                          process_path: str, process_args: str, is_git_repo: bool, is_blocking: bool,
+                          priority: int, scheduler_whitelist: list[str]) -> str:
     """Create a new single trigger in the database.
 
     Args:
@@ -288,6 +294,11 @@ def create_single_trigger(trigger_name: str, process_name: str, next_run: dateti
         process_args: The argument string of the process.
         is_git_repo: If the process_path points to a git repo.
         is_blocking: If the process should be blocking.
+        priority: The integer priority of the trigger.
+        cheduler_whitelist: A list of names of schedulers the trigger may run on.
+
+    Returns:
+        The id of the trigger that was created.
     """
     with _get_session() as session:
         trigger = SingleTrigger(
@@ -297,15 +308,19 @@ def create_single_trigger(trigger_name: str, process_name: str, next_run: dateti
             process_args = process_args,
             is_git_repo = is_git_repo,
             is_blocking = is_blocking,
-            next_run = next_run
+            next_run = next_run,
+            priority=priority,
+            scheduler_whitelist=json.dumps(scheduler_whitelist)
         )
         session.add(trigger)
         session.commit()
+        return trigger.id
 
 
+# pylint: disable=too-many-positional-arguments
 def create_scheduled_trigger(trigger_name: str, process_name: str, cron_expr: str, next_run: datetime,
                              process_path: str, process_args: str, is_git_repo: bool,
-                             is_blocking: bool) -> None:
+                             is_blocking: bool, priority: int, scheduler_whitelist: list[str]) -> str:
     """Create a new scheduled trigger in the database.
 
     Args:
@@ -317,6 +332,10 @@ def create_scheduled_trigger(trigger_name: str, process_name: str, cron_expr: st
         process_args: The argument string of the process.
         is_git_repo: If the process_path points to a git repo.
         is_blocking: If the process should be blocking.
+        priority: The integer priority of the trigger.
+        scheduler_whitelist: A list of names of schedulers the trigger may run on.
+    Returns:
+        The id of the trigger that was created.
     """
     with _get_session() as session:
         trigger = ScheduledTrigger(
@@ -327,15 +346,19 @@ def create_scheduled_trigger(trigger_name: str, process_name: str, cron_expr: st
             is_git_repo = is_git_repo,
             is_blocking = is_blocking,
             next_run = next_run,
-            cron_expr = cron_expr
+            cron_expr = cron_expr,
+            priority=priority,
+            scheduler_whitelist=json.dumps(scheduler_whitelist)
         )
         session.add(trigger)
         session.commit()
+        return trigger.id
 
 
+# pylint: disable=too-many-positional-arguments
 def create_queue_trigger(trigger_name: str, process_name: str, queue_name: str, process_path: str,
                          process_args: str, is_git_repo: bool, is_blocking: bool,
-                         min_batch_size: int) -> None:
+                         min_batch_size: int, priority: int, scheduler_whitelist: list[str]) -> str:
     """Create a new queue trigger in the database.
 
     Args:
@@ -347,6 +370,11 @@ def create_queue_trigger(trigger_name: str, process_name: str, queue_name: str,
         is_git_repo: The is_git value of the process.
         is_blocking: The is_blocking value of the process.
         min_batch_size: The minimum number of queue elements before triggering.
+        priority: The integer priority of the trigger.
+        scheduler_whitelist: A list of names of schedulers the trigger may run on.
+
+    Returns:
+        The id of the trigger that was created.
     """
     with _get_session() as session:
         trigger = QueueTrigger(
@@ -357,10 +385,13 @@ def create_queue_trigger(trigger_name: str, process_name: str, queue_name: str,
             is_git_repo = is_git_repo,
             is_blocking = is_blocking,
             queue_name = queue_name,
-            min_batch_size = min_batch_size
+            min_batch_size = min_batch_size,
+            priority=priority,
+            scheduler_whitelist=json.dumps(scheduler_whitelist)
         )
         session.add(trigger)
         session.commit()
+        return trigger.id
 
 
 def get_constant(name: str) -> Constant:
@@ -559,11 +590,11 @@ def begin_single_trigger(trigger_id: UUID | str) -> bool:
         return True
 
 
-def get_next_single_trigger() -> SingleTrigger | None:
-    """Get the single trigger that should trigger next.
+def get_pending_single_triggers() -> list[SingleTrigger]:
+    """Get all single triggers that are ready to run.
 
     Returns:
-        The next single trigger to run if any.
+        All single triggers ready to run if any.
     """
     with _get_session() as session:
         query = (
@@ -571,16 +602,15 @@ def get_next_single_trigger() -> SingleTrigger | None:
             .where(SingleTrigger.process_status == TriggerStatus.IDLE)
             .where(SingleTrigger.next_run <= datetime.now())
             .order_by(SingleTrigger.next_run)
-            .limit(1)
         )
-        return session.scalar(query)
+        return list(session.scalars(query))
 
 
-def get_next_scheduled_trigger() -> ScheduledTrigger | None:
-    """Get the scheduled trigger that should trigger next.
+def get_pending_scheduled_triggers() -> list[ScheduledTrigger]:
+    """Get all scheduled triggers that are ready to run.
 
     Returns:
-        The next scheduled trigger to run if any.
+        All scheduled triggers ready to run if any.
     """
     with _get_session() as session:
         query = (
@@ -588,9 +618,8 @@ def get_next_scheduled_trigger() -> ScheduledTrigger | None:
             .where(ScheduledTrigger.process_status == TriggerStatus.IDLE)
             .where(ScheduledTrigger.next_run <= datetime.now())
             .order_by(ScheduledTrigger.next_run)
-            .limit(1)
         )
-        return session.scalar(query)
+        return list(session.scalars(query))
 
 
 def begin_scheduled_trigger(trigger_id: UUID | str) -> bool:
@@ -625,14 +654,11 @@ def begin_scheduled_trigger(trigger_id: UUID | str) -> bool:
         return True
 
 
-def get_next_queue_trigger() -> QueueTrigger | None:
-    """Get the next queue trigger to run.
-    This functions loops through the queue triggers and checks
-    if the number of queue elements with status 'New' is above
-    the triggers min_batch_size.
+def get_pending_queue_triggers() -> list[QueueTrigger]:
+    """Get all queue triggers that are ready to run.
 
     Returns:
-        QueueTrigger | None: The next queue trigger to run if any.
+        All queue triggers that are ready to run if any.
     """
 
     with _get_session() as session:
@@ -648,9 +674,8 @@ def get_next_queue_trigger() -> QueueTrigger | None:
             select(QueueTrigger)
             .where(QueueTrigger.process_status == TriggerStatus.IDLE)
             .where(sub_query >= QueueTrigger.min_batch_size)
-            .limit(1)
         )
-        return session.scalar(query)
+        return list(session.scalars(query))
 
 
 def begin_queue_trigger(trigger_id: UUID | str) -> bool:
@@ -910,3 +935,51 @@ def delete_queue_element(element_id: UUID | str) -> None:
         q_element = session.get(QueueElement, element_id)
         session.delete(q_element)
         session.commit()
+
+
+def get_schedulers() -> tuple[Scheduler, ...]:
+    """Get Schedulers from the database"""
+    with _get_session() as session:
+        query = select(Scheduler).order_by(Scheduler.machine_name)
+        result = session.scalars(query).all()
+        return tuple(result)
+
+
+def send_ping_from_scheduler(machine_name: str) -> None:
+    """Send a ping from a running scheduler, updating the machine in the database.
+
+    Args:
+        machine_name: The machine pinging the Orchestrator
+    """
+    with _get_session() as session:
+        scheduler = session.get(Scheduler, machine_name)
+
+        if scheduler:
+            scheduler.last_update = datetime.now()
+        else:
+            scheduler = Scheduler(machine_name=machine_name, last_update=datetime.now())
+            session.add(scheduler)
+
+        session.commit()
+
+
+def start_trigger_from_machine(machine_name: str, trigger_name: str) -> None:
+    """Start a trigger from a machine running a scheduler, updating the name and time for triggers in the database.
+
+    Args:
+        machine_name: The machine starting the trigger
+        trigger_name: The trigger being started på the machine
+    """
+    with _get_session() as session:
+        scheduler = session.get(Scheduler, machine_name)
+        now = datetime.now()
+
+        if scheduler:
+            scheduler.last_update = now
+            scheduler.latest_trigger = trigger_name
+            scheduler.latest_trigger_time = now
+        else:
+            scheduler = Scheduler(machine_name=machine_name, last_update=now, latest_trigger=trigger_name, latest_trigger_time=now)
+            session.add(scheduler)
+
+        session.commit()

OpenOrchestrator/database/logs.py

@@ -4,10 +4,11 @@ from datetime import datetime
 import enum
 import uuid
 
-from sqlalchemy import String, Engine
-from sqlalchemy.orm import Mapped, DeclarativeBase, mapped_column
+from sqlalchemy import String
+from sqlalchemy.orm import Mapped, mapped_column
 
 from OpenOrchestrator.common import datetime_util
+from OpenOrchestrator.database.base import Base
 
 # All classes in this module are effectively dataclasses without methods.
 # pylint: disable=too-few-public-methods
@@ -20,10 +21,6 @@ class LogLevel(enum.Enum):
     ERROR = "Error"
 
 
-class Base(DeclarativeBase):
-    """SqlAlchemy base class for all ORM classes in this module."""
-
-
 class Log(Base):
     """A class representing log objects in the ORM."""
     __tablename__ = "Logs"
@@ -43,12 +40,3 @@ class Log(Base):
             "Message": self.log_message,
             "ID": str(self.id)
         }
-
-
-def create_tables(engine: Engine):
-    """Create all SQL tables related to ORM classes in this module.
-
-    Args:
-        engine: The SqlAlchemy connection engine used to create the tables.
-    """
-    Base.metadata.create_all(engine)

OpenOrchestrator/database/queues.py

@@ -5,10 +5,11 @@ import enum
 from typing import Optional
 import uuid
 
-from sqlalchemy import String, Engine
-from sqlalchemy.orm import Mapped, DeclarativeBase, mapped_column
+from sqlalchemy import String
+from sqlalchemy.orm import Mapped, mapped_column
 
 from OpenOrchestrator.common import datetime_util
+from OpenOrchestrator.database.base import Base
 
 # All classes in this module are effectively dataclasses without methods.
 # pylint: disable=too-few-public-methods
@@ -23,10 +24,6 @@ class QueueStatus(enum.Enum):
     ABANDONED = 'Abandoned'
 
 
-class Base(DeclarativeBase):
-    """SqlAlchemy base class for all ORM classes in this module."""
-
-
 class QueueElement(Base):
     """A class representing a queue element in the ORM."""
     __tablename__ = "Queues"
@@ -55,12 +52,3 @@ class QueueElement(Base):
             "Message": self.message,
             "Created By": self.created_by
         }
-
-
-def create_tables(engine: Engine):
-    """Create all SQL tables related to ORM classes in this module.
-
-    Args:
-        engine: The SqlAlchemy connection engine used to create the tables.
-    """
-    Base.metadata.create_all(engine)

OpenOrchestrator/database/schedulers.py

@@ -0,0 +1,32 @@
+"""Module for the Scheduler ORM class"""
+
+from datetime import datetime
+from typing import Optional
+
+from sqlalchemy import String
+from sqlalchemy.orm import Mapped, mapped_column
+
+from OpenOrchestrator.common import datetime_util
+from OpenOrchestrator.database.base import Base
+
+# All classes in this module are effectively dataclasses without methods.
+# pylint: disable=too-few-public-methods
+
+
+class Scheduler(Base):
+    """Class containing the ORM model for schedulers"""
+    __tablename__ = "Schedulers"
+
+    machine_name: Mapped[str] = mapped_column(String(100), primary_key=True)
+    last_update: Mapped[datetime] = mapped_column(onupdate=datetime.now, default=datetime.now)
+    latest_trigger: Mapped[Optional[str]] = mapped_column(String(100), nullable=True, default=None)
+    latest_trigger_time: Mapped[Optional[datetime]] = mapped_column(nullable=True, default=None)
+
+    def to_row_dict(self) -> dict[str, str]:
+        """Convert scheduler to a row dictionary for display in a table."""
+        return {
+            "Machine Name": self.machine_name,
+            "Last Connection": datetime_util.format_datetime(self.last_update),
+            "Latest Trigger": self.latest_trigger if self.latest_trigger else "None yet",
+            "Latest Trigger Time": datetime_util.format_datetime(self.latest_trigger_time) if self.latest_trigger_time else ""
+        }

OpenOrchestrator/database/triggers.py

@@ -5,10 +5,11 @@ import enum
 from typing import Optional
 import uuid
 
-from sqlalchemy import String, ForeignKey, Engine
-from sqlalchemy.orm import Mapped, DeclarativeBase, mapped_column
+from sqlalchemy import String, ForeignKey
+from sqlalchemy.orm import Mapped, mapped_column
 
 from OpenOrchestrator.common import datetime_util
+from OpenOrchestrator.database.base import Base
 
 # All classes in this module are effectively dataclasses without methods.
 # pylint: disable=too-few-public-methods
@@ -31,10 +32,6 @@ class TriggerType(enum.Enum):
     QUEUE = "Queue"
 
 
-class Base(DeclarativeBase):
-    """SqlAlchemy base class for all ORM classes in this module."""
-
-
 class Trigger(Base):
     """A base class for all triggers in the ORM."""
     __tablename__ = "Triggers"
@@ -48,6 +45,8 @@ class Trigger(Base):
     process_status: Mapped[TriggerStatus] = mapped_column(default=TriggerStatus.IDLE)
     is_git_repo: Mapped[bool]
     is_blocking: Mapped[bool]
+    scheduler_whitelist: Mapped[str] = mapped_column(String(250))
+    priority: Mapped[int] = mapped_column(default=0)
     type: Mapped[TriggerType]
 
     __mapper_args__ = {
@@ -66,7 +65,8 @@ class Trigger(Base):
             "Status": self.process_status.value,
             "Process Name": self.process_name,
             "Last Run": datetime_util.format_datetime(self.last_run, "Never"),
-            "ID": str(self.id)
+            "ID": str(self.id),
+            "Priority": str(self.priority)
         }
 
 
@@ -115,12 +115,3 @@ class QueueTrigger(Trigger):
         row_dict = super().to_row_dict()
         row_dict["Next Run"] = "N/A"
         return row_dict
-
-
-def create_tables(engine: Engine):
-    """Create all SQL tables related to ORM classes in this module.
-
-    Args:
-        engine: The SqlAlchemy connection engine used to create the tables.
-    """
-    Base.metadata.create_all(engine)