meerschaum 2.9.5__py3-none-any.whl → 3.0.0__py3-none-any.whl

meerschaum/actions/register.py

@@ -7,6 +7,9 @@ Register new Pipes. Requires the API to be running.
 """
 
 from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+
 import meerschaum as mrsm
 from meerschaum.utils.typing import SuccessTuple, Any, List, Optional, Dict
 
@@ -26,6 +29,7 @@ def register(
         'plugins'   : _register_plugins,
         'users'     : _register_users,
         'connectors': _register_connectors,
+        'tokens'    : _register_tokens,
     }
     return choose_subaction(action, options, **kw)
 
@@ -175,11 +179,11 @@ def _register_plugins(
     reload_plugins(debug=debug)
 
     repo_connector = parse_repo_keys(repository)
-    if repo_connector.__dict__.get('type', None) != 'api':
+    if repo_connector.type != 'api':
         return False, (
             "Can only upload plugins to the Meerschaum API." +
-            f"Connector '{repo_connector}' is of type " +
-            f"'{repo_connector.get('type', type(repo_connector))}'."
+            f"Connector '{repo_connector}' is " +
+            f"'{type(repo_connector)}'."
         )
 
     if len(action) == 0 or action == ['']:
@@ -272,6 +276,7 @@ def _complete_register_plugins(*args, **kw):
 def _register_users(
     action: Optional[List[str]] = None,
     mrsm_instance: Optional[str] = None,
+    scopes: Optional[List[str]] = None,
     shell: bool = False,
     debug: bool = False,
     **kw: Any
@@ -280,7 +285,7 @@ def _register_users(
     Register a new user to a Meerschaum instance.
     """
     from meerschaum.config import get_config
-    from meerschaum.config.static import STATIC_CONFIG
+    from meerschaum._internal.static import STATIC_CONFIG
     from meerschaum.connectors.parse import parse_instance_keys
     from meerschaum.utils.warnings import warn, info
     from meerschaum.core import User
@@ -333,9 +338,17 @@ def _register_users(
             )
         if len(email) == 0:
             email = None
-        user = User(username, password, email=email)
+        user = User(
+            username,
+            password,
+            email=email,
+            attributes={
+                'scopes': scopes or list(STATIC_CONFIG['tokens']['scopes']),
+            },
+            instance=instance_connector,
+        )
         info(f"Registering user '{user}' to Meerschaum instance '{instance_connector}'...")
-        result_tuple = instance_connector.register_user(user, debug=debug)
+        result_tuple = user.register(debug=debug)
         print_tuple(result_tuple)
         success[username] = result_tuple[0]
         if success[username]:
@@ -443,6 +456,130 @@ def _complete_register_connectors(
     return _complete_show_connectors(action)
 
 
+def _register_tokens(
+    mrsm_instance: Optional[str] = None,
+    name: Optional[str] = None,
+    ttl_days: Optional[int] = None,
+    scopes: Optional[List[str]] = None,
+    end: Optional[datetime] = None,
+    force: bool = False,
+    yes: bool = False,
+    noask: bool = False,
+    nopretty: bool = False,
+    debug: bool = False,
+    **kwargs: Any
+) -> mrsm.SuccessTuple:
+    """
+    Register a new long-lived access token for API access.
+    Note that omitting and end time or TTL will generate token which does not expire.
+
+    Examples:
+
+    mrsm register token --end 2032-01-01
+
+    mrsm register token --ttl-days 1000
+
+    mrsm register token --name weather-sensor
+
+    """
+    import json
+    from meerschaum.utils.schedule import parse_start_time
+    from meerschaum.core import User
+    from meerschaum.core.Token._Token import Token, _PLACEHOLDER_EXPIRATION
+    from meerschaum.connectors.parse import parse_instance_keys
+    from meerschaum.utils.prompt import yes_no, choose
+    from meerschaum.utils.formatting import make_header, print_tuple
+    from meerschaum.utils.dtypes import value_is_null
+
+    expiration = (
+        end
+        if end is not None
+        else (
+            parse_start_time(f"starting in {ttl_days} days")
+            if ttl_days
+            else _PLACEHOLDER_EXPIRATION
+        )
+    )
+
+    instance_connector = parse_instance_keys(mrsm_instance)
+    user = None
+    if instance_connector.type != 'api':
+        usernames = instance_connector.get_users(debug=debug)
+        if not usernames:
+            return False, f"No users are registered to '{instance_connector}'."
+
+        username = choose(
+            "To which user should this token be registered? Enter the number.",
+            usernames,
+        )
+    else:
+        username = getattr(instance_connector, 'username')
+
+    if not username:
+        return False, "Cannot register a token without a user."
+
+    user_id = instance_connector.get_user_id(
+        User(username, instance=mrsm_instance),
+        debug=debug,
+    )
+    if user_id is None:
+        return False, f"Cannot load ID for user '{username}'."
+
+    user = User(username, user_id=user_id, instance=mrsm_instance)
+
+    token = Token(
+        label=name,
+        expiration=expiration,
+        scopes=scopes,
+        user=user,
+        instance=mrsm_instance,
+    )
+
+    register_success, register_msg = token.register(debug=debug)
+    token_id = token.id
+    token_secret = token.secret
+    if not register_success:
+        return False, f"Failed to register token '{token.label}':\n{register_msg}"
+
+    token_model = token.to_model(refresh=True)
+    token_kwargs = dict(token_model)
+    token_kwargs['id'] = token_id
+    token_kwargs['secret'] = token_secret
+    token = Token(**token_kwargs)
+
+    if not nopretty:
+        print_tuple(
+            (
+                True,
+                (
+                    f"Registered token '{token}'.\n    "
+                    "Write down the client secret, because it won't be shown again."
+                )
+            ),
+            calm=True
+        )
+
+    msg_to_print = (
+        (
+            make_header(f"Attributes for token '{token}':") + "\n"
+            + f"    - Client ID: {token_model.id}\n"
+            + f"    - Client Secret: {token_secret}\n"
+            + "    - Expiration: "
+            + (
+                token.expiration.isoformat()
+                if not value_is_null(token.expiration)
+                else 'Does not expire.'
+            )
+            + "\n"
+            + f"    - API Key: {token.get_api_key()}\n"
+        )
+        if not nopretty
+        else json.dumps({'secret': token_secret, **dict(token_model)})
+    ) 
+    print(msg_to_print)
+    return True, f"Successfully registered token '{token.label}'."
+
+
 ### NOTE: This must be the final statement of the module.
 ###       Any subactions added below these lines will not
 ###       be added to the `help` docstring.

meerschaum/actions/reload.py

@@ -9,13 +9,30 @@ Reload the running Meerschaum instance.
 from __future__ import annotations
 from meerschaum.utils.typing import Any, SuccessTuple, List, Optional
 
+
 def reload(
-        action: Optional[List[str]] = None,
-        debug: bool = False,
-        **kw: Any
-    ) -> SuccessTuple:
+    action: Optional[List[str]] = None,
+    debug: bool = False,
+    _stop_daemons: bool = True,
+    **kw: Any
+) -> SuccessTuple:
     """
     Reload the running Meerschaum instance.
     """
     from meerschaum.utils.packages import reload_meerschaum
-    return reload_meerschaum(debug=debug)
+    from meerschaum.actions import actions
+
+    if _stop_daemons:
+        from meerschaum._internal.cli.workers import get_existing_cli_worker_indices
+        indices = get_existing_cli_worker_indices()
+        cli_action = 'restart' if indices else 'stop'
+
+        stop_daemon_success, stop_daemon_msg = actions[cli_action](['daemons'], debug=debug, **kw)
+        if not stop_daemon_success:
+            return stop_daemon_success, stop_daemon_msg
+
+    reload_success, reload_msg = reload_meerschaum(debug=debug)
+    if not reload_success:
+        return reload_success, reload_msg
+
+    return True, "Success"

meerschaum/actions/restart.py

@@ -7,6 +7,7 @@ Restart stopped jobs which have not been manually stopped.
 
 from meerschaum.utils.typing import SuccessTuple, Optional, List, Any
 
+
 def restart(
     action: Optional[List[str]] = None,
     executor_keys: Optional[str] = None,
@@ -18,6 +19,7 @@ def restart(
     from meerschaum.actions import choose_subaction
     attach_options = {
         'jobs': _restart_jobs,
+        'daemons': _restart_daemons,
     }
     return choose_subaction(action, attach_options, executor_keys=executor_keys, **kwargs)
 
@@ -111,3 +113,15 @@ def _restart_jobs(
             time.sleep(min_seconds)
 
     return check_success, check_msg
+
+
+def _restart_daemons(
+    action: Optional[List[str]] = None,
+    **kwargs
+) -> SuccessTuple:
+    """
+    Restart the CLI daemons.
+    """
+    from meerschaum.actions import actions
+    kwargs['force'] = True
+    return actions['start'](['daemons'], **kwargs)

meerschaum/actions/show.py

@@ -45,6 +45,8 @@ def show(
         'tags'       : _show_tags,
         'schedules'  : _show_schedules,
         'venvs'      : _show_venvs,
+        'tokens'     : _show_tokens,
+        'daemons'    : _show_daemons,
     }
     return choose_subaction(action, show_options, **kw)
 
@@ -113,11 +115,11 @@ def _show_help(**kw: Any) -> SuccessTuple:
 
 
 def _show_config(
-        action: Optional[List[str]] = None,
-        debug: bool = False,
-        nopretty: bool = False,
-        **kw: Any
-    ) -> SuccessTuple:
+    action: Optional[List[str]] = None,
+    debug: bool = False,
+    nopretty: bool = False,
+    **kw: Any
+) -> SuccessTuple:
     """
     Show the configuration dictionary.
     Sub-actions defined in the action list are recursive indices in the config dictionary.
@@ -204,7 +206,7 @@ def _show_version(nopretty: bool = False, **kw : Any) -> SuccessTuple:
         msg = "Meerschaum v" + version
         _print = info
     _print(msg)
-    return (True, "Success")
+    return True, "Success"
 
 
 def _show_connectors(
@@ -562,6 +564,7 @@ def _complete_show_packages(
 
     return possibilities
 
+
 def _show_jobs(
     action: Optional[List[str]] = None,
     executor_keys: Optional[str] = None,
@@ -762,7 +765,7 @@ def _show_environment(
     """
     import os
     from meerschaum.utils.formatting import pprint
-    from meerschaum.config._environment import get_env_vars
+    from meerschaum.config.environment import get_env_vars
     pprint(
         {
             env_var: os.environ[env_var]
@@ -789,34 +792,22 @@ def _show_tags(
     from meerschaum.utils.formatting import pipe_repr, UNICODE, ANSI
     from meerschaum.utils.pool import get_pool
     from meerschaum.config import get_config
+    from meerschaum.connectors.parse import parse_instance_keys
     rich_tree, rich_panel, rich_text, rich_console, rich_columns = (
         mrsm.attempt_import('rich.tree', 'rich.panel', 'rich.text', 'rich.console', 'rich.columns')
     )
-    panel = rich_panel.Panel.fit('Tags')
-    tree = rich_tree.Tree(panel)
     action = action or []
     tags = action + (tags or [])
-    pipes = mrsm.get_pipes(as_list=True, tags=tags, **kwargs)
-    if not pipes:
-        return False, f"No pipes were found with the given tags."
 
-    pool = get_pool(workers=workers)
     tag_prefix = get_config('formatting', 'pipes', 'unicode', 'icons', 'tag') if UNICODE else ''
     tag_style = get_config('formatting', 'pipes', 'ansi', 'styles', 'tags') if ANSI else None
 
-    tags_pipes = defaultdict(lambda: [])
-    gather_pipe_tags = lambda pipe: (pipe, (pipe.tags or []))
-
-    pipes_tags = dict(pool.map(gather_pipe_tags, pipes))
-
-    for pipe, tags in pipes_tags.items():
-        for tag in tags:
-            if action and tag not in action:
-                continue
-            tags_pipes[tag].append(pipe)
+    tags_pipes = mrsm.get_pipes(as_tags_dict=True, tags=tags, **kwargs)
+    if action:
+        tags_pipes = {tag: pipes for tag, pipes in tags_pipes.items() if tag in action}
 
     columns = []
-    sorted_tags = sorted([tag for tag in tags_pipes])
+    sorted_tags = sorted(list(tags_pipes))
     for tag in sorted_tags:
         _pipes = tags_pipes[tag]
         tag_text = (
@@ -904,9 +895,7 @@ def _show_schedules(
     return True, "Success"
         
 
-def _show_venvs(
-    **kwargs: Any    
-):
+def _show_venvs(**kwargs: Any) -> SuccessTuple:
     """
     Print the available virtual environments in the current MRSM_ROOT_DIR.
     """
@@ -921,12 +910,176 @@ def _show_venvs(
         for _venv in os.listdir(VIRTENV_RESOURCES_PATH)
         if venv_exists(_venv)
     ]
-    print_options(
-        venvs,
-        name = 'Venvs:',
-        **kwargs
+    print_options(venvs, name='Virtual Environments:', **kwargs)
+    return True, "Success"
+
+
+def _show_tokens(
+    action: Optional[List[str]] = None,
+    mrsm_instance: Optional[str] = None,
+    nopretty: bool = False,
+    debug: bool = False,
+    **kwargs: Any
+) -> SuccessTuple:
+    """
+    Print a table of the registered tokens on the instance.
+    """
+    import json
+    import uuid
+    from meerschaum.connectors.parse import parse_instance_keys
+    from meerschaum.utils.dtypes import value_is_null, json_serialize_value
+    from meerschaum.utils.misc import is_uuid
+    from meerschaum.utils.packages import import_rich
+    from meerschaum.utils.formatting import UNICODE, get_console
+    rich = import_rich()
+    rich_table, rich_json, rich_box = mrsm.attempt_import('rich.table', 'rich.json', 'rich.box')
+
+    conn = parse_instance_keys(mrsm_instance)
+
+    labels = [
+        label
+        for label in (action or [])
+        if not is_uuid(label)
+    ]
+    potential_token_ids = [
+        uuid.UUID(potential_id)
+        for potential_id in (action or [])
+        if is_uuid(potential_id)
+    ]
+
+    tokens = conn.get_tokens(
+        labels=(labels or None),
+        ids=(potential_token_ids or None),
+        debug=debug,
+    )
+
+    if nopretty:
+        for token in tokens:
+            print(json.dumps({
+                'id': str(token.id),
+                'label': token.label,
+                'scopes': token.scopes,
+                'expiration': (
+                    token.expiration.isoformat()
+                    if not value_is_null(token.expiration)
+                    else None
+                ),
+                'creation': (token.creation.isoformat() if not value_is_null(token.creation) else None),
+                'user': (token.user.username if token.user is not None else None),
+                'is_valid': token.is_valid,
+            }))
+        return True, "Success"
+
+    if len(tokens) == 1:
+        token = tokens[0]
+        tokens_json = json.dumps({
+            'id': str(token.id),
+            'label': token.label,
+            'scopes': token.scopes,
+            'creation': (token.creation.isoformat() if not value_is_null(token.creation) else None),
+            'expiration': (token.expiration.isoformat() if not value_is_null(token.expiration) else None),
+            'user': (token.user.username if token.user is not None else None),
+            'is_valid': token.is_valid,
+        }, default=json_serialize_value, indent=4)
+        get_console().print(rich_json.JSON(tokens_json))
+        
+        return True, "Success"
+
+    is_valid_true = (
+        "🟢"
+        if UNICODE
+        else "[+]"
+    )
+    is_valid_false = (
+        "🔴"
+        if UNICODE
+        else "[-]"
+    )
+
+    table = rich_table.Table(
+        title=f"Registered Tokens on instance '{conn}'",
+        box=rich_box.ROUNDED,
+        title_style='bold',
     )
+    table.add_column("ID")
+    table.add_column("Label")
+    table.add_column("User")
+    table.add_column("Expiration")
+    table.add_column("Valid")
+
+    for token in tokens:
+        table.add_row(
+            str(token.id),
+            token.label,
+            (token.user.username if token.user is not None else ""),
+            (
+                token.expiration.isoformat()
+                if not value_is_null(token.expiration)
+                else 'Does not expire'
+            ),
+            (is_valid_true if token.is_valid else is_valid_false),
+        )
+
+    mrsm.pprint(table)
+    return True, "Success"
+
+
+def _show_daemons() -> SuccessTuple:
+    """
+    Print information about the running CLI daemons.
+    """
+    from meerschaum._internal.cli.workers import get_existing_cli_workers
+    workers = get_existing_cli_workers()
+    if not workers:
+        return True, "No CLI daemons are running."
+
+    rich_table, rich_box, rich_text = mrsm.attempt_import('rich.table', 'rich.box', 'rich.text')
+    table = rich_table.Table(
+        rich_table.Column("Index", justify='center'),
+        rich_table.Column("Status"),
+        rich_table.Column("Locked", justify='center'),
+        rich_table.Column("Ready", justify='center'),
+        rich_table.Column("Actions"),
+        title='CLI Daemons',
+        box=rich_box.ROUNDED,
+        title_style='bold',
+    )
+
+    running_icon = mrsm.get_config('formatting', 'emoji', 'running')
+    paused_icon = mrsm.get_config('formatting', 'emoji', 'paused')
+    stopped_icon = mrsm.get_config('formatting', 'emoji', 'stopped')
+    locked_icon = mrsm.get_config('formatting', 'emoji', 'locked')
+    unlocked_icon = mrsm.get_config('formatting', 'emoji', 'unlocked')
+    status_styles = {
+        'running': 'green',
+        'stopped': 'red',
+        'paused': 'yellow',
+    }
+
+    for worker in workers:
+        status = worker.job.status
+        status_text = rich_text.Text(
+            status,
+            style=status_styles[status],
+        )
+        if not worker.is_ready():
+            ready_icon = stopped_icon
+        else:
+            ready_icon = (
+                running_icon
+                if status == 'running'
+                else paused_icon
+            )
+
+        table.add_row(
+            str(worker.ix),
+            status_text,
+            (locked_icon if worker.lock_path.exists() else unlocked_icon),
+            ready_icon,
+            str(worker.job.daemon.rotating_log.get_latest_subfile_index()),
+        )
 
+    mrsm.pprint(table)
     return True, "Success"