meerschaum 2.2.6__py3-none-any.whl → 2.3.0.dev1__py3-none-any.whl

meerschaum/actions/stop.py

@@ -50,11 +50,13 @@ def _complete_stop(
 
 def _stop_jobs(
     action: Optional[List[str]] = None,
+    executor_keys: Optional[str] = None,
     timeout_seconds: Optional[int] = None,
     noask: bool = False,
     force: bool = False,
     yes: bool = False,
     nopretty: bool = False,
+    debug: bool = False,
     **kw
 ) -> SuccessTuple:
     """
@@ -62,66 +64,93 @@ def _stop_jobs(
     
     To see running processes, run `show jobs`.
     """
+    from meerschaum.utils.jobs import (
+        get_filtered_jobs,
+        get_running_jobs,
+        get_paused_jobs,
+        get_stopped_jobs,
+        get_restart_jobs,
+    )
     from meerschaum.utils.formatting._jobs import pprint_jobs
     from meerschaum.utils.daemon import (
         get_filtered_daemons, get_running_daemons, get_stopped_daemons, get_paused_daemons,
     )
     from meerschaum.utils.warnings import warn
     from meerschaum.utils.prompt import yes_no
-    daemons = get_filtered_daemons(action, warn=(not nopretty))
-    _running_daemons = get_running_daemons(daemons)
-    _paused_daemons = get_paused_daemons(daemons)
-    _stopped_daemons = get_stopped_daemons(daemons)
-    if action and _stopped_daemons and not nopretty:
+
+    jobs = get_filtered_jobs(executor_keys, action, warn=(not nopretty))
+    running_jobs = get_running_jobs(executor_keys, jobs)
+    paused_jobs = get_paused_jobs(executor_keys, jobs)
+    restart_jobs = get_restart_jobs(executor_keys, jobs)
+    stopped_jobs = {
+        name: job
+        for name, job in get_stopped_jobs(executor_keys, jobs).items()
+        if name not in restart_jobs
+    }
+
+    jobs_to_stop = {
+        **running_jobs,
+        **paused_jobs,
+        **restart_jobs,
+    }
+
+    if action and stopped_jobs and not nopretty:
         warn(
-            f"Skipping stopped job" + ("s" if len(_stopped_daemons) > 1 else '') + " '" +
-                ("', '".join(d.daemon_id for d in _stopped_daemons)) + "'.",
-            stack = False
+            "Skipping stopped job"
+            + ("s" if len(stopped_jobs) != 1 else '')
+            + " '"
+            + ("', '".join(name for name in stopped_jobs)) + "'.",
+            stack=False,
         )
 
-    daemons_to_stop = _running_daemons + _paused_daemons
-    if not daemons_to_stop:
-        return False, "No running or paused jobs to stop."
+    if not jobs_to_stop:
+        return False, "No running, paused or restarting jobs to stop."
 
     if not action:
         if not force:
-            pprint_jobs(daemons_to_stop)
+            pprint_jobs(jobs_to_stop)
             if not yes_no(
                 "Stop the above jobs?",
                 noask=noask, yes=yes, default='n'
             ):
                 return False, "No jobs were stopped."
 
-    _quit_daemons, _kill_daemons = [], []
-    for d in daemons_to_stop:
-        quit_success, quit_msg = d.quit(timeout=timeout_seconds)
-        if quit_success:
-            _quit_daemons.append(d)
-            continue
-        else:
-            warn(
-                f"Failed to gracefully quit job '{d.daemon_id}', attempting to terminate:\n   "
-                + f"{quit_msg}",
-                stack = False,
-            )
-
-        kill_success, kill_msg = d.kill(timeout=timeout_seconds)
-        if kill_success:
-            _kill_daemons.append(d)
-            continue
-        if not nopretty:
-            warn(f"Failed to kill job '{d.daemon_id}' (PID {d.pid}):\n{kill_msg}", stack=False)
+    job_success_tuples = {}
+    for name, job in jobs_to_stop.items():
+        stop_success, stop_msg = job.stop(
+            timeout_seconds=timeout_seconds,
+            debug=debug,
+        )
+        job_success_tuples[name] = (stop_success, stop_msg)
 
+    num_success = sum(
+        (
+            1
+            for name, (stop_success, stop_msg) in job_success_tuples.items()
+            if stop_success
+        )
+    )
+    num_fail = sum(
+        (
+            1
+            for name, (stop_success, stop_msg) in job_success_tuples.items()
+            if not stop_success
+        )
+    )
+    success = num_success > 0
     msg = (
-        (("Stopped job" + ("s" if len(_quit_daemons) != 1 else '') +
-            " '" + "', '".join([d.daemon_id for d in _quit_daemons]) + "'.")
-            if _quit_daemons else '')
-        + (("\n" if _quit_daemons else "")
-           + ("Killed job" + ("s" if len(_kill_daemons) != 1 else '') +
-            " '" + "', '".join([d.daemon_id for d in _kill_daemons]) + "'.")
-            if _kill_daemons else '')
+        f"Stopped {num_success} job"
+        + ('s' if num_success != 1 else '')
+        + '.'
     )
-    return (len(_quit_daemons + _kill_daemons) > 0), msg
+    if num_fail > 0:
+        msg += (
+            f"\nFailed to stop {num_fail} job"
+            + ('s' if num_fail != 1 else '')
+            + '.'
+        )
+
+    return success, msg
 
 
 ### NOTE: This must be the final statement of the module.

meerschaum/actions/sync.py

@@ -14,9 +14,9 @@ import meerschaum as mrsm
 from meerschaum.utils.typing import SuccessTuple, Any, List, Optional, Tuple, Union
 
 def sync(
-        action: Optional[List[str]] = None,
-        **kw: Any
-    ) -> SuccessTuple:
+    action: Optional[List[str]] = None,
+    **kw: Any
+) -> SuccessTuple:
     """
     Fetch and sync data for pipes.
 

meerschaum/actions/upgrade.py

@@ -10,9 +10,9 @@ from __future__ import annotations
 from meerschaum.utils.typing import SuccessTuple, Any, List, Optional, Union
 
 def upgrade(
-        action: Optional[List[str]] = None,
-        **kw: Any
-    ) -> SuccessTuple:
+    action: Optional[List[str]] = None,
+    **kw: Any
+) -> SuccessTuple:
     """
     Upgrade Meerschaum, plugins, or packages.
     
@@ -34,13 +34,13 @@ def upgrade(
 
 
 def _upgrade_meerschaum(
-        action: Optional[List[str]] = None,
-        yes: bool = False,
-        force: bool = False,
-        noask: bool = False,
-        debug: bool = False,
-        **kw: Any
-    ) -> SuccessTuple:
+    action: Optional[List[str]] = None,
+    yes: bool = False,
+    force: bool = False,
+    noask: bool = False,
+    debug: bool = False,
+    **kw: Any
+) -> SuccessTuple:
     """
     Upgrade the current Meerschaum instance.
     Optionally specify dependency groups.
@@ -109,26 +109,21 @@ class NoVenv:
     pass
 
 def _upgrade_packages(
-        action: Optional[List[str]] = None,
-        venv: Union[str, None, NoVenv] = NoVenv,
-        yes: bool = False,
-        force: bool = False,
-        noask: bool = False,
-        debug: bool = False,
-        **kw: Any
-    ) -> SuccessTuple:
+    action: Optional[List[str]] = None,
+    venv: Union[str, None, NoVenv] = NoVenv,
+    yes: bool = False,
+    force: bool = False,
+    noask: bool = False,
+    debug: bool = False,
+    **kw: Any
+) -> SuccessTuple:
     """
     Upgrade and install dependencies.
     If provided, upgrade only a dependency group, otherwise default to `full`.
     
     Examples:
-        ```
         upgrade packages
-        ```
-    
-        ```
-        upgrade packages docs
-        ```
+        upgrade packages full
     """
     from meerschaum.utils.packages import packages, pip_install, get_prerelease_dependencies
     from meerschaum.utils.warnings import info, warn
@@ -178,14 +173,15 @@ def _upgrade_packages(
         )
     return success, msg
 
+
 def _upgrade_plugins(
-        action: Optional[List[str]] = None,
-        yes: bool = False,
-        force: bool = False,
-        noask: bool = False,
-        debug: bool = False,
-        **kw: Any
-    ) -> SuccessTuple:
+    action: Optional[List[str]] = None,
+    yes: bool = False,
+    force: bool = False,
+    noask: bool = False,
+    debug: bool = False,
+    **kw: Any
+) -> SuccessTuple:
     """
     Upgrade all installed plugins to the latest versions.
     If no plugins are specified, attempt to upgrade all,
@@ -193,12 +189,8 @@ def _upgrade_plugins(
     
     Examples:
     
-    ```
     upgrade plugins
-    ```
-    ```
-    upgrade plugins testing
-    ```
+    upgrade plugins noaa
     """
     from meerschaum.actions import actions
     from meerschaum.plugins import get_plugins_names

meerschaum/api/_events.py

@@ -19,14 +19,20 @@ from meerschaum.utils.debug import dprint
 from meerschaum.connectors.poll import retry_connect
 from meerschaum.utils.warnings import warn
 from meerschaum._internal.term.tools import is_webterm_running
+from meerschaum.utils.jobs import start_check_jobs_thread, stop_check_jobs_thread
+
+_check_jobs_thread = None
 
 @app.on_event("startup")
 async def startup():
-    conn = get_api_connector()
+    """
+    Connect to the instance database and begin monitoring jobs.
+    """
     try:
         if not no_dash:
             from meerschaum.api.dash.webterm import start_webterm
             start_webterm()
+
         connected = retry_connect(
             get_api_connector(),
             workers = get_uvicorn_config().get('workers', None),
@@ -36,18 +42,29 @@ async def startup():
         import traceback
         traceback.print_exc()
         connected = False
+
     if not connected:
         await shutdown()
         os._exit(1)
 
+    start_check_jobs_thread()
+
 
 @app.on_event("shutdown")
 async def shutdown():
+    """
+    Close the database connection and stop monitoring jobs.
+    """
     if debug:
         dprint("Closing connection...")
     if get_api_connector().type == 'sql':
         get_api_connector().engine.dispose()
 
+    stop_check_jobs_thread()
+    from meerschaum.api.routes._actions import _temp_jobs
+    for name, job in _temp_jobs.items():
+        job.delete()
+
     ### Terminate any running jobs left over.
     if 'meerschaum.api.dash' in sys.modules:
         from meerschaum.api.dash.actions import running_jobs, stop_action

meerschaum/api/_oauth2.py

@@ -30,6 +30,7 @@ class CustomOAuth2PasswordRequestForm:
         self.client_id = client_id
         self.client_secret = client_secret
 
+
 LoginManager = fastapi_login.LoginManager
 def generate_secret_key() -> str:
     """
@@ -46,5 +47,6 @@ def generate_secret_key() -> str:
 
     return secret_key.encode('utf-8')
 
+
 SECRET = generate_secret_key()
 manager = LoginManager(SECRET, token_url=endpoints['login'])

meerschaum/api/_websockets.py

@@ -23,8 +23,8 @@ sessions = {}
 @app.websocket('/dash/ws')
 @app.websocket(_websocket_endpoint)
 async def websocket_endpoint(
-        websocket: fastapi.WebSocket,
-    ):
+    websocket: fastapi.WebSocket,
+):
     """
     Communicate with the Web Interface over a websocket.
     """

meerschaum/api/dash/jobs.py

@@ -16,6 +16,7 @@ from meerschaum.api import CHECK_UPDATE
 dbc = attempt_import('dash_bootstrap_components', lazy=False, check_update=CHECK_UPDATE)
 html, dcc = import_html(), import_dcc()
 dateutil_parser = attempt_import('dateutil.parser', check_update=CHECK_UPDATE)
+from meerschaum.utils.jobs import get_jobs, Job
 from meerschaum.utils.daemon import (
     get_daemons,
     get_running_daemons,
@@ -38,13 +39,15 @@ def get_jobs_cards(state: WebState):
     Build cards and alerts lists for jobs.
     """
     daemons = get_daemons()
+    jobs = get_jobs()
     session_id = state['session-store.data'].get('session-id', None)
     is_authenticated = is_session_authenticated(session_id)
 
     alert = alert_from_success_tuple(daemons)
     cards = []
 
-    for d in daemons:
+    for name, job in jobs.items():
+        d = job.daemon
         footer_children = html.Div(
             build_process_timestamps_children(d),
             id = {'type': 'process-timestamps-div', 'index': d.daemon_id},
@@ -80,7 +83,7 @@ def get_jobs_cards(state: WebState):
         ]
 
         body_children = [
-            html.H4(html.B(d.daemon_id), className="card-title"),
+            html.H4(html.B(name), className="card-title"),
             html.Div(
                 html.P(
                     d.label,

meerschaum/api/routes/__init__.py

@@ -8,6 +8,7 @@ Import all routes from other modules in package
 
 from meerschaum.api.routes._login import *
 from meerschaum.api.routes._actions import *
+from meerschaum.api.routes._jobs import *
 from meerschaum.api.routes._connectors import *
 from meerschaum.api.routes._index import *
 from meerschaum.api.routes._misc import *

meerschaum/api/routes/_actions.py

@@ -7,65 +7,143 @@ Execute Meerschaum Actions via the API
 """
 
 from __future__ import annotations
-from meerschaum.utils.typing import SuccessTuple, Union
+import asyncio
+import traceback
+import shlex
+from functools import partial
+from datetime import datetime, timezone
 
+from fastapi import WebSocket, WebSocketDisconnect
+
+from meerschaum.utils.misc import generate_password
+from meerschaum.utils.jobs import Job
+from meerschaum.utils.warnings import warn
+from meerschaum.utils.typing import SuccessTuple, Union, List, Dict
 from meerschaum.api import (
     fastapi, app, endpoints, get_api_connector, debug, manager, private, no_auth
 )
 from meerschaum.actions import actions
 import meerschaum.core
+from meerschaum.config import get_config
+from meerschaum._internal.arguments._parse_arguments import parse_dict_to_sysargs, parse_arguments
+
 actions_endpoint = endpoints['actions']
 
+def is_user_allowed_to_execute(user) -> SuccessTuple:
+    if user is None:
+        return False, "Could not load user."
+
+    if user.type == 'admin':
+        return True, "Success"
+
+    allow_non_admin = get_config(
+        'system', 'api', 'permissions', 'actions', 'non_admin', patch=True
+    )
+    if not allow_non_admin:
+        return False, (
+            "The administrator for this server has not allowed users to perform actions.\n\n"
+            + "Please contact the system administrator, or if you are running this server, "
+            + "open the configuration file with `edit config system` "
+            + "and search for 'permissions'. "
+            + "\nUnder the keys 'api:permissions:actions', "
+            + "you can allow non-admin users to perform actions."
+        )
+
+    return True, "Success"
+
+
 @app.get(actions_endpoint, tags=['Actions'])
 def get_actions(
-        curr_user = (
-            fastapi.Depends(manager) if private else None
-        ),
-    ) -> list:
+    curr_user = (
+        fastapi.Depends(manager) if private else None
+    ),
+) -> List[str]:
+    """
+    Return a list of the available actions.
+    """
     return list(actions)
 
-@app.post(actions_endpoint + "/{action}", tags=['Actions'])
-def do_action(
-        action: str,
-        keywords: dict = fastapi.Body(...),
-        curr_user = (
-            fastapi.Depends(manager) if not no_auth else None
-        ),
-    ) -> SuccessTuple:
+
+async def notify_client(client, content: str):
+    """
+    Send a line of text to a client.
     """
-    Perform a Meerschaum action (if permissions allow it).
+    try:
+        await client.send_text(content)
+    except WebSocketDisconnect:
+        pass
 
-    Parameters
-    ----------
-    action: str :
-        The action to perform.
-        
-    keywords: dict :
-        The keywords dictionary to pass to the action.
+_temp_jobs = {}
+@app.websocket(actions_endpoint + '/ws')
+async def do_action_websocket(websocket: WebSocket):
+    """
+    Execute an action and stream the output to the client.
+    """
+    await websocket.accept()
 
-    Returns
-    -------
-    A `SuccessTuple` of success, message.
+    stop_event = asyncio.Event()
 
-    """
-    if curr_user is not None and curr_user.type != 'admin':
-        from meerschaum.config import get_config
-        allow_non_admin = get_config(
-            'system', 'api', 'permissions', 'actions', 'non_admin', patch=True
-        )
-        if not allow_non_admin:
-            return False, (
-                "The administrator for this server has not allowed users to perform actions.\n\n"
-                + "Please contact the system administrator, or if you are running this server, "
-                + "open the configuration file with `edit config system` "
-                + "and search for 'permissions'. "
-                + "\nUnder the keys 'api:permissions:actions', "
-                + "you can allow non-admin users to perform actions."
+    async def monitor_logs(job):
+        success, msg = job.start()
+        await job.monitor_logs_async(
+            partial(notify_client, websocket),
+            stop_event=stop_event,
+            stop_on_exit=True,
         )
 
-    if action not in actions:
-        return False, f"Invalid action '{action}'."
-    keywords['mrsm_instance'] = keywords.get('mrsm_instance', str(get_api_connector()))
-    _debug = keywords.get('debug', debug)
-    keywords.pop('debug', None)
-    return actions[action](debug=_debug, **keywords)
+    job = None
+    try:
+        token = await websocket.receive_text()
+        user = await manager.get_current_user(token) if not no_auth else None
+        if user is None and not no_auth:
+            raise fastapi.HTTPException(
+                status_code=401,
+                detail="Invalid credentials.",
+            )
+
+        auth_success, auth_msg = is_user_allowed_to_execute(user)
+        auth_payload = {
+            'is_authenticated': auth_success,
+            'timestamp': datetime.now(timezone.utc).isoformat(),
+        }
+        await websocket.send_json(auth_payload)
+        if not auth_success:
+            await websocket.close()
+
+        sysargs = await websocket.receive_json()
+        kwargs = parse_arguments(sysargs)
+        _ = kwargs.pop('executor_keys', None)
+        _ = kwargs.pop('shell', None)
+        sysargs = parse_dict_to_sysargs(kwargs)
+
+        job_name = '.' + generate_password(12)
+        job = Job(
+            job_name,
+            sysargs,
+            _properties={
+                'logs': {
+                    'write_timestamps': False,
+                },
+            },
+        )
+        _temp_jobs[job_name] = job
+        monitor_task = asyncio.create_task(monitor_logs(job))
+        await monitor_task
+        try:
+            await websocket.close()
+        except RuntimeError:
+            pass
+    except fastapi.HTTPException:
+        await websocket.send_text("Invalid credentials.")
+        await websocket.close()
+    except WebSocketDisconnect:
+        pass
+    except asyncio.CancelledError:
+        pass
+    except Exception:
+        warn(f"Error in logs websocket:\n{traceback.format_exc()}")
+    finally:
+        if job is not None:
+            job.delete()
+        _ = _temp_jobs.pop(job_name, None)
+        stop_event.set()