diff --git a/bluesky_queueserver/manager/plan_queue_ops.py b/bluesky_queueserver/manager/plan_queue_ops.py index fd77d325..0f82ac45 100644 --- a/bluesky_queueserver/manager/plan_queue_ops.py +++ b/bluesky_queueserver/manager/plan_queue_ops.py @@ -1,1944 +1,59 @@ -import asyncio -import copy -import json -import logging -import time as ttime +import os import uuid +from typing import Any -import redis.asyncio - -logger = logging.getLogger(__name__) - +# Use the centralized backend selection mechanism +from bluesky_queueserver.manager.plan_queue_ops_backends import get_default_backend class PlanQueueOperations: """ - The class supports operations with plan queue based on Redis. The public methods - of the class are protected with ``asyncio.Lock``. - - Parameters - ---------- - redis_host: str - Address of Redis host. - - name_prefix: str - Prefix for the names of the keys used in Redis. The prefix is used to avoid conflicts - with the keys used by other instances of Queue Server. For example, the prefix used - for unit tests should be different from the prefix used in production. If the prefix - is an empty string, then no prefix will be added (not recommended). - - Examples - -------- - - .. code-block:: python - - pq = PlanQueueOperations(prefix="qs_unit_test") # Redis located at `localhost` - await pq.start() - - # Fill queue - await pq.add_item_to_queue() - await pq.add_item_to_queue() - await pq.add_item_to_queue() - await pq.add_item_to_queue() - - # Number of plans in the queue - qsize = await pq.get_queue_size() - - # Read the queue (as a list) - queue, _ = await pq.get_queue() - - # Start the first plan (This doesn't actually execute the plan. It is just for bookkeeping.) - plan = await pq.set_next_item_as_running() - # ... - # Here place the code for executing the plan in dictionary `plan` - - # Again this only shows whether a plan was set as running. Expected to be True in - # this example. - is_running = await pq.is_item_running() - - # Assume that plan execution is completed, so move the plan to history - # This also clears the currently processed plan. - plan = await pq.set_processed_item_as_completed(exit_status="completed", run_uids=[]) - - # We are ready to start the next plan - plan = await pq.set_next_item_as_running() - - # Assume that we paused and then stopped the plan. Clear the running plan and - # push it back to the queue. Also create the respective history entry. - plan = await pq.set_processed_item_as_stopped(exit_status="stopped") - - # 'stopping' disconnects all connections. This step is not required in normal use. - await pq.stop() + Wrapper class for plan queue operations. Delegates method calls to the appropriate backend + implementation (Redis, SQLite, PostgreSQL, or Dict) while keeping the interface consistent. + + The backend selection is controlled through the PLAN_QUEUE_BACKEND environment variable, + which can be set to 'redis', 'sqlite', 'postgresql', or 'dict'. """ - def __init__(self, redis_host="localhost", name_prefix="qs_default"): - self._redis_host = redis_host - self._uid_dict = dict() - self._r_pool = None - - if not isinstance(name_prefix, str): - raise TypeError(f"Parameter 'name_prefix' should be a string: {name_prefix}") - - # The case of an empty string - if name_prefix: - name_prefix = name_prefix + "_" - - self._name_running_plan = name_prefix + "running_plan" - self._name_plan_queue = name_prefix + "plan_queue" - self._name_plan_history = name_prefix + "plan_history" - self._name_plan_queue_mode = name_prefix + "plan_queue_mode" - - # Redis is also used for storage of some additional information not related to the queue. - # The class contains only the functions for saving and retrieving the data, which is - # not used by other functions of the class. - self._name_user_group_permissions = name_prefix + "user_group_permissions" - self._name_lock_info = name_prefix + "lock_info" - self._name_autostart_mode_info = name_prefix + "autostart_mode_info" - self._name_stop_pending_info = name_prefix + "stop_pending_info" - - # The list of allowed item parameters used for parameter filtering. Filtering operation - # involves removing all parameters that are not in the list. - self._allowed_item_parameters = ( - "item_uid", - "item_type", - "name", - "args", - "kwargs", - "meta", - "user", - "user_group", - "properties", - ) - - # Plan queue UID is expected to change each time the contents of the queue is changed. - # Since `self._uid_dict` is modified each time the queue is updated, it is sufficient - # to update Plan queue UID in the functions that update `self._uid_dict`. - self._plan_queue_uid = self.new_item_uid() - # Plan history UID is expected to change each time the history is changed. - self._plan_history_uid = self.new_item_uid() - - self._lock = None - - # Settings that determine the mode of queue operation. The set of supported modes - # may be extended if additional modes are to be implemented. The mode will be saved in - # Redis, so that it is not modified between restarts of the manager. - # Loop mode: - # loop, True/False. If enabled, then each executed item (plan or instruction) - # will be placed to the back of the queue. - # ignore_failures, True/False. Run all the plans in the queue to the end - # even if some or all of the plans fail. The queue is still stopped if the user - # stops/aborts/halts a plan. - self._plan_queue_mode_default = {"loop": False, "ignore_failures": False} - self._plan_queue_mode = self.plan_queue_mode_default - - @property - def plan_queue_uid(self): - """ - Get current plan queue UID (str). Note, that the UID may be updated multiple times during - complex queue operations, so the returned UID may not represent a valid queue state. - The intended use: the changes of UID could be monitored to detect changes in the queue - without accessing the queue. If the UID is different from UID returned by - ``PlanQueueOperations.get_queue()``, then the contents of the queue changed. - """ - return self._plan_queue_uid - - @property - def plan_history_uid(self): - """ - Get current plan history UID. See notes for ``PlanQueueOperations.plan_queue_uid``. - """ - return self._plan_history_uid - - @property - def plan_queue_mode(self): - """ - Returns current plan queue mode. Plan queue mode is a dictionary with parameters - used for selection of the algorithm(s) for handling queue items. Supported parameters: - ``loop (boolean)`` enables and disables the loop mode. - """ - return self._plan_queue_mode.copy() - - @property - def plan_queue_mode_default(self): - """ - Returns the default queue mode (default settings) - """ - return self._plan_queue_mode_default.copy() - - def _validate_plan_queue_mode(self, plan_queue_mode): - """ - Validate the dictionary 'plan_queue_mode'. Check that the dictionary contains all - the required parameters and no unsupported parameters. - - Parameters - ---------- - plan_queue_mode : dict - Dictionary that contains plan queue mode. See ``self.plan_queue_mode_default``. - """ - # It is assumed that 'plan_queue_mode' will be a single-level dictionary that contains - # simple types (bool, int etc), so the following code provide better error reporting - # than schema validation. - expected_params = {"loop": bool, "ignore_failures": bool} - missing_keys = set(expected_params.keys()) - for k, v in plan_queue_mode.items(): - if k not in expected_params: - raise ValueError( - f"Unsupported plan queue mode parameter '{k}': " - f"supported parameters {list(expected_params.keys())}" - ) - missing_keys.remove(k) - key_type = expected_params.get(k) # Using [k] makes PyCharm to display annoying error - if not isinstance(v, key_type): - raise TypeError( - f"Unsupported type '{type(v)}' of the parameter '{k}': " f"expected type '{key_type}'" - ) - if missing_keys: - raise ValueError( - f"Parameters {missing_keys} are missing from 'plan_queue_mode' dictionary. " - f"The following keys are expected: {list(expected_params.keys())}" - ) - - async def _load_plan_queue_mode(self): - """ - Load plan queue mode from Redis. - """ - queue_mode = await self._r_pool.get(self._name_plan_queue_mode) - self._plan_queue_mode = json.loads(queue_mode) if queue_mode else self.plan_queue_mode_default - try: - self._validate_plan_queue_mode(self._plan_queue_mode) - except Exception as ex: - logger.error("Failed to load plan queue mode from Redis. The default mode is used: %s", ex) - self._plan_queue_mode = self.plan_queue_mode_default - - async def set_plan_queue_mode(self, plan_queue_mode, *, update=True): - """ - Set plan queue mode. The plan queue mode can be a string ``default`` or a dictionary with - parameters. See ``self.plan_queue_mode_default`` for an example of the parameter dictionary. - - Parameters - ---------- - plan_queue_mode : dict or str - The dictionary of parameters that define queue mode. If ``update`` is ``True``, then - the dictionary may contain only the parameters that need to be updated. Otherwise - ``plan_queue_mode`` dictionary must contain full valid parameter dictionary. The function - fails if the dictionary contains unsupported parameters. Calling the function with the - string value ``plan_queue_mode="default"`` will reset all the parameters to the default - values. - update : boolean (optional) - Indicates if the dictionary ``plan_queue_mode`` should be used to update mode parameters. - If ``True``, then the dictionary may contain only the parameters that should be changed. - """ - if not isinstance(plan_queue_mode, dict) and plan_queue_mode != "default": - raise TypeError( - f"Queue mode is passed using object of unsupported type '{type(plan_queue_mode)}': " - f"({plan_queue_mode}). Supported types: ('dict', 'str'), supported " - f"string value: 'default'" - ) - - if plan_queue_mode == "default": - plan_queue_mode = self.plan_queue_mode_default - elif update: - # Generate full parameter dictionary based on the existing and submitted parameters. - queue_mode = self.plan_queue_mode # Create a copy of current parameters - queue_mode.update(plan_queue_mode) - plan_queue_mode = queue_mode - - self._validate_plan_queue_mode(plan_queue_mode) - - # Prevent changes of the queue mode in the middle of queue operations. - async with self._lock: - self._plan_queue_mode = plan_queue_mode.copy() - await self._r_pool.set(self._name_plan_queue_mode, json.dumps(self._plan_queue_mode)) - - async def start(self): - """ - Create the pool and initialize the set of UIDs from the queue if it exists in the pool. - """ - if not self._r_pool: # Initialize only once - self._lock = asyncio.Lock() - async with self._lock: - try: - host = f"redis://{self._redis_host}" - self._r_pool = redis.asyncio.from_url(host, encoding="utf-8", decode_responses=True) - await self._r_pool.ping() - except Exception as ex: - error_msg = ( - f"Failed to create the Redis pool: " - f"Redis server may not be available at '{self._redis_host}'. " - f"Exception: {ex}" - ) - logger.error(error_msg) - raise OSError(error_msg) from ex - - await self._queue_clean() - await self._uid_dict_initialize() - await self._load_plan_queue_mode() - - self._plan_queue_uid = self.new_item_uid() - self._plan_history_uid = self.new_item_uid() - - async def stop(self): - """ - Close all connections in the pool. - """ - if self._r_pool: - await self._r_pool.aclose() - self._r_pool = None - - async def _queue_clean(self): - """ - Delete all the invalid queue entries (there could be some entries from failed unit tests). - """ - pq, _ = await self._get_queue() - - def verify_item(item): - # The criteria may be changed. - return "item_uid" in item - - items_to_remove = [] - for item in pq: - if not verify_item(item): - items_to_remove.append(item) - - for item in items_to_remove: - await self._remove_item(item, single=False) - - # Clean running plan info also (on the development computer it may contain garbage) - item = await self._get_running_item_info() - if item and not verify_item(item): - await self._clear_running_item_info() - - async def _delete_pool_entries(self): - """ - See ``self.delete_pool_entries()`` method. - """ - await self._r_pool.delete(self._name_running_plan) - await self._r_pool.delete(self._name_plan_queue) - await self._r_pool.delete(self._name_plan_history) - await self._r_pool.delete(self._name_plan_queue_mode) - self._uid_dict_clear() - - self._plan_queue_uid = self.new_item_uid() - self._plan_history_uid = self.new_item_uid() - - async def delete_pool_entries(self): - """ - Delete pool entries used by RE Manager. This method is mostly intended for use in testing, - but may be used for other purposes if needed. Deleting pool entries also resets plan queue mode. - """ - async with self._lock: - await self._delete_pool_entries() - - @staticmethod - def _verify_item_type(item): - """ - Check that the item (plan) is a dictionary. - """ - if not isinstance(item, dict): - raise TypeError(f"Parameter 'item' should be a dictionary: '{item}', (type '{type(item)}')") - - def _verify_item(self, item, *, ignore_uids=None): - """ - Verify that item (plan) structure is valid enough to be put in the queue. - Current checks: item is a dictionary, ``item_uid`` key is present, Plan with the UID is not in - the queue or currently running. Ignore UIDs in the list ``ignore_uids``: those UIDs are expected - to be in the dictionary. - """ - ignore_uids = ignore_uids or [] - self._verify_item_type(item) - # Verify plan UID - if "item_uid" not in item: - raise ValueError("Item does not have UID.") - uid = item["item_uid"] - if (uid not in ignore_uids) and self._is_uid_in_dict(uid): - raise RuntimeError(f"Item with UID {uid} is already in the queue") - - @staticmethod - def new_item_uid(): - """ - Generate UID for an item (plan). - """ - return str(uuid.uuid4()) - - def set_new_item_uuid(self, item): - """ - Replaces Item UID with a new one or creates a new UID. - - Parameters - ---------- - item: dict - Dictionary of item parameters. The dictionary may or may not have the key ``item_uid``. - - Returns - ------- - dict - Plan with new UID. - """ - item = copy.deepcopy(item) - self._verify_item_type(item) - item["item_uid"] = self.new_item_uid() - return item - - async def _get_index_by_uid(self, *, uid): - """ - Get index of a plan in Redis list by UID. This is inefficient operation and should - be avoided whenever possible. Raises an exception if the plan is not found. - - Parameters - ---------- - uid: str - UID of the plans to find. - - Returns - ------- - int - Index of the plan with given UID. - - Raises - ------ - IndexError - No plan is found. - """ - queue, _ = await self._get_queue() - for n, plan in enumerate(queue): - if plan["item_uid"] == uid: - return n - raise IndexError(f"No plan with UID '{uid}' was found in the list.") - - async def _get_index_by_uid_batch(self, *, uids): - """ - Batch version of ``_get_index_by_uid``. The operation is implemented efficiently - and should be used for finding indices of large number of items (in batch operations). - Returns a list of indices. Index is set to -1 for items that are not found. - - Parameters - ---------- - uids: list(str) - List of UIDs of the items in the batch to find. - - Returns - ------- - list(int) - List of indices of the items. The list has the same number of items as the list ``uids``. - Indices are set to -1 for the items that were not found in the queue. - """ - queue, _ = await self._get_queue() - - uids_set = set(uids) # Set should be more efficient for searching elements - uid_to_index = {} - for n, plan in enumerate(queue): - uid = plan["item_uid"] - if uid in uids_set: - uid_to_index[uid] = n - - indices = [None] * len(uids) - for n, uid in enumerate(uids): - indices[n] = uid_to_index.get(uid, -1) - - return indices - - # -------------------------------------------------------------------------- - # Operations with UID set - def _uid_dict_clear(self): - """ - Clear ``self._uid_dict``. - """ - self._plan_queue_uid = self.new_item_uid() - self._uid_dict.clear() - - def _is_uid_in_dict(self, uid): - """ - Checks if UID exists in ``self._uid_dict``. - """ - return uid in self._uid_dict - - def _uid_dict_add(self, item): - """ - Add UID to ``self._uid_dict``. - """ - uid = item["item_uid"] - if self._is_uid_in_dict(uid): - raise RuntimeError(f"Trying to add plan with UID '{uid}', which is already in the queue") - self._plan_queue_uid = self.new_item_uid() - self._uid_dict.update({uid: item}) - - def _uid_dict_remove(self, uid): + def __init__(self, **kwargs): """ - Remove UID from ``self._uid_dict``. - """ - if not self._is_uid_in_dict(uid): - raise RuntimeError(f"Trying to remove plan with UID '{uid}', which is not in the queue") - self._plan_queue_uid = self.new_item_uid() - self._uid_dict.pop(uid) - - def _uid_dict_update(self, item): - """ - Update a plan with UID that is already in the dictionary. - """ - uid = item["item_uid"] - if not self._is_uid_in_dict(uid): - raise RuntimeError(f"Trying to update plan with UID '{uid}', which is not in the queue") - self._plan_queue_uid = self.new_item_uid() - self._uid_dict.update({uid: item}) - - def _uid_dict_get_item(self, uid): - """ - Returns a plan with the given UID. - """ - return self._uid_dict[uid] - - async def _uid_dict_initialize(self): - """ - Initialize ``self._uid_dict`` with UIDs extracted from the plans in the queue. - """ - pq, _ = await self._get_queue() - self._uid_dict_clear() - # Go over all plans in the queue - for item in pq: - self._uid_dict_add(item) - # If plan is currently running - item = await self._get_running_item_info() - if item: - self._uid_dict_add(item) - - # ------------------------------------------------------------- - # Currently Running Plan - - async def _is_item_running(self): - """ - See ``self.is_item_running()`` method. - """ - return bool(await self._get_running_item_info()) - - async def is_item_running(self): - """ - Check if an item is set as running. True does not indicate that the plan is actually running. - - Returns - ------- - boolean - True - an item is set as running, False otherwise. - """ - async with self._lock: - return await self._is_item_running() - - async def _get_running_item_info(self): - """ - See ``self._get_running_item_info()`` method. - """ - plan = await self._r_pool.get(self._name_running_plan) - return json.loads(plan) if plan else {} - - async def get_running_item_info(self): - """ - Read info on the currently running item (plan) from Redis. - - Returns - ------- - dict - Dictionary representing currently running plan. Empty dictionary if - no plan is currently running (key value is ``{}`` or the key does not exist). - """ - async with self._lock: - return await self._get_running_item_info() - - async def _set_running_item_info(self, plan): - """ - Write info on the currently running item (plan) to Redis - - Parameters - ---------- - plan: dict - dictionary that contains plan parameters - """ - await self._r_pool.set(self._name_running_plan, json.dumps(plan)) - - async def _clear_running_item_info(self): - """ - Clear info on the currently running item (plan) in Redis. - """ - await self._set_running_item_info({}) - - # ------------------------------------------------------------- - # Plan Queue - - async def _get_queue_size(self): - """ - See ``self.get_queue_size()`` method. - """ - return await self._r_pool.llen(self._name_plan_queue) - - async def get_queue_size(self): - """ - Get the number of plans in the queue. - - Returns - ------- - int - The number of plans in the queue. - """ - async with self._lock: - return await self._get_queue_size() - - async def _get_queue(self): - """ - See ``self.get_queue()`` method. - """ - all_plans_json = await self._r_pool.lrange(self._name_plan_queue, 0, -1) - return [json.loads(_) for _ in all_plans_json], self._plan_queue_uid - - async def get_queue(self): - """ - Get the list of all items in the queue. The first element of the list is the first - item in the queue. - - Returns - ------- - list(dict) - The list of items in the queue. Each item is represented as a dictionary. - Empty list is returned if the queue is empty. - str - Plan queue UID. - """ - async with self._lock: - return await self._get_queue() - - async def _get_queue_full(self): - plan_queue, plan_queue_uid = await self._get_queue() - running_item = await self._get_running_item_info() - return plan_queue, running_item, plan_queue_uid - - async def get_queue_full(self): - """ - Get the list of all items in the queue and information on currently running item. - The first element of the list is the first item in the queue. This is 'atomic' operation, - i.e. it guarantees that all returned data represent a valid queue state and - the queue was not changed while the data was collected. - - Returns - ------- - list(dict) - The list of items in the queue. Each item is represented as a dictionary. - Empty list is returned if the queue is empty. - dict - Dictionary representing currently running plan. Empty dictionary if - no plan is currently running (key value is ``{}`` or the key does not exist). - str - Plan queue UID. - """ - async with self._lock: - return await self._get_queue_full() + Initialize the PlanQueueOperations class with the specified backend. - async def _get_item(self, *, pos=None, uid=None): - """ - See ``self.get_item()`` method. - """ - if (pos is not None) and (uid is not None): - raise ValueError("Ambiguous parameters: plan position and UID is specified") - - if uid is not None: - if not self._is_uid_in_dict(uid): - raise IndexError(f"Item with UID '{uid}' is not in the queue.") - running_item = await self._get_running_item_info() - if running_item and (uid == running_item["item_uid"]): - raise IndexError("The item with UID '{uid}' is currently running.") - item = self._uid_dict_get_item(uid) - - else: - pos = pos if pos is not None else "back" - - if pos == "back": - index = -1 - elif pos == "front": - index = 0 - elif isinstance(pos, int): - index = pos - else: - raise TypeError(f"Parameter 'pos' has incorrect type: pos={str(pos)} (type={type(pos)})") - - item_json = await self._r_pool.lindex(self._name_plan_queue, index) - if item_json is None: - raise IndexError(f"Index '{index}' is out of range (parameter pos = '{pos}')") - - item = json.loads(item_json) if item_json else {} - - return item - - async def get_item(self, *, pos=None, uid=None): - """ - Get item at a given position or with a given UID. If UID is specified, then - the position is ignored. - - Parameters - ---------- - pos: int, str or None - Position of the element ``(0, ..)`` or ``(-1, ..)``, ``front`` or ``back``. - - uid: str or None - Plan UID of the plan to be retrieved. UID always overrides position. - - Returns - ------- - dict - Dictionary of item parameters. - - Raises - ------ - TypeError - Incorrect value of ``pos`` (most likely a string different from ``front`` or ``back``) - IndexError - No element with position ``pos`` exists in the queue (index is out of range). - """ - async with self._lock: - return await self._get_item(pos=pos, uid=uid) - - async def _remove_item(self, item, single=True): - """ - Remove an item from the queue. If ``single=True`` then the exception is - raised in case of no or multiple matching plans are found in the queue. - The function is not part of user API and shouldn't be used on exception from - the other methods of the class. - - Parameters - ---------- - item: dict - Dictionary of item parameters. Must be identical to the item that is - expected to be deleted. - single: boolean - True - RuntimeError exception is raised if no or more than one matching - plan is found, the plans are removed anyway; False - no exception is - raised. - - Raises - ------ - RuntimeError - No or multiple matching plans are removed and ``single=True``. - """ - n_rem_items = await self._r_pool.lrem(self._name_plan_queue, 0, json.dumps(item)) - if (n_rem_items != 1) and single: - raise RuntimeError(f"The number of removed items is {n_rem_items}. One item is expected.") - - async def _pop_item_from_queue(self, *, pos=None, uid=None): - """ - See ``self.pop_item_from_queue()`` method - """ - - if (pos is not None) and (uid is not None): - raise ValueError("Ambiguous parameters: plan position and UID is specified") - - pos = pos if pos is not None else "back" - - if uid is not None: - if not self._is_uid_in_dict(uid): - raise IndexError(f"Plan with UID '{uid}' is not in the queue.") - running_item = await self._get_running_item_info() - if running_item and (uid == running_item["item_uid"]): - raise IndexError("Can not remove an item which is currently running.") - item = self._uid_dict_get_item(uid) - await self._remove_item(item) - elif pos == "back": - item_json = await self._r_pool.rpop(self._name_plan_queue) - if item_json is None: - raise IndexError("Queue is empty") - item = json.loads(item_json) if item_json else {} - elif pos == "front": - item_json = await self._r_pool.lpop(self._name_plan_queue) - if item_json is None: - raise IndexError("Queue is empty") - item = json.loads(item_json) if item_json else {} - elif isinstance(pos, int): - item = await self._get_item(pos=pos) - if item: - await self._remove_item(item) - else: - raise ValueError(f"Parameter 'pos' has incorrect value: pos={str(pos)} (type={type(pos)})") - - if item: - self._uid_dict_remove(item["item_uid"]) - - qsize = await self._get_queue_size() - - return item, qsize - - async def pop_item_from_queue(self, *, pos=None, uid=None): - """ - Pop a plan from the queue. Raises ``IndexError`` if plan with index ``pos`` is unavailable - or if the queue is empty. - - Parameters - ---------- - pos : int or str or None - Integer index specified position in the queue. Available string values: "front" or "back". - The range for the index is ``-qsize..qsize-1``: ``0, -qsize`` - front element of the queue, - ``-1, qsize-1`` - back element of the queue. If ``pos`` is ``None``, then the plan is popped - from the back of the queue. - uid : str or None - UID of the item to be removed - - Returns - ------- - dict or None - The last plan in the queue represented as a dictionary. - int - The size of the queue after completion of the operation. - - Raises - ------ - ValueError - Incorrect value of the parameter ``pos`` (typically unrecognized string). - IndexError - Position ``pos`` does not exist or the queue is empty. - """ - async with self._lock: - return await self._pop_item_from_queue(pos=pos, uid=uid) - - async def _pop_item_from_queue_batch(self, *, uids=None, ignore_missing=True): - """ - See ``self.pop_item_from_queue_batch()`` method - """ - - uids = uids or [] - - if not isinstance(uids, list): - raise TypeError(f"Parameter 'uids' must be a list: type(uids) = {type(uids)}") - - if not ignore_missing: - # Check if 'uids' contains only unique items - uids_set = set(uids) - if len(uids_set) != len(uids): - raise ValueError(f"The list of contains repeated UIDs ({len(uids) - len(uids_set)} UIDs)") - - # Check if all UIDs in 'uids' exist in the queue - uids_missing = [] - for uid in uids: - if not self._is_uid_in_dict(uid): - uids_missing.append(uid) - if uids_missing: - raise ValueError(f"The queue does not contain items with the following UIDs: {uids_missing}") - - items = [] - for uid in uids: - try: - item, _ = await self._pop_item_from_queue(uid=uid) - items.append(item) - except Exception as ex: - logger.debug("Failed to remove item with UID '%s' from the queue: %s", uid, ex) - - qsize = await self._get_queue_size() - return items, qsize - - async def pop_item_from_queue_batch(self, *, uids=None, ignore_missing=True): - """ - Pop a batch of items from the queue. Raises ``IndexError`` if plan with index ``pos`` is unavailable - or if the queue is empty. + The backend is determined by the environment variable `PLAN_QUEUE_BACKEND`. + If the variable is not defined, the default backend is Redis. Parameters ---------- - uids : list(str) - list of UIDs of items to be removed. The list may be empty. - ignore_missing : boolean - if the parameter is ``True`` (default), then all items from the batch that are found in - the queue are removed, if ``False``, then the function fails if the list contains repeated - entries or at least one of the items is not found in the queue. No items are removed - from the queue if the function fails. - - Returns - ------- - list(dict) - The list of items that were removed from the queue. - int - Size of the queue after completion of the operation. - - Raises - ------ - ValueError - Function failed due to missing or repeated items in the batch. + kwargs : dict + Additional arguments to pass to the backend implementation. + [...] """ - async with self._lock: - return await self._pop_item_from_queue_batch(uids=uids, ignore_missing=ignore_missing) + # Get the appropriate backend class from the centralized mechanism + BackendClass = get_default_backend() + self._backend = BackendClass(**kwargs) - def filter_item_parameters(self, item): + def __getattr__(self, name): """ - Remove parameters that are not in the list of allowed parameters. - Current parameter list includes parameters ``item_type``, ``item_uid``, - ``name``, ``args``, ``kwargs``, ``meta``, ``user``, ``user_group``. - - The list does not include ``result`` parameter, i.e. ``result`` will - be removed from the list of parameters. - + Delegate attribute access to the backend instance. + + This allows us to call methods on the backend instance directly + without needing to explicitly define them in this wrapper class. + Parameters ---------- - item : dict - dictionary of item parameters - + name : str + The name of the attribute or method to access. + Returns ------- - dict - dictionary of filtered item parameters + Any + The requested attribute or method from the backend instance. """ - return {k: w for k, w in item.items() if k in self._allowed_item_parameters} + return getattr(self._backend, name) + + async def start(self): + """Initialize the backend connection.""" + await self._backend.start() - async def _add_item_to_queue(self, item, *, pos=None, before_uid=None, after_uid=None, filter_parameters=True): - """ - See ``self.add_item_to_queue()`` method. - """ - if (pos is not None) and (before_uid is not None or after_uid is not None): - raise ValueError("Ambiguous parameters: plan position and UID is specified") - - if (before_uid is not None) and (after_uid is not None): - raise ValueError( - "Ambiguous parameters: request to insert the plan before and after the reference plan" - ) - - pos = pos if pos is not None else "back" - - if "item_uid" not in item: - item = self.set_new_item_uuid(item) - else: - self._verify_item(item) - - if filter_parameters: - item = self.filter_item_parameters(item) - - qsize0 = await self._get_queue_size() - if isinstance(pos, int): - if (pos == 0) or (pos < -qsize0): - pos = "front" - elif (pos == -1) or (pos >= qsize0): - pos = "back" - - if (before_uid is not None) or (after_uid is not None): - uid = before_uid if before_uid is not None else after_uid - before = uid == before_uid - - if not self._is_uid_in_dict(uid): - raise IndexError(f"Plan with UID '{uid}' is not in the queue.") - running_item = await self._get_running_item_info() - if running_item and (uid == running_item["item_uid"]): - if before: - raise IndexError("Can not insert a plan in the queue before a currently running plan.") - else: - # Push to the plan front of the queue (after the running plan). - qsize = await self._r_pool.lpush(self._name_plan_queue, json.dumps(item)) - else: - item_to_displace = self._uid_dict_get_item(uid) - where = "BEFORE" if (uid == before_uid) else "AFTER" - qsize = await self._r_pool.linsert( - self._name_plan_queue, where, json.dumps(item_to_displace), json.dumps(item) - ) - - elif pos == "back": - qsize = await self._r_pool.rpush(self._name_plan_queue, json.dumps(item)) - elif pos == "front": - qsize = await self._r_pool.lpush(self._name_plan_queue, json.dumps(item)) - elif isinstance(pos, int): - pos_reference = pos if (pos > 0) else (pos + 1) - - item_to_displace = await self._get_item(pos=pos_reference) - if item_to_displace: - qsize = await self._r_pool.linsert( - self._name_plan_queue, "BEFORE", json.dumps(item_to_displace), json.dumps(item) - ) - else: - raise RuntimeError(f"Could not find an existing plan at {pos}. Queue size: {qsize0}") - else: - raise ValueError(f"Parameter 'pos' has incorrect value: pos='{str(pos)}' (type={type(pos)})") - - self._uid_dict_add(item) - return item, qsize - - async def add_item_to_queue(self, item, *, pos=None, before_uid=None, after_uid=None, filter_parameters=True): - """ - Add an item to the queue. By default, the item is added to the back of the queue. - If position is integer, it is clipped to fit within the range of meaningful indices. - For the index too large or too low, the plan is pushed to the front or the back of the queue. - - Parameters - ---------- - item : dict - Item (plan or instruction) represented as a dictionary of parameters - pos : int, str or None - Integer that specifies the position index, "front" or "back". - If ``pos`` is in the range ``0..qsize`` (qsize counted before the new - item is inserted), the item is inserted to the specified position - and items at positions ``pos..qsize-1`` are shifted by one position - to the right. If ``-qsizeqsize`` or ``pos==-1``, the plan is added to the back of - the queue. If ``pos==0`` or ``pos<-qsize``, the plan is pushed to - the front of the queue. - before_uid : str or None - If UID is specified, then the item is inserted before the plan with UID. - ``before_uid`` has precedence over ``after_uid``. - after_uid : str or None - If UID is specified, then the item is inserted before the plan with UID. - filter_parameters : boolean (optional) - Remove all parameters that do not belong to the list of allowed parameters. - Default: ``True``. - - Returns - ------- - dict, int - The dictionary that contains the item that was added and the new size of the queue. - - Raises - ------ - ValueError - Incorrect value of the parameter ``pos`` (typically unrecognized string). - TypeError - Incorrect type of ``item`` (should be dict) - """ - async with self._lock: - return await self._add_item_to_queue( - item, pos=pos, before_uid=before_uid, after_uid=after_uid, filter_parameters=filter_parameters - ) - - async def _add_batch_to_queue( - self, items, *, pos=None, before_uid=None, after_uid=None, filter_parameters=True - ): - """ - See ``self.add_batch_to_queue()`` method. - """ - items_added, results = [], [] - - # Approach: attempt ot add each item to queue. In case of a failure remove all added item. - # Operation of adding items to queue is perfectly reversible. - - # Adding a batch to a specified position in the queue: - # - add the first plan by passing all positional parameters to 'self._add_item_to_queue'. - # - add the remaining plans after the first plan (it doesn't matter how position of the first - # plan was defined. - - success = True - added_item_uids = [] # List of plans with successfully added UIDs. Used for 'undo' operation. - for item in items: - try: - if not added_item_uids: - item_added, _ = await self._add_item_to_queue( - item, - pos=pos, - before_uid=before_uid, - after_uid=after_uid, - filter_parameters=filter_parameters, - ) - else: - item_added, _ = await self._add_item_to_queue( - item, after_uid=added_item_uids[-1], filter_parameters=filter_parameters - ) - added_item_uids.append(item_added["item_uid"]) - items_added.append(item_added) - results.append({"success": True, "msg": ""}) - except Exception as ex: - success = False - items_added.append(item) - msg = f"Failed to add item to queue: {ex}" - results.append({"success": False, "msg": msg}) - - # 'Undo' operation: remove all inserted items - if not success: - for uid in added_item_uids: - # The 'try-except' here is just in case anything goes wrong. Operation of removing - # items that were just added are expected to succeed. - try: - await self._pop_item_from_queue(uid=uid) - except Exception as ex: - logger.error( - "Failed to remove an item with uid='%s' after failure to add a batch of plans", ex - ) - - # Also do not return 'changed' items if adding the batch failed. - items_added = items - - qsize = await self._get_queue_size() - - # 'items_added' and 'results' ALWAYS have the same number of elements as 'items' - return items_added, results, qsize, success - - async def add_batch_to_queue( - self, items, *, pos=None, before_uid=None, after_uid=None, filter_parameters=True - ): - """ - Add a batch of item to the queue. The behavior of the function is similar - to ``add_item_to_queue`` except that it accepts the list of items to add. - The function will not add any plans to the queue if at least one of the plans - is rejected. The function returns ``success`` flag, which is ``True`` if - the batch was added and ``False`` otherwise. Success status and error messages - for each added plan can be found in ``results`` list. If the batch was added - successfully, then all ``results`` element indicate success. - - The function is not expected to raise exceptions in case of failure, but instead - report results of processing for each item in the ``results`` list. - - Parameters - ---------- - items : list(dict) - List of items (plans or instructions). Each element of the list is a dictionary - of plan parameters - pos, before_uid, after_uid, filter_parmeters - see documentation for ``add_item_to_queue`` for details. - - Returns - ------- - items_added : list(dict) - List of items that were added to queue. In case the operation fails, the list - of submitted items is returned. The list always has the same number of elements - as ``items``. - results : list(dict) - List of processing results. The list always has the same number of elements as - ``items``. Each element contains a report on processing the respective item in - the form of a dictionary with the keys ``success`` (boolean) and ``msg`` (str). - ``msg`` is always an empty string if ``success==True``. In case the batch was - added successfully, all elements are ``{"success": True, "msg": ""}``. - qsize : int - Size of the queue after the batch was added. The size will not change if the - batch is rejected. - success : bool - Indicates success of the operation of adding the batch. If ``False``, then - the batch is rejected and error messages for each item could be found in - the ``results`` list. - """ - - async with self._lock: - return await self._add_batch_to_queue( - items, pos=pos, before_uid=before_uid, after_uid=after_uid, filter_parameters=filter_parameters - ) - - async def _replace_item(self, item, *, item_uid): - """ - See ``self._replace_item()`` method - """ - # We can not replace currently running item, since it is technically not in the queue - running_item = await self._get_running_item_info() - running_item_uid = running_item["item_uid"] if running_item else None - if not self._is_uid_in_dict(item_uid): - raise RuntimeError(f"Failed to replace item: Item with UID '{item_uid}' is not in the queue") - if (running_item_uid is not None) and (running_item_uid == item_uid): - raise RuntimeError(f"Failed to replace item: Item with UID '{item_uid}' is currently running") - - if "item_uid" not in item: - item = self.set_new_item_uuid(item) - - item_to_replace = self._uid_dict_get_item(item_uid) - if item == item_to_replace: - # There is nothing to do. Consider operation as successful. - qsize = await self._get_queue_size() - return item, qsize - - # Item with 'item_uid' is expected to be in the queue, so ignore 'item_uid'. - # The concern is that the queue may contain a plan with `item["item_uid"]`, which could be - # different from 'item_uid'. Then inserting the item may violate integrity of the queue. - self._verify_item(item, ignore_uids=[item_uid]) - - # Parameters of the edited item should be verified against the list of the allowed parameters - item = self.filter_item_parameters(item) - - # Insert the new item after the old one and remove the old one. At this point it is guaranteed - # that they are not equal. - await self._r_pool.linsert(self._name_plan_queue, "AFTER", json.dumps(item_to_replace), json.dumps(item)) - - await self._remove_item(item_to_replace) - - # Update self._uid_dict - self._uid_dict_remove(item_uid) - self._uid_dict_add(item) - - # Read the actual size of the queue. - qsize = await self._get_queue_size() - - return item, qsize - - async def replace_item(self, item, *, item_uid): - """ - Replace item in the queue. Item with UID ``item_uid`` is replaced by item ``item``. The new - item may have UID which is the same or different from UID of the item being replaced. - - Parameters - ---------- - item: dict - Item (plan or instruction) represented as a dictionary of parameters. - - item_uid: str - UID of existing item in the queue that will be replaced. - - Returns - ------- - dict, int - The dictionary that contains the item that was added and the new size of the queue. - """ - async with self._lock: - return await self._replace_item(item, item_uid=item_uid) - - async def _move_item(self, *, pos=None, uid=None, pos_dest=None, before_uid=None, after_uid=None): - """ - See ``self.move_item()`` method. - """ - if (pos is None) and (uid is None): - raise ValueError("Source position or UID is not specified.") - if (pos_dest is None) and (before_uid is None) and (after_uid is None): - raise ValueError("Destination position or UID is not specified.") - - if (pos is not None) and (uid is not None): - raise ValueError("Ambiguous parameters: Both position and uid is specified for the source plan.") - if (pos_dest is not None) and (before_uid is not None or after_uid is not None): - raise ValueError("Ambiguous parameters: Both position and uid is specified for the destination plan.") - if (before_uid is not None) and (after_uid is not None): - raise ValueError("Ambiguous parameters: source should be moved 'before' and 'after' the destination.") - - queue_size = await self._get_queue_size() - - # Find the source plan - src_txt = "" - src_by_index = False # Indicates that the source is addressed by index - try: - if uid is not None: - src_txt = f"UID '{uid}'" - item_source = await self._get_item(uid=uid) - else: - src_txt = f"position {pos}" - src_by_index = True - item_source = await self._get_item(pos=pos) - except Exception as ex: - raise IndexError(f"Source plan ({src_txt}) was not found: {str(ex)}.") - - uid_source = item_source["item_uid"] - - # Find the destination plan - dest_txt, before = "", True - try: - if (before_uid is not None) or (after_uid is not None): - uid_dest = before_uid if before_uid else after_uid - before = uid_dest == before_uid - dest_txt = f"UID '{uid_dest}'" - item_dest = await self._get_item(uid=uid_dest) - else: - dest_txt = f"position {pos_dest}" - item_dest = await self._get_item(pos=pos_dest) - - # Find the index of the source in the most efficient way - src_index = pos if src_by_index else (await self._get_index_by_uid(uid=uid)) - if src_index == "front": - src_index = 0 - elif src_index == "back": - src_index = queue_size - 1 - - # Determine if the item must be inserted before or after the destination - if pos_dest == "front": - before = True - elif pos_dest == "back": - # This is one case when we need to insert the plan after the 'destination' plan. - before = False - else: - si = src_index if src_index >= 0 else queue_size + src_index - pi = pos_dest if pos_dest >= 0 else queue_size + pos_dest - before = si > pi - - except Exception as ex: - raise IndexError(f"Destination plan ({dest_txt}) was not found: {str(ex)}.") - - # Copy destination UID from the plan (we need it for the case of if addressing is positional - # so we convert it to UID, but we can do it for the case of UID addressing as well) - # In case of positional addressing 'before' is True, so the source is going to be - # inserted in place of destination. - uid_dest = item_dest["item_uid"] - - # If source and destination point to the same plan, then do nothing, - # but consider it a valid operation. - if uid_source != uid_dest: - item, _ = await self._pop_item_from_queue(uid=uid_source) - kw = {"before_uid": uid_dest} if before else {"after_uid": uid_dest} - kw.update({"item": item}) - # The item is moved 'as is'. No filtering of parameters is applied. - item, qsize = await self._add_item_to_queue(**kw, filter_parameters=False) - else: - item = item_dest - qsize = await self._get_queue_size() - return item, qsize - - async def move_item(self, *, pos=None, uid=None, pos_dest=None, before_uid=None, after_uid=None): - """ - Move existing item within the queue. Plan is moved within the queue without modification. - No parameter filtering is applied, so the ``result`` item parameter will not be removed if - present. - - Parameters - ---------- - pos: str or int - Position of the source item: positive or negative integer that specifieds the index - of the item in the queue or a string from the set {"back", "front"}. - uid: str - UID of the source item. UID overrides the position - pos_dest: str or int - Index of the new position of the item in the queue: positive or negative integer that - specifieds the index of the item in the queue or a string from the set {"back", "front"}. - before_uid: str - Insert the item before the item with the given UID. - after_uid: str - Insert the item after the item with the given UID. - - Returns - ------- - dict, int - The dictionary that contains a item that was moved and the size of the queue. - - Raises - ------ - ValueError - Error in specification of source or destination. - """ - async with self._lock: - return await self._move_item( - pos=pos, uid=uid, pos_dest=pos_dest, before_uid=before_uid, after_uid=after_uid - ) - - async def _move_batch(self, *, uids=None, pos_dest=None, before_uid=None, after_uid=None, reorder=False): - """ - See ``self.move_batch()`` method. - """ - uids = uids or [] - - if not isinstance(uids, list): - raise TypeError(f"Parameter 'uids' must be a list: type(uids) = {type(uids)}") - - # Make sure only one of the mutually exclusive parameters is not None - param_list = [pos_dest, before_uid, after_uid] - n_params = len(param_list) - param_list.count(None) - if n_params < 1: - raise ValueError( - "Destination for the batch is not specified: use parameters 'pos_dest', " - "'before_uid' or 'after_uid'" - ) - elif n_params > 1: - raise ValueError( - "The function was called with more than one mutually exclusive parameter " - "('pos_dest', 'before_uid', 'after_uid')" - ) - - # Check if 'uids' contains only unique items - uids_set = set(uids) - if len(uids_set) != len(uids): - raise ValueError(f"The list of contains repeated UIDs ({len(uids) - len(uids_set)} UIDs)") - - # Check if all UIDs in 'uids' exist in the queue - uids_missing = [] - for uid in uids: - if not self._is_uid_in_dict(uid): - uids_missing.append(uid) - if uids_missing: - raise ValueError(f"The queue does not contain items with the following UIDs: {uids_missing}") - - # Check that 'before_uid' and 'after_uid' are not in 'uids' - if (before_uid is not None) and (before_uid in uids): - raise ValueError(f"Parameter 'before_uid': item with UID '{before_uid}' is in the batch") - if (after_uid is not None) and (after_uid in uids): - raise ValueError(f"Parameter 'after_uid': item with UID '{after_uid}' is in the batch") - - # Rearrange UIDs in the list if items need to be reordered - if reorder: - indices = await self._get_index_by_uid_batch(uids=uids) - - def sorting_key(element): - return element[0] - - uids_with_indices = list(zip(indices, uids)) - uids_with_indices.sort(key=sorting_key) - uids_prepared = [_[1] for _ in uids_with_indices] - else: - uids_prepared = uids - - # Perform the 'move' operation. - last_item_uid = None - items_moved = [] - for uid in uids_prepared: - if last_item_uid is None: - # First item is moved according to specified parameters - item, _ = await self._move_item( - uid=uid, pos_dest=pos_dest, before_uid=before_uid, after_uid=after_uid - ) - else: - # Consecutive items are placed after the first item - item, _ = await self._move_item(uid=uid, after_uid=last_item_uid) - last_item_uid = uid - items_moved.append(item) - - qsize = await self._get_queue_size() - return items_moved, qsize - - async def move_batch(self, *, uids=None, pos_dest=None, before_uid=None, after_uid=None, reorder=False): - """ - Move a batch of existing items within the queue. The items are specified as a list of uids. - If at least one of the uids can not be found in the queue, the operation fails and no items - are moved. The moved items can be ordered based on the order of uids in the list (``reorder=False``) - or based on the original order in the queue (``reorder=True``). Plan is moved within the queue - without modification. No parameter filtering is applied, so the ``result`` item parameter - will not be removed if present. - - The items in the batch do not have to be contiguous. Destination items specified by ``after_uid`` - and ``before_uid`` may not belong to the batch. The parameters ``pos_dest``, ``before_uid`` and - ``after_uid`` are mutually exclusive. - - The function raises an exception with error message in case the operation fails. - - Parameters - ---------- - uids : list(str) - List of UIDs of the items in the batch. The list may not contain repeated UIDs. All UIDs - must be present in the queue. - pos_dest : str - Destination of the moved batch. Only string values ``front`` and ``back`` are allowed. - before_uid : str - Insert the batch before the item with the given UID. - after_uid : str - Insert the batch after the item with the given UID. - reorder : boolean - Arranged moved items according to the order of UIDs in the ``uids`` list (``False``) or - according to the original order of items in the queue (``True``). - - Returns - ------- - list(dict), int - The list of items that were moved and the size of the queue. The order of the items - matches the order of the items in the moved batch. Depending on the value of ``reorder`` - it may or may not match the order of ``uids``. - - Raises - ------ - ValueError - Operation could not be performed due to incorrectly specified parameters or invalid - list of ``uids``. - TypeError - Incorrect type of parameter ``uids``. - """ - async with self._lock: - return await self._move_batch( - uids=uids, pos_dest=pos_dest, before_uid=before_uid, after_uid=after_uid, reorder=reorder - ) - - async def _clear_queue(self): - """ - See ``self.clear_queue()`` method. - """ - self._plan_queue_uid = self.new_item_uid() - await self._r_pool.delete(self._name_plan_queue) - - # Remove all entries from 'self._uid_dict' except the running item. - running_item = await self._get_running_item_info() - if running_item: - uid = running_item["item_uid"] - item = self._uid_dict_get_item(uid) - self._uid_dict_clear() - self._uid_dict_add(item) - else: - self._uid_dict_clear() - - async def clear_queue(self): - """ - Remove all entries from the plan queue. Does not touch the running item (plan). - The item may be pushed back into the queue if it is stopped. - """ - async with self._lock: - await self._clear_queue() - - # ----------------------------------------------------------------------- - # Plan History - - async def _add_to_history(self, item): - """ - Add an item to history. - - Parameters - ---------- - item: dict - Item (plan) represented as a dictionary of parameters. No verifications are performed - on the plan. The function is not intended to be used outside of this class. - - Returns - ------- - int - The new size of the history. - """ - self._plan_history_uid = self.new_item_uid() - history_size = await self._r_pool.rpush(self._name_plan_history, json.dumps(item)) - return history_size - - async def _get_history_size(self): - """ - See ``self.get_history_size()`` method. - """ - return await self._r_pool.llen(self._name_plan_history) - - async def get_history_size(self): - """ - Get the number of items in the plan history. - - Returns - ------- - int - The number of plans in the history. - """ - async with self._lock: - return await self._get_history_size() - - async def _get_history(self): - """ - See ``self.get_history()`` method. - """ - all_plans_json = await self._r_pool.lrange(self._name_plan_history, 0, -1) - return [json.loads(_) for _ in all_plans_json], self._plan_history_uid - - async def get_history(self): - """ - Get the list of all items in the plan history. The first element of the list is - the oldest history entry. - - Returns - ------- - list(dict) - The list of items in the plan history. Each plan is represented as a dictionary. - Empty list is returned if the queue is empty. - str - Plan history UID - """ - async with self._lock: - return await self._get_history() - - async def _clear_history(self): - """ - See ``self.clear_history()`` method. - """ - self._plan_history_uid = self.new_item_uid() - await self._r_pool.delete(self._name_plan_history) - - async def clear_history(self): - """ - Remove all entries from the plan queue. Does not touch the running item. - The item (plan) may be pushed back into the queue if it is stopped. - """ - async with self._lock: - await self._clear_history() - - # ---------------------------------------------------------------------- - # Standard item operations during queue execution - - def _clean_item_properties(self, item): - """ - The function removes unneccessary item properties before adding the item to history or - returning it to queue. - """ - item = copy.deepcopy(item) - if "properties" in item: - p = item["properties"] - # "immediate_execution" flag is set internally by the server and should not be exposed to users - if "immediate_execution" in p: - del p["immediate_execution"] - # 'time_start' is a temporary parameter and should be removed - if "time_start" in p: - del p["time_start"] - if not p: - del item["properties"] - return item - - async def _process_next_item(self, *, item=None): - """ - See ``self.process_next_item()`` method. - """ - loop_mode = self._plan_queue_mode["loop"] - - # Read the item from the front of the queue - - immediate_execution = bool(item) - if immediate_execution: - # Generate UID if it does not exist or creates a deep copy - item = copy.deepcopy(item) if "item_uid" in item else self.set_new_item_uuid(item) - else: - item = await self._get_item(pos="front") - - item_to_return = item - if item: - item_type = item["item_type"] - if item_type == "plan": - kwargs = {"item": item} if immediate_execution else {} - item_to_return = await self._set_next_item_as_running(**kwargs) - - elif not immediate_execution: - # Items other than plans should be pushed to the back of the queue. - await self._pop_item_from_queue(pos="front") - if loop_mode: - item_to_add = self.set_new_item_uuid(item) - await self._add_item_to_queue(item_to_add) - - return item_to_return - - async def process_next_item(self, *, item=None): - """ - Process the next item in the queue. If the item is a plan, it is set as currently - running (``self.set_next_item_as_running``), otherwise it is popped from the queue. - If the queue is LOOP mode, then it the item other than plan is pushed to the back - of the queue. (Plans are pushed to the back of the queue upon successful completion.) - - If ``item`` is a plan, then the plan is set for immediate execution. If ``item`` is - an instruction, then the function does nothing. - - If an item submitted for 'immediate' execution has no UID, a new UID is created. - The returned plan is in the exact form, which is set for execution. - - For more details on processing of queue plans, see the description for - ``self.set_next_item_as_running`` method. - """ - async with self._lock: - return await self._process_next_item(item=item) - - async def _set_next_item_as_running(self, *, item=None): - """ - See ``self.set_next_item_as_running()`` method. - """ - immediate_execution = bool(item) - if immediate_execution: - # Generate UID if it does not exist or creates a deep copy - item = copy.deepcopy(item) if "item_uid" in item else self.set_new_item_uuid(item) - item.setdefault("properties", {})["immediate_execution"] = True - - # UID remains in the `self._uid_dict` after this operation. - try: - if await self._is_item_running(): - raise Exception() - - if immediate_execution: - plan = item - else: - plan = await self._get_item(pos="front") - if not plan: - raise Exception() - - if "item_type" not in plan: - raise Exception() - - if plan["item_type"] != "plan": - raise RuntimeError( - "Function 'PlanQueueOperations.set_next_item_as_running' was called for " - f"an item other than plan: {plan}" - ) - - if not immediate_execution: - # Pop plan from the front of the queue (it is the same plan as currently loaded) - await self._r_pool.lpop(self._name_plan_queue) - - # Record start time for the plan - plan.setdefault("properties", {})["time_start"] = ttime.time() - - await self._set_running_item_info(plan) - self._plan_queue_uid = self.new_item_uid() - - except RuntimeError: - raise - except Exception: - plan = {} - - return plan - - async def set_next_item_as_running(self, *, item=None): - """ - Sets the next item from the queue as 'running'. The item MUST be a plan - (e.g. not an instruction), otherwise an exception will be raised. The item is removed - from the queue. UID remains in ``self._uid_dict``, i.e. item with the same UID - may not be added to the queue while it is being executed. If ``item`` parameter - represents a plan, then the plan is set for immediate execution. - - This function can only be applied to the plans. Use ``process_next_item`` that - works correctly for any item (it calls this function if an item is a plan). - - If a plan submitted for 'immediate' execution has no UID, a new UID is created. - The returned plan is in the exact form, which is set for execution. - - Parameters - ---------- - item: dict or None - The dictionary that represents a plan submitted for immediate execution. - If ``item`` is a plan, then the queue remains intact and the plan is - set for immediate execution. If ``item=None``, then the top queue item - is removed from the queue and set for execution. - - Returns - ------- - dict - The item that was set as currently running. If another item is currently - set as 'running' or the queue is empty, then ``{}`` is returned. - - Raises - ------ - RuntimeError - The function is called for an item other than plan. - """ - async with self._lock: - return await self._set_next_item_as_running(item=item) - - async def _set_processed_item_as_completed(self, *, exit_status, run_uids, scan_ids, err_msg, err_tb): - """ - See ``self.set_processed_item_as_completed`` method. - """ - # If loop_mode is True, then add item to the back of the queue - loop_mode = self._plan_queue_mode["loop"] - - # Note: UID remains in the `self._uid_dict` after this operation - if await self._is_item_running(): - item = await self._get_running_item_info() - immediate_execution = item["properties"].get("immediate_execution", False) - item_time_start = item["properties"]["time_start"] - item_cleaned = self._clean_item_properties(item) - - if loop_mode and not immediate_execution: - item_to_add = item_cleaned.copy() - item_to_add = self.set_new_item_uuid(item_to_add) - await self._r_pool.rpush(self._name_plan_queue, json.dumps(item_to_add)) - self._uid_dict_add(item_to_add) - item_cleaned.setdefault("result", {}) - item_cleaned["result"]["exit_status"] = exit_status - item_cleaned["result"]["run_uids"] = run_uids - item_cleaned["result"]["scan_ids"] = scan_ids - item_cleaned["result"]["time_start"] = item_time_start - item_cleaned["result"]["time_stop"] = ttime.time() - item_cleaned["result"]["msg"] = err_msg - item_cleaned["result"]["traceback"] = err_tb - await self._clear_running_item_info() - if not loop_mode and not immediate_execution: - self._uid_dict_remove(item["item_uid"]) - self._plan_queue_uid = self.new_item_uid() - await self._add_to_history(item_cleaned) - else: - item_cleaned = {} - return item_cleaned - - async def set_processed_item_as_completed(self, *, exit_status, run_uids, scan_ids, err_msg, err_tb): - """ - Moves currently executed item (plan) to history and sets ``exit_status`` key. - UID is removed from ``self._uid_dict``, so a copy of the item with - the same UID may be added to the queue. - - Known ``exit_status`` values: ``"completed"`` and ``"unknown"`` (status - information was lost due to restart of RE Manager, assume that the completion - was successful and start the next plan). - - Parameters - ---------- - exit_status: str - Completion status of the plan. - run_uids: list(str) - A list of uids of completed runs. - scan_ids: list(int) - A list of scan IDs for the completed runs. - err_msg: str - Error message in case of failure. - err_tb: str - Traceback in case of failure. - - Returns - ------- - dict - The item added to the history including ``exit_status``. If another item (plan) - is currently running, then ``{}`` is returned. - """ - async with self._lock: - return await self._set_processed_item_as_completed( - exit_status=exit_status, run_uids=run_uids, scan_ids=scan_ids, err_msg=err_msg, err_tb=err_tb - ) - - async def _set_processed_item_as_stopped(self, *, exit_status, run_uids, scan_ids, err_msg, err_tb): - """ - See ``self.set_processed_item_as_stopped()`` method. - """ - # Note: UID is removed from `self._uid_dict`. - if exit_status == "stopped": - # Stopped item is considered successful, so it is not pushed back to the beginning - # of the queue, and it is added to the back of the queue in LOOP mode. - item_cleaned = await self._set_processed_item_as_completed( - exit_status=exit_status, run_uids=run_uids, scan_ids=scan_ids, err_msg=err_msg, err_tb=err_tb - ) - elif await self._is_item_running(): - item = await self._get_running_item_info() - immediate_execution = item.get("properties", {}).get("immediate_execution", False) - item_time_start = item["properties"]["time_start"] - item_cleaned = self._clean_item_properties(item) - - item_cleaned.setdefault("result", {}) - item_cleaned["result"]["exit_status"] = exit_status - item_cleaned["result"]["run_uids"] = run_uids - item_cleaned["result"]["scan_ids"] = scan_ids - item_cleaned["result"]["time_start"] = item_time_start - item_cleaned["result"]["time_stop"] = ttime.time() - item_cleaned["result"]["msg"] = err_msg - item_cleaned["result"]["traceback"] = err_tb - - await self._add_to_history(item_cleaned) - await self._clear_running_item_info() - - # Generate new UID for the item that is pushed back into the queue. - if not immediate_execution: - self._uid_dict_remove(item["item_uid"]) - # "stopped" - successful completion. Do not insert the item back in the queue. - if exit_status != "stopped": - item_pushed_to_queue = self.set_new_item_uuid(item_cleaned) - await self._add_item_to_queue(item_pushed_to_queue, pos="front", filter_parameters=False) - self._plan_queue_uid = self.new_item_uid() - else: - item_cleaned = {} - return item_cleaned - - async def set_processed_item_as_stopped(self, *, exit_status, run_uids, scan_ids, err_msg, err_tb): - """ - A stopped plan is considered successfully completed (if ``exit_status=="stopped"``) or - failed (otherwise). All items are added to history with respective ``exit_status``. - Failed items are pushed to the beginning of the queue. Item UID is removed in ``self._uid_dict``. - A new ``item_uid`` is generated for the item that is pushed back into the queue. - - Known ``exit_status`` values: ``"failed"``, ``"stopped"`` (success), ``"aborted"``, ``"halted"``. - - Parameters - ---------- - exit_status: str - Completion status of the plan. - run_uids: list(str) - A list of uids of completed runs. - scan_ids: list(int) - A list of scan IDs for the completed runs. - err_msg: str - Error message in case of failure. - err_tb: str - Traceback in case of failure. - - Returns - ------- - dict - The item (plan) added to the history including ``exit_status``. If no item (plan) - is running, then the function returns ``{}``. The item pushed back into the queue - will have different ``item_uid`` than the item added to the history. - """ - async with self._lock: - return await self._set_processed_item_as_stopped( - exit_status=exit_status, run_uids=run_uids, scan_ids=scan_ids, err_msg=err_msg, err_tb=err_tb - ) - - # ============================================================================================= - # Methods for saving and retrieving user group permissions. - - async def user_group_permissions_clear(self): - """ - Clear user group permissions saved in Redis. - """ - await self._r_pool.delete(self._name_user_group_permissions) - - async def user_group_permissions_save(self, user_group_permissions): - """ - Save user group permissions to Redis. - - Parameters - ---------- - user_group_permissions: dict - A dictionary containing user group permissions. - """ - await self._r_pool.set(self._name_user_group_permissions, json.dumps(user_group_permissions)) - - async def user_group_permissions_retrieve(self): - """ - Retreive saved user group permissions. - - Returns - ------- - dict or None - Returns dictionary with saved user group permissions or ``None`` if no permissions are saved. - """ - ugp_json = await self._r_pool.get(self._name_user_group_permissions) - return json.loads(ugp_json) if ugp_json else None - - # ============================================================================================= - # Methods for saving and retrieving lock info. - - async def lock_info_clear(self): - """ - Clear lock info saved in Redis. - """ - await self._r_pool.delete(self._name_lock_info) - - async def lock_info_save(self, lock_info): - """ - Save lock info to Redis. - - Parameters - ---------- - lock_info: dict - A dictionary containing lock info. - """ - await self._r_pool.set(self._name_lock_info, json.dumps(lock_info)) - - async def lock_info_retrieve(self): - """ - Retreive saved lock info. - - Returns - ------- - dict or None - Returns dictionary with saved lock info or ``None`` if no lock info is saved. - """ - lock_info_json = await self._r_pool.get(self._name_lock_info) - return json.loads(lock_info_json) if lock_info_json else None - - # ============================================================================================= - # Methods for saving and retrieving 'queue_stop_pending'. - - async def stop_pending_clear(self): - """ - Clear 'stop_pending' mode info info saved in Redis. - """ - await self._r_pool.delete(self._name_stop_pending_info) - - async def stop_pending_save(self, stop_pending): - """ - Save 'stop_pending' mode info to Redis. - - Parameters - ---------- - lock_info: dict - A dictionary containing lock info. - """ - await self._r_pool.set(self._name_stop_pending_info, json.dumps(stop_pending)) - - async def stop_pending_retrieve(self): - """ - Retreive saved 'stop_pending' mode info. - - Returns - ------- - dict or None - Returns dictionary with saved 'stop_pending' or ``None`` if no 'stop_pending' is saved. - """ - stop_pending_json = await self._r_pool.get(self._name_stop_pending_info) - return json.loads(stop_pending_json) if stop_pending_json else None - - # ============================================================================================= - # Methods for saving and retrieving 'autostart' mode. - - async def autostart_mode_clear(self): - """ - Clear 'autostart' mode info info saved in Redis. - """ - await self._r_pool.delete(self._name_autostart_mode_info) - - async def autostart_mode_save(self, lock_info): - """ - Save 'autostart' mode info to Redis. - - Parameters - ---------- - lock_info: dict - A dictionary containing lock info. - """ - await self._r_pool.set(self._name_autostart_mode_info, json.dumps(lock_info)) - - async def autostart_mode_retrieve(self): - """ - Retreive saved 'autostart' mode info. - - Returns - ------- - dict or None - Returns dictionary with saved lock info or ``None`` if no lock info is saved. - """ - lock_info_json = await self._r_pool.get(self._name_autostart_mode_info) - return json.loads(lock_info_json) if lock_info_json else None + async def stop(self): + """Close the backend connection.""" + await self._backend.stop() \ No newline at end of file diff --git a/bluesky_queueserver/manager/plan_queue_ops_backends/__init__.py b/bluesky_queueserver/manager/plan_queue_ops_backends/__init__.py new file mode 100644 index 00000000..250d307b --- /dev/null +++ b/bluesky_queueserver/manager/plan_queue_ops_backends/__init__.py @@ -0,0 +1,89 @@ +""" +Plan Queue Operations Backend Implementations. + +This package contains the various backend implementations for the Plan Queue Operations: +- Redis: Uses Redis as the backend storage +- SQLite: Uses SQLite database as the backend storage +- PostgreSQL: Uses PostgreSQL database as the backend storage +- Dict: Uses an in-memory dictionary (optionally persisted to file) as the backend storage +""" +import os + +# Dictionary mapping backend names to their import paths +_BACKENDS = { + "abstract": ".plan_queue_ops_abstract.AbstractPlanQueueOperations", + "redis": ".plan_queue_ops_redis.RedisPlanQueueOperations", + "sqlite": ".plan_queue_ops_sqlite.SQLitePlanQueueOperations", + "dict": ".plan_queue_ops_dict.DictPlanQueueOperations", + "postgresql": ".plan_queue_ops_postgresql.PostgreSQLPlanQueueOperations", +} + +def get_backend(backend_name=None): + """ + Dynamically import and return the requested backend class. + + Parameters + ---------- + backend_name : str, optional + Name of the backend ('redis', 'sqlite', 'dict', 'postgresql') + If None, will look for PLAN_QUEUE_BACKEND environment variable, + and default to 'redis' if not found. + + Returns + ------- + class + The requested backend class + + Raises + ------ + ValueError + If the backend name is not recognized + ImportError + If the backend cannot be imported (e.g., missing dependencies) + """ + # If no backend specified, check environment variable + if backend_name is None: + backend_name = os.getenv("PLAN_QUEUE_BACKEND", "redis").lower() + + if backend_name not in _BACKENDS: + raise ValueError(f"Unknown backend: {backend_name}. Available backends: {list(_BACKENDS.keys())}") + + import importlib + module_path, class_name = _BACKENDS[backend_name].rsplit('.', 1) + module = importlib.import_module(module_path, package=__package__) + return getattr(module, class_name) + +def get_default_backend(): + """ + Get the default backend class based on PLAN_QUEUE_BACKEND environment variable. + Defaults to Redis if no environment variable is set. + + Returns + ------- + class + The default backend class + """ + return get_backend() + +# For compatibility with old code that expects direct imports +def __getattr__(name): + """ + Support for direct attribute access like: + from plan_queue_ops_backends import RedisPlanQueueOperations + """ + # Map class names to backend names + class_to_backend = { + "AbstractPlanQueueOperations": "abstract", + "RedisPlanQueueOperations": "redis", + "SQLitePlanQueueOperations": "sqlite", + "DictPlanQueueOperations": "dict", + "PostgreSQLPlanQueueOperations": "postgresql", + } + + if name in class_to_backend: + try: + return get_backend(class_to_backend[name]) + except ImportError: + raise ImportError(f"Could not import {name} - backend dependencies may be missing") + + raise AttributeError(f"module {__name__} has no attribute {name}") \ No newline at end of file diff --git a/bluesky_queueserver/manager/plan_queue_ops_backends/config.py b/bluesky_queueserver/manager/plan_queue_ops_backends/config.py new file mode 100644 index 00000000..736e53f1 --- /dev/null +++ b/bluesky_queueserver/manager/plan_queue_ops_backends/config.py @@ -0,0 +1,94 @@ +""" +Configuration settings for Plan Queue Operation backends. + +This module contains default configuration values for all supported backends. +These settings can be overridden by passing keyword arguments to the backend constructors. +""" + +import os + +# Common configuration shared by all backends +COMMON_CONFIG = { + "name_prefix": "qs_default", # Prefix for database tables/keys +} + +# Redis backend configuration +REDIS_CONFIG = { + "host": "localhost", + "port": 6379, + "db": 0, + "password": None, + "name_prefix": COMMON_CONFIG["name_prefix"], + "decode_responses": True, # Essential for proper string handling +} + +# SQLite backend configuration +SQLITE_CONFIG = { + "database": os.path.join(os.getcwd(), "queue_storage.sqlite"), + "timeout": 5.0, # Connection timeout in seconds + "detect_types": 0, + "isolation_level": None, # Autocommit mode + "check_same_thread": False, # Allow access from multiple threads + "name_prefix": COMMON_CONFIG["name_prefix"], +} + +# PostgreSQL backend configuration +POSTGRESQL_CONFIG = { + "host": "localhost", + "port": 5432, + "database": "bluesky_queue", + "user": "postgres", + "password": "postgres", + "min_size": 1, # Minimum connection pool size + "max_size": 10, # Maximum connection pool size + "timeout": 10.0, # Connection timeout in seconds + "name_prefix": COMMON_CONFIG["name_prefix"], +} + +# Dict backend configuration (in-memory Python dictionary, optionally backed by file) +DICT_CONFIG = { + "in_memory": False, # Whether to use in-memory only or persist to file + "storage_file": None, # Path to storage file (if None, uses default based on name_prefix) + "name_prefix": COMMON_CONFIG["name_prefix"], + "autosave": True, # Whether to automatically save changes to file + "autosave_interval": 1.0, # Seconds between autosaves (if autosave is enabled) +} + +# Backend selection +DEFAULT_BACKEND = "redis" # Default backend if not specified in environment + +def get_backend_config(backend_name=None): + """ + Get configuration for the specified backend. + + Parameters + ---------- + backend_name : str, optional + Name of the backend ('redis', 'sqlite', 'postgresql', 'dict') + If None, get from PLAN_QUEUE_BACKEND environment variable or default + + Returns + ------- + dict + Configuration dictionary for the specified backend + + Raises + ------ + ValueError + If the backend name is not recognized + """ + # Determine backend name + if backend_name is None: + backend_name = os.getenv("PLAN_QUEUE_BACKEND", DEFAULT_BACKEND).lower() + + # Return appropriate config + if backend_name == "redis": + return dict(REDIS_CONFIG) + elif backend_name == "sqlite": + return dict(SQLITE_CONFIG) + elif backend_name == "postgresql": + return dict(POSTGRESQL_CONFIG) + elif backend_name == "dict": + return dict(DICT_CONFIG) + else: + raise ValueError(f"Unknown backend: {backend_name}") \ No newline at end of file diff --git a/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_abstract.py b/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_abstract.py new file mode 100644 index 00000000..75c89fb2 --- /dev/null +++ b/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_abstract.py @@ -0,0 +1,874 @@ +""" +Abstract base class for plan queue operations. + +This module defines the required interface contract that all backend implementations +must follow. This ensures consistency across different storage backends (Redis, SQLite, +PostgreSQL, Dict). +""" + +import abc +import asyncio +import copy +import uuid +from typing import Any, Dict, List, Optional, Tuple, Union + + +class AbstractPlanQueueOperations(abc.ABC): + """ + Abstract base class defining the interface for plan queue operations. + + All backend implementations must inherit from this class and implement + all abstract methods and properties. + """ + + def __init__(self): + """ + Initialize the base class. + """ + # Initialize UID dictionary + self._uid_dict = {} + self._plan_queue_uid = self.new_item_uid() + self._plan_history_uid = self.new_item_uid() + + @staticmethod + def new_item_uid() -> str: + """ + Generate a new unique ID for an item. + + Returns + ------- + str + A new unique ID + """ + return str(uuid.uuid4()) + + @property + @abc.abstractmethod + def plan_queue_uid(self) -> str: + """ + Get the current plan queue UID. + + Returns + ------- + str + The current plan queue UID + """ + pass + + @property + @abc.abstractmethod + def plan_history_uid(self) -> str: + """ + Get the current plan history UID. + + Returns + ------- + str + The current plan history UID + """ + pass + + @property + @abc.abstractmethod + def plan_queue_mode(self) -> Dict[str, Any]: + """ + Get the current plan queue mode. + + Returns + ------- + dict + The current plan queue mode + """ + pass + + @abc.abstractmethod + async def start(self) -> None: + """ + Start the plan queue operations. + + This method should initialize any necessary connections and resources. + """ + pass + + @abc.abstractmethod + async def stop(self) -> None: + """ + Stop the plan queue operations. + + This method should clean up any resources and close connections. + """ + pass + + @abc.abstractmethod + async def add_item_to_queue( + self, + item: Dict[str, Any], + *, + pos: Optional[Union[int, str]] = None, + before_uid: Optional[str] = None, + after_uid: Optional[str] = None, + filter_parameters: bool = True + ) -> Tuple[Dict[str, Any], int]: + """ + Add a single item to the queue. + + Parameters + ---------- + item : dict + Item to add to the queue + pos : int or str, optional + Position to add the item at. Can be an integer or 'front'/'back'. + If None, the item is added to the back of the queue. + before_uid : str, optional + UID of the item to add the new item before + after_uid : str, optional + UID of the item to add the new item after + filter_parameters : bool, default True + Whether to filter the item parameters + + Returns + ------- + tuple + Tuple of (item, queue_size) + + Raises + ------ + ValueError + If the position specification is ambiguous + """ + pass + + @abc.abstractmethod + async def add_batch_to_queue( + self, + items: List[Dict[str, Any]], + *, + pos: Optional[Union[int, str]] = None, + before_uid: Optional[str] = None, + after_uid: Optional[str] = None, + filter_parameters: bool = True + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], int, bool]: + """ + Add a batch of items to the queue. + + Parameters + ---------- + items : list of dict + Items to add to the queue + pos : int or str, optional + Position to add the items at. Can be an integer or 'front'/'back'. + If None, the items are added to the back of the queue. + before_uid : str, optional + UID of the item to add the new items before + after_uid : str, optional + UID of the item to add the new items after + filter_parameters : bool, default True + Whether to filter the item parameters + + Returns + ------- + tuple + Tuple of (items_added, results, queue_size, success) + + Raises + ------ + ValueError + If the position specification is ambiguous + """ + pass + + @abc.abstractmethod + async def get_queue(self) -> Tuple[List[Dict[str, Any]], str]: + """ + Get the current queue. + + Returns + ------- + tuple + Tuple of (queue, queue_uid) + """ + pass + + @abc.abstractmethod + async def get_queue_size(self) -> int: + """ + Get the current queue size. + + Returns + ------- + int + The current queue size + """ + pass + + @abc.abstractmethod + async def clear_queue(self) -> int: + """ + Clear the queue. + + Returns + ------- + int + The number of items removed from the queue + """ + pass + + @abc.abstractmethod + async def pop_item_from_queue( + self, + *, + pos: Optional[int] = None, + uid: Optional[str] = None + ) -> Dict[str, Any]: + """ + Remove an item from the queue. + + Parameters + ---------- + pos : int, optional + Position of the item to remove + uid : str, optional + UID of the item to remove + + Returns + ------- + dict + The removed item + + Raises + ------ + IndexError + If the position is invalid + ValueError + If both pos and uid are specified or neither is specified + """ + pass + + @abc.abstractmethod + async def pop_item_from_queue_batch(self, *, uids: List[str]) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + """ + Remove multiple items from the queue. + + Parameters + ---------- + uids : list of str + UIDs of the items to remove + + Returns + ------- + tuple + Tuple of (removed_items, remaining_queue) + """ + pass + + @abc.abstractmethod + async def set_plan_queue_mode( + self, + plan_queue_mode: Union[Dict[str, Any], str], + *, + update: bool = False + ) -> Dict[str, Any]: + """ + Set the plan queue mode. + + Parameters + ---------- + plan_queue_mode : dict or str + The new plan queue mode or 'default' to reset to default mode + update : bool, default False + If True, update only the specified fields in the mode dictionary + + Returns + ------- + dict + The updated plan queue mode + """ + pass + + @abc.abstractmethod + async def move_item( + self, + *, + pos: int = None, + uid: str = None, + pos_dest: int = None, + before_uid: str = None, + after_uid: str = None + ) -> Dict[str, Any]: + """ + Move an item in the queue. + + Parameters + ---------- + pos : int, optional + Position of the item to move + uid : str, optional + UID of the item to move + pos_dest : int, optional + Destination position + before_uid : str, optional + UID of the item to move the item before + after_uid : str, optional + UID of the item to move the item after + + Returns + ------- + dict + The moved item + + Raises + ------ + ValueError + If the source or destination specification is ambiguous + IndexError + If the position is invalid + """ + pass + + @abc.abstractmethod + async def swap_items( + self, + *, + pos1: int = None, + uid1: str = None, + pos2: int = None, + uid2: str = None + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """ + Swap two items in the queue. + + Parameters + ---------- + pos1 : int, optional + Position of the first item + uid1 : str, optional + UID of the first item + pos2 : int, optional + Position of the second item + uid2 : str, optional + UID of the second item + + Returns + ------- + tuple + Tuple of (item1, item2) + + Raises + ------ + ValueError + If the item specification is ambiguous + IndexError + If the position is invalid + """ + pass + + @abc.abstractmethod + async def get_item(self, *, pos: int = None, uid: str = None) -> Dict[str, Any]: + """ + Get an item from the queue without removing it. + + Parameters + ---------- + pos : int, optional + Position of the item + uid : str, optional + UID of the item + + Returns + ------- + dict + The item + + Raises + ------ + ValueError + If both pos and uid are specified or neither is specified + IndexError + If the position is invalid + """ + pass + + @abc.abstractmethod + async def update_item( + self, + item: Dict[str, Any], + *, + pos: int = None, + uid: str = None + ) -> Dict[str, Any]: + """ + Update an item in the queue. + + Parameters + ---------- + item : dict + New item data + pos : int, optional + Position of the item to update + uid : str, optional + UID of the item to update + + Returns + ------- + dict + The updated item + + Raises + ------ + ValueError + If both pos and uid are specified or neither is specified + IndexError + If the position is invalid + """ + pass + + @abc.abstractmethod + async def is_item_running(self) -> bool: + """ + Check if an item is currently running. + + Returns + ------- + bool + True if an item is running, False otherwise + """ + pass + + @abc.abstractmethod + async def process_next_item(self) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """ + Process the next item in the queue. + + This method typically pops the first item from the queue. + + Returns + ------- + tuple + Tuple of (next_item, running_item) + """ + pass + + @abc.abstractmethod + async def set_next_item_as_running(self, *, item: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + """ + Set the next item in the queue as the running item. + + Parameters + ---------- + item : dict, optional + Item to set as running. If None, the next item in the queue is used. + + Returns + ------- + dict + The item set as running + """ + pass + + @abc.abstractmethod + async def set_processed_item_as_completed( + self, + *, + exit_status: str, + run_uids: Optional[List[str]] = None, + scan_ids: Optional[List[int]] = None, + err_msg: str = "", + err_tb: str = "" + ) -> Dict[str, Any]: + """ + Mark the currently running item as completed and move it to history. + + Parameters + ---------- + exit_status : str + Exit status of the processed item + run_uids : list of str, optional + List of run UIDs + scan_ids : list of int, optional + List of scan IDs + err_msg : str, optional + Error message + err_tb : str, optional + Error traceback + + Returns + ------- + dict + The completed item + """ + pass + + @abc.abstractmethod + async def set_processed_item_as_stopped( + self, + *, + exit_status: str, + run_uids: Optional[List[str]] = None, + scan_ids: Optional[List[int]] = None, + err_msg: str = "", + err_tb: str = "" + ) -> Dict[str, Any]: + """ + Mark the currently running item as stopped and either move it to history + or push it back to the front of the queue depending on exit status. + + Parameters + ---------- + exit_status : str + Exit status of the processed item + run_uids : list of str, optional + List of run UIDs + scan_ids : list of int, optional + List of scan IDs + err_msg : str, optional + Error message + err_tb : str, optional + Error traceback + + Returns + ------- + dict + The stopped item + """ + pass + + @abc.abstractmethod + async def clear_running_item(self) -> Dict[str, Any]: + """ + Clear the currently running item. + + Returns + ------- + dict + The previously running item + """ + pass + + @abc.abstractmethod + async def get_running_item(self) -> Dict[str, Any]: + """ + Get the currently running item. + + Returns + ------- + dict + The currently running item + """ + pass + + @abc.abstractmethod + async def get_history(self) -> Tuple[List[Dict[str, Any]], str]: + """ + Get the plan history. + + Returns + ------- + tuple + Tuple of (history, history_uid) + """ + pass + + @abc.abstractmethod + async def clear_history(self) -> int: + """ + Clear the plan history. + + Returns + ------- + int + The number of items removed from history + """ + pass + + @abc.abstractmethod + async def user_group_permissions_get(self) -> Dict[str, Any]: + """ + Get the user group permissions. + + Returns + ------- + dict + User group permissions + """ + pass + + @abc.abstractmethod + async def user_group_permissions_set(self, user_group_permissions: Dict[str, Any]) -> Dict[str, Any]: + """ + Set the user group permissions. + + Parameters + ---------- + user_group_permissions : dict + User group permissions + + Returns + ------- + dict + Updated user group permissions + """ + pass + + @abc.abstractmethod + async def lock_info_get(self) -> Dict[str, Any]: + """ + Get the lock information. + + Returns + ------- + dict + Lock information + """ + pass + + @abc.abstractmethod + async def lock_info_set(self, lock_info: Dict[str, Any]) -> Dict[str, Any]: + """ + Set the lock information. + + Parameters + ---------- + lock_info : dict + Lock information + + Returns + ------- + dict + Updated lock information + """ + pass + + @abc.abstractmethod + async def autostart_mode_get(self) -> Dict[str, Any]: + """ + Get the autostart mode information. + + Returns + ------- + dict + Autostart mode information + """ + pass + + @abc.abstractmethod + async def autostart_mode_set(self, autostart_mode_info: Dict[str, Any]) -> Dict[str, Any]: + """ + Set the autostart mode information. + + Parameters + ---------- + autostart_mode_info : dict + Autostart mode information + + Returns + ------- + dict + Updated autostart mode information + """ + pass + + @abc.abstractmethod + async def stop_pending_info_get(self) -> Dict[str, Any]: + """ + Get the stop pending information. + + Returns + ------- + dict + Stop pending information + """ + pass + + @abc.abstractmethod + async def stop_pending_info_set(self, stop_pending_info: Dict[str, Any]) -> Dict[str, Any]: + """ + Set the stop pending information. + + Parameters + ---------- + stop_pending_info : dict + Stop pending information + + Returns + ------- + dict + Updated stop pending information + """ + pass + + + # Common utility methods + + async def _verify_item_type(self, item: Dict[str, Any]) -> None: + """ + Verify that item is a dictionary. + + Parameters + ---------- + item : dict + Item to verify + + Raises + ------ + TypeError + If the item is not a dictionary + """ + if not isinstance(item, dict): + raise TypeError(f"Parameter 'item' should be a dictionary, got {type(item)}") + + async def _verify_item(self, item: Dict[str, Any], *, ignore_uids: Optional[List[str]] = None) -> None: + """ + Verify that item structure is valid enough to be put in the queue. + + Parameters + ---------- + item : dict + Item to verify + ignore_uids : list of str, optional + List of UIDs to ignore when checking for duplicates + + Raises + ------ + ValueError + If the item is invalid + RuntimeError + If the item UID is already in the queue + """ + # Verify item is a dictionary + await self._verify_item_type(item) + + # Ensure ignore_uids is a list + ignore_uids = ignore_uids or [] + + # Check for UID + if "item_uid" not in item: + raise ValueError("Item does not have UID.") + + # Check for valid name + if "name" not in item or not item["name"]: + raise ValueError("Item must have a valid name") + + # Check if UID is already in the queue + uid = item["item_uid"] + if uid not in ignore_uids and self._is_uid_in_dict(uid): + raise RuntimeError(f"Item with UID {uid} is already in the queue") + + async def set_new_item_uuid(self, item: Dict[str, Any]) -> Dict[str, Any]: + """ + Set a new UUID for an item. + + Parameters + ---------- + item : dict + Item to set UUID for + + Returns + ------- + dict + Copy of the item with a new UUID + """ + # Create a deep copy to avoid modifying the original + item_copy = copy.deepcopy(item) + item_copy["item_uid"] = self.new_item_uid() + return item_copy + + async def _clean_item_properties(self, item: Dict[str, Any]) -> Dict[str, Any]: + """ + Clean the item properties. + + Parameters + ---------- + item : dict + Item to clean + + Returns + ------- + dict + Cleaned item + """ + # Create a deep copy to avoid modifying the original + cleaned_item = copy.deepcopy(item) + + # If there's a properties section, clean it + if "properties" in cleaned_item: + # Create a clean properties dictionary with only permitted fields + allowed_properties = [ + "time_submit", + "time_start", + "immediate_execution" + ] + + # Extract only allowed properties + permitted_properties = { + k: v for k, v in cleaned_item["properties"].items() + if k in allowed_properties + } + + # Replace with cleaned properties + cleaned_item["properties"] = permitted_properties + + return cleaned_item + + # -------------------------------------------------------------------------- + # Operations with UID set + def _uid_dict_clear(self): + """ + Clear ``self._uid_dict``. + """ + self._plan_queue_uid = self.new_item_uid() + self._uid_dict.clear() + + def _is_uid_in_dict(self, uid): + """ + Checks if UID exists in ``self._uid_dict``. + """ + return uid in self._uid_dict + + def _uid_dict_add(self, item): + """ + Add UID to ``self._uid_dict``. + """ + uid = item["item_uid"] + if self._is_uid_in_dict(uid): + raise RuntimeError(f"Trying to add plan with UID '{uid}', which is already in the queue") + self._plan_queue_uid = self.new_item_uid() + self._uid_dict.update({uid: item}) + + def _uid_dict_remove(self, uid): + """ + Remove UID from ``self._uid_dict``. + """ + if not self._is_uid_in_dict(uid): + raise RuntimeError(f"Trying to remove plan with UID '{uid}', which is not in the queue") + self._plan_queue_uid = self.new_item_uid() + self._uid_dict.pop(uid) + + def _uid_dict_update(self, item): + """ + Update a plan with UID that is already in the dictionary. + """ + uid = item["item_uid"] + if not self._is_uid_in_dict(uid): + raise RuntimeError(f"Trying to update plan with UID '{uid}', which is not in the queue") + self._plan_queue_uid = self.new_item_uid() + self._uid_dict.update({uid: item}) + + def _uid_dict_get_item(self, uid): + """ + Returns a plan with the given UID. + """ + return self._uid_dict[uid] + + async def _uid_dict_initialize(self): + """ + Initialize ``self._uid_dict`` with UIDs extracted from the plans in the queue. + """ + pq, _ = await self.get_queue() # Changed from _get_queue to get_queue + self._uid_dict_clear() + # Go over all plans in the queue + for item in pq: + self._uid_dict_add(item) + # If plan is currently running + item = await self.get_running_item() # Changed from _get_running_item_info + if item: + self._uid_dict_add(item) diff --git a/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_dict.py b/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_dict.py new file mode 100644 index 00000000..ebc59aeb --- /dev/null +++ b/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_dict.py @@ -0,0 +1,458 @@ +import asyncio +import aiofiles +import os +import json +import logging +import time as ttime +import copy +from typing import Any, Dict, List, Optional, Tuple, Union +from bluesky_queueserver.manager.plan_queue_ops_backends.plan_queue_ops_abstract import AbstractPlanQueueOperations + +logger = logging.getLogger(__name__) + +class DictPlanQueueOperations(AbstractPlanQueueOperations): + """ + Backend implementation for plan queue operations using an in-memory Python dictionary. + Data is saved and loaded to/from storage using JSON serialization with aiofiles. + """ + + def __init__(self, storage_file=None, in_memory=False, name_prefix="qs_default"): + """ + Initialize the DictPlanQueueOperations class. + """ + # Initialize the parent class + super().__init__() + + self._in_memory = in_memory + self._name_prefix = name_prefix # Store name_prefix for consistency with other backends + self._uid_dict = dict() # Initialize UID dictionary for tracking + + if not in_memory: + if not storage_file: + # Use name_prefix in auto-generated filename + storage_file = os.getenv( + "BACKEND_DB_PATH", + os.path.join(os.getcwd(), f"{name_prefix}_plan_queue_storage.json") + ) + self._storage_file = storage_file + else: + self._storage_file = None # No file is used for in-memory mode + + self._data = { + "plan_queue": [], + "plan_history": [], + "running_plan": None, + "plan_queue_mode": {"loop": False, "ignore_failures": False}, + "user_group_permissions": None, + "lock_info": None, + "autostart_mode_info": None, + "stop_pending_info": None, + } + self._plan_queue_uid = self.new_item_uid() + self._plan_history_uid = self.new_item_uid() + + # The list of allowed item parameters used for parameter filtering + self._allowed_item_parameters = ( + "item_uid", + "item_type", + "name", + "args", + "kwargs", + "meta", + "user", + "user_group", + "properties", + ) + + self._lock = asyncio.Lock() + + # UID Dictionary Operations + def _uid_dict_add(self, item): + """Add an item to the UID dictionary.""" + if "item_uid" in item: + self._uid_dict[item["item_uid"]] = item + + def _uid_dict_remove(self, uid): + """Remove an item from the UID dictionary.""" + if uid in self._uid_dict: + del self._uid_dict[uid] + + def _uid_dict_clear(self): + """Clear the UID dictionary.""" + self._uid_dict.clear() + + def _uid_dict_get_item(self, uid): + """Get an item from the UID dictionary.""" + return self._uid_dict.get(uid) + + def _is_uid_in_dict(self, uid): + """Check if a UID exists in the dictionary.""" + return uid in self._uid_dict + + async def _verify_item_type(self, item): + """Verify that item is a dictionary.""" + if not isinstance(item, dict): + raise TypeError(f"Parameter 'item' should be a dictionary, got {type(item)}") + + async def _verify_item(self, item, *, ignore_uids=None): + """ + Verify that item structure is valid enough to be put in the queue. + """ + ignore_uids = ignore_uids or [] + await self._verify_item_type(item) + + # First check for UID (original behavior that tests expect) + if "item_uid" not in item: + raise ValueError("Item does not have UID.") + + # Then check if UID is already in use + uid = item["item_uid"] + if (uid not in ignore_uids) and self._is_uid_in_dict(uid): + raise RuntimeError(f"Item with UID {uid} is already in the queue") + + # Only then check name (this won't be reached if UID check fails) + if "name" not in item or not item["name"]: + raise ValueError("Item must have a valid name") + + async def set_new_item_uuid(self, item): + """ + Set a new UUID for an item. + + Parameters + ---------- + item : dict + Item to set UUID for + + Returns + ------- + dict + Copy of the item with a new UUID + """ + item_copy = copy.deepcopy(item) + item_copy["item_uid"] = self.new_item_uid() + return item_copy + + async def _save_to_storage(self): + """Save the data to storage file if not in memory mode.""" + if not self._in_memory and self._storage_file: + try: + async with aiofiles.open(self._storage_file, 'w') as f: + await f.write(json.dumps(self._data)) + except Exception as ex: + logger.error(f"Failed to save data to storage file: {ex}") + + async def _load_from_storage(self): + """Load the data from storage file if it exists.""" + if not self._in_memory and self._storage_file and os.path.exists(self._storage_file): + try: + async with aiofiles.open(self._storage_file, 'r') as f: + content = await f.read() + if content: + self._data = json.loads(content) + except Exception as ex: + logger.error(f"Failed to load data from storage file: {ex}") + + def filter_item_parameters(self, item): + """ + Remove parameters that are not in the list of allowed parameters. + + Parameters + ---------- + item : dict + dictionary of item parameters + + Returns + ------- + dict + dictionary of filtered item parameters + """ + return {k: v for k, v in item.items() if k in self._allowed_item_parameters} + + # -------------------------------------------------------------------------- + # Queue Operations + async def add_item_to_queue(self, item, *, pos=None, before_uid=None, after_uid=None, filter_parameters=True): + """Add a single item to the queue.""" + async with self._lock: + if (pos is not None) and (before_uid is not None or after_uid is not None): + raise ValueError("Ambiguous parameters: plan position and UID is specified.") + + if (before_uid is not None) and (after_uid is not None): + raise ValueError( + "Ambiguous parameters: request to insert the plan before and after the reference plan." + ) + + pos = pos if pos is not None else "back" + + if "item_uid" not in item: + item = await self.set_new_item_uuid(item) # Added await here + else: + await self._verify_item(item) + + if filter_parameters: + item = self.filter_item_parameters(item) + + qsize = len(self._data["plan_queue"]) + + if isinstance(pos, int): + if (pos <= 0) or (pos < -qsize): + pos = "front" + elif (pos >= qsize) or (pos == -1): + pos = "back" + + if (before_uid is not None) or (after_uid is not None): + uid = before_uid if before_uid is not None else after_uid + before = uid == before_uid + + if not self._is_uid_in_dict(uid): + raise IndexError(f"Item with UID '{uid}' not found in the queue.") + + # Find the item with the specified UID + idx = next((i for i, q in enumerate(self._data["plan_queue"]) if q["item_uid"] == uid), None) + if idx is None: + raise IndexError(f"Item with UID '{uid}' not found in the queue.") + + insert_pos = idx if before else idx + 1 + self._data["plan_queue"].insert(insert_pos, item) + elif pos == "front": + self._data["plan_queue"].insert(0, item) + elif pos == "back": + self._data["plan_queue"].append(item) + else: # pos is an integer + self._data["plan_queue"].insert(pos, item) + + # Add item to UID dictionary + self._uid_dict_add(item) + self._plan_queue_uid = self.new_item_uid() + + await self._save_to_storage() + return item, len(self._data["plan_queue"]) + + async def add_batch_to_queue(self, items, *, pos=None, before_uid=None, after_uid=None, filter_parameters=True): + """Add a batch of items to the queue.""" + async with self._lock: + items_added, results = [], [] + + # Approach: attempt to add each item to queue. In case of a failure, remove all added items. + success = True + added_item_uids = [] # List of successfully added UIDs (used for rollback) + + for item in items: + try: + if not added_item_uids: + item_added, _ = await self.add_item_to_queue( + item, + pos=pos, + before_uid=before_uid, + after_uid=after_uid, + filter_parameters=filter_parameters, + ) + else: + item_added, _ = await self.add_item_to_queue( + item, after_uid=added_item_uids[-1], filter_parameters=filter_parameters + ) + added_item_uids.append(item_added["item_uid"]) + items_added.append(item_added) + results.append({"success": True, "msg": ""}) + except Exception as ex: + success = False + items_added.append(item) + msg = f"Failed to add item to queue: {ex}" + results.append({"success": False, "msg": msg}) + break # Stop processing more items if one fails + + # 'Undo' operation: remove all inserted items if there was a failure + if not success: + for uid in added_item_uids: + try: + await self.pop_item_from_queue(uid=uid) + except Exception as ex: + logger.error( + f"Failed to remove an item with UID '{uid}' after failure to add a batch of plans: {ex}" + ) + # Return the original items if the batch operation fails + items_added = items + + qsize = len(self._data["plan_queue"]) + return items_added, results, qsize, success + + async def set_processed_item_as_completed(self, *, exit_status, run_uids=None, scan_ids=None, err_msg="", err_tb=""): + """Mark the currently running item as completed and move it to history.""" + async with self._lock: + # Default empty lists for run_uids and scan_ids + run_uids = run_uids or [] + scan_ids = scan_ids or [] + + # Check if loop mode is enabled + loop_enabled = self._data["plan_queue_mode"].get("loop", False) + + running_item = self._data["running_plan"] + if running_item: + immediate_execution = running_item.get("properties", {}).get("immediate_execution", False) + item_time_start = running_item.get("properties", {}).get("time_start", ttime.time()) + item_cleaned = await self._clean_item_properties(running_item) # Added await here + + if loop_enabled and not immediate_execution: + # Add a copy of the item back to the queue + item_to_add = copy.deepcopy(item_cleaned) + item_to_add = await self.set_new_item_uuid(item_to_add) # Added await here + self._data["plan_queue"].append(item_to_add) + self._uid_dict_add(item_to_add) + + # Add result information + item_cleaned.setdefault("result", {}) + item_cleaned["result"]["exit_status"] = exit_status + item_cleaned["result"]["run_uids"] = run_uids + item_cleaned["result"]["scan_ids"] = scan_ids + item_cleaned["result"]["time_start"] = item_time_start + item_cleaned["result"]["time_stop"] = ttime.time() + item_cleaned["result"]["msg"] = err_msg + item_cleaned["result"]["traceback"] = err_tb + + # Clear running item + self._data["running_plan"] = None + + # Remove from UID dictionary if not in loop mode and not immediate execution + if not loop_enabled and not immediate_execution: + self._uid_dict_remove(running_item["item_uid"]) + + # Add to history + self._data["plan_history"].append(item_cleaned) + self._plan_history_uid = self.new_item_uid() + self._plan_queue_uid = self.new_item_uid() + + await self._save_to_storage() + return item_cleaned + else: + return {} + + async def set_processed_item_as_stopped(self, *, exit_status, run_uids=None, scan_ids=None, err_msg="", err_tb=""): + """ + Mark the currently running item as stopped and either move it to history + or push it back to the front of the queue depending on exit status. + """ + async with self._lock: + run_uids = run_uids or [] + scan_ids = scan_ids or [] + + if exit_status == "stopped": + # Stopped item is considered successful, so it is not pushed back to the queue + return await self.set_processed_item_as_completed( + exit_status=exit_status, run_uids=run_uids, scan_ids=scan_ids, err_msg=err_msg, err_tb=err_tb + ) + + running_item = self._data["running_plan"] + if running_item: + immediate_execution = running_item.get("properties", {}).get("immediate_execution", False) + item_time_start = running_item.get("properties", {}).get("time_start", ttime.time()) + item_cleaned = await self._clean_item_properties(running_item) # Added await here + + # Add result information + item_cleaned.setdefault("result", {}) + item_cleaned["result"]["exit_status"] = exit_status + item_cleaned["result"]["run_uids"] = run_uids + item_cleaned["result"]["scan_ids"] = scan_ids + item_cleaned["result"]["time_start"] = item_time_start + item_cleaned["result"]["time_stop"] = ttime.time() + item_cleaned["result"]["msg"] = err_msg + item_cleaned["result"]["traceback"] = err_tb + + # Add to history + self._data["plan_history"].append(item_cleaned) + self._plan_history_uid = self.new_item_uid() + + # Clear running item + self._data["running_plan"] = None + + # Generate new UID for the item that is pushed back into the queue + if not immediate_execution: + self._uid_dict_remove(running_item["item_uid"]) + # Only push back to queue for non-stopped statuses + if exit_status != "stopped": + item_copy = copy.deepcopy(item_cleaned) + if "result" in item_copy: + del item_copy["result"] + item_pushed_to_queue = await self.set_new_item_uuid(item_copy) # Added await here + self._data["plan_queue"].insert(0, item_pushed_to_queue) + self._uid_dict_add(item_pushed_to_queue) + + self._plan_queue_uid = self.new_item_uid() + await self._save_to_storage() + return item_cleaned + else: + return {} + + async def set_plan_queue_mode( + self, + plan_queue_mode: Union[Dict[str, Any], str], + *, + update: bool = False + ) -> Dict[str, Any]: + """ + Set the plan queue mode. + + Parameters + ---------- + plan_queue_mode : dict or str + The new plan queue mode or 'default' to reset to default mode + update : bool, default False + If True, update only the specified fields in the mode dictionary + + Returns + ------- + dict + The updated plan queue mode + """ + async with self._lock: + # Default mode for reference + default_mode = {"loop": False, "ignore_failures": False} + + # Reset to default if requested + if isinstance(plan_queue_mode, str) and plan_queue_mode.lower() == "default": + self._data["plan_queue_mode"] = copy.deepcopy(default_mode) + elif not isinstance(plan_queue_mode, dict): + raise TypeError(f"Parameter 'plan_queue_mode' must be a dictionary or 'default': {plan_queue_mode}") + else: + # Create a new mode or update the current one + if update: + # Verify keys in plan_queue_mode + for key in plan_queue_mode: + if key not in default_mode: + raise ValueError(f"Unsupported plan queue mode parameter '{key}'") + # Update mode + self._data["plan_queue_mode"].update(plan_queue_mode) + else: + # Verify that all required keys are present + missing_keys = set(default_mode.keys()) - set(plan_queue_mode.keys()) + if missing_keys: + raise ValueError(f"Parameters {missing_keys} are missing") + + # Verify keys in plan_queue_mode + for key in plan_queue_mode: + if key not in default_mode: + raise ValueError(f"Unsupported plan queue mode parameter '{key}'") + + self._data["plan_queue_mode"] = copy.deepcopy(plan_queue_mode) + + # Verify value types + for key, value in self._data["plan_queue_mode"].items(): + if key == "loop" and not isinstance(value, bool): + raise TypeError(f"Unsupported type {type(value)} of the parameter 'loop'") + elif key == "ignore_failures" and not isinstance(value, bool): + raise TypeError(f"Unsupported type {type(value)} of the parameter 'ignore_failures'") + + # Save to storage + await self._save_to_storage() + + return copy.deepcopy(self._data["plan_queue_mode"]) + + @property + def plan_queue_mode(self) -> Dict[str, Any]: + """ + Get the current plan queue mode. + + Returns + ------- + dict + The current plan queue mode + """ + return copy.deepcopy(self._data["plan_queue_mode"]) \ No newline at end of file diff --git a/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_postgresql.py b/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_postgresql.py new file mode 100644 index 00000000..86d637ae --- /dev/null +++ b/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_postgresql.py @@ -0,0 +1,1733 @@ +import asyncio +import json +import logging +import time as ttime +import copy +import asyncpg +from typing import Any, Dict, List, Optional, Tuple, Union + +from bluesky_queueserver.manager.plan_queue_ops_backends.plan_queue_ops_abstract import AbstractPlanQueueOperations + + +logger = logging.getLogger(__name__) + +class PostgreSQLPlanQueueOperations(AbstractPlanQueueOperations): + """ + Backend implementation for plan queue operations using PostgreSQL database. + + This class provides persistent storage of queue items, history, and other state + information using a PostgreSQL database. It implements all the methods required + by AbstractPlanQueueOperations and utilizes UIDOperations for consistent handling + of unique identifiers. + + Parameters + ---------- + host: str + PostgreSQL host address + port: int + PostgreSQL port + database: str + PostgreSQL database name + user: str + PostgreSQL username + password: str + PostgreSQL password + name_prefix: str + Prefix for table names to avoid conflicts with other instances + """ + + def __init__( + self, + host: str = "localhost", + port: int = 5432, + database: str = "bluesky_queue", + user: str = "postgres", + password: str = "postgres", + name_prefix: str = "qs_default" + ): + # Initialize parent class + super().__init__() + + # Add these lines to track UIDs consistently + self._plan_queue_uid = self.new_item_uid() + self._plan_history_uid = self.new_item_uid() + + self._connection_params = { + "host": host, + "port": port, + "database": database, + "user": user, + "password": password + } + self._conn = None + self._uid_dict = dict() + + if not isinstance(name_prefix, str): + raise TypeError(f"Parameter 'name_prefix' should be a string: {name_prefix}") + + # Add underscore to prefix if it's not empty + if name_prefix: + name_prefix = name_prefix + "_" + + # Define table names with prefix + self._table_plan_queue = name_prefix + "plan_queue" + self._table_plan_history = name_prefix + "plan_history" + self._table_running_plan = name_prefix + "running_plan" + self._table_lock_info = name_prefix + "lock_info" + self._table_autostart_mode_info = name_prefix + "autostart_mode_info" + self._table_stop_pending_info = name_prefix + "stop_pending_info" + self._table_user_group_permissions = name_prefix + "user_group_permissions" + self._table_plan_queue_mode = name_prefix + "plan_queue_mode" + + # Define default queue mode + self._plan_queue_mode_default = {"loop": False, "ignore_failures": False} + self._plan_queue_mode = self._plan_queue_mode_default + + # The list of allowed item parameters for filtering + self._allowed_item_parameters = ( + "item_uid", + "item_type", + "name", + "args", + "kwargs", + "meta", + "user", + "user_group", + "properties", + ) + + self._lock = asyncio.Lock() + + async def _initialize_tables(self): + """ + Initialize database tables if they don't exist. + """ + async with self._conn.transaction(): + # Create plan queue table + await self._conn.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_plan_queue} ( + id SERIAL PRIMARY KEY, + position INTEGER, + item JSONB + ) + """) + + # Create plan history table + await self._conn.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_plan_history} ( + id SERIAL PRIMARY KEY, + item JSONB, + timestamp TIMESTAMPTZ DEFAULT NOW() + ) + """) + + # Create running plan table - stores at most one record + await self._conn.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_running_plan} ( + id SERIAL PRIMARY KEY, + item JSONB + ) + """) + + # Create lock info table + await self._conn.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_lock_info} ( + id SERIAL PRIMARY KEY, + info JSONB + ) + """) + + # Create autostart mode info table + await self._conn.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_autostart_mode_info} ( + id SERIAL PRIMARY KEY, + info JSONB + ) + """) + + # Create stop pending info table + await self._conn.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_stop_pending_info} ( + id SERIAL PRIMARY KEY, + info JSONB + ) + """) + + # Create user group permissions table + await self._conn.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_user_group_permissions} ( + id SERIAL PRIMARY KEY, + info JSONB + ) + """) + + # Create index on item_uid for efficient lookup + await self._conn.execute(f""" + CREATE INDEX IF NOT EXISTS idx_{self._table_plan_queue}_item_uid + ON {self._table_plan_queue} ((item->>'item_uid')) + """) + + def filter_item_parameters(self, item: Dict[str, Any]) -> Dict[str, Any]: + """ + Remove parameters that are not in the list of allowed parameters. + + Parameters + ---------- + item : dict + Dictionary of item parameters + + Returns + ------- + dict + Dictionary of filtered item parameters + """ + return {k: v for k, v in item.items() if k in self._allowed_item_parameters} + + # -------------------------------------------------------------------------- + # Backend Management + async def start(self) -> None: + """ + Start the PostgreSQL backend connection. + """ + if not self._conn: + try: + self._conn = await asyncpg.connect(**self._connection_params) + await self._initialize_tables() + await self._uid_dict_initialize() + except Exception as ex: + logger.error(f"Failed to connect to PostgreSQL: {ex}") + if self._conn: + await self._conn.close() + self._conn = None + raise + + async def stop(self) -> None: + """ + Stop the PostgreSQL backend connection. + """ + if self._conn: + await self._conn.close() + self._conn = None + + async def reset(self) -> None: + """ + Reset the backend to its initial state. + """ + async with self._lock: + async with self._conn.transaction(): + await self._conn.execute(f"TRUNCATE TABLE {self._table_plan_queue}") + await self._conn.execute(f"TRUNCATE TABLE {self._table_plan_history}") + await self._conn.execute(f"TRUNCATE TABLE {self._table_running_plan}") + await self._conn.execute(f"TRUNCATE TABLE {self._table_lock_info}") + await self._conn.execute(f"TRUNCATE TABLE {self._table_autostart_mode_info}") + await self._conn.execute(f"TRUNCATE TABLE {self._table_stop_pending_info}") + await self._conn.execute(f"TRUNCATE TABLE {self._table_user_group_permissions}") + + self._uid_dict_clear() + self._plan_queue_mode = self._plan_queue_mode_default + + async def delete_pool_entries(self) -> None: + """ + Delete all pool entries used by this backend. + """ + await self.reset() + + # -------------------------------------------------------------------------- + # UID Operations + async def _uid_dict_initialize(self) -> None: + """ + Initialize the UID dictionary from the current state in PostgreSQL. + """ + self._uid_dict_clear() + + # Add all items from plan queue + rows = await self._conn.fetch(f"SELECT item FROM {self._table_plan_queue} ORDER BY position") + for row in rows: + item = json.loads(row['item']) + if "item_uid" in item: + self._uid_dict_add(item) + + # Add running item if it exists + row = await self._conn.fetchrow(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + if row: + item = json.loads(row['item']) + if "item_uid" in item: + self._uid_dict_add(item) + + # -------------------------------------------------------------------------- + # Queue Operations + async def add_item_to_queue( + self, + item: Dict[str, Any], + *, + pos: Optional[Union[int, str]] = None, + before_uid: Optional[str] = None, + after_uid: Optional[str] = None, + filter_parameters: bool = True + ) -> Tuple[Dict[str, Any], int]: + """ + Add an item to the queue. + + Parameters + ---------- + item : dict + Item to add to the queue + pos : int or str, optional + Position in the queue. Can be an integer (0-based index) or string ('front' or 'back') + before_uid : str, optional + UID of the item before which to place the new item + after_uid : str, optional + UID of the item after which to place the new item + filter_parameters : bool, default True + Whether to filter item parameters + + Returns + ------- + tuple + (item, queue_size) - the added item and the new queue size + """ + async with self._lock: + if (pos is not None) and (before_uid is not None or after_uid is not None): + raise ValueError("Ambiguous parameters: plan position and UID is specified.") + + if (before_uid is not None) and (after_uid is not None): + raise ValueError( + "Ambiguous parameters: request to insert the plan before and after the reference plan." + ) + + pos = pos if pos is not None else "back" + + if "item_uid" not in item: + item = await self.set_new_item_uuid(item) + else: + await self._verify_item(item) + + if filter_parameters: + item = self.filter_item_parameters(item) + + async with self._conn.transaction(): + # Get current queue size + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + + # Determine the position + position = None + if isinstance(pos, int): + if pos <= 0: + position = 0 + elif pos >= qsize: + position = qsize + else: + position = pos + elif pos == "front": + position = 0 + elif pos == "back": + position = qsize + elif before_uid or after_uid: + # Find the index of the item with the specified UID + ref_uid = before_uid if before_uid else after_uid + ref_position = await self._conn.fetchval( + f"SELECT position FROM {self._table_plan_queue} WHERE item->>'item_uid' = $1", + ref_uid + ) + if ref_position is None: + raise IndexError(f"Item with UID '{ref_uid}' not found in the queue.") + + position = ref_position if before_uid else ref_position + 1 + else: + raise ValueError(f"Invalid value for 'pos': {pos}") + + # Shift positions for items that come after the insertion point + await self._conn.execute( + f"UPDATE {self._table_plan_queue} SET position = position + 1 WHERE position >= $1", + position + ) + + # Insert the new item + item_json = json.dumps(item) + await self._conn.execute( + f"INSERT INTO {self._table_plan_queue} (position, item) VALUES ($1, $2)", + position, item_json + ) + + # Add item to UID dictionary + self._uid_dict_add(item) + + # Update the plan queue UID after modifications + self._plan_queue_uid = self.new_item_uid() + + # Get the new queue size + new_qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + + return item, new_qsize + + async def add_batch_to_queue( + self, + items: List[Dict[str, Any]], + *, + pos: Optional[Union[int, str]] = None, + before_uid: Optional[str] = None, + after_uid: Optional[str] = None, + filter_parameters: bool = True + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], int, bool]: + """ + Add a batch of items to the queue. + + Parameters + ---------- + items : list + List of items to add to the queue + pos : int or str, optional + Position in the queue. Can be an integer (0-based index) or string ('front' or 'back') + before_uid : str, optional + UID of the item before which to place the new items + after_uid : str, optional + UID of the item after which to place the new items + filter_parameters : bool, default True + Whether to filter item parameters + + Returns + ------- + tuple + (items_added, results, queue_size, success) - added items, results of each operation, + new queue size, and overall success flag + """ + async with self._lock: + items_added = [] + results = [] + success = True + added_item_uids = [] # For rollback in case of failure + + try: + async with self._conn.transaction(): + for item in items: + try: + if not added_item_uids: + # First item is placed according to parameters + item_added, _ = await self.add_item_to_queue( + item, + pos=pos, + before_uid=before_uid, + after_uid=after_uid, + filter_parameters=filter_parameters + ) + else: + # Subsequent items are placed after the previous one + item_added, _ = await self.add_item_to_queue( + item, + after_uid=added_item_uids[-1], + filter_parameters=filter_parameters + ) + + added_item_uids.append(item_added["item_uid"]) + items_added.append(item_added) + results.append({"success": True, "msg": ""}) + except Exception as ex: + success = False + items_added.append(item) + msg = f"Failed to add item to queue: {ex}" + results.append({"success": False, "msg": msg}) + # Raise to trigger transaction rollback + raise + + # Get the new queue size + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + return items_added, results, qsize, success + + except Exception: + # Transaction rollback has already happened in PostgreSQL + # Reset items_added to original items for consistent return value + return items, results, await self.get_queue_size(), False + + async def pop_item_from_queue( + self, + *, + pos: Optional[Union[int, str]] = None, + uid: Optional[str] = None + ) -> Tuple[Optional[Dict[str, Any]], int]: + """ + Remove and return an item from the queue. + + Parameters + ---------- + pos : int or str, optional + Position in the queue. Can be an integer (0-based index) or string ('front' or 'back') + uid : str, optional + UID of the item to remove + + Returns + ------- + tuple + (item, queue_size) - the removed item and the new queue size + """ + async with self._lock: + if (pos is not None) and (uid is not None): + raise ValueError("Ambiguous parameters: plan position and UID is specified") + + pos = pos if pos is not None else "front" + + async with self._conn.transaction(): + item = None + + if uid is not None: + # Check if item with this UID exists + if not self._is_uid_in_dict(uid): + raise IndexError(f"Item with UID '{uid}' is not in the queue.") + + # Check if the item is not currently running + row = await self._conn.fetchrow( + f"SELECT item FROM {self._table_running_plan} WHERE item->>'item_uid' = $1", + uid + ) + if row: + raise IndexError(f"Cannot remove an item which is currently running.") + + # Get the item and its position + row = await self._conn.fetchrow( + f"SELECT position, item FROM {self._table_plan_queue} WHERE item->>'item_uid' = $1", + uid + ) + if not row: + raise IndexError(f"Item with UID '{uid}' not found in the queue.") + + position = row['position'] + item = json.loads(row['item']) + + # Delete the item + await self._conn.execute( + f"DELETE FROM {self._table_plan_queue} WHERE item->>'item_uid' = $1", + uid + ) + + # Remove from UID dictionary + self._uid_dict_remove(uid) + + elif pos == "front": + row = await self._conn.fetchrow( + f"SELECT item FROM {self._table_plan_queue} ORDER BY position LIMIT 1" + ) + if not row: + return None, 0 + + item = json.loads(row['item']) + + # Delete the item + await self._conn.execute( + f"DELETE FROM {self._table_plan_queue} WHERE position = 0" + ) + + # Remove from UID dictionary + self._uid_dict_remove(item["item_uid"]) + + elif pos == "back": + # Find the position of the last item + max_pos = await self._conn.fetchval( + f"SELECT MAX(position) FROM {self._table_plan_queue}" + ) + if max_pos is None: + return None, 0 + + row = await self._conn.fetchrow( + f"SELECT item FROM {self._table_plan_queue} WHERE position = $1", + max_pos + ) + if not row: + return None, 0 + + item = json.loads(row['item']) + + # Delete the item + await self._conn.execute( + f"DELETE FROM {self._table_plan_queue} WHERE position = $1", + max_pos + ) + + # Remove from UID dictionary + self._uid_dict_remove(item["item_uid"]) + + elif isinstance(pos, int): + # Adjust negative indices + if pos < 0: + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + pos = qsize + pos + + if pos < 0: + raise IndexError(f"Position {pos} out of range") + + row = await self._conn.fetchrow( + f"SELECT item FROM {self._table_plan_queue} WHERE position = $1", + pos + ) + if not row: + raise IndexError(f"Position {pos} out of range") + + item = json.loads(row['item']) + + # Delete the item + await self._conn.execute( + f"DELETE FROM {self._table_plan_queue} WHERE position = $1", + pos + ) + + # Remove from UID dictionary + self._uid_dict_remove(item["item_uid"]) + else: + raise ValueError(f"Invalid value for 'pos': {pos}") + + # Reposition remaining items to eliminate gaps + await self._conn.execute(f""" + WITH ranked AS ( + SELECT id, ROW_NUMBER() OVER (ORDER BY position) - 1 AS new_position + FROM {self._table_plan_queue} + ) + UPDATE {self._table_plan_queue} SET position = ranked.new_position + FROM ranked WHERE {self._table_plan_queue}.id = ranked.id + """) + + # Get the new queue size + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + + return item, qsize + + async def pop_item_from_queue_batch( + self, + *, + uids: Optional[List[str]] = None, + ignore_missing: bool = True + ) -> Tuple[List[Dict[str, Any]], int]: + """ + Remove and return multiple items from the queue. + + Parameters + ---------- + uids : list, optional + List of UIDs of the items to remove + ignore_missing : bool, default True + Whether to ignore missing items + + Returns + ------- + tuple + (items, queue_size) - the removed items and the new queue size + """ + async with self._lock: + uids = uids or [] + + if not isinstance(uids, list): + raise TypeError(f"Parameter 'uids' must be a list: type(uids) = {type(uids)}") + + if not ignore_missing: + # Check if 'uids' contains only unique items + uids_set = set(uids) + if len(uids_set) != len(uids): + raise ValueError(f"The list contains repeated UIDs ({len(uids) - len(uids_set)} UIDs)") + + # Check if all UIDs in 'uids' exist in the queue + uids_missing = [] + for uid in uids: + if not self._is_uid_in_dict(uid): + uids_missing.append(uid) + if uids_missing: + raise ValueError(f"The queue does not contain items with the following UIDs: {uids_missing}") + + items = [] + async with self._conn.transaction(): + for uid in uids: + try: + item, _ = await self.pop_item_from_queue(uid=uid) + items.append(item) + except Exception as ex: + if not ignore_missing: + raise + logger.debug(f"Failed to remove item with UID '{uid}' from the queue: {ex}") + + # Get the new queue size + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + + return items, qsize + + async def clear_queue(self) -> None: + """Clear the plan queue.""" + async with self._lock: + async with self._conn.transaction(): + # Get the running item's UID if it exists + row = await self._conn.fetchrow(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + running_uid = None + if row: + running_item = json.loads(row['item']) + running_uid = running_item.get('item_uid') + + # Clear UID dictionary but preserve running item + if running_uid: + running_item = self._uid_dict_get_item(running_uid) + self._uid_dict_clear() + if running_item: + self._uid_dict_add(running_item) + else: + self._uid_dict_clear() + + # Clear the queue table + await self._conn.execute(f"TRUNCATE TABLE {self._table_plan_queue}") + + async def get_queue_size(self) -> int: + """Get the size of the queue.""" + return await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + + async def get_queue_full(self) -> Tuple[List[Dict[str, Any]], Dict[str, Any], str]: + """Retrieve the full queue.""" + async with self._lock: + # Get all items from queue ordered by position + rows = await self._conn.fetch(f"SELECT item FROM {self._table_plan_queue} ORDER BY position") + queue = [json.loads(row['item']) for row in rows] + + # Get running item if exists + row = await self._conn.fetchrow(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + running_item = json.loads(row['item']) if row else {} + + # Generate a unique identifier for this queue state + queue_uid = self.new_item_uid() + + return queue, running_item, queue_uid + + async def move_item( + self, + *, + pos: Optional[Union[int, str]] = None, + uid: Optional[str] = None, + pos_dest: Optional[Union[int, str]] = None, + before_uid: Optional[str] = None, + after_uid: Optional[str] = None + ) -> Tuple[Dict[str, Any], int]: + """ + Move an item to a new position in the queue. + + Parameters + ---------- + pos : int or str, optional + Source position in the queue + uid : str, optional + UID of the item to move + pos_dest : int or str, optional + Destination position + before_uid : str, optional + UID of the item before which to place the moved item + after_uid : str, optional + UID of the item after which to place the moved item + + Returns + ------- + tuple + (item, queue_size) - the moved item and the new queue size + """ + async with self._lock: + if (pos is None) and (uid is None): + raise ValueError("Source position or UID is not specified.") + if (pos_dest is None) and (before_uid is None) and (after_uid is None): + raise ValueError("Destination position or UID is not specified.") + + if (pos is not None) and (uid is not None): + raise ValueError("Ambiguous parameters: Both position and uid is specified for the source plan.") + if (pos_dest is not None) and (before_uid is not None or after_uid is not None): + raise ValueError("Ambiguous parameters: Both position and uid is specified for the destination plan.") + if (before_uid is not None) and (after_uid is not None): + raise ValueError("Ambiguous parameters: source should be moved 'before' and 'after' the destination.") + + async with self._conn.transaction(): + # Find the source item and its position + source_pos = None + item = None + + if uid is not None: + row = await self._conn.fetchrow( + f"SELECT position, item FROM {self._table_plan_queue} WHERE item->>'item_uid' = $1", + uid + ) + if not row: + raise IndexError(f"Item with UID '{uid}' not found in the queue.") + source_pos = row['position'] + item = json.loads(row['item']) + else: + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + if pos == "front": + source_pos = 0 + elif pos == "back": + source_pos = qsize - 1 + elif isinstance(pos, int): + if pos < 0: + source_pos = qsize + pos + else: + source_pos = pos + + if source_pos < 0 or source_pos >= qsize: + raise IndexError(f"Position {pos} is out of range.") + else: + raise ValueError(f"Invalid value for 'pos': {pos}") + + row = await self._conn.fetchrow( + f"SELECT item FROM {self._table_plan_queue} WHERE position = $1", + source_pos + ) + if not row: + raise IndexError(f"No item found at position {source_pos}.") + item = json.loads(row['item']) + + # Find the destination position + dest_pos = None + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + + if before_uid is not None or after_uid is not None: + dest_uid = before_uid if before_uid is not None else after_uid + before = dest_uid == before_uid + + row = await self._conn.fetchrow( + f"SELECT position FROM {self._table_plan_queue} WHERE item->>'item_uid' = $1", + dest_uid + ) + if not row: + raise IndexError(f"Item with UID '{dest_uid}' not found in the queue.") + + dest_pos = row['position'] + if not before: # If after_uid, increment the dest_pos + dest_pos += 1 + else: + if pos_dest == "front": + dest_pos = 0 + elif pos_dest == "back": + dest_pos = qsize + elif isinstance(pos_dest, int): + if pos_dest < 0: + dest_pos = qsize + pos_dest + else: + dest_pos = pos_dest + + if dest_pos < 0 or dest_pos > qsize: + raise IndexError(f"Position {pos_dest} is out of range.") + else: + raise ValueError(f"Invalid value for 'pos_dest': {pos_dest}") + + # If source_pos is the same as dest_pos, no need to move + if source_pos == dest_pos: + return item, qsize + + # If source is before destination, adjust dest_pos since array length will change + if source_pos < dest_pos: + dest_pos -= 1 + + # Remove item from source position + await self._conn.execute( + f"DELETE FROM {self._table_plan_queue} WHERE position = $1", + source_pos + ) + + # Reposition remaining items to eliminate gaps + await self._conn.execute(f""" + WITH ranked AS ( + SELECT id, ROW_NUMBER() OVER (ORDER BY position) - 1 AS new_position + FROM {self._table_plan_queue} + ) + UPDATE {self._table_plan_queue} SET position = ranked.new_position + FROM ranked WHERE {self._table_plan_queue}.id = ranked.id + """) + + # Make space at destination position + await self._conn.execute( + f"UPDATE {self._table_plan_queue} SET position = position + 1 WHERE position >= $1", + dest_pos + ) + + # Insert item at destination position + item_json = json.dumps(item) + await self._conn.execute( + f"INSERT INTO {self._table_plan_queue} (position, item) VALUES ($1, $2)", + dest_pos, item_json + ) + + # Get the new queue size + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + + return item, qsize + + async def move_batch( + self, + *, + uids: Optional[List[str]] = None, + pos_dest: Optional[Union[int, str]] = None, + before_uid: Optional[str] = None, + after_uid: Optional[str] = None, + reorder: bool = False + ) -> Tuple[List[Dict[str, Any]], int]: + """ + Move a batch of items within the queue. + + Parameters + ---------- + uids : list, optional + List of UIDs of the items to move + pos_dest : int or str, optional + Destination position for the first item + before_uid : str, optional + UID of the item before which to place the first moved item + after_uid : str, optional + UID of the item after which to place the first moved item + reorder : bool, default False + Whether to reorder items according to their current order in the queue + + Returns + ------- + tuple + (items, queue_size) - the moved items and the new queue size + """ + async with self._lock: + uids = uids or [] + + if not isinstance(uids, list): + raise TypeError(f"Parameter 'uids' must be a list: type(uids) = {type(uids)}") + + # Make sure only one of the mutually exclusive parameters is not None + param_list = [pos_dest, before_uid, after_uid] + n_params = len(param_list) - param_list.count(None) + if n_params < 1: + raise ValueError( + "Destination for the batch is not specified: use parameters 'pos_dest', " + "'before_uid' or 'after_uid'" + ) + elif n_params > 1: + raise ValueError( + "The function was called with more than one mutually exclusive parameter " + "('pos_dest', 'before_uid', 'after_uid')" + ) + + # Check if 'uids' contains only unique items + uids_set = set(uids) + if len(uids_set) != len(uids): + raise ValueError(f"The list contains repeated UIDs ({len(uids) - len(uids_set)} UIDs)") + + # Check if all UIDs in 'uids' exist in the queue + async with self._conn.transaction(): + for uid in uids: + count = await self._conn.fetchval( + f"SELECT COUNT(*) FROM {self._table_plan_queue} WHERE item->>'item_uid' = $1", + uid + ) + if count == 0: + raise ValueError(f"The queue does not contain an item with UID: {uid}") + + # Check that 'before_uid' and 'after_uid' are not in 'uids' + if (before_uid is not None) and (before_uid in uids): + raise ValueError(f"Parameter 'before_uid': item with UID '{before_uid}' is in the batch") + if (after_uid is not None) and (after_uid in uids): + raise ValueError(f"Parameter 'after_uid': item with UID '{after_uid}' is in the batch") + + # If reorder is True, arrange UIDs based on their current positions in the queue + if reorder: + uids_with_positions = [] + for uid in uids: + position = await self._conn.fetchval( + f"SELECT position FROM {self._table_plan_queue} WHERE item->>'item_uid' = $1", + uid + ) + uids_with_positions.append((position, uid)) + + uids_with_positions.sort(key=lambda x: x[0]) + uids_prepared = [pair[1] for pair in uids_with_positions] + else: + uids_prepared = uids + + # Perform the 'move' operation + last_item_uid = None + items_moved = [] + + for uid in uids_prepared: + if last_item_uid is None: + # First item is moved according to specified parameters + item, _ = await self.move_item( + uid=uid, pos_dest=pos_dest, before_uid=before_uid, after_uid=after_uid + ) + else: + # Consecutive items are placed after the previous item + item, _ = await self.move_item(uid=uid, after_uid=last_item_uid) + + last_item_uid = uid + items_moved.append(item) + + # Get the new queue size + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + + return items_moved, qsize + + async def replace_item( + self, + item: Dict[str, Any], + *, + item_uid: str + ) -> Tuple[Dict[str, Any], int]: + """ + Replace an existing item in the queue with a new item. + + Parameters + ---------- + item : dict + The new item to replace the existing one + item_uid : str + UID of the item to replace + + Returns + ------- + tuple + (old_item, queue_size) - the replaced item and the new queue size + """ + async with self._lock: + # Verify the new item, ignoring the UID of the item being replaced + await self._verify_item(item, ignore_uids=[item_uid]) + + async with self._conn.transaction(): + # Find the item to replace + row = await self._conn.fetchrow( + f"SELECT position, item FROM {self._table_plan_queue} WHERE item->>'item_uid' = $1", + item_uid + ) + if not row: + raise ValueError(f"Item with UID '{item_uid}' not found in the queue.") + + old_item = json.loads(row['item']) + position = row['position'] + + # Replace the item + item_json = json.dumps(item) + await self._conn.execute( + f"UPDATE {self._table_plan_queue} SET item = $1 WHERE position = $2", + item_json, position + ) + + # Update UID dictionary + self._uid_dict_remove(item_uid) + self._uid_dict_add(item) + + # Get the new queue size + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + + return old_item, qsize + + # -------------------------------------------------------------------------- + # History Management + async def clear_history(self) -> None: + """Clear the plan history.""" + async with self._lock: + async with self._conn.transaction(): + await self._conn.execute(f"TRUNCATE TABLE {self._table_plan_history}") + + async def get_history(self) -> Tuple[List[Dict[str, Any]], str]: + """Retrieve the full history.""" + rows = await self._conn.fetch( + f"SELECT item FROM {self._table_plan_history} ORDER BY id DESC" + ) + history = [json.loads(row['item']) for row in rows] + history_uid = self.new_item_uid() # Generate a unique ID for this history state + return history, self._plan_history_uid + + async def get_history_size(self) -> int: + """Get the size of the history.""" + return await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_history}") + + # -------------------------------------------------------------------------- + # Running Item Management + async def get_running_item_info(self) -> Dict[str, Any]: + """Get information about the currently running item.""" + row = await self._conn.fetchrow(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + return json.loads(row['item']) if row else {} + + async def _clean_item_properties(self, item: Dict[str, Any]) -> Dict[str, Any]: + """ + Clean unnecessary item properties. + + Parameters + ---------- + item : dict + The item to clean + + Returns + ------- + dict + The cleaned item + """ + item = copy.deepcopy(item) + if "properties" in item: + p = item["properties"] + if "immediate_execution" in p: + del p["immediate_execution"] + if "time_start" in p: + del p["time_start"] + if not p: + del item["properties"] + return item + + async def set_next_item_as_running(self, *, item: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + """ + Set the next item from the queue as running. + + Parameters + ---------- + item : dict, optional + Item for immediate execution + + Returns + ------- + dict + The item that is now running + """ + async with self._lock: + async with self._conn.transaction(): + # Check if an item is already running + row = await self._conn.fetchrow(f"SELECT COUNT(*) FROM {self._table_running_plan}") + if row and row[0] > 0: + raise RuntimeError("An item is already running") + + immediate_execution = bool(item) + plan = None + + if immediate_execution: + # Generate UID if it doesn't exist + plan = copy.deepcopy(item) + if "item_uid" not in plan: + plan = await self.set_new_item_uuid(plan) + + plan.setdefault("properties", {})["immediate_execution"] = True + else: + # Get the first item from the queue + row = await self._conn.fetchrow( + f"SELECT item FROM {self._table_plan_queue} ORDER BY position LIMIT 1" + ) + if not row: + return {} # Queue is empty + + plan = json.loads(row['item']) + + # Remove the plan from the queue + await self._conn.execute( + f"DELETE FROM {self._table_plan_queue} WHERE position = 0" + ) + + # Reposition remaining items + await self._conn.execute(f""" + WITH ranked AS ( + SELECT id, ROW_NUMBER() OVER (ORDER BY position) - 1 AS new_position + FROM {self._table_plan_queue} + ) + UPDATE {self._table_plan_queue} SET position = ranked.new_position + FROM ranked WHERE {self._table_plan_queue}.id = ranked.id + """) + + # Verify that the item is a plan + if "item_type" not in plan: + raise ValueError("Item does not have 'item_type'") + + if plan["item_type"] != "plan": + raise RuntimeError( + "Function 'set_next_item_as_running' was called for " + f"an item other than plan: {plan}" + ) + + # Record start time + plan.setdefault("properties", {})["time_start"] = ttime.time() + + # Set as running plan + plan_json = json.dumps(plan) + await self._conn.execute( + f"INSERT INTO {self._table_running_plan} (item) VALUES ($1)", + plan_json + ) + + # Add to UID dictionary + self._uid_dict_add(plan) + + return plan + + async def set_processed_item_as_completed( + self, + *, + exit_status: str, + run_uids: List[str], + scan_ids: List[int], + err_msg: str = "", + err_tb: str = "" + ) -> Dict[str, Any]: + """ + Mark the currently running item as completed and move it to history. + + Parameters + ---------- + exit_status : str + Exit status of the item + run_uids : list + List of run UIDs + scan_ids : list + List of scan IDs + err_msg : str, default "" + Error message if any + err_tb : str, default "" + Error traceback if any + + Returns + ------- + dict + The completed item + """ + async with self._lock: + # Check if loop mode is enabled in queue mode + row = await self._conn.fetchrow(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + if not row: + return {} + + running_item = json.loads(row['item']) + immediate_execution = running_item.get("properties", {}).get("immediate_execution", False) + item_time_start = running_item.get("properties", {}).get("time_start", ttime.time()) + + if loop_enabled and not immediate_execution: + # Add a copy of the item back to the queue + item_to_add = copy.deepcopy(item_cleaned) + item_to_add = await self.set_new_item_uuid(item_to_add) # Add await here + + row = await self._conn.fetchrow( + f"SELECT info FROM {self._table_autostart_mode_info} LIMIT 1" + ) + if row: + info = json.loads(row['info']) + loop_enabled = info.get("loop", False) + + async with self._conn.transaction(): + # Clean item properties + item_cleaned = await self._clean_item_properties(running_item) + + if loop_enabled and not immediate_execution: + # Add a copy of the item back to the queue + item_to_add = copy.deepcopy(item_cleaned) + item_to_add = await self.set_new_item_uuid(item_to_add) + + # Get the current size of the queue + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + + # Add to the end of the queue + item_json = json.dumps(item_to_add) + await self._conn.execute( + f"INSERT INTO {self._table_plan_queue} (position, item) VALUES ($1, $2)", + qsize, item_json + ) + + # Add to UID dictionary + self._uid_dict_add(item_to_add) + + # Add result information to the item + item_cleaned.setdefault("result", {}) + item_cleaned["result"]["exit_status"] = exit_status + item_cleaned["result"]["run_uids"] = run_uids + item_cleaned["result"]["scan_ids"] = scan_ids + item_cleaned["result"]["time_start"] = item_time_start + item_cleaned["result"]["time_stop"] = ttime.time() + item_cleaned["result"]["msg"] = err_msg + item_cleaned["result"]["traceback"] = err_tb + + # Add to history + item_json = json.dumps(item_cleaned) + await self._conn.execute( + f"INSERT INTO {self._table_plan_history} (item) VALUES ($1)", + item_json + ) + + # Remove from running + await self._conn.execute(f"TRUNCATE TABLE {self._table_running_plan}") + + # Remove from UID dictionary if not in loop mode + if not loop_enabled and not immediate_execution: + self._uid_dict_remove(running_item["item_uid"]) + + # Update UIDs after modifications + self._plan_queue_uid = self.new_item_uid() + self._plan_history_uid = self.new_item_uid() + + return item_cleaned + + async def set_processed_item_as_stopped( + self, + *, + exit_status: str, + run_uids: Optional[List[str]] = None, + scan_ids: Optional[List[int]] = None, + err_msg: str = "", + err_tb: str = "" + ) -> Dict[str, Any]: + """ + Mark the currently running item as stopped and move it to history. + If exit_status is not "stopped", also push the item back to the queue. + + Parameters + ---------- + exit_status : str + Exit status of the item + run_uids : list, optional + List of run UIDs + scan_ids : list, optional + List of scan IDs + err_msg : str, default "" + Error message if any + err_tb : str, default "" + Error traceback if any + + Returns + ------- + dict + The stopped item + """ + async with self._lock: + run_uids = run_uids or [] + scan_ids = scan_ids or [] + + # If the status is "stopped", treat it as a completed item + if exit_status == "stopped": + return await self.set_processed_item_as_completed( + exit_status=exit_status, + run_uids=run_uids, + scan_ids=scan_ids, + err_msg=err_msg, + err_tb=err_tb + ) + + row = await self._conn.fetchrow(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + if not row: + return {} + + running_item = json.loads(row['item']) + immediate_execution = running_item.get("properties", {}).get("immediate_execution", False) + item_time_start = running_item.get("properties", {}).get("time_start", ttime.time()) + + async with self._conn.transaction(): + # Clean item properties + item_cleaned = await self._clean_item_properties(running_item) + + # Add result information to the item + item_cleaned.setdefault("result", {}) + item_cleaned["result"]["exit_status"] = exit_status + item_cleaned["result"]["run_uids"] = run_uids + item_cleaned["result"]["scan_ids"] = scan_ids + item_cleaned["result"]["time_start"] = item_time_start + item_cleaned["result"]["time_stop"] = ttime.time() + item_cleaned["result"]["msg"] = err_msg + item_cleaned["result"]["traceback"] = err_tb + + # Add to history + item_json = json.dumps(item_cleaned) + await self._conn.execute( + f"INSERT INTO {self._table_plan_history} (item) VALUES ($1)", + item_json + ) + + # Push back to the queue with new UID if not immediate execution and status != stopped + if not immediate_execution and exit_status != "stopped": + # Create a copy without result information + item_copy = copy.deepcopy(item_cleaned) + if "result" in item_copy: + del item_copy["result"] + + # Set a new UID + item_copy = await self.set_new_item_uuid(item_copy) + + # Insert at the front of the queue + await self._conn.execute( + f"UPDATE {self._table_plan_queue} SET position = position + 1" + ) + + item_json = json.dumps(item_copy) + await self._conn.execute( + f"INSERT INTO {self._table_plan_queue} (position, item) VALUES (0, $1)", + item_json + ) + + # Add to UID dictionary + self._uid_dict_add(item_copy) + + # Remove from running + await self._conn.execute(f"TRUNCATE TABLE {self._table_running_plan}") + + # Remove original item from UID dictionary + self._uid_dict_remove(running_item["item_uid"]) + + return item_cleaned + + # -------------------------------------------------------------------------- + # Lock Management + async def lock_info_save(self, lock_info: Dict[str, Any]) -> None: + """Save lock information.""" + async with self._lock: + async with self._conn.transaction(): + # Clear existing lock info + await self._conn.execute(f"TRUNCATE TABLE {self._table_lock_info}") + + # Save new lock info + lock_json = json.dumps(lock_info) + await self._conn.execute( + f"INSERT INTO {self._table_lock_info} (info) VALUES ($1)", + lock_json + ) + + async def lock_info_retrieve(self) -> Optional[Dict[str, Any]]: + """Retrieve lock information.""" + row = await self._conn.fetchrow(f"SELECT info FROM {self._table_lock_info} LIMIT 1") + return json.loads(row['info']) if row else None + + async def lock_info_clear(self) -> None: + """Clear lock information.""" + await self._conn.execute(f"TRUNCATE TABLE {self._table_lock_info}") + + # -------------------------------------------------------------------------- + # Autostart Mode Management + async def autostart_mode_save(self, autostart_mode: Dict[str, Any]) -> None: + """Save autostart mode information.""" + async with self._lock: + async with self._conn.transaction(): + # Clear existing autostart mode info + await self._conn.execute(f"TRUNCATE TABLE {self._table_autostart_mode_info}") + + # Save new autostart mode info + mode_json = json.dumps(autostart_mode) + await self._conn.execute( + f"INSERT INTO {self._table_autostart_mode_info} (info) VALUES ($1)", + mode_json + ) + + async def autostart_mode_retrieve(self) -> bool: + """ + Retrieve autostart mode information. + + Returns + ------- + bool + True if autostart mode is enabled, False otherwise + """ + row = await self._conn.fetchrow(f"SELECT info FROM {self._table_autostart_mode_info} LIMIT 1") + if row: + info = json.loads(row['info']) + return info.get("enabled", False) + return False + + async def autostart_mode_clear(self) -> None: + """Clear autostart mode information.""" + await self._conn.execute(f"TRUNCATE TABLE {self._table_autostart_mode_info}") + + # -------------------------------------------------------------------------- + # Stop Pending State Management + async def stop_pending_save(self, stop_pending: Dict[str, Any]) -> None: + """Save stop pending information.""" + async with self._lock: + async with self._conn.transaction(): + # Clear existing stop pending info + await self._conn.execute(f"TRUNCATE TABLE {self._table_stop_pending_info}") + + # Save new stop pending info + pending_json = json.dumps(stop_pending) + await self._conn.execute( + f"INSERT INTO {self._table_stop_pending_info} (info) VALUES ($1)", + pending_json + ) + + async def stop_pending_retrieve(self) -> Optional[Dict[str, Any]]: + """Retrieve stop pending information.""" + row = await self._conn.fetchrow(f"SELECT info FROM {self._table_stop_pending_info} LIMIT 1") + return json.loads(row['info']) if row else None + + async def stop_pending_clear(self) -> None: + """Clear stop pending information.""" + await self._conn.execute(f"TRUNCATE TABLE {self._table_stop_pending_info}") + + # -------------------------------------------------------------------------- + # User Group Permissions + async def user_group_permissions_save(self, user_group_permissions: Dict[str, Any]) -> None: + """Save user group permissions.""" + async with self._lock: + async with self._conn.transaction(): + # Clear existing permissions + await self._conn.execute(f"TRUNCATE TABLE {self._table_user_group_permissions}") + + # Save new permissions + perms_json = json.dumps(user_group_permissions) + await self._conn.execute( + f"INSERT INTO {self._table_user_group_permissions} (info) VALUES ($1)", + perms_json + ) + + async def user_group_permissions_retrieve(self) -> Optional[Dict[str, Any]]: + """Retrieve user group permissions.""" + row = await self._conn.fetchrow(f"SELECT info FROM {self._table_user_group_permissions} LIMIT 1") + return json.loads(row['info']) if row else None + + async def user_group_permissions_clear(self) -> None: + """Clear user group permissions.""" + await self._conn.execute(f"TRUNCATE TABLE {self._table_user_group_permissions}") + + # -------------------------------------------------------------------------- + # Utility Methods + async def get_queue_state(self) -> Dict[str, Any]: + """ + Get the overall queue state. + + Returns + ------- + dict + Dictionary containing the queue state, including queue size, history size, + and information about the currently running item. + """ + queue_size = await self.get_queue_size() + history_size = await self.get_history_size() + running_item = await self.get_running_item_info() + + return { + "queue_size": queue_size, + "history_size": history_size, + "running_item": running_item, + } + + + async def set_plan_queue_mode( + self, + plan_queue_mode: Union[Dict[str, Any], str], + *, + update: bool = False + ) -> Dict[str, Any]: + """ + Set the plan queue mode. + + Parameters + ---------- + plan_queue_mode : dict or str + The new plan queue mode or 'default' to reset to default mode + update : bool, default False + If True, update only the specified fields in the mode dictionary + + Returns + ------- + dict + The updated plan queue mode + """ + async with self._lock: + # Create a table for plan queue mode if it doesn't exist + await self._conn.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_plan_queue_mode} ( + id SERIAL PRIMARY KEY, + info JSONB + ) + """) + + # Get current mode from database or use default + row = await self._conn.fetchrow(f"SELECT info FROM {self._table_plan_queue_mode} LIMIT 1") + current_mode = json.loads(row['info']) if row else copy.deepcopy(self._plan_queue_mode_default) + + # Reset to default if requested + if isinstance(plan_queue_mode, str) and plan_queue_mode.lower() == "default": + new_mode = copy.deepcopy(self._plan_queue_mode_default) + elif not isinstance(plan_queue_mode, dict): + raise TypeError(f"Parameter 'plan_queue_mode' must be a dictionary or 'default': {plan_queue_mode}") + else: + # Create a new mode or update the current one + if update: + new_mode = copy.deepcopy(current_mode) + # Verify keys in plan_queue_mode + for key in plan_queue_mode: + if key not in self._plan_queue_mode_default: + raise ValueError(f"Unsupported plan queue mode parameter '{key}'") + # Update mode + new_mode.update(plan_queue_mode) + else: + # Verify that all required keys are present + missing_keys = set(self._plan_queue_mode_default.keys()) - set(plan_queue_mode.keys()) + if missing_keys: + raise ValueError(f"Parameters {missing_keys} are missing") + + # Verify keys in plan_queue_mode + for key in plan_queue_mode: + if key not in self._plan_queue_mode_default: + raise ValueError(f"Unsupported plan queue mode parameter '{key}'") + + new_mode = copy.deepcopy(plan_queue_mode) + + # Verify value types + for key, value in new_mode.items(): + if key == "loop" and not isinstance(value, bool): + raise TypeError(f"Unsupported type {type(value)} of the parameter 'loop'") + elif key == "ignore_failures" and not isinstance(value, bool): + raise TypeError(f"Unsupported type {type(value)} of the parameter 'ignore_failures'") + + # Save the new mode to the database + async with self._conn.transaction(): + await self._conn.execute(f"TRUNCATE TABLE {self._table_plan_queue_mode}") + mode_json = json.dumps(new_mode) + await self._conn.execute( + f"INSERT INTO {self._table_plan_queue_mode} (info) VALUES ($1)", + mode_json + ) + + # Update instance variable + self._plan_queue_mode = new_mode + + return copy.deepcopy(new_mode) + + @property + def plan_queue_mode(self) -> Dict[str, Any]: + """ + Get the current plan queue mode. + + Returns + ------- + dict + The current plan queue mode + """ + return copy.deepcopy(self._plan_queue_mode) + + @property + def plan_queue_mode_default(self) -> Dict[str, Any]: + """ + Get the default plan queue mode. + + Returns + ------- + dict + The default plan queue mode + """ + return copy.deepcopy(self._plan_queue_mode_default) + + @property + def plan_queue_uid(self) -> str: + """ + Get the current plan queue UID. + + Returns + ------- + str + The current plan queue UID + """ + # Generate a new UID for each request, or implement a tracking mechanism + return self.new_item_uid() + + @property + def plan_history_uid(self) -> str: + """ + Get the current plan history UID. + + Returns + ------- + str + The current plan history UID + """ + # Generate a new UID for each request, or implement a tracking mechanism + return self.new_item_uid() + + async def get_item( + self, + *, + pos: Optional[Union[int, str]] = None, + uid: Optional[str] = None + ) -> Dict[str, Any]: + """ + Get an item from the queue without removing it. + + Parameters + ---------- + pos : int or str, optional + Position in the queue. Can be an integer (0-based index) or string ('front' or 'back') + uid : str, optional + UID of the item to get + + Returns + ------- + dict + The requested item + """ + async with self._lock: + if (pos is None) and (uid is None): + raise ValueError("Position or UID must be specified") + if (pos is not None) and (uid is not None): + raise ValueError("Ambiguous parameters: both position and UID are specified") + + # Get queue size + qsize = await self._conn.fetchval(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + + if qsize == 0: + raise IndexError("The queue is empty") + + if uid is not None: + # Get by UID + row = await self._conn.fetchrow( + f"SELECT item FROM {self._table_plan_queue} WHERE item->>'item_uid' = $1", + uid + ) + if not row: + raise RuntimeError(f"Item with UID '{uid}' is not in the queue") + + return json.loads(row['item']) + else: + # Get by position + if pos == "front": + position = 0 + elif pos == "back": + position = qsize - 1 + elif isinstance(pos, int): + if pos < 0: + position = qsize + pos + else: + position = pos + + if position < 0 or position >= qsize: + raise IndexError(f"Position {pos} is out of range") + else: + raise ValueError(f"Parameter 'pos' has incorrect value: {pos}") + + row = await self._conn.fetchrow( + f"SELECT item FROM {self._table_plan_queue} WHERE position = $1", + position + ) + if not row: + raise IndexError(f"No item found at position {position}") + + return json.loads(row['item']) + + async def process_next_item( + self, + *, + item: Optional[Dict[str, Any]] = None + ) -> Dict[str, Any]: + """ + Process the next item in the queue or a specified item. + + Parameters + ---------- + item : dict, optional + Item for immediate execution, if not provided, use the next item from the queue + + Returns + ------- + dict + The processed item + """ + # This is essentially a wrapper around set_next_item_as_running + # which seems to already be implemented in the class + return await self.set_next_item_as_running(item=item) + + @property + def plan_queue_uid(self) -> str: + """ + Get the current plan queue UID. + + Returns + ------- + str + The current plan queue UID + """ + return self._plan_queue_uid + + @property + def plan_history_uid(self) -> str: + """ + Get the current plan history UID. + + Returns + ------- + str + The current plan history UID + """ + return self._plan_history_uid \ No newline at end of file diff --git a/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_redis.py b/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_redis.py new file mode 100644 index 00000000..f07806d9 --- /dev/null +++ b/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_redis.py @@ -0,0 +1,2662 @@ +import asyncio +import copy +import json +import logging +import time as ttime +from typing import Any, Dict, List, Optional, Tuple, Union + +from bluesky_queueserver.manager.plan_queue_ops_backends.plan_queue_ops_abstract import AbstractPlanQueueOperations + +import redis.asyncio + + +logger = logging.getLogger(__name__) + + +class PlanQueueOperations: + """ + The class supports operations with a plan queue based on either Redis or SQLite. The backend + can be selected using the `PLAN_QUEUE_BACKEND` environment variable or by defaulting to Redis + if the variable is not set. The public methods of the class are protected with ``asyncio.Lock``. + + Parameters + ---------- + redis_host : str + Address of the Redis host. This parameter is only relevant if the backend is set to "redis". + name_prefix : str + Prefix for the names of the keys used in Redis or SQLite. The prefix is used to avoid conflicts + with the keys used by other instances of the Queue Server. For example, the prefix used + for unit tests should be different from the prefix used in production. If the prefix + is an empty string, then no prefix will be added (not recommended). + + Backend Selection + ----------------- + The backend can be selected using the `PLAN_QUEUE_BACKEND` environment variable: + - Set `PLAN_QUEUE_BACKEND=redis` to use Redis as the backend. + - Set `PLAN_QUEUE_BACKEND=sqlite` to use SQLite as the backend. + - If the environment variable is not set, the backend defaults to Redis. + + SQLite Database Path + --------------------- + When using SQLite as the backend, the database file path can be specified using the + `PLAN_QUEUE_SQLITE_PATH` environment variable: + - Set `PLAN_QUEUE_SQLITE_PATH=/path/to/database.db` to specify the SQLite database file path. + - If the environment variable is not set, the default database file `plan_queue.db` will be used + in the current working directory. + + Notes + ----- + - The backend is responsible for storing the plan queue, history, and other related data. + - The class supports both Redis and SQLite backends, and the backend-specific logic is encapsulated + in helper methods (e.g., `_backend_set`, `_backend_get`, etc.). + - The `name_prefix` ensures that multiple instances of the Queue Server can operate without + interfering with each other's data. + """ + + def __init__(self, redis_host="localhost", name_prefix="qs_default"): + """ + Initialize the PlanQueueOperations class. + + Parameters + ---------- + redis_host : str + Address of the Redis host. This parameter is only relevant if the backend is set to "redis". + name_prefix : str + Prefix for the names of the keys used in Redis or SQLite. The prefix is used to avoid conflicts + with the keys used by other instances of the Queue Server. For example, the prefix used + for unit tests should be different from the prefix used in production. If the prefix + is an empty string, then no prefix will be added (not recommended). + + Backend Selection + ----------------- + The backend can be selected using the `PLAN_QUEUE_BACKEND` environment variable: + - Set `PLAN_QUEUE_BACKEND=redis` to use Redis as the backend. + - Set `PLAN_QUEUE_BACKEND=sqlite` to use SQLite as the backend. + - If the environment variable is not set, the backend defaults to Redis. + + Notes + ----- + - The backend is responsible for storing the plan queue, history, and other related data. + - The class supports both Redis and SQLite backends, and the backend-specific logic is encapsulated + in helper methods (e.g., `_backend_set`, `_backend_get`, etc.). + - The `name_prefix` ensures that multiple instances of the Queue Server can operate without + interfering with each other's data. + """ + self._redis_host = redis_host + self._uid_dict = dict() + self._r_pool = None + self._sqlite_conn = None + + # Check for environment variable to set the backend + self._backend = os.getenv("PLAN_QUEUE_BACKEND", "redis").lower() # Default to "redis" + # self._backend = os.getenv("PLAN_QUEUE_BACKEND", "sqlite").lower() # DELETE + if self._backend not in ["redis", "sqlite"]: + raise ValueError(f"Invalid backend specified: {self._backend}. Must be 'redis' or 'sqlite'.") + + if not isinstance(name_prefix, str): + raise TypeError(f"Parameter 'name_prefix' should be a string: {name_prefix}") + + # The case of an empty string + if name_prefix: + name_prefix = name_prefix + "_" + + self._name_running_plan = name_prefix + "running_plan" + self._name_plan_queue = name_prefix + "plan_queue" + self._name_plan_history = name_prefix + "plan_history" + self._name_plan_queue_mode = name_prefix + "plan_queue_mode" + + # Redis is also used for storage of some additional information not related to the queue. + # The class contains only the functions for saving and retrieving the data, which is + # not used by other functions of the class. + self._name_user_group_permissions = name_prefix + "user_group_permissions" + self._name_lock_info = name_prefix + "lock_info" + self._name_autostart_mode_info = name_prefix + "autostart_mode_info" + self._name_stop_pending_info = name_prefix + "stop_pending_info" + + # The list of allowed item parameters used for parameter filtering. Filtering operation + # involves removing all parameters that are not in the list. + self._allowed_item_parameters = ( + "item_uid", + "item_type", + "name", + "args", + "kwargs", + "meta", + "user", + "user_group", + "properties", + ) + + # Plan queue UID is expected to change each time the contents of the queue is changed. + # Since `self._uid_dict` is modified each time the queue is updated, it is sufficient + # to update Plan queue UID in the functions that update `self._uid_dict`. + self._plan_queue_uid = self.new_item_uid() + # Plan history UID is expected to change each time the history is changed. + self._plan_history_uid = self.new_item_uid() + + self._lock = None + + # Settings that determine the mode of queue operation. The set of supported modes + # may be extended if additional modes are to be implemented. The mode will be saved in + # Redis, so that it is not modified between restarts of the manager. + # Loop mode: + # loop, True/False. If enabled, then each executed item (plan or instruction) + # will be placed to the back of the queue. + # ignore_failures, True/False. Run all the plans in the queue to the end + # even if some or all of the plans fail. The queue is still stopped if the user + # stops/aborts/halts a plan. + self._plan_queue_mode_default = {"loop": False, "ignore_failures": False} + self._plan_queue_mode = self.plan_queue_mode_default + + @property + def plan_queue_uid(self): + """ + Get current plan queue UID (str). Note, that the UID may be updated multiple times during + complex queue operations, so the returned UID may not represent a valid queue state. + The intended use: the changes of UID could be monitored to detect changes in the queue + without accessing the queue. If the UID is different from UID returned by + ``PlanQueueOperations.get_queue()``, then the contents of the queue changed. + """ + return self._plan_queue_uid + + @property + def plan_history_uid(self): + """ + Get current plan history UID. See notes for ``PlanQueueOperations.plan_queue_uid``. + """ + return self._plan_history_uid + + @property + def plan_queue_mode(self): + """ + Returns current plan queue mode. Plan queue mode is a dictionary with parameters + used for selection of the algorithm(s) for handling queue items. Supported parameters: + ``loop (boolean)`` enables and disables the loop mode. + """ + return self._plan_queue_mode.copy() + + @property + def plan_queue_mode_default(self): + """ + Returns the default queue mode (default settings) + """ + return self._plan_queue_mode_default.copy() + + def _validate_plan_queue_mode(self, plan_queue_mode): + """ + Validate the dictionary 'plan_queue_mode'. Check that the dictionary contains all + the required parameters and no unsupported parameters. + + Parameters + ---------- + plan_queue_mode : dict + Dictionary that contains plan queue mode. See ``self.plan_queue_mode_default``. + """ + # It is assumed that 'plan_queue_mode' will be a single-level dictionary that contains + # simple types (bool, int etc), so the following code provide better error reporting + # than schema validation. + expected_params = {"loop": bool, "ignore_failures": bool} + missing_keys = set(expected_params.keys()) + for k, v in plan_queue_mode.items(): + if k not in expected_params: + raise ValueError( + f"Unsupported plan queue mode parameter '{k}': " + f"supported parameters {list(expected_params.keys())}" + ) + missing_keys.remove(k) + key_type = expected_params.get(k) # Using [k] makes PyCharm to display annoying error + if not isinstance(v, key_type): + raise TypeError( + f"Unsupported type '{type(v)}' of the parameter '{k}': " f"expected type '{key_type}'" + ) + if missing_keys: + raise ValueError( + f"Parameters {missing_keys} are missing from 'plan_queue_mode' dictionary. " + f"The following keys are expected: {list(expected_params.keys())}" + ) + + async def _load_plan_queue_mode(self): + """ + Load plan queue mode from Redis. + """ + queue_mode = await self._backend_get(self._name_plan_queue_mode) + self._plan_queue_mode = json.loads(queue_mode) if queue_mode else self.plan_queue_mode_default + try: + self._validate_plan_queue_mode(self._plan_queue_mode) + except Exception as ex: + logger.error("Failed to load plan queue mode from Redis. The default mode is used: %s", ex) + self._plan_queue_mode = self.plan_queue_mode_default + + async def set_plan_queue_mode(self, plan_queue_mode, *, update=True): + """ + Set plan queue mode. The plan queue mode can be a string ``default`` or a dictionary with + parameters. See ``self.plan_queue_mode_default`` for an example of the parameter dictionary. + + Parameters + ---------- + plan_queue_mode : dict or str + The dictionary of parameters that define queue mode. If ``update`` is ``True``, then + the dictionary may contain only the parameters that need to be updated. Otherwise + ``plan_queue_mode`` dictionary must contain full valid parameter dictionary. The function + fails if the dictionary contains unsupported parameters. Calling the function with the + string value ``plan_queue_mode="default"`` will reset all the parameters to the default + values. + update : boolean (optional) + Indicates if the dictionary ``plan_queue_mode`` should be used to update mode parameters. + If ``True``, then the dictionary may contain only the parameters that should be changed. + """ + if not isinstance(plan_queue_mode, dict) and plan_queue_mode != "default": + raise TypeError( + f"Queue mode is passed using object of unsupported type '{type(plan_queue_mode)}': " + f"({plan_queue_mode}). Supported types: ('dict', 'str'), supported " + f"string value: 'default'" + ) + + if plan_queue_mode == "default": + plan_queue_mode = self.plan_queue_mode_default + elif update: + # Generate full parameter dictionary based on the existing and submitted parameters. + queue_mode = self.plan_queue_mode # Create a copy of current parameters + queue_mode.update(plan_queue_mode) + plan_queue_mode = queue_mode + + self._validate_plan_queue_mode(plan_queue_mode) + + # Prevent changes of the queue mode in the middle of queue operations. + async with self._lock: + self._plan_queue_mode = plan_queue_mode.copy() + await self._backend_set(self._name_plan_queue_mode, json.dumps(self._plan_queue_mode)) + + async def start(self): + """ + Create the connection pool (Redis or SQLite) and initialize the set of UIDs from the queue + if it exists in the backend. + """ + if self._backend == "redis": + if not self._r_pool: # Initialize only once + self._lock = asyncio.Lock() + async with self._lock: + try: + host = f"redis://{self._redis_host}" + self._r_pool = redis.asyncio.from_url(host, encoding="utf-8", decode_responses=True) + await self._backend_ping() + except Exception as ex: + error_msg = ( + f"Failed to create the Redis pool: " + f"Redis server may not be available at '{self._redis_host}'. " + f"Exception: {ex}" + ) + logger.error(error_msg) + raise OSError(error_msg) from ex + + await self._queue_clean() + await self._uid_dict_initialize() + await self._load_plan_queue_mode() + + self._plan_queue_uid = self.new_item_uid() + self._plan_history_uid = self.new_item_uid() + + + elif self._backend == "sqlite": + if not self._sqlite_conn: # Initialize only once + self._lock = asyncio.Lock() + async with self._lock: + try: + # Default to a persistent SQLite database file + sqlite_db_path = os.getenv("PLAN_QUEUE_SQLITE_PATH", "plan_queue.db") + self._sqlite_conn = await aiosqlite.connect(sqlite_db_path) + await self._backend_ping() + + # Initialize the SQLite database + await self._initialize_sqlite_database() + + except Exception as ex: + error_msg = ( + f"Failed to create the SQLite connection: " + f"SQLite database may not be accessible. Exception: {ex}" + ) + logger.error(error_msg) + raise OSError(error_msg) from ex + + await self._queue_clean() + await self._uid_dict_initialize() + await self._load_plan_queue_mode() + + self._plan_queue_uid = self.new_item_uid() + self._plan_history_uid = self.new_item_uid() + + async def stop(self): + """ + Close all connections in the backend (Redis or SQLite). + """ + if self._backend == "redis": + if self._r_pool: + await self._backend_aclose() + self._r_pool = None + elif self._backend == "sqlite": + if self._sqlite_conn: + await self._backend_aclose() + self._sqlite_conn = None + + async def _queue_clean(self): + """ + Delete all the invalid queue entries (there could be some entries from failed unit tests). + """ + pq, _ = await self._get_queue() + + def verify_item(item): + # The criteria may be changed. + return "item_uid" in item + + items_to_remove = [] + for item in pq: + if not verify_item(item): + items_to_remove.append(item) + + for item in items_to_remove: + await self._remove_item(item, single=False) + + # Clean running plan info also (on the development computer it may contain garbage) + item = await self._get_running_item_info() + if item and not verify_item(item): + await self._clear_running_item_info() + + async def _delete_pool_entries(self): + """ + See ``self.delete_pool_entries()`` method. + """ + await self._backend_delete(self._name_running_plan) + await self._backend_delete(self._name_plan_queue) + await self._backend_delete(self._name_plan_history) + await self._backend_delete(self._name_plan_queue_mode) + self._uid_dict_clear() + + self._plan_queue_uid = self.new_item_uid() + self._plan_history_uid = self.new_item_uid() + + async def delete_pool_entries(self): + """ + Delete pool entries used by RE Manager. This method is mostly intended for use in testing, + but may be used for other purposes if needed. Deleting pool entries also resets plan queue mode. + """ + async with self._lock: + await self._delete_pool_entries() + + # @staticmethod + # def _verify_item_type(item): + # """ + # Check that the item (plan) is a dictionary. + # """ + # if not isinstance(item, dict): + # raise TypeError(f"Parameter 'item' should be a dictionary: '{item}', (type '{type(item)}')") + + # def _verify_item(self, item, *, ignore_uids=None): + # """ + # Verify that item (plan) structure is valid enough to be put in the queue. + # Current checks: item is a dictionary, ``item_uid`` key is present, Plan with the UID is not in + # the queue or currently running. Ignore UIDs in the list ``ignore_uids``: those UIDs are expected + # to be in the dictionary. + # """ + # ignore_uids = ignore_uids or [] + # self._verify_item_type(item) + # # Verify plan UID + # if "item_uid" not in item: + # raise ValueError("Item does not have UID.") + # uid = item["item_uid"] + # if (uid not in ignore_uids) and self._is_uid_in_dict(uid): + # raise RuntimeError(f"Item with UID {uid} is already in the queue") + + # @staticmethod + # def new_item_uid(): + # """ + # Generate UID for an item (plan). + # """ + # return str(uuid.uuid4()) + + # def set_new_item_uuid(self, item): + # """ + # Replaces Item UID with a new one or creates a new UID. + + # Parameters + # ---------- + # item: dict + # Dictionary of item parameters. The dictionary may or may not have the key ``item_uid``. + + # Returns + # ------- + # dict + # Plan with new UID. + # """ + # item = copy.deepcopy(item) + # self._verify_item_type(item) + # item["item_uid"] = self.new_item_uid() + # return item + + async def _get_index_by_uid(self, *, uid): + """ + Get index of a plan in Redis list by UID. This is inefficient operation and should + be avoided whenever possible. Raises an exception if the plan is not found. + + Parameters + ---------- + uid: str + UID of the plans to find. + + Returns + ------- + int + Index of the plan with given UID. + + Raises + ------ + IndexError + No plan is found. + """ + queue, _ = await self._get_queue() + for n, plan in enumerate(queue): + if plan["item_uid"] == uid: + return n + raise IndexError(f"No plan with UID '{uid}' was found in the list.") + + async def _get_index_by_uid_batch(self, *, uids): + """ + Batch version of ``_get_index_by_uid``. The operation is implemented efficiently + and should be used for finding indices of large number of items (in batch operations). + Returns a list of indices. Index is set to -1 for items that are not found. + + Parameters + ---------- + uids: list(str) + List of UIDs of the items in the batch to find. + + Returns + ------- + list(int) + List of indices of the items. The list has the same number of items as the list ``uids``. + Indices are set to -1 for the items that were not found in the queue. + """ + queue, _ = await self._get_queue() + + uids_set = set(uids) # Set should be more efficient for searching elements + uid_to_index = {} + for n, plan in enumerate(queue): + uid = plan["item_uid"] + if uid in uids_set: + uid_to_index[uid] = n + + indices = [None] * len(uids) + for n, uid in enumerate(uids): + indices[n] = uid_to_index.get(uid, -1) + + return indices + + # -------------------------------------------------------------------------- + # Operations with UID set + # def _uid_dict_clear(self): + # """ + # Clear ``self._uid_dict``. + # """ + # self._plan_queue_uid = self.new_item_uid() + # self._uid_dict.clear() + + # def _is_uid_in_dict(self, uid): + # """ + # Checks if UID exists in ``self._uid_dict``. + # """ + # return uid in self._uid_dict + + # def _uid_dict_add(self, item): + # """ + # Add UID to ``self._uid_dict``. + # """ + # uid = item["item_uid"] + # if self._is_uid_in_dict(uid): + # raise RuntimeError(f"Trying to add plan with UID '{uid}', which is already in the queue") + # self._plan_queue_uid = self.new_item_uid() + # self._uid_dict.update({uid: item}) + + # def _uid_dict_remove(self, uid): + # """ + # Remove UID from ``self._uid_dict``. + # """ + # if not self._is_uid_in_dict(uid): + # raise RuntimeError(f"Trying to remove plan with UID '{uid}', which is not in the queue") + # self._plan_queue_uid = self.new_item_uid() + # self._uid_dict.pop(uid) + + # def _uid_dict_update(self, item): + # """ + # Update a plan with UID that is already in the dictionary. + # """ + # uid = item["item_uid"] + # if not self._is_uid_in_dict(uid): + # raise RuntimeError(f"Trying to update plan with UID '{uid}', which is not in the queue") + # self._plan_queue_uid = self.new_item_uid() + # self._uid_dict.update({uid: item}) + + # def _uid_dict_get_item(self, uid): + # """ + # Returns a plan with the given UID. + # """ + # return self._uid_dict[uid] + + # async def _uid_dict_initialize(self): + # """ + # Initialize ``self._uid_dict`` with UIDs extracted from the plans in the queue. + # """ + # pq, _ = await self._get_queue() + # self._uid_dict_clear() + # # Go over all plans in the queue + # for item in pq: + # self._uid_dict_add(item) + # # If plan is currently running + # item = await self._get_running_item_info() + # if item: + # self._uid_dict_add(item) + + # ------------------------------------------------------------- + # Currently Running Plan + + async def _is_item_running(self): + """ + See ``self.is_item_running()`` method. + """ + return bool(await self._get_running_item_info()) + + async def is_item_running(self): + """ + Check if an item is set as running. True does not indicate that the plan is actually running. + + Returns + ------- + boolean + True - an item is set as running, False otherwise. + """ + async with self._lock: + return await self._is_item_running() + + async def _get_running_item_info(self): + """ + See ``self._get_running_item_info()`` method. + """ + plan = await self._backend_get(self._name_running_plan) + return json.loads(plan) if plan else {} + + async def get_running_item_info(self): + """ + Read info on the currently running item (plan) from Redis. + + Returns + ------- + dict + Dictionary representing currently running plan. Empty dictionary if + no plan is currently running (key value is ``{}`` or the key does not exist). + """ + async with self._lock: + return await self._get_running_item_info() + + async def _set_running_item_info(self, plan): + """ + Write info on the currently running item (plan) to Redis + + Parameters + ---------- + plan: dict + dictionary that contains plan parameters + """ + await self._backend_set(self._name_running_plan, json.dumps(plan)) + + async def _clear_running_item_info(self): + """ + Clear info on the currently running item (plan) in Redis. + """ + await self._set_running_item_info({}) + + # ------------------------------------------------------------- + # Plan Queue + + async def _get_queue_size(self): + """ + See ``self.get_queue_size()`` method. + """ + return await self._backend_llen(self._name_plan_queue) + + async def get_queue_size(self): + """ + Get the number of plans in the queue. + + Returns + ------- + int + The number of plans in the queue. + """ + async with self._lock: + return await self._get_queue_size() + + async def _get_queue(self): + """ + See ``self.get_queue()`` method. + """ + all_plans_json = await self._backend_lrange(self._name_plan_queue, 0, -1) + return [json.loads(_) for _ in all_plans_json], self._plan_queue_uid + + async def get_queue(self): + """ + Get the list of all items in the queue. The first element of the list is the first + item in the queue. + + Returns + ------- + list(dict) + The list of items in the queue. Each item is represented as a dictionary. + Empty list is returned if the queue is empty. + str + Plan queue UID. + """ + async with self._lock: + return await self._get_queue() + + async def _get_queue_full(self): + plan_queue, plan_queue_uid = await self._get_queue() + running_item = await self._get_running_item_info() + return plan_queue, running_item, plan_queue_uid + + async def get_queue_full(self): + """ + Get the list of all items in the queue and information on currently running item. + The first element of the list is the first item in the queue. This is 'atomic' operation, + i.e. it guarantees that all returned data represent a valid queue state and + the queue was not changed while the data was collected. + + Returns + ------- + list(dict) + The list of items in the queue. Each item is represented as a dictionary. + Empty list is returned if the queue is empty. + dict + Dictionary representing currently running plan. Empty dictionary if + no plan is currently running (key value is ``{}`` or the key does not exist). + str + Plan queue UID. + """ + async with self._lock: + return await self._get_queue_full() + + async def _get_item(self, *, pos=None, uid=None): + """ + See ``self.get_item()`` method. + """ + if (pos is not None) and (uid is not None): + raise ValueError("Ambiguous parameters: plan position and UID is specified") + + if uid is not None: + if not self._is_uid_in_dict(uid): + raise IndexError(f"Item with UID '{uid}' is not in the queue.") + running_item = await self._get_running_item_info() + if running_item and (uid == running_item["item_uid"]): + raise IndexError("The item with UID '{uid}' is currently running.") + item = self._uid_dict_get_item(uid) + + else: + pos = pos if pos is not None else "back" + + if pos == "back": + index = -1 + elif pos == "front": + index = 0 + elif isinstance(pos, int): + index = pos + else: + raise TypeError(f"Parameter 'pos' has incorrect type: pos={str(pos)} (type={type(pos)})") + + item_json = await self._backend_lindex(self._name_plan_queue, index) + if item_json is None: + raise IndexError(f"Index '{index}' is out of range (parameter pos = '{pos}')") + + item = json.loads(item_json) if item_json else {} + + return item + + async def get_item(self, *, pos=None, uid=None): + """ + Get item at a given position or with a given UID. If UID is specified, then + the position is ignored. + + Parameters + ---------- + pos: int, str or None + Position of the element ``(0, ..)`` or ``(-1, ..)``, ``front`` or ``back``. + + uid: str or None + Plan UID of the plan to be retrieved. UID always overrides position. + + Returns + ------- + dict + Dictionary of item parameters. + + Raises + ------ + TypeError + Incorrect value of ``pos`` (most likely a string different from ``front`` or ``back``) + IndexError + No element with position ``pos`` exists in the queue (index is out of range). + """ + async with self._lock: + return await self._get_item(pos=pos, uid=uid) + + async def _remove_item(self, item, single=True): + """ + Remove an item from the queue. If ``single=True`` then the exception is + raised in case of no or multiple matching plans are found in the queue. + The function is not part of user API and shouldn't be used on exception from + the other methods of the class. + + Parameters + ---------- + item: dict + Dictionary of item parameters. Must be identical to the item that is + expected to be deleted. + single: boolean + True - RuntimeError exception is raised if no or more than one matching + plan is found, the plans are removed anyway; False - no exception is + raised. + + Raises + ------ + RuntimeError + No or multiple matching plans are removed and ``single=True``. + """ + n_rem_items = await self._backend_lrem(self._name_plan_queue, 0, json.dumps(item)) + if (n_rem_items != 1) and single: + raise RuntimeError(f"The number of removed items is {n_rem_items}. One item is expected.") + + async def _pop_item_from_queue(self, *, pos=None, uid=None): + """ + See ``self.pop_item_from_queue()`` method + """ + + if (pos is not None) and (uid is not None): + raise ValueError("Ambiguous parameters: plan position and UID is specified") + + pos = pos if pos is not None else "back" + + if uid is not None: + if not self._is_uid_in_dict(uid): + raise IndexError(f"Plan with UID '{uid}' is not in the queue.") + running_item = await self._get_running_item_info() + if running_item and (uid == running_item["item_uid"]): + raise IndexError("Can not remove an item which is currently running.") + item = self._uid_dict_get_item(uid) + await self._remove_item(item) + elif pos == "back": + item_json = await self._backend_rpop(self._name_plan_queue) + if item_json is None: + raise IndexError("Queue is empty") + item = json.loads(item_json) if item_json else {} + elif pos == "front": + item_json = await self._backend_lpop(self._name_plan_queue) + if item_json is None: + raise IndexError("Queue is empty") + item = json.loads(item_json) if item_json else {} + elif isinstance(pos, int): + item = await self._get_item(pos=pos) + if item: + await self._remove_item(item) + else: + raise ValueError(f"Parameter 'pos' has incorrect value: pos={str(pos)} (type={type(pos)})") + + if item: + self._uid_dict_remove(item["item_uid"]) + + qsize = await self._get_queue_size() + + return item, qsize + + async def pop_item_from_queue(self, *, pos=None, uid=None): + """ + Pop a plan from the queue. Raises ``IndexError`` if plan with index ``pos`` is unavailable + or if the queue is empty. + + Parameters + ---------- + pos : int or str or None + Integer index specified position in the queue. Available string values: "front" or "back". + The range for the index is ``-qsize..qsize-1``: ``0, -qsize`` - front element of the queue, + ``-1, qsize-1`` - back element of the queue. If ``pos`` is ``None``, then the plan is popped + from the back of the queue. + uid : str or None + UID of the item to be removed + + Returns + ------- + dict or None + The last plan in the queue represented as a dictionary. + int + The size of the queue after completion of the operation. + + Raises + ------ + ValueError + Incorrect value of the parameter ``pos`` (typically unrecognized string). + IndexError + Position ``pos`` does not exist or the queue is empty. + """ + async with self._lock: + return await self._pop_item_from_queue(pos=pos, uid=uid) + + async def _pop_item_from_queue_batch(self, *, uids=None, ignore_missing=True): + """ + See ``self.pop_item_from_queue_batch()`` method + """ + + uids = uids or [] + + if not isinstance(uids, list): + raise TypeError(f"Parameter 'uids' must be a list: type(uids) = {type(uids)}") + + if not ignore_missing: + # Check if 'uids' contains only unique items + uids_set = set(uids) + if len(uids_set) != len(uids): + raise ValueError(f"The list of contains repeated UIDs ({len(uids) - len(uids_set)} UIDs)") + + # Check if all UIDs in 'uids' exist in the queue + uids_missing = [] + for uid in uids: + if not self._is_uid_in_dict(uid): + uids_missing.append(uid) + if uids_missing: + raise ValueError(f"The queue does not contain items with the following UIDs: {uids_missing}") + + items = [] + for uid in uids: + try: + item, _ = await self._pop_item_from_queue(uid=uid) + items.append(item) + except Exception as ex: + logger.debug("Failed to remove item with UID '%s' from the queue: %s", uid, ex) + + qsize = await self._get_queue_size() + return items, qsize + + async def pop_item_from_queue_batch(self, *, uids=None, ignore_missing=True): + """ + Pop a batch of items from the queue. Raises ``IndexError`` if plan with index ``pos`` is unavailable + or if the queue is empty. + + Parameters + ---------- + uids : list(str) + list of UIDs of items to be removed. The list may be empty. + ignore_missing : boolean + if the parameter is ``True`` (default), then all items from the batch that are found in + the queue are removed, if ``False``, then the function fails if the list contains repeated + entries or at least one of the items is not found in the queue. No items are removed + from the queue if the function fails. + + Returns + ------- + list(dict) + The list of items that were removed from the queue. + int + Size of the queue after completion of the operation. + + Raises + ------ + ValueError + Function failed due to missing or repeated items in the batch. + """ + async with self._lock: + return await self._pop_item_from_queue_batch(uids=uids, ignore_missing=ignore_missing) + + def filter_item_parameters(self, item): + """ + Remove parameters that are not in the list of allowed parameters. + Current parameter list includes parameters ``item_type``, ``item_uid``, + ``name``, ``args``, ``kwargs``, ``meta``, ``user``, ``user_group``. + + The list does not include ``result`` parameter, i.e. ``result`` will + be removed from the list of parameters. + + Parameters + ---------- + item : dict + dictionary of item parameters + + Returns + ------- + dict + dictionary of filtered item parameters + """ + return {k: w for k, w in item.items() if k in self._allowed_item_parameters} + + async def _add_item_to_queue(self, item, *, pos=None, before_uid=None, after_uid=None, filter_parameters=True): + """ + See ``self.add_item_to_queue()`` method. + """ + if (pos is not None) and (before_uid is not None or after_uid is not None): + raise ValueError("Ambiguous parameters: plan position and UID is specified") + + if (before_uid is not None) and (after_uid is not None): + raise ValueError( + "Ambiguous parameters: request to insert the plan before and after the reference plan" + ) + + pos = pos if pos is not None else "back" + + if "item_uid" not in item: + item = await self.set_new_item_uuid(item) + else: + await self._verify_item(item) + + if filter_parameters: + item = self.filter_item_parameters(item) + + qsize0 = await self._get_queue_size() + if isinstance(pos, int): + if (pos == 0) or (pos < -qsize0): + pos = "front" + elif (pos == -1) or (pos >= qsize0): + pos = "back" + + if (before_uid is not None) or (after_uid is not None): + uid = before_uid if before_uid is not None else after_uid + before = uid == before_uid + + if not self._is_uid_in_dict(uid): + raise IndexError(f"Plan with UID '{uid}' is not in the queue.") + running_item = await self._get_running_item_info() + if running_item and (uid == running_item["item_uid"]): + if before: + raise IndexError("Can not insert a plan in the queue before a currently running plan.") + else: + # Push to the plan front of the queue (after the running plan). + qsize = await self._backend_lpush(self._name_plan_queue, json.dumps(item)) + else: + item_to_displace = self._uid_dict_get_item(uid) + where = "BEFORE" if (uid == before_uid) else "AFTER" + qsize = await self._backend_linsert( + self._name_plan_queue, where, json.dumps(item_to_displace), json.dumps(item) + ) + + elif pos == "back": + qsize = await self._backend_rpush(self._name_plan_queue, json.dumps(item)) + elif pos == "front": + qsize = await self._backend_lpush(self._name_plan_queue, json.dumps(item)) + elif isinstance(pos, int): + pos_reference = pos if (pos > 0) else (pos + 1) + + item_to_displace = await self._get_item(pos=pos_reference) + if item_to_displace: + qsize = await self._backend_linsert( + self._name_plan_queue, "BEFORE", json.dumps(item_to_displace), json.dumps(item) + ) + else: + raise RuntimeError(f"Could not find an existing plan at {pos}. Queue size: {qsize0}") + else: + raise ValueError(f"Parameter 'pos' has incorrect value: pos='{str(pos)}' (type={type(pos)})") + + self._uid_dict_add(item) + return item, qsize + + async def add_item_to_queue(self, item, *, pos=None, before_uid=None, after_uid=None, filter_parameters=True): + """ + Add an item to the queue. By default, the item is added to the back of the queue. + If position is integer, it is clipped to fit within the range of meaningful indices. + For the index too large or too low, the plan is pushed to the front or the back of the queue. + + Parameters + ---------- + item : dict + Item (plan or instruction) represented as a dictionary of parameters + pos : int, str or None + Integer that specifies the position index, "front" or "back". + If ``pos`` is in the range ``0..qsize`` (qsize counted before the new + item is inserted), the item is inserted to the specified position + and items at positions ``pos..qsize-1`` are shifted by one position + to the right. If ``-qsizeqsize`` or ``pos==-1``, the plan is added to the back of + the queue. If ``pos==0`` or ``pos<-qsize``, the plan is pushed to + the front of the queue. + before_uid : str or None + If UID is specified, then the item is inserted before the plan with UID. + ``before_uid`` has precedence over ``after_uid``. + after_uid : str or None + If UID is specified, then the item is inserted before the plan with UID. + filter_parameters : boolean (optional) + Remove all parameters that do not belong to the list of allowed parameters. + Default: ``True``. + + Returns + ------- + dict, int + The dictionary that contains the item that was added and the new size of the queue. + + Raises + ------ + ValueError + Incorrect value of the parameter ``pos`` (typically unrecognized string). + TypeError + Incorrect type of ``item`` (should be dict) + """ + async with self._lock: + return await self._add_item_to_queue( + item, pos=pos, before_uid=before_uid, after_uid=after_uid, filter_parameters=filter_parameters + ) + + async def _add_batch_to_queue( + self, items, *, pos=None, before_uid=None, after_uid=None, filter_parameters=True + ): + """ + See ``self.add_batch_to_queue()`` method. + """ + items_added, results = [], [] + + # Approach: attempt ot add each item to queue. In case of a failure remove all added item. + # Operation of adding items to queue is perfectly reversible. + + # Adding a batch to a specified position in the queue: + # - add the first plan by passing all positional parameters to 'self._add_item_to_queue'. + # - add the remaining plans after the first plan (it doesn't matter how position of the first + # plan was defined. + + success = True + added_item_uids = [] # List of plans with successfully added UIDs. Used for 'undo' operation. + for item in items: + try: + if not added_item_uids: + item_added, _ = await self._add_item_to_queue( + item, + pos=pos, + before_uid=before_uid, + after_uid=after_uid, + filter_parameters=filter_parameters, + ) + else: + item_added, _ = await self._add_item_to_queue( + item, after_uid=added_item_uids[-1], filter_parameters=filter_parameters + ) + added_item_uids.append(item_added["item_uid"]) + items_added.append(item_added) + results.append({"success": True, "msg": ""}) + except Exception as ex: + success = False + items_added.append(item) + msg = f"Failed to add item to queue: {ex}" + results.append({"success": False, "msg": msg}) + + # 'Undo' operation: remove all inserted items + if not success: + for uid in added_item_uids: + # The 'try-except' here is just in case anything goes wrong. Operation of removing + # items that were just added are expected to succeed. + try: + await self._pop_item_from_queue(uid=uid) + except Exception as ex: + logger.error( + "Failed to remove an item with uid='%s' after failure to add a batch of plans", ex + ) + + # Also do not return 'changed' items if adding the batch failed. + items_added = items + + qsize = await self._get_queue_size() + + # 'items_added' and 'results' ALWAYS have the same number of elements as 'items' + return items_added, results, qsize, success + + async def add_batch_to_queue( + self, items, *, pos=None, before_uid=None, after_uid=None, filter_parameters=True + ): + """ + Add a batch of item to the queue. The behavior of the function is similar + to ``add_item_to_queue`` except that it accepts the list of items to add. + The function will not add any plans to the queue if at least one of the plans + is rejected. The function returns ``success`` flag, which is ``True`` if + the batch was added and ``False`` otherwise. Success status and error messages + for each added plan can be found in ``results`` list. If the batch was added + successfully, then all ``results`` element indicate success. + + The function is not expected to raise exceptions in case of failure, but instead + report results of processing for each item in the ``results`` list. + + Parameters + ---------- + items : list(dict) + List of items (plans or instructions). Each element of the list is a dictionary + of plan parameters + pos, before_uid, after_uid, filter_parmeters + see documentation for ``add_item_to_queue`` for details. + + Returns + ------- + items_added : list(dict) + List of items that were added to queue. In case the operation fails, the list + of submitted items is returned. The list always has the same number of elements + as ``items``. + results : list(dict) + List of processing results. The list always has the same number of elements as + ``items``. Each element contains a report on processing the respective item in + the form of a dictionary with the keys ``success`` (boolean) and ``msg`` (str). + ``msg`` is always an empty string if ``success==True``. In case the batch was + added successfully, all elements are ``{"success": True, "msg": ""}``. + qsize : int + Size of the queue after the batch was added. The size will not change if the + batch is rejected. + success : bool + Indicates success of the operation of adding the batch. If ``False``, then + the batch is rejected and error messages for each item could be found in + the ``results`` list. + """ + + async with self._lock: + return await self._add_batch_to_queue( + items, pos=pos, before_uid=before_uid, after_uid=after_uid, filter_parameters=filter_parameters + ) + + async def _replace_item(self, item, *, item_uid): + """ + See ``self._replace_item()`` method + """ + # We can not replace currently running item, since it is technically not in the queue + running_item = await self._get_running_item_info() + running_item_uid = running_item["item_uid"] if running_item else None + if not self._is_uid_in_dict(item_uid): + raise RuntimeError(f"Failed to replace item: Item with UID '{item_uid}' is not in the queue") + if (running_item_uid is not None) and (running_item_uid == item_uid): + raise RuntimeError(f"Failed to replace item: Item with UID '{item_uid}' is currently running") + + if "item_uid" not in item: + item = await self.set_new_item_uuid(item) + + item_to_replace = self._uid_dict_get_item(item_uid) + if item == item_to_replace: + # There is nothing to do. Consider operation as successful. + qsize = await self._get_queue_size() + return item, qsize + + # Item with 'item_uid' is expected to be in the queue, so ignore 'item_uid'. + # The concern is that the queue may contain a plan with `item["item_uid"]`, which could be + # different from 'item_uid'. Then inserting the item may violate integrity of the queue. + await self._verify_item(item, ignore_uids=[item_uid]) + + # Parameters of the edited item should be verified against the list of the allowed parameters + item = self.filter_item_parameters(item) + + # Insert the new item after the old one and remove the old one. At this point it is guaranteed + # that they are not equal. + await self._backend_linsert(self._name_plan_queue, "AFTER", json.dumps(item_to_replace), json.dumps(item)) + + await self._remove_item(item_to_replace) + + # Update self._uid_dict + self._uid_dict_remove(item_uid) + self._uid_dict_add(item) + + # Read the actual size of the queue. + qsize = await self._get_queue_size() + + return item, qsize + + async def replace_item(self, item, *, item_uid): + """ + Replace item in the queue. Item with UID ``item_uid`` is replaced by item ``item``. The new + item may have UID which is the same or different from UID of the item being replaced. + + Parameters + ---------- + item: dict + Item (plan or instruction) represented as a dictionary of parameters. + + item_uid: str + UID of existing item in the queue that will be replaced. + + Returns + ------- + dict, int + The dictionary that contains the item that was added and the new size of the queue. + """ + async with self._lock: + return await self._replace_item(item, item_uid=item_uid) + + async def _move_item(self, *, pos=None, uid=None, pos_dest=None, before_uid=None, after_uid=None): + """ + See ``self.move_item()`` method. + """ + if (pos is None) and (uid is None): + raise ValueError("Source position or UID is not specified.") + if (pos_dest is None) and (before_uid is None) and (after_uid is None): + raise ValueError("Destination position or UID is not specified.") + + if (pos is not None) and (uid is not None): + raise ValueError("Ambiguous parameters: Both position and uid is specified for the source plan.") + if (pos_dest is not None) and (before_uid is not None or after_uid is not None): + raise ValueError("Ambiguous parameters: Both position and uid is specified for the destination plan.") + if (before_uid is not None) and (after_uid is not None): + raise ValueError("Ambiguous parameters: source should be moved 'before' and 'after' the destination.") + + queue_size = await self._get_queue_size() + + # Find the source plan + src_txt = "" + src_by_index = False # Indicates that the source is addressed by index + try: + if uid is not None: + src_txt = f"UID '{uid}'" + item_source = await self._get_item(uid=uid) + else: + src_txt = f"position {pos}" + src_by_index = True + item_source = await self._get_item(pos=pos) + except Exception as ex: + raise IndexError(f"Source plan ({src_txt}) was not found: {str(ex)}.") + + uid_source = item_source["item_uid"] + + # Find the destination plan + dest_txt, before = "", True + try: + if (before_uid is not None) or (after_uid is not None): + uid_dest = before_uid if before_uid else after_uid + before = uid_dest == before_uid + dest_txt = f"UID '{uid_dest}'" + item_dest = await self._get_item(uid=uid_dest) + else: + dest_txt = f"position {pos_dest}" + item_dest = await self._get_item(pos=pos_dest) + + # Find the index of the source in the most efficient way + src_index = pos if src_by_index else (await self._get_index_by_uid(uid=uid)) + if src_index == "front": + src_index = 0 + elif src_index == "back": + src_index = queue_size - 1 + + # Determine if the item must be inserted before or after the destination + if pos_dest == "front": + before = True + elif pos_dest == "back": + # This is one case when we need to insert the plan after the 'destination' plan. + before = False + else: + si = src_index if src_index >= 0 else queue_size + src_index + pi = pos_dest if pos_dest >= 0 else queue_size + pos_dest + before = si > pi + + except Exception as ex: + raise IndexError(f"Destination plan ({dest_txt}) was not found: {str(ex)}.") + + # Copy destination UID from the plan (we need it for the case of if addressing is positional + # so we convert it to UID, but we can do it for the case of UID addressing as well) + # In case of positional addressing 'before' is True, so the source is going to be + # inserted in place of destination. + uid_dest = item_dest["item_uid"] + + # If source and destination point to the same plan, then do nothing, + # but consider it a valid operation. + if uid_source != uid_dest: + item, _ = await self._pop_item_from_queue(uid=uid_source) + kw = {"before_uid": uid_dest} if before else {"after_uid": uid_dest} + kw.update({"item": item}) + # The item is moved 'as is'. No filtering of parameters is applied. + item, qsize = await self._add_item_to_queue(**kw, filter_parameters=False) + else: + item = item_dest + qsize = await self._get_queue_size() + return item, qsize + + async def move_item(self, *, pos=None, uid=None, pos_dest=None, before_uid=None, after_uid=None): + """ + Move existing item within the queue. Plan is moved within the queue without modification. + No parameter filtering is applied, so the ``result`` item parameter will not be removed if + present. + + Parameters + ---------- + pos: str or int + Position of the source item: positive or negative integer that specifieds the index + of the item in the queue or a string from the set {"back", "front"}. + uid: str + UID of the source item. UID overrides the position + pos_dest: str or int + Index of the new position of the item in the queue: positive or negative integer that + specifieds the index of the item in the queue or a string from the set {"back", "front"}. + before_uid: str + Insert the item before the item with the given UID. + after_uid: str + Insert the item after the item with the given UID. + + Returns + ------- + dict, int + The dictionary that contains a item that was moved and the size of the queue. + + Raises + ------ + ValueError + Error in specification of source or destination. + """ + async with self._lock: + return await self._move_item( + pos=pos, uid=uid, pos_dest=pos_dest, before_uid=before_uid, after_uid=after_uid + ) + + async def _move_batch(self, *, uids=None, pos_dest=None, before_uid=None, after_uid=None, reorder=False): + """ + See ``self.move_batch()`` method. + """ + uids = uids or [] + + if not isinstance(uids, list): + raise TypeError(f"Parameter 'uids' must be a list: type(uids) = {type(uids)}") + + # Make sure only one of the mutually exclusive parameters is not None + param_list = [pos_dest, before_uid, after_uid] + n_params = len(param_list) - param_list.count(None) + if n_params < 1: + raise ValueError( + "Destination for the batch is not specified: use parameters 'pos_dest', " + "'before_uid' or 'after_uid'" + ) + elif n_params > 1: + raise ValueError( + "The function was called with more than one mutually exclusive parameter " + "('pos_dest', 'before_uid', 'after_uid')" + ) + + # Check if 'uids' contains only unique items + uids_set = set(uids) + if len(uids_set) != len(uids): + raise ValueError(f"The list of contains repeated UIDs ({len(uids) - len(uids_set)} UIDs)") + + # Check if all UIDs in 'uids' exist in the queue + uids_missing = [] + for uid in uids: + if not self._is_uid_in_dict(uid): + uids_missing.append(uid) + if uids_missing: + raise ValueError(f"The queue does not contain items with the following UIDs: {uids_missing}") + + # Check that 'before_uid' and 'after_uid' are not in 'uids' + if (before_uid is not None) and (before_uid in uids): + raise ValueError(f"Parameter 'before_uid': item with UID '{before_uid}' is in the batch") + if (after_uid is not None) and (after_uid in uids): + raise ValueError(f"Parameter 'after_uid': item with UID '{after_uid}' is in the batch") + + # Rearrange UIDs in the list if items need to be reordered + if reorder: + indices = await self._get_index_by_uid_batch(uids=uids) + + def sorting_key(element): + return element[0] + + uids_with_indices = list(zip(indices, uids)) + uids_with_indices.sort(key=sorting_key) + uids_prepared = [_[1] for _ in uids_with_indices] + else: + uids_prepared = uids + + # Perform the 'move' operation. + last_item_uid = None + items_moved = [] + for uid in uids_prepared: + if last_item_uid is None: + # First item is moved according to specified parameters + item, _ = await self._move_item( + uid=uid, pos_dest=pos_dest, before_uid=before_uid, after_uid=after_uid + ) + else: + # Consecutive items are placed after the first item + item, _ = await self._move_item(uid=uid, after_uid=last_item_uid) + last_item_uid = uid + items_moved.append(item) + + qsize = await self._get_queue_size() + return items_moved, qsize + + async def move_batch(self, *, uids=None, pos_dest=None, before_uid=None, after_uid=None, reorder=False): + """ + Move a batch of existing items within the queue. The items are specified as a list of uids. + If at least one of the uids can not be found in the queue, the operation fails and no items + are moved. The moved items can be ordered based on the order of uids in the list (``reorder=False``) + or based on the original order in the queue (``reorder=True``). Plan is moved within the queue + without modification. No parameter filtering is applied, so the ``result`` item parameter + will not be removed if present. + + The items in the batch do not have to be contiguous. Destination items specified by ``after_uid`` + and ``before_uid`` may not belong to the batch. The parameters ``pos_dest``, ``before_uid`` and + ``after_uid`` are mutually exclusive. + + The function raises an exception with error message in case the operation fails. + + Parameters + ---------- + uids : list(str) + List of UIDs of the items in the batch. The list may not contain repeated UIDs. All UIDs + must be present in the queue. + pos_dest : str + Destination of the moved batch. Only string values ``front`` and ``back`` are allowed. + before_uid : str + Insert the batch before the item with the given UID. + after_uid : str + Insert the batch after the item with the given UID. + reorder : boolean + Arranged moved items according to the order of UIDs in the ``uids`` list (``False``) or + according to the original order of items in the queue (``True``). + + Returns + ------- + list(dict), int + The list of items that were moved and the size of the queue. The order of the items + matches the order of the items in the moved batch. Depending on the value of ``reorder`` + it may or may not match the order of ``uids``. + + Raises + ------ + ValueError + Operation could not be performed due to incorrectly specified parameters or invalid + list of ``uids``. + TypeError + Incorrect type of parameter ``uids``. + """ + async with self._lock: + return await self._move_batch( + uids=uids, pos_dest=pos_dest, before_uid=before_uid, after_uid=after_uid, reorder=reorder + ) + + async def _clear_queue(self): + """ + See ``self.clear_queue()`` method. + """ + self._plan_queue_uid = self.new_item_uid() + await self._backend_delete(self._name_plan_queue) + + # Remove all entries from 'self._uid_dict' except the running item. + running_item = await self._get_running_item_info() + if running_item: + uid = running_item["item_uid"] + item = self._uid_dict_get_item(uid) + self._uid_dict_clear() + self._uid_dict_add(item) + else: + self._uid_dict_clear() + + async def clear_queue(self): + """ + Remove all entries from the plan queue. Does not touch the running item (plan). + The item may be pushed back into the queue if it is stopped. + """ + async with self._lock: + await self._clear_queue() + + # ----------------------------------------------------------------------- + # Plan History + + async def _add_to_history(self, item): + """ + Add an item to history. + + Parameters + ---------- + item: dict + Item (plan) represented as a dictionary of parameters. No verifications are performed + on the plan. The function is not intended to be used outside of this class. + + Returns + ------- + int + The new size of the history. + """ + self._plan_history_uid = self.new_item_uid() + history_size = await self._backend_rpush(self._name_plan_history, json.dumps(item)) + return history_size + + async def _get_history_size(self): + """ + See ``self.get_history_size()`` method. + """ + return await self._backend_llen(self._name_plan_history) + + async def get_history_size(self): + """ + Get the number of items in the plan history. + + Returns + ------- + int + The number of plans in the history. + """ + async with self._lock: + return await self._get_history_size() + + async def _get_history(self): + """ + See ``self.get_history()`` method. + """ + all_plans_json = await self._backend_lrange(self._name_plan_history, 0, -1) + return [json.loads(_) for _ in all_plans_json], self._plan_history_uid + + async def get_history(self): + """ + Get the list of all items in the plan history. The first element of the list is + the oldest history entry. + + Returns + ------- + list(dict) + The list of items in the plan history. Each plan is represented as a dictionary. + Empty list is returned if the queue is empty. + str + Plan history UID + """ + async with self._lock: + return await self._get_history() + + async def _clear_history(self): + """ + See ``self.clear_history()`` method. + """ + self._plan_history_uid = self.new_item_uid() + await self._backend_delete(self._name_plan_history) + + async def clear_history(self): + """ + Remove all entries from the plan queue. Does not touch the running item. + The item (plan) may be pushed back into the queue if it is stopped. + """ + async with self._lock: + await self._clear_history() + + # ---------------------------------------------------------------------- + # Standard item operations during queue execution + + # def _clean_item_properties(self, item): + # """ + # The function removes unneccessary item properties before adding the item to history or + # returning it to queue. + # """ + # item = copy.deepcopy(item) + # if "properties" in item: + # p = item["properties"] + # # "immediate_execution" flag is set internally by the server and should not be exposed to users + # if "immediate_execution" in p: + # del p["immediate_execution"] + # # 'time_start' is a temporary parameter and should be removed + # if "time_start" in p: + # del p["time_start"] + # if not p: + # del item["properties"] + # return item + + async def _process_next_item(self, *, item=None): + """ + See ``self.process_next_item()`` method. + """ + loop_mode = self._plan_queue_mode["loop"] + + # Read the item from the front of the queue + + immediate_execution = bool(item) + if immediate_execution: + # Generate UID if it does not exist or creates a deep copy + item = copy.deepcopy(item) if "item_uid" in item else await self.set_new_item_uuid(item) + else: + item = await self._get_item(pos="front") + + item_to_return = item + if item: + item_type = item["item_type"] + if item_type == "plan": + kwargs = {"item": item} if immediate_execution else {} + item_to_return = await self._set_next_item_as_running(**kwargs) + + elif not immediate_execution: + # Items other than plans should be pushed to the back of the queue. + await self._pop_item_from_queue(pos="front") + if loop_mode: + item_to_add = await self.set_new_item_uuid(item) + await self._add_item_to_queue(item_to_add) + + return item_to_return + + async def process_next_item(self, *, item=None): + """ + Process the next item in the queue. If the item is a plan, it is set as currently + running (``self.set_next_item_as_running``), otherwise it is popped from the queue. + If the queue is LOOP mode, then it the item other than plan is pushed to the back + of the queue. (Plans are pushed to the back of the queue upon successful completion.) + + If ``item`` is a plan, then the plan is set for immediate execution. If ``item`` is + an instruction, then the function does nothing. + + If an item submitted for 'immediate' execution has no UID, a new UID is created. + The returned plan is in the exact form, which is set for execution. + + For more details on processing of queue plans, see the description for + ``self.set_next_item_as_running`` method. + """ + async with self._lock: + return await self._process_next_item(item=item) + + async def _set_next_item_as_running(self, *, item=None): + """ + See ``self.set_next_item_as_running()`` method. + """ + immediate_execution = bool(item) + if immediate_execution: + # Generate UID if it does not exist or creates a deep copy + item = copy.deepcopy(item) if "item_uid" in item else await self.set_new_item_uuid(item) + item.setdefault("properties", {})["immediate_execution"] = True + + # UID remains in the `self._uid_dict` after this operation. + try: + if await self._is_item_running(): + raise Exception() + + if immediate_execution: + plan = item + else: + plan = await self._get_item(pos="front") + if not plan: + raise Exception() + + if "item_type" not in plan: + raise Exception() + + if plan["item_type"] != "plan": + raise RuntimeError( + "Function 'PlanQueueOperations.set_next_item_as_running' was called for " + f"an item other than plan: {plan}" + ) + + if not immediate_execution: + # Pop plan from the front of the queue (it is the same plan as currently loaded) + await self._backend_lpop(self._name_plan_queue) + + # Record start time for the plan + plan.setdefault("properties", {})["time_start"] = ttime.time() + + await self._set_running_item_info(plan) + self._plan_queue_uid = self.new_item_uid() + + except RuntimeError: + raise + except Exception: + plan = {} + + return plan + + async def set_next_item_as_running(self, *, item=None): + """ + Sets the next item from the queue as 'running'. The item MUST be a plan + (e.g. not an instruction), otherwise an exception will be raised. The item is removed + from the queue. UID remains in ``self._uid_dict``, i.e. item with the same UID + may not be added to the queue while it is being executed. If ``item`` parameter + represents a plan, then the plan is set for immediate execution. + + This function can only be applied to the plans. Use ``process_next_item`` that + works correctly for any item (it calls this function if an item is a plan). + + If a plan submitted for 'immediate' execution has no UID, a new UID is created. + The returned plan is in the exact form, which is set for execution. + + Parameters + ---------- + item: dict or None + The dictionary that represents a plan submitted for immediate execution. + If ``item`` is a plan, then the queue remains intact and the plan is + set for immediate execution. If ``item=None``, then the top queue item + is removed from the queue and set for execution. + + Returns + ------- + dict + The item that was set as currently running. If another item is currently + set as 'running' or the queue is empty, then ``{}`` is returned. + + Raises + ------ + RuntimeError + The function is called for an item other than plan. + """ + async with self._lock: + return await self._set_next_item_as_running(item=item) + + async def _set_processed_item_as_completed(self, *, exit_status, run_uids, scan_ids, err_msg, err_tb): + """ + See ``self.set_processed_item_as_completed`` method. + """ + # If loop_mode is True, then add item to the back of the queue + loop_mode = self._plan_queue_mode["loop"] + + # Note: UID remains in the `self._uid_dict` after this operation + if await self._is_item_running(): + item = await self._get_running_item_info() + immediate_execution = item["properties"].get("immediate_execution", False) + item_time_start = item["properties"]["time_start"] + item_cleaned = self._clean_item_properties(item) + + if loop_mode and not immediate_execution: + item_to_add = item_cleaned.copy() + item_to_add = await self.set_new_item_uuid(item_to_add) + await self._backend_rpush(self._name_plan_queue, json.dumps(item_to_add)) + self._uid_dict_add(item_to_add) + item_cleaned.setdefault("result", {}) + item_cleaned["result"]["exit_status"] = exit_status + item_cleaned["result"]["run_uids"] = run_uids + item_cleaned["result"]["scan_ids"] = scan_ids + item_cleaned["result"]["time_start"] = item_time_start + item_cleaned["result"]["time_stop"] = ttime.time() + item_cleaned["result"]["msg"] = err_msg + item_cleaned["result"]["traceback"] = err_tb + await self._clear_running_item_info() + if not loop_mode and not immediate_execution: + self._uid_dict_remove(item["item_uid"]) + self._plan_queue_uid = self.new_item_uid() + await self._add_to_history(item_cleaned) + else: + item_cleaned = {} + return item_cleaned + + async def set_processed_item_as_completed(self, *, exit_status, run_uids, scan_ids, err_msg, err_tb): + """ + Moves currently executed item (plan) to history and sets ``exit_status`` key. + UID is removed from ``self._uid_dict``, so a copy of the item with + the same UID may be added to the queue. + + Known ``exit_status`` values: ``"completed"`` and ``"unknown"`` (status + information was lost due to restart of RE Manager, assume that the completion + was successful and start the next plan). + + Parameters + ---------- + exit_status: str + Completion status of the plan. + run_uids: list(str) + A list of uids of completed runs. + scan_ids: list(int) + A list of scan IDs for the completed runs. + err_msg: str + Error message in case of failure. + err_tb: str + Traceback in case of failure. + + Returns + ------- + dict + The item added to the history including ``exit_status``. If another item (plan) + is currently running, then ``{}`` is returned. + """ + async with self._lock: + return await self._set_processed_item_as_completed( + exit_status=exit_status, run_uids=run_uids, scan_ids=scan_ids, err_msg=err_msg, err_tb=err_tb + ) + + async def _set_processed_item_as_stopped(self, *, exit_status, run_uids, scan_ids, err_msg, err_tb): + """ + See ``self.set_processed_item_as_stopped()`` method. + """ + # Note: UID is removed from `self._uid_dict`. + if exit_status == "stopped": + # Stopped item is considered successful, so it is not pushed back to the beginning + # of the queue, and it is added to the back of the queue in LOOP mode. + item_cleaned = await self._set_processed_item_as_completed( + exit_status=exit_status, run_uids=run_uids, scan_ids=scan_ids, err_msg=err_msg, err_tb=err_tb + ) + elif await self._is_item_running(): + item = await self._get_running_item_info() + immediate_execution = item.get("properties", {}).get("immediate_execution", False) + item_time_start = item["properties"]["time_start"] + item_cleaned = self._clean_item_properties(item) + + item_cleaned.setdefault("result", {}) + item_cleaned["result"]["exit_status"] = exit_status + item_cleaned["result"]["run_uids"] = run_uids + item_cleaned["result"]["scan_ids"] = scan_ids + item_cleaned["result"]["time_start"] = item_time_start + item_cleaned["result"]["time_stop"] = ttime.time() + item_cleaned["result"]["msg"] = err_msg + item_cleaned["result"]["traceback"] = err_tb + + await self._add_to_history(item_cleaned) + await self._clear_running_item_info() + + # Generate new UID for the item that is pushed back into the queue. + if not immediate_execution: + self._uid_dict_remove(item["item_uid"]) + # "stopped" - successful completion. Do not insert the item back in the queue. + if exit_status != "stopped": + item_pushed_to_queue = await self.set_new_item_uuid(item_cleaned) + await self._add_item_to_queue(item_pushed_to_queue, pos="front", filter_parameters=False) + self._plan_queue_uid = self.new_item_uid() + else: + item_cleaned = {} + return item_cleaned + + async def set_processed_item_as_stopped(self, *, exit_status, run_uids, scan_ids, err_msg, err_tb): + """ + A stopped plan is considered successfully completed (if ``exit_status=="stopped"``) or + failed (otherwise). All items are added to history with respective ``exit_status``. + Failed items are pushed to the beginning of the queue. Item UID is removed in ``self._uid_dict``. + A new ``item_uid`` is generated for the item that is pushed back into the queue. + + Known ``exit_status`` values: ``"failed"``, ``"stopped"`` (success), ``"aborted"``, ``"halted"``. + + Parameters + ---------- + exit_status: str + Completion status of the plan. + run_uids: list(str) + A list of uids of completed runs. + scan_ids: list(int) + A list of scan IDs for the completed runs. + err_msg: str + Error message in case of failure. + err_tb: str + Traceback in case of failure. + + Returns + ------- + dict + The item (plan) added to the history including ``exit_status``. If no item (plan) + is running, then the function returns ``{}``. The item pushed back into the queue + will have different ``item_uid`` than the item added to the history. + """ + async with self._lock: + return await self._set_processed_item_as_stopped( + exit_status=exit_status, run_uids=run_uids, scan_ids=scan_ids, err_msg=err_msg, err_tb=err_tb + ) + + # ============================================================================================= + # Methods for saving and retrieving user group permissions. + + async def user_group_permissions_clear(self): + """ + Clear user group permissions saved in Redis. + """ + await self._backend_delete(self._name_user_group_permissions) + + async def user_group_permissions_save(self, user_group_permissions): + """ + Save user group permissions to Redis. + + Parameters + ---------- + user_group_permissions: dict + A dictionary containing user group permissions. + """ + await self._backend_set(self._name_user_group_permissions, json.dumps(user_group_permissions)) + + async def user_group_permissions_retrieve(self): + """ + Retreive saved user group permissions. + + Returns + ------- + dict or None + Returns dictionary with saved user group permissions or ``None`` if no permissions are saved. + """ + ugp_json = await self._backend_get(self._name_user_group_permissions) + return json.loads(ugp_json) if ugp_json else None + + # ============================================================================================= + # Methods for saving and retrieving lock info. + + async def lock_info_clear(self): + """ + Clear lock info saved in Redis. + """ + await self._backend_delete(self._name_lock_info) + + async def lock_info_save(self, lock_info): + """ + Save lock info to Redis. + + Parameters + ---------- + lock_info: dict + A dictionary containing lock info. + """ + await self._backend_set(self._name_lock_info, json.dumps(lock_info)) + + async def lock_info_retrieve(self): + """ + Retreive saved lock info. + + Returns + ------- + dict or None + Returns dictionary with saved lock info or ``None`` if no lock info is saved. + """ + lock_info_json = await self._backend_get(self._name_lock_info) + return json.loads(lock_info_json) if lock_info_json else None + + # ============================================================================================= + # Methods for saving and retrieving 'queue_stop_pending'. + + async def stop_pending_clear(self): + """ + Clear 'stop_pending' mode info info saved in Redis. + """ + await self._backend_delete(self._name_stop_pending_info) + + async def stop_pending_save(self, stop_pending): + """ + Save 'stop_pending' mode info to Redis. + + Parameters + ---------- + lock_info: dict + A dictionary containing lock info. + """ + await self._backend_set(self._name_stop_pending_info, json.dumps(stop_pending)) + + async def stop_pending_retrieve(self): + """ + Retreive saved 'stop_pending' mode info. + + Returns + ------- + dict or None + Returns dictionary with saved 'stop_pending' or ``None`` if no 'stop_pending' is saved. + """ + stop_pending_json = await self._backend_get(self._name_stop_pending_info) + return json.loads(stop_pending_json) if stop_pending_json else None + + # ============================================================================================= + # Methods for saving and retrieving 'autostart' mode. + + async def autostart_mode_clear(self): + """ + Clear 'autostart' mode info info saved in Redis. + """ + await self._backend_delete(self._name_autostart_mode_info) + + async def autostart_mode_save(self, lock_info): + """ + Save 'autostart' mode info to Redis. + + Parameters + ---------- + lock_info: dict + A dictionary containing lock info. + """ + await self._backend_set(self._name_autostart_mode_info, json.dumps(lock_info)) + + async def autostart_mode_retrieve(self): + """ + Retreive saved 'autostart' mode info. + + Returns + ------- + dict or None + Returns dictionary with saved lock info or ``None`` if no lock info is saved. + """ + lock_info_json = await self._backend_get(self._name_autostart_mode_info) + return json.loads(lock_info_json) if lock_info_json else None + + # ============================================================================================= + # Proper initialization of the SQLite database + # This method is called when the SQLite backend is used. + + async def _initialize_sqlite_database(self): + """ + Initialize the SQLite database by creating required tables if they do not exist. + """ + async with self._sqlite_conn.cursor() as cursor: + # Create the key-value store table + await cursor.execute(""" + CREATE TABLE IF NOT EXISTS kv_store ( + key TEXT PRIMARY KEY, + value TEXT + ) + """) + + # Create the plan queue table with a UNIQUE constraint on item_uid + await cursor.execute(""" + CREATE TABLE IF NOT EXISTS plan_queue ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT, + item_uid TEXT UNIQUE + ) + """) + + # Create the plan history table with a UNIQUE constraint on item_uid + await cursor.execute(""" + CREATE TABLE IF NOT EXISTS plan_history ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT, + item_uid TEXT UNIQUE + ) + """) + + await self._sqlite_conn.commit() + + # ============================================================================================= + # Abstract helper methods for both redis and sqlite backends + # These methods encapsulate the backend-specific logic. + # This allows the rest of the code to remain agnostic + # of the backend being used. + + async def _backend_set(self, key, value): + """ + Save a key-value pair to the backend (Redis or SQLite). + + Parameters + ---------- + key : str + The key to save the data under. + value : str + The data to save (must be a JSON-encoded string). + + Returns + ------- + None + """ + # Redis Backend + if self._backend == "redis": + await self._r_pool.set(key, value) + + # SQLite Backend + elif self._backend == "sqlite": + value_json = json.dumps(value) + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(""" + CREATE TABLE IF NOT EXISTS kv_store ( + key TEXT PRIMARY KEY, + value TEXT UNIQUE + ) + """) + await cursor.execute( + "INSERT OR REPLACE INTO kv_store (key, value) VALUES (?, ?)", (key, value_json) + ) + await self._sqlite_conn.commit() + + + async def _backend_get(self, key): + """ + Retrieve a value from the backend (Redis or SQLite). + + Parameters + ---------- + key : str + The key to retrieve the data from. + + Returns + ------- + str or None + The value as a JSON-encoded string, or `None` if the key does not exist. + """ + # Redis Backend + if self._backend == "redis": + # Retrieve the value from Redis + value = await self._r_pool.get(key) + return value # Redis returns a string or None + + # SQLite Backend + elif self._backend == "sqlite": + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(""" + CREATE TABLE IF NOT EXISTS kv_store ( + key TEXT PRIMARY KEY, + value TEXT UNIQUE + ) + """) + # Retrieve the value from the SQLite table + await cursor.execute("SELECT value FROM kv_store WHERE key=?", (key,)) + row = await cursor.fetchone() + value = row[0] if row else None # Extract the value if the row exists, otherwise return None + return value + + + async def _backend_delete(self, key): + """ + Delete a key from the backend (Redis or SQLite). + + Parameters + ---------- + key : str + The key to delete. + """ + # Redis Backend + if self._backend == "redis": + await self._r_pool.delete(key) + + # SQLite Backend + elif self._backend == "sqlite": + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(""" + CREATE TABLE IF NOT EXISTS kv_store ( + key TEXT PRIMARY KEY, + value TEXT UNIQUE + ) + """) + # Delete the key from the SQLite table + await cursor.execute("DELETE FROM kv_store WHERE key=?", (key,)) + await self._sqlite_conn.commit() + + async def _backend_rpush(self, key, value): + """ + Append a value to the end of a list in the backend (Redis or SQLite). + + Parameters + ---------- + key : str + The key representing the list. + value : str + The value to append (must be a JSON-encoded string). + + Returns + ------- + int + The size of the list after the operation. + """ + # Redis Backend + if self._backend == "redis": + list_size = await self._r_pool.rpush(key, value) + return list_size + + # SQLite Backend + elif self._backend == "sqlite": + # Extract item_uid from the JSON-encoded value + item = json.loads(value) + item_uid = item.get("item_uid") + if not item_uid: + raise ValueError("Item does not contain 'item_uid'.") + + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {key} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT, + item_uid TEXT UNIQUE + ) + """) + # Insert the value into the SQLite table + try: + await cursor.execute(f"INSERT INTO {key} (value, item_uid) VALUES (?, ?)", (value, item_uid)) + except aiosqlite.IntegrityError: + raise RuntimeError(f"Item with UID '{item_uid}' already exists in the queue.") + + # Get the size of the list (number of rows in the table) + await cursor.execute(f"SELECT COUNT(*) FROM {key}") + row = await cursor.fetchone() + list_size = row[0] if row else 0 + + await self._sqlite_conn.commit() + return list_size + + + async def _backend_lpush(self, key, value): + """ + Prepend a value to the beginning of a list in the backend (Redis or SQLite). + + Parameters + ---------- + key : str + The key representing the list. + value : dict + The value to prepend (will be serialized to JSON). + + Returns + ------- + int + The size of the list after the operation. + """ + # Redis Backend + if self._backend == "redis": + # Prepend the value to the Redis list and return the list size + list_size = await self._r_pool.lpush(key, value) + return list_size + + # SQLite Backend + elif self._backend == "sqlite": + # Extract item_uid from the JSON-encoded value + item = json.loads(value) + item_uid = item.get("item_uid") + if not item_uid: + raise ValueError("Item does not contain 'item_uid'.") + + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {key} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT UNIQUE + ) + """) + # Shift all existing IDs to maintain order + await cursor.execute(f"UPDATE {key} SET id = id + 1") + # Insert the new value at the beginning (id = 1) + await cursor.execute(f"INSERT INTO {key} (id, value) VALUES (1, ?)", (value)) + # Get the size of the list (number of rows in the table) + await cursor.execute(f"SELECT COUNT(*) FROM {key}") + row = await cursor.fetchone() + list_size = row[0] if row else 0 + await self._sqlite_conn.commit() + return list_size + + async def _backend_rpop(self, key): + """ + Remove and return the last element of a list in the backend (Redis or SQLite). + + Parameters + ---------- + key : str + The key representing the list. + + Returns + ------- + dict or None + The last element of the list as a dictionary, or `None` if the list is empty. + """ + # Redis Backend + if self._backend == "redis": + # Remove and return the last element from the Redis list + value = await self._r_pool.rpop(key) + return value + + # SQLite Backend + elif self._backend == "sqlite": + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {key} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT UNIQUE + ) + """) + # Get the last element from the SQLite table + await cursor.execute(f"SELECT id, value FROM {key} ORDER BY id DESC LIMIT 1") + row = await cursor.fetchone() + if row: + value = row[1] + # Remove the last element from the SQLite table + await cursor.execute(f"DELETE FROM {key} WHERE id = ?", (row[0],)) + else: + value = None + await self._sqlite_conn.commit() + + # return json.loads(value_json) if value_json else None + return value if value else None + + async def _backend_lpop(self, key): + """ + Remove and return the first element of a list in the backend (Redis or SQLite). + + Parameters + ---------- + key : str + The key representing the list. + + Returns + ------- + dict or None + The first element of the list as a dictionary, or `None` if the list is empty. + """ + # Redis Backend + if self._backend == "redis": + # Remove and return the first element from the Redis list + value = await self._r_pool.lpop(key) + return value + + # SQLite Backend + elif self._backend == "sqlite": + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {key} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT UNIQUE + ) + """) + # Get the first element from the SQLite table + await cursor.execute(f"SELECT id, value FROM {key} ORDER BY id ASC LIMIT 1") + row = await cursor.fetchone() + if row: + value = row[1] + # Remove the first element from the SQLite table + await cursor.execute(f"DELETE FROM {key} WHERE id = ?", (row[0],)) + # Reorder the IDs to maintain sequential order + await cursor.execute(f"UPDATE {key} SET id = id - 1 WHERE id > ?", (row[0],)) + else: + value = None + await self._sqlite_conn.commit() + + # return json.loads(value_json) if value_json else None + return value if value else None + + async def _backend_llen(self, key): + """ + Get the length of a list in the backend (Redis or SQLite). + + Parameters + ---------- + key : str + The key representing the list. + + Returns + ------- + int + The length of the list. + """ + # Redis Backend + if self._backend == "redis": + # Get the length of the list in Redis + list_length = await self._r_pool.llen(key) + return list_length + + # SQLite Backend + elif self._backend == "sqlite": + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {key} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT UNIQUE + ) + """) + # Get the length of the list (number of rows in the table) + await cursor.execute(f"SELECT COUNT(*) FROM {key}") + row = await cursor.fetchone() + list_length = row[0] if row else 0 + return list_length + + async def _backend_lrange(self, key, start, stop): + """ + Get a range of elements from a list in the backend (Redis or SQLite). + + Parameters + ---------- + key : str + The key representing the list. + start : int + The starting index of the range (inclusive). + stop : int + The ending index of the range (inclusive). Use -1 to indicate the end of the list. + + Returns + ------- + list + A list of elements in the specified range, where each element is a string. + """ + # Redis Backend + if self._backend == "redis": + # Get the range of elements from the Redis list + elements = await self._r_pool.lrange(key, start, stop) + return elements # Redis returns a list of strings + + # SQLite Backend + elif self._backend == "sqlite": + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {key} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT UNIQUE + ) + """) + # Adjust the stop index for SQLite (inclusive range) + if stop == -1: + stop = float('inf') # SQLite doesn't have a concept of -1 for the end of the list + # Get the range of elements from the SQLite table + await cursor.execute(f""" + SELECT value FROM {key} + WHERE id >= ? AND id <= ? + ORDER BY id ASC + """, (start + 1, stop + 1)) # SQLite IDs are 1-based + rows = await cursor.fetchall() + elements = [row[0] for row in rows] # Extract the string values from the rows + return elements + + async def _backend_linsert(self, key, where, pivot, value): + """ + Insert a value into a list in the backend (Redis or SQLite) before or after a pivot element. + + Parameters + ---------- + key : str + The key representing the list. + where : str + Either "BEFORE" or "AFTER", indicating where to insert the value relative to the pivot. + pivot : str + The pivot element (must be a JSON-encoded string). + value : str + The value to insert (must be a JSON-encoded string). + + Returns + ------- + int + The size of the list after the operation. + """ + # Redis Backend + if self._backend == "redis": + result = await self._r_pool.linsert(key, where, pivot, value) + return result + + # SQLite Backend + elif self._backend == "sqlite": + # Extract item_uid from the JSON-encoded value + item = json.loads(value) + item_uid = item.get("item_uid") + if not item_uid: + raise ValueError("Item does not contain 'item_uid'.") + + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {key} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT, + item_uid TEXT UNIQUE + ) + """) + # Check for duplicate item_uid + await cursor.execute(f"SELECT id FROM {key} WHERE item_uid = ?", (item_uid,)) + row = await cursor.fetchone() + if row: + raise RuntimeError(f"Item with UID '{item_uid}' already exists in the queue.") + + # Find the ID of the pivot element + await cursor.execute(f"SELECT id FROM {key} WHERE value = ?", (pivot,)) + row = await cursor.fetchone() + if not row: + raise ValueError(f"Pivot element not found in the list: {pivot}") + pivot_id = row[0] + + # Determine the insertion position + if where == "BEFORE": + insert_position = pivot_id + elif where == "AFTER": + insert_position = pivot_id + 1 + else: + raise ValueError(f"Invalid value for 'where': {where}. Must be 'BEFORE' or 'AFTER'.") + + # Shift IDs to make space for the new element + await cursor.execute(f"UPDATE {key} SET id = id + 1 WHERE id >= ?", (insert_position,)) + # Insert the new value + await cursor.execute(f"INSERT INTO {key} (id, value, item_uid) VALUES (?, ?, ?)", + (insert_position, value, item_uid)) + + # Get the size of the list + await cursor.execute(f"SELECT COUNT(*) FROM {key}") + row = await cursor.fetchone() + result = row[0] if row else 0 + + await self._sqlite_conn.commit() + return result + + async def _backend_lrem(self, key, count, value): + """ + Remove elements from a list in the backend (Redis or SQLite). + + Parameters + ---------- + key : str + The key representing the list. + count : int + The number of matching elements to remove: + - count > 0: Remove the first `count` occurrences of `value`. + - count == 0: Remove all occurrences of `value`. + - count < 0: Remove the last `abs(count)` occurrences of `value`. + value : dict + The value to remove (will be serialized to JSON). + + Returns + ------- + int + The number of removed elements. + """ + # Redis Backend + if self._backend == "redis": + # Remove elements from the Redis list + removed_count = await self._r_pool.lrem(key, count, value) + return removed_count + + # SQLite Backend + elif self._backend == "sqlite": + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {key} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT UNIQUE + ) + """) + + if count > 0: + # Remove the first `count` occurrences + await cursor.execute(f""" + DELETE FROM {key} + WHERE id IN ( + SELECT id FROM {key} WHERE value = ? ORDER BY id ASC LIMIT ? + ) + """, (value, count)) + elif count < 0: + # Remove the last `abs(count)` occurrences + await cursor.execute(f""" + DELETE FROM {key} + WHERE id IN ( + SELECT id FROM {key} WHERE value = ? ORDER BY id DESC LIMIT ? + ) + """, (value, abs(count))) + else: + # Remove all occurrences + await cursor.execute(f"DELETE FROM {key} WHERE value = ?", (value,)) + + # Get the number of rows affected + removed_count = cursor.rowcount + + await self._sqlite_conn.commit() + return removed_count + + async def _backend_lindex(self, key, index): + """ + Get an element at a specific index from a list in the backend (Redis or SQLite). + + Parameters + ---------- + key : str + The key representing the list. + index : int + The index of the element to retrieve. Negative indices are supported. + + Returns + ------- + dict or None + The element at the specified index as a dictionary, or `None` if the index is out of range. + """ + # Redis Backend + if self._backend == "redis": + # Get the element at the specified index from the Redis list + value = await self._r_pool.lindex(key, index) + return value + + # SQLite Backend + elif self._backend == "sqlite": + async with self._sqlite_conn.cursor() as cursor: + # Ensure the table exists + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {key} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT UNIQUE + ) + """) + if index >= 0: + # Positive index + await cursor.execute(f""" + SELECT value FROM {key} + WHERE id = ? + """, (index + 1,)) + else: + # Negative index + await cursor.execute(f""" + SELECT value FROM {key} + WHERE id = ( + SELECT id FROM {key} + ORDER BY id DESC + LIMIT 1 OFFSET ? + ) + """, (-index - 1,)) + row = await cursor.fetchone() + value = row[0] if row else None + + return value if value else None + + async def _backend_ping(self): + """ + Ping the backend (Redis or SQLite) to check if the connection is alive. + + Returns + ------- + bool + True if the backend is reachable, False otherwise. + """ + # Redis Backend + if self._backend == "redis": + try: + # Ping the Redis server + await self._r_pool.ping() + return True + except Exception as ex: + logger.error(f"Redis ping failed: {ex}") + return False + + # SQLite Backend + elif self._backend == "sqlite": + try: + # Execute a simple query to check the SQLite connection + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute("SELECT 1") + return True + except Exception as ex: + logger.error(f"SQLite ping failed: {ex}") + return False + + async def _backend_aclose(self): + """ + Close the backend connection (Redis or SQLite). + + This method ensures that the connection to the backend is properly closed. + + Returns + ------- + None + """ + # Redis Backend + if self._backend == "redis": + if self._r_pool: + await self._r_pool.aclose() + self._r_pool = None + + # SQLite Backend + elif self._backend == "sqlite": + if self._sqlite_conn: + await self._sqlite_conn.close() + self._sqlite_conn = None \ No newline at end of file diff --git a/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_sqlite.py b/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_sqlite.py new file mode 100644 index 00000000..769a04c7 --- /dev/null +++ b/bluesky_queueserver/manager/plan_queue_ops_backends/plan_queue_ops_sqlite.py @@ -0,0 +1,1953 @@ +import asyncio +import aiosqlite +import json +import logging +import os +import time as ttime +import copy +from typing import Any, Dict, List, Optional, Tuple, Union + +from bluesky_queueserver.manager.plan_queue_ops_backends.plan_queue_ops_abstract import AbstractPlanQueueOperations + + +logger = logging.getLogger(__name__) + +class SQLitePlanQueueOperations(AbstractPlanQueueOperations): + """ + Backend implementation for plan queue operations using SQLite database. + + This class provides persistent storage of queue items, history, and other state + information using a SQLite database. It implements all the methods required + by AbstractPlanQueueOperations and utilizes UIDOperations for consistent handling + of unique identifiers. + + Parameters + ---------- + sqlite_db_path : str, optional + Path to the SQLite database file. Default is determined by environment variable + or a default path in the current working directory. + name_prefix : str, optional + Prefix for table names to avoid conflicts with other instances. Default is "qs_default". + """ + def __init__( + self, + sqlite_db_path: Optional[str] = None, + name_prefix: str = "qs_default" + ) -> None: + """ + Initialize SQLitePlanQueueOperations. + + Parameters + ---------- + sqlite_db_path : str, optional + Path to the SQLite database file + name_prefix : str, default "qs_default" + Prefix used for table names + """ + # Initialize the parent class first + super().__init__() + + # Set database path + if not sqlite_db_path: + sqlite_db_path = os.getenv("BACKEND_DB_PATH", os.path.join(os.getcwd(), "bluesky_queue.sqlite")) + + self._sqlite_db_path = sqlite_db_path + self._sqlite_conn = None + self._name_prefix = name_prefix + + # Ensure consistent naming pattern with Redis implementation + if name_prefix: + name_prefix = name_prefix + "_" + + # Table names + self._table_queue = name_prefix + "plan_queue" + self._table_running = name_prefix + "running_plan" + self._table_history = name_prefix + "plan_history" + self._table_plan_queue_mode = name_prefix + "plan_queue_mode" + self._table_user_group_permissions = name_prefix + "user_group_permissions" + self._table_lock_info = name_prefix + "lock_info" + self._table_stop_pending = name_prefix + "stop_pending_info" + self._table_autostart_mode = name_prefix + "autostart_mode_info" + + # Match variable names with code expectations + self._table_plan_queue = self._table_queue + self._table_running_plan = self._table_running + self._table_plan_history = self._table_history + self._table_stop_pending_info = self._table_stop_pending + self._table_autostart_mode_info = self._table_autostart_mode + + # Create lock for thread safety + self._lock = asyncio.Lock() + + async def _initialize_tables(self) -> None: + """ + Initialize database tables if they don't exist. + """ + async with self._sqlite_conn.cursor() as cursor: + # Create plan queue table + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_plan_queue} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + position INTEGER, + item TEXT + ) + """) + + # Create plan history table + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_plan_history} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + item TEXT, + timestamp DATETIME DEFAULT CURRENT_TIMESTAMP + ) + """) + + # Create running plan table + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_running_plan} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + item TEXT + ) + """) + + # Create lock info table + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_lock_info} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + info TEXT + ) + """) + + # Create autostart mode info table + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_autostart_mode_info} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + info TEXT + ) + """) + + # Create stop pending info table + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_stop_pending_info} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + info TEXT + ) + """) + + # Create user group permissions table + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_user_group_permissions} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + info TEXT + ) + """) + + # Create plan queue mode table + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_plan_queue_mode} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + info TEXT + ) + """) + + # Create indexes for efficient lookup + await cursor.execute(f"CREATE INDEX IF NOT EXISTS idx_position ON {self._table_plan_queue} (position)") + + await self._sqlite_conn.commit() + + # -------------------------------------------------------------------------- + # Backend Management + async def start(self) -> None: + """ + Start the SQLite backend connection and initialize tables. + """ + # Only establish connection if none exists + if self._sqlite_conn is None: + try: + # Connect to SQLite database + self._sqlite_conn = await aiosqlite.connect(self._sqlite_db_path) + + # Create tables + await self._initialize_tables() + + # Initialize UID dictionary from existing data + await self._uid_dict_initialize() + + # Initialize plan queue mode if not exists + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue_mode}") + count = await cursor.fetchone() + if count[0] == 0: + # Set default mode if none exists + await cursor.execute( + f"INSERT INTO {self._table_plan_queue_mode} (info) VALUES (?)", + (json.dumps(self._plan_queue_mode_default),) + ) + await self._sqlite_conn.commit() + else: + # Load existing mode + await cursor.execute(f"SELECT info FROM {self._table_plan_queue_mode} LIMIT 1") + row = await cursor.fetchone() + if row: + self._plan_queue_mode = json.loads(row[0]) + + except Exception as ex: + logger.error(f"Failed to connect to SQLite database: {ex}") + if self._sqlite_conn: + await self._sqlite_conn.close() + self._sqlite_conn = None + raise + + async def stop(self) -> None: + """ + Stop the SQLite backend connection. + """ + if self._sqlite_conn: + await self._sqlite_conn.close() + self._sqlite_conn = None + + async def reset(self) -> None: + """ + Reset the backend to its initial state by clearing all data. + """ + async with self._lock: + if self._sqlite_conn: + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"DELETE FROM {self._table_plan_queue}") + await cursor.execute(f"DELETE FROM {self._table_plan_history}") + await cursor.execute(f"DELETE FROM {self._table_running_plan}") + await cursor.execute(f"DELETE FROM {self._table_lock_info}") + await cursor.execute(f"DELETE FROM {self._table_autostart_mode_info}") + await cursor.execute(f"DELETE FROM {self._table_stop_pending_info}") + await cursor.execute(f"DELETE FROM {self._table_user_group_permissions}") + + # Reset plan queue mode to default + await cursor.execute(f"DELETE FROM {self._table_plan_queue_mode}") + default_mode = {"loop": False, "ignore_failures": False} + await cursor.execute( + f"INSERT INTO {self._table_plan_queue_mode} (info) VALUES (?)", + (json.dumps(default_mode),) + ) + + await self._sqlite_conn.commit() + self._uid_dict_clear() + + async def delete_pool_entries(self) -> None: + """ + Delete all pool entries used by this backend. + """ + await self.reset() + + # -------------------------------------------------------------------------- + # UID Operations + async def _uid_dict_initialize(self) -> None: + """ + Initialize the UID dictionary from the current state in SQLite. + """ + self._uid_dict_clear() + + # Add all items from plan queue + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT item FROM {self._table_plan_queue} ORDER BY position") + rows = await cursor.fetchall() + for row in rows: + item = json.loads(row[0]) + if "item_uid" in item: + self._uid_dict_add(item) + + # Add running item if it exists + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + row = await cursor.fetchone() + if row: + item = json.loads(row[0]) + if "item_uid" in item: + self._uid_dict_add(item) + + # -------------------------------------------------------------------------- + # Queue Operations + async def add_item_to_queue( + self, + item: Dict[str, Any], + *, + pos: Optional[Union[int, str]] = None, + before_uid: Optional[str] = None, + after_uid: Optional[str] = None, + filter_parameters: bool = True + ) -> Tuple[Dict[str, Any], int]: + """ + Add an item to the queue. + + Parameters + ---------- + item : dict + Item to add to the queue + pos : int or str, optional + Position in the queue. Can be an integer (0-based index) or string ('front' or 'back') + before_uid : str, optional + UID of the item before which to place the new item + after_uid : str, optional + UID of the item after which to place the new item + filter_parameters : bool, default True + Whether to filter item parameters + + Returns + ------- + tuple + (item, queue_size) - the added item and the new queue size + """ + async with self._lock: + # FIXED: Add validation for required fields + if "name" not in item or not item["name"]: + raise ValueError("Item must have a valid name") + + if (pos is not None) and (before_uid is not None or after_uid is not None): + raise ValueError("Ambiguous parameters: plan position and UID is specified.") + + if (before_uid is not None) and (after_uid is not None): + raise ValueError( + "Ambiguous parameters: request to insert the plan before and after the reference plan." + ) + + pos = pos if pos is not None else "back" + + if "item_uid" not in item: + item = await self.set_new_item_uuid(item) + else: + await self._verify_item(item) + + if filter_parameters: + item = self.filter_item_parameters(item) + + # Get current queue size + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + qsize = (await cursor.fetchone())[0] + + # Determine the position + position = None + if isinstance(pos, int): + if pos <= 0: + position = 0 + elif pos >= qsize: + position = qsize + else: + position = pos + elif pos == "front": + position = 0 + elif pos == "back": + position = qsize + elif before_uid or after_uid: + # Find the index of the item with the specified UID + ref_uid = before_uid if before_uid else after_uid + await cursor.execute( + f"SELECT position FROM {self._table_plan_queue} WHERE json_extract(item, '$.item_uid') = ?", + (ref_uid,) + ) + row = await cursor.fetchone() + if not row: + raise IndexError(f"Item with UID '{ref_uid}' not found in the queue.") + + position = row[0] if before_uid else row[0] + 1 + else: + raise ValueError(f"Invalid value for 'pos': {pos}") + + # Shift positions for items that come after the insertion point + await cursor.execute( + f"UPDATE {self._table_plan_queue} SET position = position + 1 WHERE position >= ?", + (position,) + ) + + # Insert the new item + item_json = json.dumps(item) + await cursor.execute( + f"INSERT INTO {self._table_plan_queue} (position, item) VALUES (?, ?)", + (position, item_json) + ) + + # Add item to UID dictionary + self._uid_dict_add(item) + + # Get the new queue size + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + new_qsize = (await cursor.fetchone())[0] + + await self._sqlite_conn.commit() + return item, new_qsize + + async def add_batch_to_queue( + self, + items: List[Dict[str, Any]], + *, + pos: Optional[Union[int, str]] = None, + before_uid: Optional[str] = None, + after_uid: Optional[str] = None, + filter_parameters: bool = True + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], int, bool]: + """ + Add a batch of items to the queue. + + Parameters + ---------- + items : list + List of items to add to the queue + pos : int or str, optional + Position in the queue. Can be an integer (0-based index) or string ('front' or 'back') + before_uid : str, optional + UID of the item before which to place the new items + after_uid : str, optional + UID of the item after which to place the new items + filter_parameters : bool, default True + Whether to filter item parameters + + Returns + ------- + tuple + (items_added, results, queue_size, success) - added items, results of each operation, + new queue size, and overall success flag + """ + async with self._lock: + items_added = [] + results = [] + success = True + added_item_uids = [] # For rollback in case of failure + + try: + # Begin transaction + await self._sqlite_conn.execute("BEGIN TRANSACTION") + + for item in items: + try: + if not added_item_uids: + # First item is placed according to parameters + item_added, _ = await self.add_item_to_queue( + item, + pos=pos, + before_uid=before_uid, + after_uid=after_uid, + filter_parameters=filter_parameters, + ) + else: + # Subsequent items are placed after the previous one + item_added, _ = await self.add_item_to_queue( + item, + after_uid=added_item_uids[-1], + filter_parameters=filter_parameters + ) + + added_item_uids.append(item_added["item_uid"]) + items_added.append(item_added) + results.append({"success": True, "msg": ""}) + except Exception as ex: + success = False + items_added.append(item) + msg = f"Failed to add item to queue: {ex}" + results.append({"success": False, "msg": msg}) + # No raise here - this was causing the test to hang + break # Stop processing more items if one fails + + # Handle transaction based on overall success + if success: + await self._sqlite_conn.commit() + else: + await self._sqlite_conn.rollback() + items_added = items # Reset to original items for consistent return value + + # Get the new queue size - always needs to be awaited + qsize = 0 + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + row = await cursor.fetchone() + if row: + qsize = row[0] + + return items_added, results, qsize, success + + except Exception: + # Ensure transaction is rolled back in case of error + await self._sqlite_conn.rollback() + queue_size = await self.get_queue_size() # Ensure we await this + return items, results, queue_size, False + + async def pop_item_from_queue( + self, + *, + pos: Optional[Union[int, str]] = None, + uid: Optional[str] = None + ) -> Tuple[Optional[Dict[str, Any]], int]: + """ + Remove and return an item from the queue. + + Parameters + ---------- + pos : int or str, optional + Position in the queue. Can be an integer (0-based index) or string ('front' or 'back') + uid : str, optional + UID of the item to remove + + Returns + ------- + tuple + (item, queue_size) - the removed item and the new queue size + """ + async with self._lock: + if (pos is not None) and (uid is not None): + raise ValueError("Ambiguous parameters: plan position and UID is specified") + + pos = pos if pos is not None else "front" + + async with self._sqlite_conn.cursor() as cursor: + item = None + + if uid is not None: + # Check if item with this UID exists + if not self._is_uid_in_dict(uid): + raise IndexError(f"Item with UID '{uid}' is not in the queue.") + + # Check if the item is not currently running + await cursor.execute( + f"SELECT item FROM {self._table_running_plan} WHERE json_extract(item, '$.item_uid') = ?", + (uid,) + ) + row = await cursor.fetchone() + if row: + raise IndexError(f"Cannot remove an item which is currently running.") + + # Get the item and its position + await cursor.execute( + f"SELECT position, item FROM {self._table_plan_queue} WHERE json_extract(item, '$.item_uid') = ?", + (uid,) + ) + row = await cursor.fetchone() + if not row: + raise IndexError(f"Item with UID '{uid}' not found in the queue.") + + position = row[0] + item = json.loads(row[1]) + + # Delete the item + await cursor.execute( + f"DELETE FROM {self._table_plan_queue} WHERE json_extract(item, '$.item_uid') = ?", + (uid,) + ) + + # Remove from UID dictionary + self._uid_dict_remove(uid) + + elif pos == "front": + await cursor.execute( + f"SELECT item FROM {self._table_plan_queue} ORDER BY position LIMIT 1" + ) + row = await cursor.fetchone() + if not row: + return None, 0 + + item = json.loads(row[0]) + + # Delete the item + await cursor.execute( + f"DELETE FROM {self._table_plan_queue} WHERE position = 0" + ) + + # Remove from UID dictionary + self._uid_dict_remove(item["item_uid"]) + + elif pos == "back": + # Find the position of the last item + await cursor.execute( + f"SELECT MAX(position) FROM {self._table_plan_queue}" + ) + row = await cursor.fetchone() + if not row or row[0] is None: + return None, 0 + + max_pos = row[0] + + await cursor.execute( + f"SELECT item FROM {self._table_plan_queue} WHERE position = ?", + (max_pos,) + ) + row = await cursor.fetchone() + if not row: + return None, 0 + + item = json.loads(row[0]) + + # Delete the item + await cursor.execute( + f"DELETE FROM {self._table_plan_queue} WHERE position = ?", + (max_pos,) + ) + + # Remove from UID dictionary + self._uid_dict_remove(item["item_uid"]) + + elif isinstance(pos, int): + # Adjust negative indices + if pos < 0: + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + qsize = (await cursor.fetchone())[0] + pos = qsize + pos + + if pos < 0: + raise IndexError(f"Position {pos} out of range") + + await cursor.execute( + f"SELECT item FROM {self._table_plan_queue} WHERE position = ?", + (pos,) + ) + row = await cursor.fetchone() + if not row: + raise IndexError(f"Position {pos} out of range") + + item = json.loads(row[0]) + + # Delete the item + await cursor.execute( + f"DELETE FROM {self._table_plan_queue} WHERE position = ?", + (pos,) + ) + + # Remove from UID dictionary + self._uid_dict_remove(item["item_uid"]) + else: + raise ValueError(f"Invalid value for 'pos': {pos}") + + # Reposition remaining items to eliminate gaps + await cursor.execute(f""" + UPDATE {self._table_plan_queue} + SET position = ( + SELECT COUNT(*) - 1 + FROM {self._table_plan_queue} AS t2 + WHERE t2.position <= {self._table_plan_queue}.position + AND t2.id != {self._table_plan_queue}.id + ) + """) + + # Get the new queue size + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + qsize = (await cursor.fetchone())[0] + + await self._sqlite_conn.commit() + return item, qsize + + async def pop_item_from_queue_batch( + self, + *, + uids: Optional[List[str]] = None, + ignore_missing: bool = True + ) -> Tuple[List[Dict[str, Any]], int]: + """ + Remove and return multiple items from the queue. + + Parameters + ---------- + uids : list, optional + List of UIDs of the items to remove + ignore_missing : bool, default True + Whether to ignore missing items + + Returns + ------- + tuple + (items, queue_size) - the removed items and the new queue size + """ + async with self._lock: + uids = uids or [] + + if not isinstance(uids, list): + raise TypeError(f"Parameter 'uids' must be a list: type(uids) = {type(uids)}") + + if not ignore_missing: + # Check if 'uids' contains only unique items + uids_set = set(uids) + if len(uids_set) != len(uids): + raise ValueError(f"The list contains repeated UIDs ({len(uids) - len(uids_set)} UIDs)") + + # Check if all UIDs in 'uids' exist in the queue + uids_missing = [] + for uid in uids: + if not self._is_uid_in_dict(uid): + uids_missing.append(uid) + if uids_missing: + raise ValueError(f"The queue does not contain items with the following UIDs: {uids_missing}") + + # Begin transaction + await self._sqlite_conn.execute("BEGIN TRANSACTION") + + items = [] + try: + for uid in uids: + try: + item, _ = await self.pop_item_from_queue(uid=uid) + if item: # Add null check + items.append(item) + except Exception as ex: + if not ignore_missing: + await self._sqlite_conn.rollback() # Ensure rollback on exception + raise + logger.debug(f"Failed to remove item with UID '{uid}' from the queue: {ex}") + + await self._sqlite_conn.commit() + + # Get the new queue size + qsize = 0 # Default value in case of empty result + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + row = await cursor.fetchone() + if row: + qsize = row[0] + + return items, qsize + except Exception as ex: + # Make sure the rollback is awaited + await self._sqlite_conn.rollback() + logger.error(f"Error in pop_item_from_queue_batch: {ex}") + raise + + async def clear_queue(self) -> None: + """Clear the plan queue.""" + async with self._lock: + async with self._sqlite_conn.cursor() as cursor: + # Get the running item's UID if it exists + await cursor.execute(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + row = await cursor.fetchone() + running_uid = None + if row: + running_item = json.loads(row[0]) + running_uid = running_item.get('item_uid') + + # FIXED: Complete reset of UID dictionary + self._uid_dict_clear() + + # Re-add just the running item if it exists + if running_uid and row: + running_item = json.loads(row[0]) + self._uid_dict_add(running_item) + + # Clear the queue table + await cursor.execute(f"DELETE FROM {self._table_plan_queue}") + + await self._sqlite_conn.commit() + + async def get_queue_size(self) -> int: + """ + Get the size of the queue. + + Returns + ------- + int + The number of items in the queue + """ + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + return (await cursor.fetchone())[0] + + async def get_queue_full(self) -> Tuple[List[Dict[str, Any]], Dict[str, Any], str]: + """ + Retrieve the full queue. + + Returns + ------- + tuple + (queue_items, running_item, queue_uid) - all items in the queue, the currently running item, + and a unique identifier for this queue state + """ + async with self._lock: + # Get all items from queue ordered by position + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT item FROM {self._table_plan_queue} ORDER BY position") + rows = await cursor.fetchall() + queue = [json.loads(row[0]) for row in rows] + + # Get running item if exists + await cursor.execute(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + row = await cursor.fetchone() + running_item = json.loads(row[0]) if row else {} + + # Generate a unique identifier for this queue state + queue_uid = self.new_item_uid() + + return queue, running_item, queue_uid + + async def move_item( + self, + *, + pos: Optional[Union[int, str]] = None, + uid: Optional[str] = None, + pos_dest: Optional[Union[int, str]] = None, + before_uid: Optional[str] = None, + after_uid: Optional[str] = None + ) -> Tuple[Dict[str, Any], int]: + """ + Move an item to a new position in the queue. + + Parameters + ---------- + pos : int or str, optional + Source position in the queue + uid : str, optional + UID of the item to move + pos_dest : int or str, optional + Destination position + before_uid : str, optional + UID of the item before which to place the moved item + after_uid : str, optional + UID of the item after which to place the moved item + + Returns + ------- + tuple + (item, queue_size) - the moved item and the new queue size + """ + async with self._lock: + if (pos is None) and (uid is None): + raise ValueError("Source position or UID is not specified.") + if (pos_dest is None) and (before_uid is None) and (after_uid is None): + raise ValueError("Destination position or UID is not specified.") + + if (pos is not None) and (uid is not None): + raise ValueError("Ambiguous parameters: Both position and uid is specified for the source plan.") + if (pos_dest is not None) and (before_uid is not None or after_uid is not None): + raise ValueError("Ambiguous parameters: Both position and uid is specified for the destination plan.") + if (before_uid is not None) and (after_uid is not None): + raise ValueError("Ambiguous parameters: source should be moved 'before' and 'after' the destination.") + + # Begin transaction + await self._sqlite_conn.execute("BEGIN TRANSACTION") + + try: + async with self._sqlite_conn.cursor() as cursor: + # Find the source item and its position + source_pos = None + item = None + + if uid is not None: + await cursor.execute( + f"SELECT position, item FROM {self._table_plan_queue} WHERE json_extract(item, '$.item_uid') = ?", + (uid,) + ) + row = await cursor.fetchone() + if not row: + raise IndexError(f"Item with UID '{uid}' not found in the queue.") + source_pos = row[0] + item = json.loads(row[1]) + else: + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + qsize = (await cursor.fetchone())[0] + if pos == "front": + source_pos = 0 + elif pos == "back": + source_pos = qsize - 1 + elif isinstance(pos, int): + if pos < 0: + source_pos = qsize + pos + else: + source_pos = pos + + if source_pos < 0 or source_pos >= qsize: + raise IndexError(f"Position {pos} is out of range.") + else: + raise ValueError(f"Invalid value for 'pos': {pos}") + + await cursor.execute( + f"SELECT item FROM {self._table_plan_queue} WHERE position = ?", + (source_pos,) + ) + row = await cursor.fetchone() + if not row: + raise IndexError(f"No item found at position {source_pos}.") + item = json.loads(row[0]) + + # Find the destination position + dest_pos = None + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + qsize = (await cursor.fetchone())[0] + + if before_uid is not None or after_uid is not None: + dest_uid = before_uid if before_uid is not None else after_uid + before = dest_uid == before_uid + + await cursor.execute( + f"SELECT position FROM {self._table_plan_queue} WHERE json_extract(item, '$.item_uid') = ?", + (dest_uid,) + ) + row = await cursor.fetchone() + if not row: + raise IndexError(f"Item with UID '{dest_uid}' not found in the queue.") + + dest_pos = row[0] + if not before: # If after_uid, increment the dest_pos + dest_pos += 1 + else: + if pos_dest == "front": + dest_pos = 0 + elif pos_dest == "back": + dest_pos = qsize + elif isinstance(pos_dest, int): + if pos_dest < 0: + dest_pos = qsize + pos_dest + else: + dest_pos = pos_dest + + if dest_pos < 0 or dest_pos > qsize: + raise IndexError(f"Position {pos_dest} is out of range.") + else: + raise ValueError(f"Invalid value for 'pos_dest': {pos_dest}") + + # If source_pos is the same as dest_pos, no need to move + if source_pos == dest_pos: + await self._sqlite_conn.rollback() # No changes needed + return item, qsize + + # If source is before destination, adjust dest_pos since array length will change + if source_pos < dest_pos: + dest_pos -= 1 + + # Remove item from source position + await cursor.execute( + f"DELETE FROM {self._table_plan_queue} WHERE position = ?", + (source_pos,) + ) + + # Reposition remaining items to eliminate gaps + await cursor.execute(f""" + UPDATE {self._table_plan_queue} + SET position = ( + SELECT COUNT(*) - 1 + FROM {self._table_plan_queue} AS t2 + WHERE t2.position <= {self._table_plan_queue}.position + AND t2.id != {self._table_plan_queue}.id + ) + """) + + # Make space at destination position + await cursor.execute( + f"UPDATE {self._table_plan_queue} SET position = position + 1 WHERE position >= ?", + (dest_pos,) + ) + + # Insert item at destination position + item_json = json.dumps(item) + await cursor.execute( + f"INSERT INTO {self._table_plan_queue} (position, item) VALUES (?, ?)", + (dest_pos, item_json) + ) + + # Get the new queue size + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + qsize = (await cursor.fetchone())[0] + + await self._sqlite_conn.commit() + return item, qsize + except Exception: + await self._sqlite_conn.rollback() + raise + + async def move_batch( + self, + *, + uids: Optional[List[str]] = None, + pos_dest: Optional[Union[int, str]] = None, + before_uid: Optional[str] = None, + after_uid: Optional[str] = None, + reorder: bool = False + ) -> Tuple[List[Dict[str, Any]], int]: + """ + Move a batch of items within the queue. + + Parameters + ---------- + uids : list, optional + List of UIDs of the items to move + pos_dest : int or str, optional + Destination position for the first item + before_uid : str, optional + UID of the item before which to place the first moved item + after_uid : str, optional + UID of the item after which to place the first moved item + reorder : bool, default False + Whether to reorder items according to their current order in the queue + + Returns + ------- + tuple + (items, queue_size) - the moved items and the new queue size + """ + async with self._lock: + uids = uids or [] + + if not isinstance(uids, list): + raise TypeError(f"Parameter 'uids' must be a list: type(uids) = {type(uids)}") + + # Make sure only one of the mutually exclusive parameters is not None + param_list = [pos_dest, before_uid, after_uid] + n_params = len(param_list) - param_list.count(None) + if n_params < 1: + raise ValueError( + "Destination for the batch is not specified: use parameters 'pos_dest', " + "'before_uid' or 'after_uid'" + ) + elif n_params > 1: + raise ValueError( + "The function was called with more than one mutually exclusive parameter " + "('pos_dest', 'before_uid', 'after_uid')" + ) + + # Check if 'uids' contains only unique items + uids_set = set(uids) + if len(uids_set) != len(uids): + raise ValueError(f"The list contains repeated UIDs ({len(uids) - len(uids_set)} UIDs)") + + # Begin transaction + await self._sqlite_conn.execute("BEGIN TRANSACTION") + + try: + async with self._sqlite_conn.cursor() as cursor: + # Check if all UIDs in 'uids' exist in the queue + for uid in uids: + await cursor.execute( + f"SELECT COUNT(*) FROM {self._table_plan_queue} WHERE json_extract(item, '$.item_uid') = ?", + (uid,) + ) + count = (await cursor.fetchone())[0] + if count == 0: + raise ValueError(f"The queue does not contain an item with UID: {uid}") + + # Check that 'before_uid' and 'after_uid' are not in 'uids' + if (before_uid is not None) and (before_uid in uids): + raise ValueError(f"Parameter 'before_uid': item with UID '{before_uid}' is in the batch") + if (after_uid is not None) and (after_uid in uids): + raise ValueError(f"Parameter 'after_uid': item with UID '{after_uid}' is in the batch") + + # If reorder is True, arrange UIDs based on their current positions in the queue + if reorder: + uids_with_positions = [] + for uid in uids: + await cursor.execute( + f"SELECT position FROM {self._table_plan_queue} WHERE json_extract(item, '$.item_uid') = ?", + (uid,) + ) + position = (await cursor.fetchone())[0] + uids_with_positions.append((position, uid)) + + uids_with_positions.sort(key=lambda x: x[0]) + uids_prepared = [pair[1] for pair in uids_with_positions] + else: + uids_prepared = uids + + # Perform the 'move' operation + last_item_uid = None + items_moved = [] + + for uid in uids_prepared: + try: + if last_item_uid is None: + # First item is moved according to specified parameters + item, _ = await self.move_item( + uid=uid, pos_dest=pos_dest, before_uid=before_uid, after_uid=after_uid + ) + else: + # Consecutive items are placed after the previous item + item, _ = await self.move_item(uid=uid, after_uid=last_item_uid) + + last_item_uid = uid + items_moved.append(item) + except Exception as ex: + await self._sqlite_conn.rollback() # Ensure rollback is awaited + raise RuntimeError(f"Error moving item {uid}: {str(ex)}") from ex + + # Get the new queue size + qsize = 0 + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + row = await cursor.fetchone() + if row: + qsize = row[0] + + await self._sqlite_conn.commit() + return items_moved, qsize + + except Exception as ex: + # Ensure transaction is rolled back in case of error + await self._sqlite_conn.rollback() + raise + + async def replace_item( + self, + item: Dict[str, Any], + *, + item_uid: str + ) -> Tuple[Dict[str, Any], int]: + """ + Replace an existing item in the queue with a new item. + + Parameters + ---------- + item : dict + The new item to replace the existing one + item_uid : str + UID of the item to replace + + Returns + ------- + tuple + (old_item, queue_size) - the replaced item and the new queue size + """ + async with self._lock: + # Verify the new item, ignoring the UID of the item being replaced + await self._verify_item(item, ignore_uids=[item_uid]) + + async with self._sqlite_conn.cursor() as cursor: + # Find the item to replace + await cursor.execute( + f"SELECT position, item FROM {self._table_plan_queue} WHERE json_extract(item, '$.item_uid') = ?", + (item_uid,) + ) + row = await cursor.fetchone() + if not row: + raise ValueError(f"Item with UID '{item_uid}' not found in the queue.") + + old_item = json.loads(row[1]) + position = row[0] + + # Replace the item + item_json = json.dumps(item) + await cursor.execute( + f"UPDATE {self._table_plan_queue} SET item = ? WHERE position = ?", + (item_json, position) + ) + + # Update UID dictionary + self._uid_dict_remove(item_uid) + self._uid_dict_add(item) + + # Get the new queue size + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + qsize = (await cursor.fetchone())[0] + + await self._sqlite_conn.commit() + return old_item, qsize + + # -------------------------------------------------------------------------- + # History Management + async def clear_history(self) -> None: + """Clear the plan history.""" + async with self._lock: + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"DELETE FROM {self._table_plan_history}") + await self._sqlite_conn.commit() + + async def get_history(self) -> Tuple[List[Dict[str, Any]], str]: + """ + Retrieve the full history. + + Returns + ------- + tuple + (history_items, history_uid) - all items in the history and a unique identifier for this history state + """ + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute( + f"SELECT item FROM {self._table_plan_history} ORDER BY id DESC" + ) + rows = await cursor.fetchall() + history = [json.loads(row[0]) for row in rows] + history_uid = self.new_item_uid() # Generate a unique ID for this history state + return history, history_uid + + async def get_history_size(self) -> int: + """ + Get the size of the history. + + Returns + ------- + int + The number of items in the history + """ + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_history}") + return (await cursor.fetchone())[0] + + # -------------------------------------------------------------------------- + # Running Item Management + async def get_running_item_info(self) -> Dict[str, Any]: + """ + Get information about the currently running item. + + Returns + ------- + dict + Information about the currently running item, or an empty dict if no item is running + """ + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + row = await cursor.fetchone() + return json.loads(row[0]) if row else {} + + async def _clean_item_properties(self, item: Dict[str, Any]) -> Dict[str, Any]: + """Clean and set default item properties.""" + item = copy.deepcopy(item) + + # Add default values for required fields if missing + if "args" not in item: + item["args"] = [] + if "kwargs" not in item: + item["kwargs"] = {} + if "item_type" not in item: + item["item_type"] = "plan" + if "user" not in item: + item["user"] = None + if "user_group" not in item: + item["user_group"] = None + + # Clean existing properties if present + if "properties" in item: + p = item["properties"] + # ONLY remove these specific properties + if "immediate_execution" in p: + del p["immediate_execution"] + if "time_start" in p: + del p["time_start"] + if not p: + del item["properties"] + + return item + + async def set_next_item_as_running(self, *, item: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + """ + Set the next item from the queue as running. + + Parameters + ---------- + item : dict, optional + Item for immediate execution + + Returns + ------- + dict + The item that is now running + """ + async with self._lock: + async with self._sqlite_conn.cursor() as cursor: + # Check if an item is already running + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_running_plan}") + count = (await cursor.fetchone())[0] + if count > 0: + raise RuntimeError("An item is already running") + + immediate_execution = bool(item) + plan = None + + if immediate_execution: + # Generate UID if it doesn't exist + plan = copy.deepcopy(item) + if "item_uid" not in plan: + plan = await self.set_new_item_uuid(plan) + + plan.setdefault("properties", {})["immediate_execution"] = True + else: + # Get the first item from the queue + await cursor.execute( + f"SELECT item FROM {self._table_plan_queue} ORDER BY position LIMIT 1" + ) + row = await cursor.fetchone() + if not row: + return {} # Queue is empty + + plan = json.loads(row[0]) + + # Remove the plan from the queue + await cursor.execute( + f"DELETE FROM {self._table_plan_queue} WHERE position = 0" + ) + + # Reposition remaining items + await cursor.execute(f""" + UPDATE {self._table_plan_queue} + SET position = ( + SELECT COUNT(*) - 1 + FROM {self._table_plan_queue} AS t2 + WHERE t2.position <= {self._table_plan_queue}.position + AND t2.id != {self._table_plan_queue}.id + ) + """) + + # Verify that the item is a plan + if "item_type" not in plan: + raise ValueError("Item does not have 'item_type'") + + if plan["item_type"] != "plan": + raise RuntimeError( + "Function 'set_next_item_as_running' was called for " + f"an item other than plan: {plan}" + ) + + # Record start time + plan.setdefault("properties", {})["time_start"] = ttime.time() + + # Set as running plan + plan_json = json.dumps(plan) + await cursor.execute( + f"INSERT INTO {self._table_running_plan} (item) VALUES (?)", + (plan_json,) + ) + + # Add to UID dictionary + self._uid_dict_add(plan) + + await self._sqlite_conn.commit() + return plan + + async def set_processed_item_as_completed( + self, + *, + exit_status: str, + run_uids: List[str], + scan_ids: List[int], + err_msg: str = "", + err_tb: str = "" + ) -> Dict[str, Any]: + """ + Mark the currently running item as completed and move it to history. + + Parameters + ---------- + exit_status : str + Exit status of the item + run_uids : list + List of run UIDs + scan_ids : list + List of scan IDs + err_msg : str, default "" + Error message if any + err_tb : str, default "" + Error traceback if any + + Returns + ------- + dict + The completed item + """ + async with self._lock: + # Check if loop mode is enabled in queue mode + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + row = await cursor.fetchone() + if not row: + return {} + + running_item = json.loads(row[0]) + immediate_execution = running_item.get("properties", {}).get("immediate_execution", False) + item_time_start = running_item.get("properties", {}).get("time_start", ttime.time()) + + # Check if loop mode is enabled + loop_enabled = False + await cursor.execute( + f"SELECT info FROM {self._table_plan_queue_mode} LIMIT 1" + ) + row = await cursor.fetchone() + if row: + mode = json.loads(row[0]) + loop_enabled = mode.get("loop", False) + + # Clean item properties + item_cleaned = await self._clean_item_properties(running_item) + + if loop_enabled and not immediate_execution: + # Add a copy of the item back to the queue + item_to_add = copy.deepcopy(item_cleaned) + item_to_add = await self.set_new_item_uuid(item_to_add) + + # Get the current size of the queue + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + qsize = (await cursor.fetchone())[0] + + # Add to the end of the queue + item_json = json.dumps(item_to_add) + await cursor.execute( + f"INSERT INTO {self._table_plan_queue} (position, item) VALUES (?, ?)", + (qsize, item_json) + ) + + # Add to UID dictionary + self._uid_dict_add(item_to_add) + + # Add result information to the item + item_cleaned.setdefault("result", {}) + item_cleaned["result"]["exit_status"] = exit_status + item_cleaned["result"]["run_uids"] = run_uids + item_cleaned["result"]["scan_ids"] = scan_ids + item_cleaned["result"]["time_start"] = item_time_start + item_cleaned["result"]["time_stop"] = ttime.time() + item_cleaned["result"]["msg"] = err_msg + item_cleaned["result"]["traceback"] = err_tb + + # Add to history + item_json = json.dumps(item_cleaned) + await cursor.execute( + f"INSERT INTO {self._table_plan_history} (item) VALUES (?)", + (item_json,) + ) + + # Remove from running + await cursor.execute(f"DELETE FROM {self._table_running_plan}") + + # Remove from UID dictionary if not in loop mode + if not loop_enabled and not immediate_execution: + self._uid_dict_remove(running_item["item_uid"]) + + await self._sqlite_conn.commit() + return item_cleaned + + async def set_processed_item_as_stopped( + self, + *, + exit_status: str, + run_uids: Optional[List[str]] = None, + scan_ids: Optional[List[int]] = None, + err_msg: str = "", + err_tb: str = "" + ) -> Dict[str, Any]: + """ + Mark the currently running item as stopped and move it to history. + If exit_status is not "stopped", also push the item back to the queue. + + Parameters + ---------- + exit_status : str + Exit status of the item + run_uids : list, optional + List of run UIDs + scan_ids : list, optional + List of scan IDs + err_msg : str, default "" + Error message if any + err_tb : str, default "" + Error traceback if any + + Returns + ------- + dict + The stopped item + """ + async with self._lock: + run_uids = run_uids or [] + scan_ids = scan_ids or [] + + # If the status is "stopped", treat it as a completed item + if exit_status == "stopped": + return await self.set_processed_item_as_completed( + exit_status=exit_status, + run_uids=run_uids, + scan_ids=scan_ids, + err_msg=err_msg, + err_tb=err_tb + ) + + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT item FROM {self._table_running_plan} LIMIT 1") + row = await cursor.fetchone() + if not row: + return {} + + running_item = json.loads(row[0]) + immediate_execution = running_item.get("properties", {}).get("immediate_execution", False) + item_time_start = running_item.get("properties", {}).get("time_start", ttime.time()) + + # Clean item properties + item_cleaned = await self._clean_item_properties(running_item) + + # Add result information to the item + item_cleaned.setdefault("result", {}) + item_cleaned["result"]["exit_status"] = exit_status + item_cleaned["result"]["run_uids"] = run_uids + item_cleaned["result"]["scan_ids"] = scan_ids + item_cleaned["result"]["time_start"] = item_time_start + item_cleaned["result"]["time_stop"] = ttime.time() + item_cleaned["result"]["msg"] = err_msg + item_cleaned["result"]["traceback"] = err_tb + + # Add to history + item_json = json.dumps(item_cleaned) + await cursor.execute( + f"INSERT INTO {self._table_plan_history} (item) VALUES (?)", + (item_json,) + ) + + # Push back to the queue with new UID if not immediate execution and status != stopped + if not immediate_execution and exit_status != "stopped": + # Create a copy without result information + item_copy = copy.deepcopy(item_cleaned) + if "result" in item_copy: + del item_copy["result"] + + # Set a new UID + item_copy = await self.set_new_item_uuid(item_copy) + + # Insert at the front of the queue + await cursor.execute( + f"UPDATE {self._table_plan_queue} SET position = position + 1" + ) + + item_json = json.dumps(item_copy) + await cursor.execute( + f"INSERT INTO {self._table_plan_queue} (position, item) VALUES (0, ?)", + (item_json,) + ) + + # Add to UID dictionary + self._uid_dict_add(item_copy) + + # Remove from running + await cursor.execute(f"DELETE FROM {self._table_running_plan}") + + # Remove original item from UID dictionary + self._uid_dict_remove(running_item["item_uid"]) + + await self._sqlite_conn.commit() + return item_cleaned + + # -------------------------------------------------------------------------- + # Lock Management + async def lock_info_save(self, lock_info: Dict[str, Any]) -> None: + """ + Save lock information. + + Parameters + ---------- + lock_info : dict + Lock information to save + """ + async with self._lock: + async with self._sqlite_conn.cursor() as cursor: + # Clear existing lock info + await cursor.execute(f"DELETE FROM {self._table_lock_info}") + + # Save new lock info + lock_json = json.dumps(lock_info) + await cursor.execute( + f"INSERT INTO {self._table_lock_info} (info) VALUES (?)", + (lock_json,) + ) + await self._sqlite_conn.commit() + + async def lock_info_retrieve(self) -> Optional[Dict[str, Any]]: + """ + Retrieve lock information. + + Returns + ------- + dict or None + The saved lock information, or None if no information exists + """ + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT info FROM {self._table_lock_info} LIMIT 1") + row = await cursor.fetchone() + return json.loads(row[0]) if row else None + + async def lock_info_clear(self) -> None: + """Clear lock information.""" + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"DELETE FROM {self._table_lock_info}") + await self._sqlite_conn.commit() + + # -------------------------------------------------------------------------- + # Autostart Mode Management + async def autostart_mode_save(self, autostart_mode: Dict[str, Any]) -> None: + """ + Save autostart mode information. + + Parameters + ---------- + autostart_mode : dict + Autostart mode information to save + """ + async with self._lock: + async with self._sqlite_conn.cursor() as cursor: + # Clear existing autostart mode info + await cursor.execute(f"DELETE FROM {self._table_autostart_mode_info}") + + # Save new autostart mode info + mode_json = json.dumps(autostart_mode) + await cursor.execute( + f"INSERT INTO {self._table_autostart_mode_info} (info) VALUES (?)", + (mode_json,) + ) + await self._sqlite_conn.commit() + + async def autostart_mode_retrieve(self) -> bool: + """ + Retrieve autostart mode information. + + Returns + ------- + bool + True if autostart mode is enabled, False otherwise + """ + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT info FROM {self._table_autostart_mode_info} LIMIT 1") + row = await cursor.fetchone() + if row: + info = json.loads(row[0]) + return info.get("enabled", False) + return False + + async def autostart_mode_clear(self) -> None: + """Clear autostart mode information.""" + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"DELETE FROM {self._table_autostart_mode_info}") + await self._sqlite_conn.commit() + + # -------------------------------------------------------------------------- + # Stop Pending State Management + async def stop_pending_save(self, stop_pending: Dict[str, Any]) -> None: + """ + Save stop pending information. + + Parameters + ---------- + stop_pending : dict + Stop pending information to save + """ + async with self._lock: + async with self._sqlite_conn.cursor() as cursor: + # Clear existing stop pending info + await cursor.execute(f"DELETE FROM {self._table_stop_pending_info}") + + # Save new stop pending info + pending_json = json.dumps(stop_pending) + await cursor.execute( + f"INSERT INTO {self._table_stop_pending_info} (info) VALUES (?)", + (pending_json,) + ) + await self._sqlite_conn.commit() + + async def stop_pending_retrieve(self) -> Optional[Dict[str, Any]]: + """ + Retrieve stop pending information. + + Returns + ------- + dict or None + The saved stop pending information, or None if no information exists + """ + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT info FROM {self._table_stop_pending_info} LIMIT 1") + row = await cursor.fetchone() + return json.loads(row[0]) if row else None + + async def stop_pending_clear(self) -> None: + """Clear stop pending information.""" + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"DELETE FROM {self._table_stop_pending_info}") + await self._sqlite_conn.commit() + + + # -------------------------------------------------------------------------- + # User Group Permissions + async def user_group_permissions_save(self, user_group_permissions: Dict[str, Any]) -> None: + """ + Save user group permissions. + + Parameters + ---------- + user_group_permissions : dict + User group permissions to save + """ + async with self._lock: + async with self._sqlite_conn.cursor() as cursor: + # Clear existing permissions + await cursor.execute(f"DELETE FROM {self._table_user_group_permissions}") + + # Save new permissions + perms_json = json.dumps(user_group_permissions) + await cursor.execute( + f"INSERT INTO {self._table_user_group_permissions} (info) VALUES (?)", + (perms_json,) + ) + await self._sqlite_conn.commit() + + async def user_group_permissions_retrieve(self) -> Optional[Dict[str, Any]]: + """ + Retrieve user group permissions. + + Returns + ------- + dict or None + The saved user group permissions, or None if no permissions exist + """ + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT info FROM {self._table_user_group_permissions} LIMIT 1") + row = await cursor.fetchone() + return json.loads(row[0]) if row else None + + async def user_group_permissions_clear(self) -> None: + """Clear user group permissions.""" + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"DELETE FROM {self._table_user_group_permissions}") + await self._sqlite_conn.commit() + + async def get_queue_state(self) -> Dict[str, Any]: + """ + Get the overall queue state. + + Returns + ------- + dict + Dictionary containing the queue state, including queue size, history size, + and information about the currently running item. + """ + queue_size = await self.get_queue_size() + history_size = await self.get_history_size() + running_item = await self.get_running_item_info() + + return { + "queue_size": queue_size, + "history_size": history_size, + "running_item": running_item, + } + + async def get_item( + self, + *, + pos: Optional[Union[int, str]] = None, + uid: Optional[str] = None + ) -> Dict[str, Any]: + """ + Get an item from the queue without removing it. + + Parameters + ---------- + pos : int or str, optional + Position in the queue. Can be an integer (0-based index) or string ('front' or 'back') + uid : str, optional + UID of the item to get + + Returns + ------- + dict + The requested item + """ + if (pos is None) and (uid is None): + raise ValueError("Position or UID must be specified") + if (pos is not None) and (uid is not None): + raise ValueError("Ambiguous parameters: both position and UID are specified") + + # Extra validation could be added here + if uid is not None: + await self._verify_item_exists(uid) # If you implement this method + + async with self._sqlite_conn.cursor() as cursor: + if uid is not None: + # Get by UID + await cursor.execute( + f"SELECT item FROM {self._table_plan_queue} WHERE json_extract(item, '$.item_uid') = ?", + (uid,) + ) + row = await cursor.fetchone() + if not row: + raise RuntimeError(f"Item with UID '{uid}' not found in the queue") + return json.loads(row[0]) + else: + # Get by position + await cursor.execute(f"SELECT COUNT(*) FROM {self._table_plan_queue}") + qsize = (await cursor.fetchone())[0] + + if qsize == 0: + raise IndexError("The queue is empty") + + if pos == "front": + position = 0 + elif pos == "back": + position = qsize - 1 + elif isinstance(pos, int): + if pos < 0: + position = qsize + pos + else: + position = pos + + if position < 0 or position >= qsize: + raise IndexError(f"Position {pos} is out of range") + else: + raise ValueError(f"Invalid value for 'pos': {pos}") + + await cursor.execute( + f"SELECT item FROM {self._table_plan_queue} WHERE position = ?", + (position,) + ) + row = await cursor.fetchone() + if not row: + raise IndexError(f"No item found at position {position}") + return json.loads(row[0]) + + async def process_next_item( + self, + *, + item: Optional[Dict[str, Any]] = None + ) -> Dict[str, Any]: + """ + Process the next item in the queue or a specified item. + + Parameters + ---------- + item : dict, optional + Item for immediate execution, if not provided, use the next item from the queue + + Returns + ------- + dict + The processed item + """ + # This method is essentially a wrapper around set_next_item_as_running + return await self.set_next_item_as_running(item=item) + + async def set_plan_queue_mode( + self, + plan_queue_mode: Union[Dict[str, Any], str], + *, + update: bool = False + ) -> Dict[str, Any]: + """ + Set the plan queue mode. + + Parameters + ---------- + plan_queue_mode : dict or str + The new plan queue mode or 'default' to reset to default mode + update : bool, default False + If True, update only the specified fields in the mode dictionary + + Returns + ------- + dict + The updated plan queue mode + """ + async with self._lock: + # Create a table for plan queue mode if it doesn't exist + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table_plan_queue_mode} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + info TEXT + ) + """) + + # Get current mode from database or use default + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"SELECT info FROM {self._table_plan_queue_mode} LIMIT 1") + row = await cursor.fetchone() + current_mode = json.loads(row[0]) if row else copy.deepcopy(self._plan_queue_mode_default) + + # Reset to default if requested + if isinstance(plan_queue_mode, str) and plan_queue_mode.lower() == "default": + new_mode = copy.deepcopy(self._plan_queue_mode_default) + elif not isinstance(plan_queue_mode, dict): + raise TypeError(f"Parameter 'plan_queue_mode' must be a dictionary or 'default': {plan_queue_mode}") + else: + # Create a new mode or update the current one + if update: + new_mode = copy.deepcopy(current_mode) + # Verify keys in plan_queue_mode + for key in plan_queue_mode: + if key not in self._plan_queue_mode_default: + raise ValueError(f"Unsupported plan queue mode parameter '{key}'") + # Update mode + new_mode.update(plan_queue_mode) + else: + # Verify that all required keys are present + missing_keys = set(self._plan_queue_mode_default.keys()) - set(plan_queue_mode.keys()) + if missing_keys: + raise ValueError(f"Parameters {missing_keys} are missing") + + # Verify keys in plan_queue_mode + for key in plan_queue_mode: + if key not in self._plan_queue_mode_default: + raise ValueError(f"Unsupported plan queue mode parameter '{key}'") + + new_mode = copy.deepcopy(plan_queue_mode) + + # Verify value types + for key, value in new_mode.items(): + if key == "loop" and not isinstance(value, bool): + raise TypeError(f"Unsupported type {type(value)} of the parameter 'loop'") + elif key == "ignore_failures" and not isinstance(value, bool): + raise TypeError(f"Unsupported type {type(value)} of the parameter 'ignore_failures'") + + # Save the new mode to the database + async with self._sqlite_conn.cursor() as cursor: + await cursor.execute(f"DELETE FROM {self._table_plan_queue_mode}") + mode_json = json.dumps(new_mode) + await cursor.execute( + f"INSERT INTO {self._table_plan_queue_mode} (info) VALUES (?)", + (mode_json,) + ) + + await self._sqlite_conn.commit() + + # Update instance variable + self._plan_queue_mode = new_mode + + return copy.deepcopy(self._plan_queue_mode) + + @property + def plan_queue_mode(self) -> Dict[str, Any]: + """ + Get the current plan queue mode. + + Returns + ------- + dict + The current plan queue mode + """ + return copy.deepcopy(self._plan_queue_mode) + + @property + def plan_queue_mode_default(self) -> Dict[str, Any]: + """ + Get the default plan queue mode. + + Returns + ------- + dict + The default plan queue mode + """ + return copy.deepcopy(self._plan_queue_mode_default) + + @property + def plan_queue_uid(self) -> str: + """ + Get the current plan queue UID. + + Returns + ------- + str + The current plan queue UID + """ + return self.new_item_uid() + + @property + def plan_history_uid(self) -> str: + """ + Get the current plan history UID. + + Returns + ------- + str + The current plan history UID + """ + return self.new_item_uid() + + # -------------------------------------------------------------------------- + # method not included in the abstract (interface) class, but necessary for testing + async def get_queue(self) -> Tuple[List[Dict[str, Any]], str]: + """ + Retrieve the queue contents and a unique queue ID. + + Returns + ------- + tuple + (queue_items, queue_uid) - all items in the queue and a unique queue ID + """ + queue, _, queue_uid = await self.get_queue_full() + return queue, queue_uid \ No newline at end of file diff --git a/bluesky_queueserver/manager/tests/test_plan_queue_ops.py b/bluesky_queueserver/manager/tests/test_plan_queue_ops.py index 021d9991..797fe167 100644 --- a/bluesky_queueserver/manager/tests/test_plan_queue_ops.py +++ b/bluesky_queueserver/manager/tests/test_plan_queue_ops.py @@ -57,14 +57,27 @@ def test_pq_start_stop(): async def testing(): pq = PlanQueueOperations(name_prefix=_test_redis_name_prefix) await pq.start() - await pq._r_pool.ping() + + # Use _backend_ping to check if the backend is reachable + assert await pq._backend_ping() is True + await pq.stop() - assert pq._r_pool is None + # Ensure the backend is properly closed + if pq._backend == "redis": + assert pq._r_pool is None + elif pq._backend == "sqlite": + assert pq._sqlite_conn is None + + # Restart and check again await pq.start() - await pq._r_pool.ping() + assert await pq._backend_ping() is True await pq.stop() - assert pq._r_pool is None + + if pq._backend == "redis": + assert pq._r_pool is None + elif pq._backend == "sqlite": + assert pq._sqlite_conn is None asyncio.run(testing()) @@ -2523,4 +2536,4 @@ async def testing(): await pq.autostart_mode_clear() assert await pq.autostart_mode_retrieve() is None - asyncio.run(testing()) + asyncio.run(testing()) \ No newline at end of file diff --git a/bluesky_queueserver/manager/tests/test_plan_queue_ops_interface_contract.py b/bluesky_queueserver/manager/tests/test_plan_queue_ops_interface_contract.py new file mode 100644 index 00000000..3d5eca7f --- /dev/null +++ b/bluesky_queueserver/manager/tests/test_plan_queue_ops_interface_contract.py @@ -0,0 +1,255 @@ +import pytest +import inspect +import asyncio +from typing import get_type_hints, get_args, get_origin, Union + +from bluesky_queueserver.manager.plan_queue_ops_abstract import AbstractPlanQueueOperations +from bluesky_queueserver.manager.plan_queue_ops_sqlite import SQLitePlanQueueOperations +from bluesky_queueserver.manager.plan_queue_ops_redis import RedisPlanQueueOperations +from bluesky_queueserver.manager.plan_queue_ops_dict import DictPlanQueueOperations +from bluesky_queueserver.manager.manager import RunEngineManager, LockInfo, MState + + +class TestPlanQueueOpsContract: + """Tests to ensure implementations properly fulfill the abstract interface contract.""" + + @pytest.mark.parametrize("implementation_class", [ + SQLitePlanQueueOperations, + RedisPlanQueueOperations, + ]) + def test_implements_all_abstract_methods(self, implementation_class): + """Verify that implementations provide all required abstract methods.""" + # Get all abstract methods from the base class + abstract_methods = { + name: method for name, method in inspect.getmembers(AbstractPlanQueueOperations) + if not name.startswith('_') or name in ('_verify_item', '_clean_item_properties') # Include specific protected methods + if inspect.isabstract(method) or (hasattr(method, '__isabstractmethod__') and method.__isabstractmethod__) + } + + # Check each abstract method is implemented + for method_name, abstract_method in abstract_methods.items(): + assert hasattr(implementation_class, method_name), f"Method '{method_name}' not implemented" + + impl_method = getattr(implementation_class, method_name) + assert callable(impl_method), f"'{method_name}' is not callable in implementation" + + # Check signature compatibility + abstract_sig = inspect.signature(abstract_method) + impl_sig = inspect.signature(impl_method) + + # Check required parameters + for param_name in abstract_sig.parameters: + if param_name == 'self': + continue + assert param_name in impl_sig.parameters, f"Parameter '{param_name}' missing in implementation of '{method_name}'" + + # Check return type annotations match + abstract_return = abstract_sig.return_annotation + impl_return = impl_sig.return_annotation + if abstract_return is not inspect.Signature.empty and impl_return is not inspect.Signature.empty: + assert impl_return == abstract_return or issubclass(impl_return, abstract_return), \ + f"Return type mismatch in '{method_name}': expected {abstract_return}, got {impl_return}" + + @pytest.mark.parametrize("implementation_class", [ + SQLitePlanQueueOperations, + RedisPlanQueueOperations, + ]) + def test_properties_implemented(self, implementation_class): + """Check that required properties are implemented.""" + required_properties = [ + "plan_queue_uid", + "plan_history_uid", + "plan_queue_mode" + ] + + impl = implementation_class() + for prop_name in required_properties: + assert hasattr(impl, prop_name), f"Property '{prop_name}' missing in implementation" + + @pytest.mark.asyncio + @pytest.mark.parametrize("implementation_class", [ + SQLitePlanQueueOperations, + RedisPlanQueueOperations, + ]) + async def test_async_interface_compatibility(self, implementation_class): + """Test async methods for interface compatibility.""" + impl = implementation_class() + await impl.start() + + try: + # Test a few key operations to ensure async compatibility + await impl.clear_queue() + + # Add an item + item = { + "name": "count", + "args": [["det1"]], + "kwargs": {"num": 5}, + "item_type": "plan" + } + item = await impl.set_new_item_uuid(item) + await impl.add_item_to_queue(item) + + # Verify the queue has the item + queue, _ = await impl.get_queue() + assert len(queue) == 1 + assert queue[0]["name"] == "count" + + # Test setting queue mode + await impl.set_plan_queue_mode({"loop": True}) + assert impl.plan_queue_mode["loop"] is True + + finally: + await impl.stop() + + def test_run_engine_manager_interface_usage(self): + """ + Analyze RunEngineManager to find all methods it calls on the queue interface. + This helps identify any missing or incompatible interface requirements. + """ + # Get source code of RunEngineManager + source = inspect.getsource(RunEngineManager) + + # Extract all calls to self._plan_queue.X + import re + plan_queue_calls = re.findall(r'self\._plan_queue\.(\w+)', source) + unique_calls = set(plan_queue_calls) + + # Print the set of methods that RunEngineManager expects from the interface + for method_name in sorted(unique_calls): + # Check if this method exists in the abstract class + assert hasattr(AbstractPlanQueueOperations, method_name), \ + f"Method '{method_name}' called by RunEngineManager but not defined in AbstractPlanQueueOperations" + + +class TestManagerIntegration: + """Integration tests with RunEngineManager.""" + + @pytest.mark.parametrize("implementation_class", [ + SQLitePlanQueueOperations, + RedisPlanQueueOperations, + ]) + def test_manager_initialization_with_implementation(self, implementation_class, monkeypatch): + """Test that RunEngineManager can initialize properly with each implementation.""" + # Mock necessary dependencies for RunEngineManager + monkeypatch.setattr("multiprocessing.Process.__init__", lambda *args, **kwargs: None) + monkeypatch.setattr("bluesky_queueserver.manager.manager.PlanQueueOperations", implementation_class) + + # Create minimal config for initialization + config = { + "zmq_addr": "tcp://*:60615", + "redis_addr": "localhost", + "redis_name_prefix": "test_prefix", + "lock_key_emergency": None, + "user_group_permissions_reload": "NEVER", + "use_ipython_kernel": False, + } + + # Create a manager instance + manager = RunEngineManager( + conn_watchdog=object(), # Fake connection + conn_worker=object(), # Fake connection + config=config, + number_of_restarts=1 + ) + + # Assert that the manager can be instantiated without errors + assert isinstance(manager, RunEngineManager) + + +class TestTypeAnnotationCompliance: + """Tests to ensure type annotations match between abstract class and implementations.""" + + @pytest.mark.parametrize("implementation_class", [ + SQLitePlanQueueOperations, + RedisPlanQueueOperations, + ]) + def test_method_type_annotations_match(self, implementation_class): + """Verify type annotations in implementations match the abstract class.""" + abstract_methods = { + name: method for name, method in inspect.getmembers(AbstractPlanQueueOperations, inspect.isfunction) + if not name.startswith('__') # Exclude dunder methods + } + + for method_name, abstract_method in abstract_methods.items(): + if not hasattr(implementation_class, method_name): + continue # Skip methods not implemented (should be caught by other tests) + + impl_method = getattr(implementation_class, method_name) + + # Get type hints for abstract method and implementation + abstract_hints = get_type_hints(abstract_method) + impl_hints = get_type_hints(impl_method) + + # Check parameters have compatible type annotations + for param_name, abstract_type in abstract_hints.items(): + if param_name == 'return': + continue # Check return type separately + + if param_name in impl_hints: + impl_type = impl_hints[param_name] + # Check if types are compatible + assert self._is_type_compatible(impl_type, abstract_type), \ + f"Parameter '{param_name}' in '{method_name}' has incompatible type: " \ + f"abstract={abstract_type}, implementation={impl_type}" + + @staticmethod + def _is_type_compatible(impl_type, abstract_type): + """Check if implementation type is compatible with abstract type.""" + # Check for Union types + if get_origin(abstract_type) is Union: + return impl_type in get_args(abstract_type) or any( + TestTypeAnnotationCompliance._is_type_compatible(impl_type, arg_type) + for arg_type in get_args(abstract_type) + ) + + # Check for inheritance relationships + try: + return issubclass(impl_type, abstract_type) + except TypeError: + # Handle non-class types + return impl_type == abstract_type + + +class TestAsyncPropertyBehavior: + """Test correct async/await behavior in implementations.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("implementation_class", [ + SQLitePlanQueueOperations, + RedisPlanQueueOperations, + ]) + async def test_async_methods_return_awaitable(self, implementation_class): + """Verify that async methods return objects that can be awaited.""" + impl = implementation_class() + await impl.start() + + try: + # Test a few key methods + methods_to_test = [ + 'clear_queue', + 'add_item_to_queue', + 'get_queue', + 'set_plan_queue_mode' + ] + + for method_name in methods_to_test: + method = getattr(impl, method_name) + if asyncio.iscoroutinefunction(method): + # Call with minimal args - this might fail but we just want to verify the return type + try: + if method_name == 'add_item_to_queue': + result = method({"name": "test", "item_type": "plan"}) + elif method_name == 'set_plan_queue_mode': + result = method({}) + else: + result = method() + + # Verify the result is awaitable + assert asyncio.iscoroutine(result), f"Method {method_name} doesn't return an awaitable" + except: + # We don't care about execution errors, just the interface + pass + + finally: + await impl.stop() \ No newline at end of file diff --git a/bluesky_queueserver/manager/tests/test_plan_queue_ops_redis.py b/bluesky_queueserver/manager/tests/test_plan_queue_ops_redis.py new file mode 100644 index 00000000..9b6293d0 --- /dev/null +++ b/bluesky_queueserver/manager/tests/test_plan_queue_ops_redis.py @@ -0,0 +1,2526 @@ +import asyncio +import copy +import json +import pprint +import re +import time as ttime + +import pytest + +from bluesky_queueserver.manager.plan_queue_ops import PlanQueueOperations + +from .common import _test_redis_name_prefix + +errmsg_wrong_plan_type = "Parameter 'item' should be a dictionary" + + +class PQ: + """ + The class that creates and initializes an instance of plan queue for testing and performs necessary + cleanup afterwards. Intended for use with ``async with``. + """ + + def __init__(self): + self._pq = None + + async def __aenter__(self): + """ + Returns initialized instance of plan queue. + """ + self._pq = PlanQueueOperations(name_prefix=_test_redis_name_prefix) + await self._pq.start() + # Clear any pool entries + await self._pq.delete_pool_entries() + await self._pq.user_group_permissions_clear() + await self._pq.lock_info_clear() + await self._pq.autostart_mode_clear() + await self._pq.stop_pending_clear() + return self._pq + + async def __aexit__(self, exc_t, exc_v, exc_tb): + """ + Cleanup after test. + """ + # Don't leave any test entries in the pool + await self._pq.delete_pool_entries() + await self._pq.user_group_permissions_clear() + await self._pq.lock_info_clear() + await self._pq.autostart_mode_clear() + await self._pq.stop_pending_clear() + + +def test_pq_start_stop(): + """ + Test for the ``PlanQueueOperations.start()`` and ``PlanQueueOperations.stop()`` + """ + + async def testing(): + pq = PlanQueueOperations(name_prefix=_test_redis_name_prefix) + await pq.start() + await pq._r_pool.ping() + await pq.stop() + assert pq._r_pool is None + + await pq.start() + await pq._r_pool.ping() + await pq.stop() + assert pq._r_pool is None + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("item_in, item_out", [ + ({"name": "plan1", "item_uid": "abcde"}, {"name": "plan1", "item_uid": "abcde"}), + ({"user": "user1", "user_group": "group1"}, {"user": "user1", "user_group": "group1"}), + ({"args": [1, 2], "kwargs": {"a": 1, "b": 2}}, {"args": [1, 2], "kwargs": {"a": 1, "b": 2}}), + ({"item_type": "plan", "meta": {"md1": 1, "md2": 2}}, {"item_type": "plan", "meta": {"md1": 1, "md2": 2}}), + ({"name": "plan1", "result": {}}, {"name": "plan1"}), +]) +# fmt: on +def test_filter_item_parameters(item_in, item_out): + """ + Tests for ``filter_item_parameters``. + """ + + async def testing(): + async with PQ() as pq: + item = pq.filter_item_parameters(item_in) + assert item == item_out + + asyncio.run(testing()) + + +def test_running_plan_info(): + """ + Basic test for the following methods: + `PlanQueueOperations.is_item_running()` + `PlanQueueOperations.get_running_item_info()` + `PlanQueueOperations.delete_pool_entries()` + """ + + async def testing(): + async with PQ() as pq: + assert await pq.get_running_item_info() == {} + assert await pq.is_item_running() is False + + some_plan = {"some_key": "some_value"} + await pq._set_running_item_info(some_plan) + assert await pq.get_running_item_info() == some_plan + + assert await pq.is_item_running() is True + + await pq._clear_running_item_info() + assert await pq.get_running_item_info() == {} + assert await pq.is_item_running() is False + + await pq._set_running_item_info(some_plan) + await pq.delete_pool_entries() + assert await pq.get_running_item_info() == {} + assert await pq.is_item_running() is False + + asyncio.run(testing()) + + +def test_redis_name_prefix(): + """ + Test that the prefix is correctly appended to the name of the redis key (test with one key). + """ + pq = PlanQueueOperations(name_prefix=_test_redis_name_prefix) + assert pq._name_plan_queue == _test_redis_name_prefix + "_plan_queue" + + pq = PlanQueueOperations(name_prefix="") + assert pq._name_plan_queue == "plan_queue" + + +# fmt: off +@pytest.mark.parametrize("plan_running, plans, result_running, result_plans", [ + ({"testing": 1}, [{"testing": 2}, {"item_uid": "ab", "name": "nm"}, {"testing": 2}], + {}, [{"item_uid": "ab", "name": "nm"}]), + ({"testing": 1}, [{"testing": 2}, {"item_uid": "ab", "name": "nm"}, {"testing": 3}], + {}, [{"item_uid": "ab", "name": "nm"}]), + ({"item_uid": "a"}, [{"item_uid": "a1"}, {"item_uid": "a2"}, {"item_uid": "a3"}], + {"item_uid": "a"}, [{"item_uid": "a1"}, {"item_uid": "a2"}, {"item_uid": "a3"}]), +]) +# fmt: on +def test_queue_clean(plan_running, plans, result_running, result_plans): + """ + Test for ``_queue_clean()`` method + """ + + async def testing(): + async with PQ() as pq: + await pq._set_running_item_info(plan_running) + for plan in plans: + await pq._r_pool.rpush(pq._name_plan_queue, json.dumps(plan)) + + assert await pq.get_running_item_info() == plan_running + plan_queue, _ = await pq.get_queue() + assert plan_queue == plans + + await pq._queue_clean() + + assert await pq.get_running_item_info() == result_running + plan_queue, _ = await pq.get_queue() + assert plan_queue == result_plans + + asyncio.run(testing()) + + +def test_set_plan_queue_mode_1(): + """ + ``set_plan_queue_mode``: basic functionality. + """ + + async def testing(): + async with PQ() as pq: + # Initially plan queue mode must be default queue mode + assert pq.plan_queue_mode == pq.plan_queue_mode_default + assert pq.plan_queue_mode["loop"] is False + assert pq.plan_queue_mode["ignore_failures"] is False + + # The properties are expecte to return copies + assert pq.plan_queue_mode is not pq._plan_queue_mode + assert pq.plan_queue_mode_default is not pq._plan_queue_mode_default + + mode_expected = pq.plan_queue_mode + assert mode_expected == dict(loop=False, ignore_failures=False) + + mode_new = {"loop": True} + mode_expected.update(mode_new) + await pq.set_plan_queue_mode(mode_new, update=True) + assert pq._plan_queue_mode is not mode_new # Verify that 'set' operation performs copy + assert pq.plan_queue_mode == mode_expected + + mode_new = {"ignore_failures": True} + mode_expected.update(mode_new) + await pq.set_plan_queue_mode(mode_new, update=True) + assert pq._plan_queue_mode is not mode_new # Verify that 'set' operation performs copy + assert pq.plan_queue_mode == mode_expected + + mode_new = {"loop": False, "ignore_failures": False} + mode_expected.update(mode_new) + await pq.set_plan_queue_mode(mode_new, update=False) + assert pq._plan_queue_mode is not mode_new # Verify that 'set' operation performs copy + assert pq.plan_queue_mode == mode_expected + + for update in (False, True): + print(f"update={update}") + with pytest.raises(ValueError, match="Unsupported plan queue mode parameter 'nonexisting_key'"): + await pq.set_plan_queue_mode({"nonexisting_key": True}, update=update) + + await pq.set_plan_queue_mode({}, update=True) + assert pq.plan_queue_mode == mode_expected + + with pytest.raises(ValueError, match="Parameters .* are missing"): + await pq.set_plan_queue_mode({}, update=False) + + mode_new = {"loop": True} + with pytest.raises(ValueError, match="Parameters {'ignore_failures'} are missing"): + await pq.set_plan_queue_mode(mode_new, update=False) + + with pytest.raises(TypeError, match="Unsupported type .* of the parameter 'loop'"): + await pq.set_plan_queue_mode({"loop": 10}, update=True) + + with pytest.raises(TypeError, match="Unsupported type .* of the parameter 'ignore_failures'"): + await pq.set_plan_queue_mode({"ignore_failures": []}, update=True) + + assert pq.plan_queue_mode == mode_expected + + # Verify that the queue mode is saved to Redis + queue_mode = {"loop": True, "ignore_failures": True} + await pq.set_plan_queue_mode(queue_mode, update=False) + + pq2 = PlanQueueOperations(name_prefix=_test_redis_name_prefix) + await pq2.start() + assert pq2.plan_queue_mode == queue_mode + + # Set queue mode to default + assert pq.plan_queue_mode != pq.plan_queue_mode_default + await pq.set_plan_queue_mode("default", update=update) + assert pq.plan_queue_mode == pq.plan_queue_mode_default + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("plan, result", + [({"a": 10}, True), + ([10, 20], False), + (50, False), + ("abc", False)]) +# fmt: on +def test_verify_item_type(plan, result): + async def testing(): + async with PQ() as pq: + if result: + pq._verify_item_type(plan) + else: + with pytest.raises(TypeError, match=errmsg_wrong_plan_type): + pq._verify_item_type(plan) + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize( + "plan, f_kwargs, result, errmsg", + [({"a": 10}, {}, False, "Item does not have UID"), + ([10, 20], {}, False, errmsg_wrong_plan_type), + ({"item_uid": "one"}, {}, True, ""), + ({"item_uid": "two"}, {}, False, "Item with UID .+ is already in the queue"), + ({"item_uid": "three"}, {}, False, "Item with UID .+ is already in the queue"), + ({"item_uid": "two"}, {"ignore_uids": None}, False, "Item with UID .+ is already in the queue"), + ({"item_uid": "two"}, {"ignore_uids": ["two"]}, True, ""), + ({"item_uid": "two"}, {"ignore_uids": ["two", "three"]}, True, ""), + ({"item_uid": "two"}, {"ignore_uids": ["one", "three"]}, False, "Item with UID .+ is already in the queue"), + ]) +# fmt: on +def test_verify_item(plan, f_kwargs, result, errmsg): + """ + Tests for method ``_verify_item()``. + """ + # Set two existiing plans and then set one of them as running + existing_plans = [{"item_type": "plan", "item_uid": "two"}, {"item_type": "plan", "item_uid": "three"}] + + async def testing(): + async with PQ() as pq: + + async def set_plans(): + # Add plan to queue + for plan in existing_plans: + await pq.add_item_to_queue(plan) + # Set one plan as currently running + await pq.set_next_item_as_running() + + # Verify that setup is correct + assert await pq.is_item_running() is True + assert await pq.get_queue_size() == 1 + + await set_plans() + + if result: + await pq._verify_item(plan, **f_kwargs) + else: + with pytest.raises(Exception, match=errmsg): + await pq._verify_item(plan, **f_kwargs) + + asyncio.run(testing()) + + +def test_new_item_uid(): + """ + Smoke test for the method ``new_item_uid()``. + """ + assert isinstance(PlanQueueOperations.new_item_uid(), str) + + +# fmt: off +@pytest.mark.parametrize("plan", [ + {"name": "a"}, + {"item_uid": "some_uid", "name": "a"}, +]) +# fmt: on +def test_set_new_item_uuid(plan): + """ + Basic test for the method ``set_new_item_uuid()``. + """ + + async def testing(): + async with PQ() as pq: + uid = plan.get("item_uid", None) + + # The function is supposed to create or replace UID + new_plan = await pq.set_new_item_uuid(plan) + + assert "item_uid" in new_plan + assert isinstance(new_plan["item_uid"], str) + assert new_plan["item_uid"] != uid + + asyncio.run(testing()) + + +def test_get_index_by_uid_1(): + """ + Test for ``_get_index_by_uid()`` + """ + plans = [ + {"item_uid": "a", "name": "name_a"}, + {"item_uid": "b", "name": "name_b"}, + {"item_uid": "c", "name": "name_c"}, + ] + + async def testing(): + async with PQ() as pq: + for plan in plans: + await pq.add_item_to_queue(plan) + + assert await pq._get_index_by_uid(uid="b") == 1 + + with pytest.raises(IndexError, match="No plan with UID 'nonexistent'"): + assert await pq._get_index_by_uid(uid="nonexistent") + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("sequence, indices_expected", [ + (["a", "b", "c"], [0, 1, 2]), + (["a", "b"], [0, 1]), + (["c", "b"], [2, 1]), + (["c", "d", "b"], [2, -1, 1]), +]) +# fmt: on +def test_get_index_by_uid_batch_1(sequence, indices_expected): + """ + Test for ``_get_index_by_uid_batch()`` + """ + plans = [ + {"item_uid": "a", "name": "name_a"}, + {"item_uid": "b", "name": "name_b"}, + {"item_uid": "c", "name": "name_c"}, + ] + + async def testing(): + async with PQ() as pq: + for plan in plans: + await pq.add_item_to_queue(plan) + + indices = await pq._get_index_by_uid_batch(uids=sequence) + assert indices == indices_expected + + asyncio.run(testing()) + + +def test_uid_dict_1(): + """ + Basic test for functions associated with `_uid_dict` + """ + plan_a = {"item_uid": "a", "name": "name_a"} + plan_b = {"item_uid": "b", "name": "name_b"} + plan_c = {"item_uid": "c", "name": "name_c"} + + plan_b_updated = {"item_uid": "b", "name": "name_b_updated"} + + async def testing(): + async with PQ() as pq: + pq._uid_dict_add(plan_a) + pq._uid_dict_add(plan_b) + + assert pq._is_uid_in_dict(plan_a["item_uid"]) is True + assert pq._is_uid_in_dict(plan_b["item_uid"]) is True + assert pq._is_uid_in_dict(plan_c["item_uid"]) is False + + assert pq._uid_dict_get_item(plan_b["item_uid"]) == plan_b + pq._uid_dict_update(plan_b_updated) + assert pq._uid_dict_get_item(plan_b["item_uid"]) == plan_b_updated + + pq._uid_dict_remove(plan_a["item_uid"]) + assert pq._is_uid_in_dict(plan_a["item_uid"]) is False + assert pq._is_uid_in_dict(plan_b["item_uid"]) is True + + pq._uid_dict_clear() + assert pq._is_uid_in_dict(plan_a["item_uid"]) is False + assert pq._is_uid_in_dict(plan_b["item_uid"]) is False + + asyncio.run(testing()) + + +def test_uid_dict_2(): + """ + Test if functions changing `pq._uid_dict` are also updating `pq.plan_queue_uid`. + """ + plan_a = {"item_uid": "a", "name": "name_a"} + plan_a_updated = {"item_uid": "a", "name": "name_a_updated"} + + async def testing(): + async with PQ() as pq: + pq_uid = pq.plan_queue_uid + + pq._uid_dict_add(plan_a) + + assert pq.plan_queue_uid != pq_uid + pq_uid = pq.plan_queue_uid + + pq._uid_dict_update(plan_a_updated) + + assert pq.plan_queue_uid != pq_uid + pq_uid = pq.plan_queue_uid + + pq._uid_dict_remove(plan_a_updated["item_uid"]) + + assert pq.plan_queue_uid != pq_uid + pq_uid = pq.plan_queue_uid + + pq._uid_dict_clear() + + assert pq.plan_queue_uid != pq_uid + + asyncio.run(testing()) + + +def test_uid_dict_3_initialize(): + """ + Basic test for functions associated with ``_uid_dict_initialize()`` + """ + + async def testing(): + async with PQ() as pq: + await pq.add_item_to_queue({"name": "a"}) + await pq.add_item_to_queue({"name": "b"}) + await pq.add_item_to_queue({"name": "c"}) + plans, _ = await pq.get_queue() + uid_dict = {_["item_uid"]: _ for _ in plans} + + pq_uid = pq.plan_queue_uid + + await pq._uid_dict_initialize() + + assert pq._uid_dict == uid_dict + assert pq.plan_queue_uid != pq_uid + + asyncio.run(testing()) + + +def test_uid_dict_4_failing(): + """ + Failing cases for functions associated with `_uid_dict` + """ + plan_a = {"item_uid": "a", "name": "name_a"} + plan_b = {"item_uid": "b", "name": "name_b"} + plan_c = {"item_uid": "c", "name": "name_c"} + + async def testing(): + async with PQ() as pq: + pq._uid_dict_add(plan_a) + pq._uid_dict_add(plan_b) + + # Add plan with UID that already exists + pq_uid = pq.plan_queue_uid + with pytest.raises(RuntimeError, match=f"'{plan_a['item_uid']}', which is already in the queue"): + pq._uid_dict_add(plan_a) + assert pq.plan_queue_uid == pq_uid + + assert len(pq._uid_dict) == 2 + + # Remove plan with UID does not exist exists + with pytest.raises(RuntimeError, match=f"'{plan_c['item_uid']}', which is not in the queue"): + pq._uid_dict_remove(plan_c["item_uid"]) + assert pq.plan_queue_uid == pq_uid + + assert len(pq._uid_dict) == 2 + + # Update plan with UID does not exist exists + with pytest.raises(RuntimeError, match=f"'{plan_c['item_uid']}', which is not in the queue"): + pq._uid_dict_update(plan_c) + assert pq.plan_queue_uid == pq_uid + + assert len(pq._uid_dict) == 2 + + asyncio.run(testing()) + + +def test_remove_item(): + """ + Basic test for functions associated with ``_remove_plan()`` + """ + + async def testing(): + async with PQ() as pq: + plan_list = [{"name": "a"}, {"name": "b"}, {"name": "c"}] + for plan in plan_list: + await pq.add_item_to_queue(plan) + + plans, _ = await pq.get_queue() + plan_to_remove = [_ for _ in plans if _["name"] == "b"][0] + + # Remove one plan + await pq._remove_item(plan_to_remove) + plans, _ = await pq.get_queue() + assert len(plans) == 2 + + # Add a copy of a plan (queue is not supposed to have copies in real life) + plan_to_add = plans[0] + await pq._r_pool.lpush(pq._name_plan_queue, json.dumps(plan_to_add)) + # Now remove both plans + await pq._remove_item(plan_to_add, single=False) # Allow deleting multiple or no plans + assert await pq.get_queue_size() == 1 + + # Delete the plan again (the plan is not in the queue, but it shouldn't raise an exception) + await pq._remove_item(plan_to_add, single=False) # Allow deleting multiple or no plans + assert await pq.get_queue_size() == 1 + + with pytest.raises(RuntimeError, match="One item is expected"): + await pq._remove_item(plan_to_add) + assert await pq.get_queue_size() == 1 + + # Now add 'plan_to_add' twice (create two copies) + await pq._r_pool.lpush(pq._name_plan_queue, json.dumps(plan_to_add)) + await pq._r_pool.lpush(pq._name_plan_queue, json.dumps(plan_to_add)) + assert await pq.get_queue_size() == 3 + # Attempt to delete two copies + with pytest.raises(RuntimeError, match="One item is expected"): + await pq._remove_item(plan_to_add) + # Exception is raised, but both copies are deleted + assert await pq.get_queue_size() == 1 + + asyncio.run(testing()) + + +def test_get_queue_full_1(): + """ + Basic test for the functions ``PlanQueueOperations.get_queue()`` and + ``PlanQueueOperations.get_queue_full()`` + """ + + async def testing(): + async with PQ() as pq: + plans = [ + {"item_type": "plan", "item_uid": "one", "name": "a"}, + {"item_type": "plan", "item_uid": "two", "name": "b"}, + {"item_type": "plan", "item_uid": "three", "name": "c"}, + ] + + for p in plans: + await pq.add_item_to_queue(p) + await pq.set_next_item_as_running() + + pq_uid = pq.plan_queue_uid + queue1, uid1 = await pq.get_queue() + running_item1 = await pq.get_running_item_info() + queue2, running_item2, uid2 = await pq.get_queue_full() + + assert queue1 == plans[1:] + assert queue2 == plans[1:] + assert {k: v for k, v in running_item1.items() if k != "properties"} == plans[0] + assert running_item2 == running_item1 + + assert isinstance(running_item1["properties"]["time_start"], float) + assert ttime.time() - running_item1["properties"]["time_start"] < 10 + + assert uid1 == pq_uid + assert uid2 == pq_uid + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("params, name", [ + ({"pos": "front"}, "a"), + ({"pos": "back"}, "c"), + ({"pos": 0}, "a"), + ({"pos": 1}, "b"), + ({"pos": 2}, "c"), + ({"pos": 3}, None), # Index out of range + ({"pos": -1}, "c"), + ({"pos": -2}, "b"), + ({"pos": -3}, "a"), + ({"pos": -4}, None), # Index out of range + ({"uid": "one"}, "a"), + ({"uid": "two"}, "b"), + ({"uid": "nonexistent"}, None), +]) +# fmt: on +def test_get_item_1(params, name): + """ + Basic test for the function ``PlanQueueOperations.get_item()`` + """ + + async def testing(): + async with PQ() as pq: + pq_uid = pq.plan_queue_uid + await pq.add_item_to_queue({"item_uid": "one", "name": "a"}) + await pq.add_item_to_queue({"item_uid": "two", "name": "b"}) + await pq.add_item_to_queue({"item_uid": "three", "name": "c"}) + assert await pq.get_queue_size() == 3 + assert pq.plan_queue_uid != pq_uid + + if name is not None: + plan = await pq.get_item(**params) + assert plan["name"] == name + else: + msg = "Index .* is out of range" if "pos" in params else "is not in the queue" + with pytest.raises(IndexError, match=msg): + await pq.get_item(**params) + + asyncio.run(testing()) + + +def test_get_item_2_fail(): + """ + Basic test for the function ``PlanQueueOperations.get_item()``. + Attempt to retrieve a running plan. + """ + + async def testing(): + async with PQ() as pq: + await pq.add_item_to_queue({"item_type": "plan", "item_uid": "one", "name": "a"}) + await pq.add_item_to_queue({"item_type": "plan", "item_uid": "two", "name": "b"}) + await pq.add_item_to_queue({"item_type": "plan", "item_uid": "three", "name": "c"}) + assert await pq.get_queue_size() == 3 + + pq_uid = pq.plan_queue_uid + await pq.set_next_item_as_running() + assert await pq.get_queue_size() == 2 + assert pq.plan_queue_uid != pq_uid + + with pytest.raises(IndexError, match="is currently running"): + await pq.get_item(uid="one") + + # Ambiguous parameters (position and UID is passed) + with pytest.raises(ValueError, match="Ambiguous parameters"): + await pq.get_item(pos=5, uid="abc") + + asyncio.run(testing()) + + +def test_add_item_to_queue_1(): + """ + Basic test for the function ``PlanQueueOperations.add_item_to_queue()`` + """ + + async def testing(): + async with PQ() as pq: + + async def add_plan(plan, n, **kwargs): + plan_added, qsize = await pq.add_item_to_queue(plan, **kwargs) + assert plan_added["name"] == plan["name"], f"plan: {plan}" + assert qsize == n, f"plan: {plan}" + + await add_plan({"name": "a"}, 1) + await add_plan({"name": "b"}, 2) + await add_plan({"name": "c"}, 3, pos="back") + await add_plan({"name": "d"}, 4, pos="front") + await add_plan({"name": "e"}, 5, pos=0) # front + await add_plan({"name": "f"}, 6, pos=5) # back (index == queue size) + await add_plan({"name": "g"}, 7, pos=5) # previous to last + await add_plan({"name": "h"}, 8, pos=-1) # last + await add_plan({"name": "i"}, 9, pos=3) # arbitrary index + await add_plan({"name": "j"}, 10, pos=100) # back (index some large number) + await add_plan({"name": "k"}, 11, pos=-11) # front (precisely) + await add_plan({"name": "l"}, 12, pos=-11) # 2nd element + await add_plan({"name": "m"}, 13, pos=-100) # front (index some large negative number) + await add_plan({"name": "n"}, 14, pos=-2) # previous to lasst + + assert await pq.get_queue_size() == 14 + + plans, _ = await pq.get_queue() + name_sequence = [_["name"] for _ in plans] + assert name_sequence == ["m", "k", "l", "e", "d", "a", "i", "b", "c", "g", "f", "h", "n", "j"] + + await pq.clear_queue() + + asyncio.run(testing()) + + +def test_add_item_to_queue_2(): + """ + Basic test for the function ``PlanQueueOperations.add_item_to_queue()`` + """ + + async def testing(): + async with PQ() as pq: + + async def add_plan(plan, n, **kwargs): + plan_added, qsize = await pq.add_item_to_queue(plan, **kwargs) + assert plan_added["name"] == plan["name"], f"plan: {plan}" + assert qsize == n, f"plan: {plan}" + + await add_plan({"item_type": "plan", "name": "a"}, 1) + await add_plan({"item_type": "plan", "name": "b"}, 2) + await add_plan({"item_type": "plan", "name": "c"}, 3, pos="back") + + plan_queue, _ = await pq.get_queue() + displaced_uid = plan_queue[1]["item_uid"] + + await add_plan({"item_type": "plan", "name": "d"}, 4, before_uid=displaced_uid) + await add_plan({"item_type": "plan", "name": "e"}, 5, after_uid=displaced_uid) + + # This reduces the number of elements in the queue by one + await pq.set_next_item_as_running() + + displaced_uid = plan_queue[0]["item_uid"] + await add_plan({"name": "f"}, 5, after_uid=displaced_uid) + + with pytest.raises( + IndexError, match="Can not insert a plan in the queue before a currently running plan" + ): + await add_plan({"name": "g"}, 5, before_uid=displaced_uid) + + with pytest.raises(IndexError, match="is not in the queue"): + await add_plan({"name": "h"}, 5, before_uid="nonexistent_uid") + + assert await pq.get_queue_size() == 5 + + plans, _ = await pq.get_queue() + name_sequence = [_["name"] for _ in plans] + assert name_sequence == ["f", "d", "b", "e", "c"] + + await pq.clear_queue() + + asyncio.run(testing()) + + +@pytest.mark.parametrize("filter_params", [False, True]) +def test_add_item_to_queue_3(filter_params): + """ + Test if parameter filtering works as expected with `add_item_to_queue` function. + """ + + async def testing(): + async with PQ() as pq: + # Parameter 'result' should be removed if filtering is enabled + plan1 = {"item_type": "plan", "name": "a", "item_uid": "1"} + plan2 = plan1.copy() + plan2["result"] = {} + + plan_added, qsize = await pq.add_item_to_queue(plan2, filter_parameters=filter_params) + + if filter_params: + assert plan_added == plan1 + else: + assert plan_added == plan2 + + asyncio.run(testing()) + + +def test_add_item_to_queue_4_fail(): + """ + Failing tests for the function ``PlanQueueOperations.add_item_to_queue()`` + """ + + async def testing(): + async with PQ() as pq: + pq_uid = pq.plan_queue_uid + with pytest.raises(ValueError, match="Parameter 'pos' has incorrect value"): + await pq.add_item_to_queue({"name": "a"}, pos="something") + assert pq.plan_queue_uid == pq_uid + + with pytest.raises(TypeError, match=errmsg_wrong_plan_type): + await pq.add_item_to_queue("plan_is_not_string") + assert pq.plan_queue_uid == pq_uid + + # Duplicate plan UID + plan = {"item_uid": "abc", "name": "a"} + await pq.add_item_to_queue(plan) + pq_uid = pq.plan_queue_uid + with pytest.raises(RuntimeError, match="Item with UID .+ is already in the queue"): + await pq.add_item_to_queue(plan) + assert pq.plan_queue_uid == pq_uid + + # Ambiguous parameters (position and UID is passed) + with pytest.raises(ValueError, match="Ambiguous parameters"): + await pq.add_item_to_queue({"name": "abc"}, pos=5, before_uid="abc") + assert pq.plan_queue_uid == pq_uid + + # Ambiguous parameters ('before_uid' and 'after_uid' is specified) + with pytest.raises(ValueError, match="Ambiguous parameters"): + await pq.add_item_to_queue({"name": "abc"}, before_uid="abc", after_uid="abc") + assert pq.plan_queue_uid == pq_uid + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("batch_params, queue_seq, batch_seq, expected_seq", [ + ({}, "", "", ""), # Add an empty batch + ({}, "", "567", "567"), + ({"pos": "front"}, "", "567", "567"), + ({"pos": "back"}, "", "567", "567"), + ({}, "1234", "567", "1234567"), + ({"pos": "front"}, "1234", "567", "5671234"), + ({"pos": "back"}, "1234", "567", "1234567"), + ({"pos": 0}, "1234", "567", "5671234"), + ({"pos": 1}, "1234", "567", "1567234"), + ({"pos": 2}, "1234", "567", "1256734"), + ({"pos": 4}, "1234", "567", "1234567"), + ({"pos": 100}, "1234", "567", "1234567"), + ({"pos": -1}, "1234", "567", "1234567"), + ({"pos": -2}, "1234", "567", "1235674"), + ({"pos": -3}, "1234", "567", "1256734"), + ({"pos": -4}, "1234", "567", "1567234"), + ({"pos": -5}, "1234", "567", "5671234"), + ({"pos": -100}, "1234", "567", "5671234"), + ({"before_uid": "1"}, "1234", "567", "5671234"), + ({"before_uid": "2"}, "1234", "567", "1567234"), + ({"before_uid": "3"}, "1234", "567", "1256734"), + ({"after_uid": "1"}, "1234", "567", "1567234"), + ({"after_uid": "2"}, "1234", "567", "1256734"), + ({"after_uid": "4"}, "1234", "567", "1234567"), +]) +# fmt: on +def test_add_batch_to_queue_1(batch_params, queue_seq, batch_seq, expected_seq): + """ + Basic test for the function ``PlanQueueOperations.add_batch_to_queue()`` + """ + + async def testing(): + async with PQ() as pq: + + async def add_plan(plan, n, **kwargs): + plan_added, qsize = await pq.add_item_to_queue(plan, **kwargs) + assert plan_added["name"] == plan["name"], f"plan: {plan}" + assert qsize == n, f"plan: {plan}" + + # Create the queue with plans + for n, p_name in enumerate(queue_seq): + await add_plan({"name": p_name, "item_uid": f"{p_name}{p_name}"}, n + 1) + + def fix_uid(uid): + return f"{uid}{uid}" + + if "before_uid" in batch_params: + batch_params["before_uid"] = fix_uid(batch_params["before_uid"]) + if "after_uid" in batch_params: + batch_params["after_uid"] = fix_uid(batch_params["after_uid"]) + + items = [] + for p_name in batch_seq: + items.append({"name": p_name, "item_uid": p_name + p_name}) + items_added, results, qsize, success = await pq.add_batch_to_queue(items, **batch_params) + assert success is True, pprint.pformat(results) + assert qsize == len(queue_seq) + len(batch_seq) + assert await pq.get_queue_size() == len(queue_seq) + len(batch_seq) + assert len(items_added) == len(items) + assert len(results) == len(items) + + # Verify that the results are set correctly (success) + for res in results: + assert res["success"] is True, pprint.pformat(results) + assert res["msg"] == "", pprint.pformat(results) + + # Verify the sequence of items in the queue + queue, _ = await pq.get_queue() + queue_sequence = [_["name"] for _ in queue] + queue_sequence = "".join(queue_sequence) + assert queue_sequence == expected_seq + + await pq.clear_queue() + + asyncio.run(testing()) + + +@pytest.mark.parametrize("filter_params", [False, True, None]) +def test_add_batch_to_queue_2(filter_params): + """ + Test if parameter filtering works as expected with `add_batch_to_queue` function. + """ + + async def testing(): + async with PQ() as pq: + # Parameter 'result' should be removed if filtering is enabled. + params = {"filter_parameters": filter_params} if (filter_params is not None) else {} + do_filtering = True if (filter_params is None) else filter_params + + items = [{"name": _, "result": {}} for _ in ("a", "b", "c")] + items_added, results, qsize, success = await pq.add_batch_to_queue(items, **params) + assert success is True + assert qsize == len(items) + assert len(items_added) == len(items) + assert len(results) == len(items) + + queue, _ = await pq.get_queue() + + assert [_["name"] for _ in queue] == [_["name"] for _ in items] + for queue_item in queue: + if do_filtering: + assert "result" not in queue_item, pprint.pformat(queue_item) + else: + assert "result" in queue_item, pprint.pformat(queue_item) + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("params, queue_seq, batch_seq, err_msgs", [ + ({}, "abcd", "eaf", ["", "Item with UID .+ is already in the queue", ""]), + ({}, "abcd", "ead", [""] + ["Item with UID .+ is already in the queue"] * 2), + ({"before_uid": "unknown"}, "abcd", "efg", ["Plan with UID .* is not in the queue"] * 3), + ({"after_uid": "unknown"}, "abcd", "efg", ["Plan with UID .* is not in the queue"] * 3), + ({"before_uid": "unknown", "after_uid": "unknown"}, "abcd", "efg", ["Ambiguous parameters"] * 3), + ({"pos": "front", "after_uid": "unknown"}, "abcd", "efg", ["Ambiguous parameters"] * 3), +]) +# fmt: on +def test_add_batch_to_queue_3_fail(params, queue_seq, batch_seq, err_msgs): + """ + Failing cases for the function ``PlanQueueOperations.add_batch_to_queue()`` + """ + + async def testing(): + async with PQ() as pq: + + async def add_plan(plan, n, **kwargs): + plan_added, qsize = await pq.add_item_to_queue(plan, **kwargs) + assert plan_added["name"] == plan["name"], f"plan: {plan}" + assert qsize == n, f"plan: {plan}" + + # Create the queue with plans + for n, p_name in enumerate(queue_seq): + await add_plan({"name": p_name, "item_uid": p_name + p_name}, n + 1) + + items = [] + for p_name in batch_seq: + items.append({"name": p_name, "item_uid": p_name + p_name}) + items_added, results, qsize, success = await pq.add_batch_to_queue(items, **params) + + assert success is False + assert qsize == len(queue_seq) + assert len(items_added) == len(items) + assert len(results) == len(items) + + # The returned item list is expected to be identical to the original list + assert items_added == items, pprint.pformat(items_added) + + for res, msg in zip(results, err_msgs): + if not msg: + assert res["success"] is True, pprint.pformat(res) + assert res["msg"] == "", pprint.pformat(res) + else: + assert res["success"] is False, pprint.pformat(res) + assert re.search(msg, res["msg"]), pprint.pformat(res) + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("replace_uid", [False, True]) +# fmt: on +def test_replace_item_1(replace_uid): + """ + Basic functionality of ``PlanQueueOperations.replace_item()`` function. + """ + + async def testing(): + async with PQ() as pq: + plans = [{"name": "a"}, {"name": "b"}, {"name": "c"}] + plans_added = [None] * len(plans) + qsizes = [None] * len(plans) + + for n, plan in enumerate(plans): + plans_added[n], qsizes[n] = await pq.add_item_to_queue(plan) + + assert qsizes == [1, 2, 3] + assert await pq.get_queue_size() == 3 + + # Change name, but keep UID + plan_names_new = ["d", "e", "f"] + for n in range(len(plans_added)): + plan = plans_added[n].copy() + plan["name"] = plan_names_new[n] + + uid_to_replace = plan["item_uid"] + if replace_uid: + plan["item_uid"] = pq.new_item_uid() # Generate new UID + + pq_uid = pq.plan_queue_uid + plan_new, qsize = await pq.replace_item(plan, item_uid=uid_to_replace) + assert plan_new["name"] == plan["name"] + assert plan_new["item_uid"] == plan["item_uid"] + assert pq.plan_queue_uid != pq_uid + + assert await pq.get_queue_size() == 3 + plans_added[n] = plan_new + + # Make sure that the plan can be correctly extracted by uid + assert await pq.get_item(uid=plan["item_uid"]) == plan + assert await pq.get_queue_size() == 3 + + # Initialize '_uid_dict' and see if the plan can still be extracted using correct UID. + await pq._uid_dict_initialize() + assert await pq.get_item(uid=plan["item_uid"]) == plan + assert await pq.get_queue_size() == 3 + + asyncio.run(testing()) + + +def test_replace_item_2(): + """ + ``PlanQueueOperations.replace_item()`` function: not UID in the plan - random UID is assigned. + """ + + async def testing(): + async with PQ() as pq: + plans = [{"name": "a"}, {"name": "b"}, {"name": "c"}] + plans_added = [None] * len(plans) + qsizes = [None] * len(plans) + + for n, plan in enumerate(plans): + plans_added[n], qsizes[n] = await pq.add_item_to_queue(plan) + + assert qsizes == [1, 2, 3] + assert await pq.get_queue_size() == 3 + + new_name = "h" + plan_new = {"name": new_name} # No UID in the plan. It should still be inserted + pq_uid = pq.plan_queue_uid + plan_replaced, qsize = await pq.replace_item(plan_new, item_uid=plans_added[1]["item_uid"]) + assert pq.plan_queue_uid != pq_uid + + new_uid = plan_replaced["item_uid"] + assert new_uid != plans_added[1]["item_uid"] + assert plan_replaced["name"] == plan_new["name"] + + plan = await pq.get_item(uid=new_uid) + assert plan["item_uid"] == new_uid + assert plan["name"] == new_name + + # Initialize '_uid_dict' and see if the plan can still be extracted using correct UID. + await pq._uid_dict_initialize() + plan = await pq.get_item(uid=new_uid) + assert plan["item_uid"] == new_uid + assert plan["name"] == new_name + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("mode", ["exact_copy", "change_uid", "change_plan"]) +# fmt: on +def test_replace_item_3(mode): + """ + ``PlanQueueOperations.replace_item()`` function: make sure that filtering + is applied to the parameters of edited plan if the plan was changed + """ + + async def testing(): + async with PQ() as pq: + plans = [{"name": "a", "result": {}}, {"name": "b"}, {"name": "c"}] + + for n, plan in enumerate(plans): + await pq.add_item_to_queue(plan, filter_parameters=False) + + assert await pq.get_queue_size() == 3 + + queue, _ = await pq.get_queue() + plan0_before = queue[0] + assert "result" in plan0_before + + if mode == "exact_copy": + # Replace the plan with the exact copy of it. Filtering is not be applied. + plan_new = plan0_before + plan_replaced, qsize = await pq.replace_item(plan_new, item_uid=plan_new["item_uid"]) + + queue, _ = await pq.get_queue() + plan0_after = queue[0] + + assert plan_replaced == plan0_after + assert "result" in plan_replaced, pprint.pformat(plan_replaced) + + elif mode == "change_uid": + # Modify plan UID. Filtering is applied + plan_new = plan0_before + item_uid_to_replace = plan_new["item_uid"] + plan_new = await pq.set_new_item_uuid(plan_new) + plan_replaced, qsize = await pq.replace_item(plan_new, item_uid=item_uid_to_replace) + + queue, _ = await pq.get_queue() + plan0_after = queue[0] + + assert plan_replaced == plan0_after + assert "result" not in plan_replaced, pprint.pformat(plan_replaced) + + elif mode == "change_plan": + # Modify plan, keep UID the same. Filtering is applied + plan_new = plan0_before + new_name = "h" + plan_new["name"] = new_name + plan_replaced, qsize = await pq.replace_item(plan_new, item_uid=plan_new["item_uid"]) + + queue, _ = await pq.get_queue() + plan0_after = queue[0] + + assert plan_replaced == plan0_after + assert "result" not in plan_replaced, pprint.pformat(plan_replaced) + + else: + raise Exception(f"Test mode '{mode}' does not exist.") + + asyncio.run(testing()) + + +def test_replace_item_4_failing(): + """ + ``PlanQueueOperations.replace_item()`` - failing cases + """ + + async def testing(): + async with PQ() as pq: + plans = [ + {"item_type": "plan", "name": "a"}, + {"item_type": "plan", "name": "b"}, + {"item_type": "plan", "name": "c"}, + ] + plans_added = [None] * len(plans) + qsizes = [None] * len(plans) + + for n, plan in enumerate(plans): + plans_added[n], qsizes[n] = await pq.add_item_to_queue(plan) + + assert qsizes == [1, 2, 3] + assert await pq.get_queue_size() == 3 + + # Set the first item as 'running' + running_plan = await pq.set_next_item_as_running() + assert {k: v for k, v in running_plan.items() if k != "properties"} == plans_added[0] + + queue, _ = await pq.get_queue() + running_item_info = await pq.get_running_item_info() + + pq_uid = pq.plan_queue_uid + + plan_new = {"name": "h"} # No UID in the plan. It should still be inserted + # Case: attempt to replace a plan that is currently running + with pytest.raises( + RuntimeError, match="Failed to replace item: Item with UID .* is currently running" + ): + await pq.replace_item(plan_new, item_uid=running_plan["item_uid"]) + assert pq.plan_queue_uid == pq_uid + + # Case: attempt to replace a plan that is not in queue + with pytest.raises(RuntimeError, match="Failed to replace item: Item with UID .* is not in the queue"): + await pq.replace_item(plan_new, item_uid="uid-that-does-not-exist") + assert pq.plan_queue_uid == pq_uid + + # Case: attempt to replace a plan with another plan that already exists in the queue + plan = plans_added[1] + with pytest.raises(RuntimeError, match="Item with UID .* is already in the queue"): + await pq.replace_item(plan, item_uid=plans_added[2]["item_uid"]) + assert pq.plan_queue_uid == pq_uid + + # Case: attempt to replace a plan with currently running plan + with pytest.raises(RuntimeError, match="Item with UID .* is already in the queue"): + await pq.replace_item(running_plan, item_uid=plans_added[2]["item_uid"]) + assert pq.plan_queue_uid == pq_uid + + # Make sure that the queue did not change during the test + plan_queue, _ = await pq.get_queue() + assert plan_queue == queue + assert await pq.get_running_item_info() == running_item_info + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("params, src, order, success, pquid_changed, msg", [ + ({"pos": 1, "pos_dest": 1}, 1, "abcde", True, False, ""), + ({"pos": "front", "pos_dest": "front"}, 0, "abcde", True, False, ""), + ({"pos": "back", "pos_dest": "back"}, 4, "abcde", True, False, ""), + ({"pos": "front", "pos_dest": "back"}, 0, "bcdea", True, True, ""), + ({"pos": "back", "pos_dest": "front"}, 4, "eabcd", True, True, ""), + ({"pos": 1, "pos_dest": 2}, 1, "acbde", True, True, ""), + ({"pos": 2, "pos_dest": 1}, 2, "acbde", True, True, ""), + ({"pos": 0, "pos_dest": 4}, 0, "bcdea", True, True, ""), + ({"pos": 4, "pos_dest": 0}, 4, "eabcd", True, True, ""), + ({"pos": 3, "pos_dest": "front"}, 3, "dabce", True, True, ""), + ({"pos": 2, "pos_dest": "back"}, 2, "abdec", True, True, ""), + ({"pos": -1, "pos_dest": 0}, 4, "eabcd", True, True, ""), + ({"pos": -1, "pos_dest": 1}, 4, "aebcd", True, True, ""), + ({"pos": -2, "pos_dest": 0}, 3, "dabce", True, True, ""), + ({"pos": -2, "pos_dest": 1}, 3, "adbce", True, True, ""), + ({"pos": -1, "pos_dest": -5}, 4, "eabcd", True, True, ""), + ({"pos": -1, "pos_dest": -4}, 4, "aebcd", True, True, ""), + ({"pos": -2, "pos_dest": -5}, 3, "dabce", True, True, ""), + ({"pos": -2, "pos_dest": -4}, 3, "adbce", True, True, ""), + ({"pos": -2, "pos_dest": -3}, 3, "abdce", True, True, ""), + ({"pos": -4, "pos_dest": -2}, 1, "acdbe", True, True, ""), + ({"pos": -3, "pos_dest": -2}, 2, "abdce", True, True, ""), + ({"pos": -2, "pos_dest": -2}, 3, "abcde", True, False, ""), + ({"uid": "p3", "after_uid": "p3"}, 2, "abcde", True, False, ""), + ({"uid": "p1", "before_uid": "p2"}, 0, "abcde", True, True, ""), + ({"uid": "p1", "after_uid": "p2"}, 0, "bacde", True, True, ""), + ({"uid": "p2", "pos_dest": "front"}, 1, "bacde", True, True, ""), + ({"uid": "p2", "pos_dest": "back"}, 1, "acdeb", True, True, ""), + ({"uid": "p1", "pos_dest": "front"}, 0, "abcde", True, False, ""), + ({"uid": "p5", "pos_dest": "back"}, 4, "abcde", True, False, ""), + ({"pos": 1, "after_uid": "p4"}, 1, "acdbe", True, True, ""), + ({"pos": "front", "after_uid": "p4"}, 0, "bcdae", True, True, ""), + ({"pos": 3, "after_uid": "p1"}, 3, "adbce", True, True, ""), + ({"pos": "back", "after_uid": "p1"}, 4, "aebcd", True, True, ""), + ({"pos": 1, "before_uid": "p4"}, 1, "acbde", True, True, ""), + ({"pos": "front", "before_uid": "p4"}, 0, "bcade", True, True, ""), + ({"pos": 3, "before_uid": "p1"}, 3, "dabce", True, True, ""), + ({"pos": "back", "before_uid": "p1"}, 4, "eabcd", True, True, ""), + ({"pos": "back", "after_uid": "p5"}, 4, "abcde", True, False, ""), + ({"pos": "front", "before_uid": "p1"}, 0, "abcde", True, False, ""), + ({"pos": 50, "before_uid": "p1"}, 0, "", False, False, r"Source plan \(position 50\) was not found"), + ({"uid": "abc", "before_uid": "p1"}, 0, "", False, False, r"Source plan \(UID 'abc'\) was not found"), + ({"pos": 3, "pos_dest": 50}, 0, "", False, False, r"Destination plan \(position 50\) was not found"), + ({"uid": "p1", "before_uid": "abc"}, 0, "", False, False, r"Destination plan \(UID 'abc'\) was not found"), + ({"before_uid": "p1"}, 0, "", False, None, r"Source position or UID is not specified"), + ({"pos": 3}, 0, "", False, False, r"Destination position or UID is not specified"), + ({"pos": 1, "uid": "p1", "before_uid": "p4"}, 1, "", False, False, "Ambiguous parameters"), + ({"pos": 1, "pos_dest": 4, "before_uid": "p4"}, 1, "", False, False, "Ambiguous parameters"), + ({"pos": 1, "after_uid": "p4", "before_uid": "p4"}, 1, "", False, False, "Ambiguous parameters"), +]) +# fmt: on +def test_move_item_1(params, src, order, success, pquid_changed, msg): + """ + Basic tests for ``move_item()``. + """ + + async def testing(): + async with PQ() as pq: + plans = [ + {"item_uid": "p1", "name": "a"}, + {"item_uid": "p2", "name": "b"}, + {"item_uid": "p3", "name": "c"}, + {"item_uid": "p4", "name": "d"}, + {"item_uid": "p5", "name": "e"}, + ] + + for plan in plans: + await pq.add_item_to_queue(plan) + + assert await pq.get_queue_size() == len(plans) + pq_uid = pq.plan_queue_uid + + if success: + plan, qsize = await pq.move_item(**params) + assert qsize == len(plans) + assert plan["name"] == plans[src]["name"] + + queue, _ = await pq.get_queue() + names = [_["name"] for _ in queue] + names = "".join(names) + assert names == order + + if pquid_changed: + assert pq.plan_queue_uid != pq_uid + else: + assert pq.plan_queue_uid == pq_uid + + else: + with pytest.raises(Exception, match=msg): + await pq.move_item(**params) + assert pq.plan_queue_uid == pq_uid + + asyncio.run(testing()) + + +def test_move_item_2(): + """ + ``move_item``: test if the item is moved 'as is', i.e. no parameter filtering is applied to it. + """ + + async def testing(): + async with PQ() as pq: + plans = [ + {"item_uid": "p1", "name": "a", "result": {}}, + {"item_uid": "p2", "name": "b"}, + {"item_uid": "p3", "name": "c"}, + ] + + for plan in plans: + await pq.add_item_to_queue(plan, filter_parameters=False) + + queue, _ = await pq.get_queue() + assert await pq.get_queue_size() == 3 + assert "result" in queue[0] + + uid_src = queue[0]["item_uid"] + uid_dest = queue[1]["item_uid"] + + await pq.move_item(uid=uid_src, after_uid=uid_dest) + + queue, _ = await pq.get_queue() + assert await pq.get_queue_size() == 3 + # Make sure that the items were rearranged + assert queue[0]["item_uid"] == uid_dest, pprint.pformat(queue) + assert queue[1]["item_uid"] == uid_src, pprint.pformat(queue) + # Make sure that the 'result' parameter was not removed + assert "result" in queue[1] + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("batch_params, queue_seq, selection_seq, batch_seq, expected_seq, success, msg", [ + ({"pos_dest": "front"}, "0123456", "23", "23", "2301456", True, ""), + ({"before_uid": "0"}, "0123456", "23", "23", "2301456", True, ""), + ({"pos_dest": "back"}, "0123456", "23", "23", "0145623", True, ""), + ({"after_uid": "6"}, "0123456", "23", "23", "0145623", True, ""), + ({"before_uid": "5"}, "0123456", "23", "23", "0142356", True, ""), + ({"after_uid": "5"}, "0123456", "23", "23", "0145236", True, ""), + + ({"pos_dest": 0}, "0123456", "23", "23", "2301456", True, ""), + ({"pos_dest": -7}, "0123456", "23", "23", "2301456", True, ""), + ({"pos_dest": 1}, "0123456", "23", "23", "0231456", True, ""), + ({"pos_dest": -6}, "0123456", "23", "23", "0231456", True, ""), + ({"pos_dest": 5}, "0123456", "23", "23", "0145236", True, ""), + ({"pos_dest": -2}, "0123456", "23", "23", "0145236", True, ""), + ({"pos_dest": 6}, "0123456", "23", "23", "0145623", True, ""), + ({"pos_dest": -1}, "0123456", "23", "23", "0145623", True, ""), + + # Controlling the order of moved items + ({"after_uid": "5"}, "0123456", "023", "023", "1450236", True, ""), + ({"after_uid": "5"}, "0123456", "203", "203", "1452036", True, ""), + ({"after_uid": "5"}, "0123456", "302", "302", "1453026", True, ""), + ({"after_uid": "5", "reorder": False}, "0123456", "302", "302", "1453026", True, ""), + ({"after_uid": "5", "reorder": True}, "0123456", "023", "023", "1450236", True, ""), + ({"after_uid": "5", "reorder": True}, "0123456", "203", "023", "1450236", True, ""), + ({"after_uid": "5", "reorder": True}, "0123456", "302", "023", "1450236", True, ""), + # Empty list of UIDS + ({"pos_dest": "front"}, "0123456", "", "", "0123456", True, ""), + ({"pos_dest": "front"}, "", "", "", "", True, ""), + # Move the batch which is already in the front or back to front or back of the queue + # (nothing is done, but operation is still successful) + ({"pos_dest": "front"}, "0123456", "01", "01", "0123456", True, ""), + ({"pos_dest": "back"}, "0123456", "56", "56", "0123456", True, ""), + # Same, but change the order of moved items + ({"pos_dest": "front"}, "0123456", "10", "10", "1023456", True, ""), + ({"pos_dest": "back"}, "0123456", "65", "65", "0123465", True, ""), + # Failing cases + ({}, "0123456", "23", "23", "0123456", False, "Destination for the batch is not specified"), + ({"pos_dest": "front", "before_uid": "5"}, "0123456", "23", "23", "0123456", False, + "more than one mutually exclusive parameter"), + ({"after_uid": "3"}, "0123456", "023", "023", "0123456", False, "item with UID '33' is in the batch"), + ({"before_uid": "3"}, "0123456", "023", "023", "0123456", False, "item with UID '33' is in the batch"), + ({"after_uid": "5"}, "0123456", "093", "093", "0123456", False, + re.escape("The queue does not contain items with the following UIDs: ['99']")), + ({"after_uid": "5"}, "0123456", "07893", "07893", "0123456", False, + re.escape("The queue does not contain items with the following UIDs: ['77', '88', '99']")), + ({"after_uid": "5"}, "0123456", "0223", "0223", "0123456", False, + re.escape("The list of contains repeated UIDs (1 UIDs)")), +]) +# fmt: on +def test_move_batch_1(batch_params, queue_seq, selection_seq, batch_seq, expected_seq, success, msg): + async def testing(): + async with PQ() as pq: + + async def add_plan(plan, n, **kwargs): + plan_added, qsize = await pq.add_item_to_queue(plan, **kwargs) + assert plan_added["name"] == plan["name"], f"plan: {plan}" + assert qsize == n, f"plan: {plan}" + + def name_to_uid(uid): + return f"{uid}{uid}" + + # Create the queue with plans + for n, p_name in enumerate(queue_seq): + await add_plan({"name": p_name, "item_uid": f"{name_to_uid(p_name)}"}, n + 1) + qsize = await pq.get_queue_size() + assert qsize == len(queue_seq) + + if "before_uid" in batch_params: + batch_params["before_uid"] = name_to_uid(batch_params["before_uid"]) + if "after_uid" in batch_params: + batch_params["after_uid"] = name_to_uid(batch_params["after_uid"]) + + uids = [name_to_uid(_) for _ in selection_seq] + if success: + items_moved, qsize_after = await pq.move_batch(uids=uids, **batch_params) + items_seq = "".join([_["name"] for _ in items_moved]) + assert items_seq == batch_seq + assert qsize_after == qsize + else: + with pytest.raises(Exception, match=msg): + await pq.move_batch(uids=uids, **batch_params) + + # Make sure that the queue is in the correct state + queue, _ = await pq.get_queue() + seq = "".join([_["name"] for _ in queue]) + assert seq == expected_seq + assert len(queue) == qsize + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("pos, name", [ + ("front", "a"), + ("back", "c"), + (0, "a"), + (1, "b"), + (2, "c"), + (3, None), # Index out of range + (-1, "c"), + (-2, "b"), + (-3, "a"), + (-4, None) # Index out of range +]) +# fmt: on +def test_pop_item_from_queue_1(pos, name): + """ + Basic test for the function ``PlanQueueOperations.pop_item_from_queue()`` + """ + + async def testing(): + async with PQ() as pq: + await pq.add_item_to_queue({"name": "a"}) + await pq.add_item_to_queue({"name": "b"}) + await pq.add_item_to_queue({"name": "c"}) + assert await pq.get_queue_size() == 3 + + pq_uid = pq.plan_queue_uid + + if name is not None: + plan, qsize = await pq.pop_item_from_queue(pos=pos) + assert plan["name"] == name + assert qsize == 2 + assert await pq.get_queue_size() == 2 + # Push the plan back to the queue (proves that UID is removed from '_uid_dict') + await pq.add_item_to_queue(plan) + assert await pq.get_queue_size() == 3 + assert pq.plan_queue_uid != pq_uid + else: + with pytest.raises(IndexError, match="Index .* is out of range"): + await pq.pop_item_from_queue(pos=pos) + assert pq.plan_queue_uid == pq_uid + + asyncio.run(testing()) + + +@pytest.mark.parametrize("pos", ["front", "back", 0, 1, -1]) +def test_pop_item_from_queue_2(pos): + """ + Test for the function ``PlanQueueOperations.pop_item_from_queue()``: + the case of empty queue. + """ + + async def testing(): + async with PQ() as pq: + assert await pq.get_queue_size() == 0 + pq_uid = pq.plan_queue_uid + with pytest.raises(IndexError, match="Index .* is out of range|Queue is empty"): + await pq.pop_item_from_queue(pos=pos) + assert pq.plan_queue_uid == pq_uid + + asyncio.run(testing()) + + +def test_pop_item_from_queue_3(): + """ + Pop plans by UID. + """ + + async def testing(): + async with PQ() as pq: + await pq.add_item_to_queue({"item_type": "plan", "name": "a"}) + await pq.add_item_to_queue({"item_type": "plan", "name": "b"}) + await pq.add_item_to_queue({"item_type": "plan", "name": "c"}) + assert await pq.get_queue_size() == 3 + + plans, _ = await pq.get_queue() + assert len(plans) == 3 + plan_to_remove = [_ for _ in plans if _["name"] == "b"][0] + + # Remove one plan + pq_uid = pq.plan_queue_uid + await pq.pop_item_from_queue(uid=plan_to_remove["item_uid"]) + assert await pq.get_queue_size() == 2 + assert pq.plan_queue_uid != pq_uid + + # Attempt to remove the plan again. This should raise an exception. + pq_uid = pq.plan_queue_uid + with pytest.raises( + IndexError, match=f"Plan with UID '{plan_to_remove['item_uid']}' " f"is not in the queue" + ): + await pq.pop_item_from_queue(uid=plan_to_remove["item_uid"]) + assert await pq.get_queue_size() == 2 + assert pq.plan_queue_uid == pq_uid + + # Attempt to remove the plan that is running. This should raise an exception. + await pq.set_next_item_as_running() + assert await pq.get_queue_size() == 1 + pq_uid = pq.plan_queue_uid + with pytest.raises(IndexError, match="Can not remove an item which is currently running"): + await pq.pop_item_from_queue(uid=plans[0]["item_uid"]) + assert await pq.get_queue_size() == 1 + assert pq.plan_queue_uid == pq_uid + + asyncio.run(testing()) + + +def test_pop_item_from_queue_4_fail(): + """ + Failing tests for the function ``PlanQueueOperations.pop_item_from_queue()`` + """ + + async def testing(): + async with PQ() as pq: + pq_uid = pq.plan_queue_uid + with pytest.raises(ValueError, match="Parameter 'pos' has incorrect value"): + await pq.pop_item_from_queue(pos="something") + assert pq.plan_queue_uid == pq_uid + + # Ambiguous parameters (position and UID is passed) + with pytest.raises(ValueError, match="Ambiguous parameters"): + await pq.pop_item_from_queue(pos=5, uid="abc") + assert pq.plan_queue_uid == pq_uid + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("batch_params, queue_seq, selection_seq, batch_seq, expected_seq, success, msg", [ + ({}, "0123456", "", "", "0123456", True, ""), + ({}, "0123456", "23", "23", "01456", True, ""), + ({}, "0123456", "32", "32", "01456", True, ""), + ({}, "0123456", "06", "06", "12345", True, ""), + ({}, "0123456", "283", "23", "01456", True, ""), + ({}, "0123456", "2893", "23", "01456", True, ""), + ({}, "0123456", "2443", "243", "0156", True, ""), + ({"ignore_missing": True}, "0123456", "2443", "243", "0156", True, ""), + ({"ignore_missing": True}, "0123456", "283", "23", "01456", True, ""), + ({"ignore_missing": False}, "0123456", "2443", "", "0123456", False, "The list of contains repeated UIDs"), + ({"ignore_missing": False}, "0123456", "283", "", "0123456", False, + "The queue does not contain items with the following UIDs"), + ({"ignore_missing": False}, "0123456", "2883", "", "0123456", False, "The list of contains repeated UIDs"), + ({}, "0123456", "", "", "0123456", True, ""), + ({}, "", "", "", "", True, ""), + ({}, "", "23", "", "", True, ""), + ({"ignore_missing": False}, "", "", "", "", True, ""), + ({"ignore_missing": False}, "", "23", "", "", False, + "The queue does not contain items with the following UIDs"), +]) +# fmt: on +def test_pop_items_from_queue_batch_1( + batch_params, queue_seq, selection_seq, batch_seq, expected_seq, success, msg +): + """ + Tests for ``pop_items_from_queue_batch``. + """ + + async def testing(): + async with PQ() as pq: + + async def add_plan(plan, n, **kwargs): + plan_added, qsize = await pq.add_item_to_queue(plan, **kwargs) + assert plan_added["name"] == plan["name"], f"plan: {plan}" + assert qsize == n, f"plan: {plan}" + + def name_to_uid(uid): + return f"{uid}{uid}" + + # Create the queue with plans + for n, p_name in enumerate(queue_seq): + await add_plan({"name": p_name, "item_uid": f"{name_to_uid(p_name)}"}, n + 1) + qsize = await pq.get_queue_size() + assert qsize == len(queue_seq) + + uids = [name_to_uid(_) for _ in selection_seq] + if success: + items_removed, qsize_after = await pq.pop_item_from_queue_batch(uids=uids, **batch_params) + items_seq = "".join([_["name"] for _ in items_removed]) + assert items_seq == batch_seq + assert qsize_after == len(expected_seq) + else: + with pytest.raises(Exception, match=msg): + await pq.pop_item_from_queue_batch(uids=uids, **batch_params) + + # Make sure that the queue is in the correct state + queue, _ = await pq.get_queue() + seq = "".join([_["name"] for _ in queue]) + assert seq == expected_seq + assert len(queue) == len(expected_seq) + + asyncio.run(testing()) + + +def test_clear_queue(): + """ + Test for ``PlanQueueOperations.clear_queue`` function + """ + + async def testing(): + async with PQ() as pq: + await pq.add_item_to_queue({"item_type": "plan", "name": "a"}) + await pq.add_item_to_queue({"item_type": "plan", "name": "b"}) + await pq.add_item_to_queue({"item_type": "plan", "name": "c"}) + # Set one of 3 plans as running (removes it from the queue) + await pq.set_next_item_as_running() + + assert await pq.get_queue_size() == 2 + assert len(pq._uid_dict) == 3 + + pq_uid = pq.plan_queue_uid + + # Clears the queue only (doesn't touch the running plan) + await pq.clear_queue() + + assert await pq.get_queue_size() == 0 + assert len(pq._uid_dict) == 1 + assert pq.plan_queue_uid != pq_uid + + with pytest.raises(ValueError, match="Parameter 'pos' has incorrect value"): + await pq.pop_item_from_queue(pos="something") + + asyncio.run(testing()) + + +def test_add_to_history_functions(): + """ + Test for ``PlanQueueOperations._add_to_history()`` method. + """ + + async def testing(): + async with PQ() as pq: + assert await pq.get_history_size() == 0 + + plans = [{"name": "a"}, {"name": "b"}, {"name": "c"}] + ph_uid = pq.plan_history_uid + for plan in plans: + await pq._add_to_history(plan) + assert await pq.get_history_size() == 3 + assert pq.plan_history_uid != ph_uid + + ph_uid = pq.plan_history_uid + plan_history, plan_history_uid_1 = await pq.get_history() + assert pq.plan_history_uid == plan_history_uid_1 + assert pq.plan_history_uid == ph_uid + + assert len(plan_history) == 3 + assert plan_history == plans + + ph_uid = pq.plan_history_uid + await pq.clear_history() + assert pq.plan_history_uid != ph_uid + + plan_history, _ = await pq.get_history() + assert plan_history == [] + + asyncio.run(testing()) + + +@pytest.mark.parametrize("immediate_execution", [False, True]) +@pytest.mark.parametrize("func", ["process_next_item", "set_next_item_as_running"]) +@pytest.mark.parametrize("loop_mode", [False, True]) +def test_process_next_item_1(func, loop_mode, immediate_execution): + """ + Test for ``PlanQueueOperations.process_next_item()`` and + ``PlanQueueOperations.set_next_item_as_running()`` functions. + Those functions are interchangeable for plans, but ``process_next_item`` works + for items and plans. + """ + + async def testing(): + async with PQ() as pq: + await pq.set_plan_queue_mode({"loop": loop_mode}) + + # Apply to empty queue + assert await pq.get_queue_size() == 0 + assert await pq.is_item_running() is False + assert await pq.set_next_item_as_running() == {} + assert await pq.get_queue_size() == 0 + assert await pq.is_item_running() is False + + # Apply to a queue with several plans + p1, _ = await pq.add_item_to_queue({"item_type": "plan", "name": "a"}) + p2, _ = await pq.add_item_to_queue({"item_type": "plan", "name": "b"}) + p3, _ = await pq.add_item_to_queue({"item_type": "plan", "name": "c"}) + + # Plan used for immediate execution + plan4 = {"item_type": "plan", "name": "d"} + + if func == "process_next_item": + f = pq.process_next_item + elif func == "set_next_item_as_running": + f = pq.set_next_item_as_running + else: + assert False, f"Unknown test case: '{func}'" + + pq_uid1 = pq.plan_queue_uid + queue_1, _ = await pq.get_queue() + + # If the next plan is set as running, it is removed from the queue and + # the queue UID is changed. If a plan is set for immediate execution, + # then the queue and the queue UID remain unchanged. + params = {"item": plan4} if immediate_execution else {} + running_plan = await f(**params) + queue_2, _ = await pq.get_queue() + + assert running_plan != {} + + if immediate_execution: + assert running_plan["name"] == "d" + assert "properties" in running_plan + assert running_plan["properties"]["immediate_execution"] is True + assert pq.plan_queue_uid != pq_uid1 + assert await pq.get_queue_size() == 3 + assert len(pq._uid_dict) == 3 + assert queue_1 == queue_2 + else: + assert running_plan["name"] == "a" + assert pq.plan_queue_uid != pq_uid1 + assert await pq.get_queue_size() == 2 + assert len(pq._uid_dict) == 3 + + # Apply if a plan is already running + pq_uid1 = pq.plan_queue_uid + assert await f() == {} + assert pq.plan_queue_uid == pq_uid1 + assert await pq.get_queue_size() == (3 if immediate_execution else 2) + assert len(pq._uid_dict) == 3 + + asyncio.run(testing()) + + +@pytest.mark.parametrize("immediate_execution", [False, True]) +@pytest.mark.parametrize("loop_mode", [False, True]) +def test_process_next_item_2(loop_mode, immediate_execution): + """ + Test for ``PlanQueueOperations.process_next_item()`` and + ``PlanQueueOperations.set_next_item_as_running()`` functions. + Those functions are interchangeable for plans, but ``process_next_item`` works + for items and plans. + + Test with 'instruction' item. + """ + + async def testing(): + async with PQ() as pq: + await pq.set_plan_queue_mode({"loop": loop_mode}) + + # Apply to empty queue + assert await pq.get_queue_size() == 0 + assert await pq.is_item_running() is False + assert await pq.set_next_item_as_running() == {} + assert await pq.get_queue_size() == 0 + assert await pq.is_item_running() is False + + # Apply to a queue with several plans + p0, _ = await pq.add_item_to_queue({"item_type": "instruction", "name": "a"}) + p1, _ = await pq.add_item_to_queue({"item_type": "plan", "name": "b"}) + p2, _ = await pq.add_item_to_queue({"item_type": "plan", "name": "c"}) + + item4 = {"item_type": "instruction", "name": "d"} + + pq_uid = pq.plan_queue_uid + queue_1 = await pq.get_queue() + + params = {"item": item4} if immediate_execution else {} + running_item = await pq.process_next_item(**params) + queue_2 = await pq.get_queue() + + if immediate_execution: + assert {_: running_item[_] for _ in ("name", "item_type")} == item4 + assert pq.plan_queue_uid == pq_uid + assert await pq.get_queue_size() == 3 + assert len(pq._uid_dict) == 3 + assert queue_2 == queue_1 + else: + assert running_item == p0 + assert pq.plan_queue_uid != pq_uid + assert await pq.get_queue_size() == 3 if loop_mode else 2 + assert len(pq._uid_dict) == 3 if loop_mode else 2 + + queue, _ = await pq.get_queue() + assert queue[0] == p1 + assert queue[1] == p2 + if loop_mode: + assert queue[2]["name"] == p0["name"] + assert queue[2]["item_uid"] != p0["name"] + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("has_uid", [False, True]) +@pytest.mark.parametrize("item_type", ["plan", "instruction"]) +# fmt: on +def test_process_next_item_3(item_type, has_uid): + """ + Verify if the function ``process_next_item`` is assigning new UID to plans, but not instructions. + """ + + async def testing(): + async with PQ() as pq: + # Apply to a queue with several plans + await pq.add_item_to_queue({"item_type": "instruction", "name": "a"}) + await pq.add_item_to_queue({"item_type": "plan", "name": "b"}) + await pq.add_item_to_queue({"item_type": "plan", "name": "c"}) + + item4 = {"item_type": item_type, "name": "d"} + if has_uid: + item4 = await pq.set_new_item_uuid(item4) + + running_item = await pq.process_next_item(item=item4) + + if has_uid: + assert running_item["item_uid"] == item4["item_uid"] + else: + assert "item_uid" in running_item + assert isinstance(running_item["item_uid"], str) + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("immediate_execution", [False, True]) +# fmt: on +def test_process_next_item_4_fail(immediate_execution): + """ + Try using 'set_next_item_as_running' with an instruction. Exception should be raised. + The queue should remain unchanged. + """ + + async def testing(): + async with PQ() as pq: + # Apply to a queue with several plans + await pq.add_item_to_queue({"item_type": "instruction", "name": "a"}) + await pq.add_item_to_queue({"item_type": "plan", "name": "b"}) + await pq.add_item_to_queue({"item_type": "plan", "name": "c"}) + + item4 = {"item_type": "instruction", "name": "d"} + + pq_uid = pq.plan_queue_uid + queue_1, _ = await pq.get_queue() + + params = {"item": item4} if immediate_execution else {} + + with pytest.raises(RuntimeError, match="was called for an item other than plan"): + await pq.set_next_item_as_running(**params) + + queue_2, _ = await pq.get_queue() # The queue should remain untouched + + assert pq.plan_queue_uid == pq_uid + assert await pq.get_queue_size() == 3 + assert len(pq._uid_dict) == 3 + assert queue_2 == queue_1 + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("item_in, item_out", [ + ({"name": "count"}, {"name": "count"}), + ({"name": "count", "properties": {}}, {"name": "count"}), + ({"name": "count", "properties": {"nonexisting": 10}}, {"name": "count", "properties": {"nonexisting": 10}}), + ({"name": "count", "properties": {"immediate_execution": True}}, {"name": "count"}), +]) +# fmt: on +def test_clean_item_properties_1(item_in, item_out): + """ + Basic test for `_clean_item_properties` function. + """ + + async def testing(): + async with PQ() as pq: + item_cleaned = pq._clean_item_properties(item_in) + assert item_cleaned == item_out + + asyncio.run(testing()) + + +def test_set_processed_item_as_completed_1(): + """ + Test for ``PlanQueueOperations.set_processed_item_as_completed()`` function. + The function moves currently running plan to history. + """ + + plan_uids = [1, 2, 3] + plans = [ + {"item_type": "plan", "item_uid": plan_uids[0], "name": "a"}, + {"item_type": "plan", "item_uid": plan_uids[1], "name": "b"}, + {"item_type": "plan", "item_uid": plan_uids[2], "name": "c"}, + ] + plans_run_uids = [["abc1"], ["abc2", "abc3"], []] + plans_scan_ids = [[1], [100, 101], []] + + def add_status_to_plans(plans, run_uids, scan_ids, exit_status): + plans = copy.deepcopy(plans) + plans_modified = [] + for plan, run_uid, scan_id in zip(plans, run_uids, scan_ids): + plan.setdefault("result", {}) + plan["result"]["exit_status"] = exit_status + plan["result"]["run_uids"] = run_uid + plan["result"]["scan_ids"] = scan_id + plan["result"]["msg"] = "" + plan["result"]["traceback"] = "" + plans_modified.append(plan) + return plans_modified + + def add_msg_to_plan_history(plan_history, run_uids, msg, tb): + plan_history = copy.deepcopy(plan_history) + for p in plan_history: + if p["item_uid"] in run_uids: + p["result"]["msg"] = msg + p["result"]["traceback"] = tb + return plan_history + + def check_plan_history(plan_history, plan_history_expected): + ph = copy.deepcopy(plan_history) + for _ in ph: + del _["result"]["time_start"] + del _["result"]["time_stop"] + assert ph == plan_history_expected + for _ in plan_history: + assert isinstance(_["result"]["time_start"], float) + assert isinstance(_["result"]["time_stop"], float) + assert _["result"]["time_start"] < _["result"]["time_stop"] + + async def testing(): + async with PQ() as pq: + for plan in plans: + await pq.add_item_to_queue(plan) + + # No plan is running + pq_uid = pq.plan_queue_uid + ph_uid = pq.plan_history_uid + plan = await pq.set_processed_item_as_completed( + exit_status="completed", + run_uids=plans_run_uids[0], + scan_ids=plans_scan_ids[0], + err_msg="", + err_tb="", + ) + assert plan == {} + assert pq.plan_queue_uid == pq_uid + assert pq.plan_history_uid == ph_uid + + # Execute the first plan + await pq.set_next_item_as_running() + pq_uid = pq.plan_queue_uid + plan = await pq.set_processed_item_as_completed( + exit_status="completed", + run_uids=plans_run_uids[0], + scan_ids=plans_scan_ids[0], + err_msg="Test message", + err_tb="Traceback", + ) + assert pq.plan_queue_uid != pq_uid + assert pq.plan_history_uid != ph_uid + + assert await pq.get_queue_size() == 2 + assert await pq.get_history_size() == 1 + assert plan["name"] == plans[0]["name"] + assert plan["result"]["exit_status"] == "completed" + assert plan["result"]["run_uids"] == plans_run_uids[0] + assert plan["result"]["scan_ids"] == plans_scan_ids[0] + + plan_history, _ = await pq.get_history() + plan_history_expected = add_status_to_plans( + plans[0:1], plans_run_uids[0:1], plans_scan_ids[0:1], "completed" + ) + plan_history_expected = add_msg_to_plan_history( + plan_history_expected, [plan_uids[0]], "Test message", "Traceback" + ) + print(plan_history) + print(plan_history_expected) + check_plan_history(plan_history, plan_history_expected) + + # Execute the second plan + await pq.set_next_item_as_running() + plan = await pq.set_processed_item_as_completed( + exit_status="completed", + run_uids=plans_run_uids[1], + scan_ids=plans_scan_ids[1], + err_msg="", + err_tb="", + ) + + assert await pq.get_queue_size() == 1 + assert await pq.get_history_size() == 2 + assert plan["name"] == plans[1]["name"] + assert plan["result"]["exit_status"] == "completed" + assert plan["result"]["run_uids"] == plans_run_uids[1] + assert plan["result"]["scan_ids"] == plans_scan_ids[1] + + plan_history, _ = await pq.get_history() + plan_history_expected = add_status_to_plans( + plans[0:2], plans_run_uids[0:2], plans_scan_ids[0:2], "completed" + ) + plan_history_expected = add_msg_to_plan_history( + plan_history_expected, [plan_uids[0]], "Test message", "Traceback" + ) + check_plan_history(plan_history, plan_history_expected) + + asyncio.run(testing()) + + +def test_set_processed_item_as_completed_2(): + """ + Test for ``PlanQueueOperations.set_processed_item_as_completed()`` function. + Similar test as the previous one, but with LOOP mode ENABLED. + """ + + plan_uids = [1, 2, 3] + plans = [ + {"item_type": "plan", "item_uid": plan_uids[0], "name": "a"}, + {"item_type": "plan", "item_uid": plan_uids[1], "name": "b"}, + {"item_type": "plan", "item_uid": plan_uids[2], "name": "c"}, + ] + plans_run_uids = [["abc1"], ["abc2", "abc3"], []] + plans_scan_ids = [[1], [100, 101], []] + + async def testing(): + async with PQ() as pq: + for plan in plans: + await pq.add_item_to_queue(plan) + + # Enable LOOP mode: the completed item will be pushed to the back of the queue + await pq.set_plan_queue_mode({"loop": True}) + + # No plan is running + pq_uid = pq.plan_queue_uid + ph_uid = pq.plan_history_uid + plan = await pq.set_processed_item_as_completed( + exit_status="completed", + run_uids=plans_run_uids[0], + scan_ids=plans_scan_ids[0], + err_msg="", + err_tb="", + ) + assert plan == {} + assert pq.plan_queue_uid == pq_uid + assert pq.plan_history_uid == ph_uid + + queue, _ = await pq.get_queue() + assert [_["name"] for _ in queue] == ["a", "b", "c"] + + # Execute the first plan + await pq.set_next_item_as_running() + pq_uid = pq.plan_queue_uid + plan = await pq.set_processed_item_as_completed( + exit_status="completed", + run_uids=plans_run_uids[0], + scan_ids=plans_scan_ids[0], + err_msg="", + err_tb="", + ) + assert pq.plan_queue_uid != pq_uid + assert pq.plan_history_uid != ph_uid + + queue, _ = await pq.get_queue() + assert [_["name"] for _ in queue] == ["b", "c", "a"] + + assert await pq.get_queue_size() == 3 + assert await pq.get_history_size() == 1 + assert plan["name"] == plans[0]["name"] + assert plan["result"]["exit_status"] == "completed" + assert plan["result"]["run_uids"] == plans_run_uids[0] + assert plan["result"]["scan_ids"] == plans_scan_ids[0] + assert plan["result"]["msg"] == "" + assert plan["result"]["traceback"] == "" + assert plan["result"]["time_stop"] > plan["result"]["time_start"] + + # Execute the second plan + await pq.set_next_item_as_running() + plan = await pq.set_processed_item_as_completed( + exit_status="unknown", + run_uids=plans_run_uids[1], + scan_ids=plans_scan_ids[1], + err_msg="Unknown exit status", + err_tb="Some traceback", + ) + + queue, _ = await pq.get_queue() + assert [_["name"] for _ in queue] == ["c", "a", "b"] + + assert await pq.get_queue_size() == 3 + assert await pq.get_history_size() == 2 + assert plan["name"] == plans[1]["name"] + assert plan["result"]["exit_status"] == "unknown" + assert plan["result"]["run_uids"] == plans_run_uids[1] + assert plan["result"]["scan_ids"] == plans_scan_ids[1] + assert plan["result"]["msg"] == "Unknown exit status" + assert plan["result"]["traceback"] == "Some traceback" + assert plan["result"]["time_stop"] > plan["result"]["time_start"] + + asyncio.run(testing()) + + +def test_set_processed_item_as_stopped_1(): + """ + Test for ``PlanQueueOperations.set_processed_item_as_stopped()`` function. + The function pushes running plan back to the queue unless ``exit_status=="stopped"`` + and saves it in history as well. + + Typically execution of single-run plans result in no UIDs, but in this test we still + assign UIDs to test functionality. + """ + plan_uids = [1, 2, 3] + plans = [ + {"item_type": "plan", "item_uid": plan_uids[0], "name": "a"}, + {"item_type": "plan", "item_uid": plan_uids[1], "name": "b"}, + {"item_type": "plan", "item_uid": plan_uids[2], "name": "c"}, + ] + plans_run_uids = [["abc1"], ["abc2", "abc3"], []] + plans_scan_ids = [[1], [100, 101], []] + + def add_status_to_plans(plans, run_uids, scan_ids, exit_status): + plans = copy.deepcopy(plans) + + if isinstance(exit_status, list): + assert len(exit_status) == len(plans) + else: + exit_status = [exit_status] * len(plans) + + plans_modified = [] + for plan, run_uid, scan_id, es in zip(plans, run_uids, scan_ids, exit_status): + plan.setdefault("result", {}) + plan["result"]["exit_status"] = es + plan["result"]["run_uids"] = run_uid + plan["result"]["scan_ids"] = scan_id + plan["result"]["msg"] = "" + plan["result"]["traceback"] = "" + plans_modified.append(plan) + return plans_modified + + def add_msg_to_plan_history(plan_history, run_uids, msg, tb): + plan_history = copy.deepcopy(plan_history) + for p in plan_history: + if p["item_uid"] in run_uids: + p["result"]["msg"] = msg + p["result"]["traceback"] = tb + return plan_history + + def check_plan_history(plan_history, plan_history_expected): + ph = copy.deepcopy(plan_history) + for _ in ph: + del _["result"]["time_start"] + del _["result"]["time_stop"] + assert ph == plan_history_expected + for _ in plan_history: + assert isinstance(_["result"]["time_start"], float) + assert isinstance(_["result"]["time_stop"], float) + assert _["result"]["time_start"] < _["result"]["time_stop"] + + async def testing(): + async with PQ() as pq: + for plan in plans: + await pq.add_item_to_queue(plan) + + # No plan is running + pq_uid = pq.plan_queue_uid + ph_uid = pq.plan_history_uid + plan = await pq.set_processed_item_as_stopped( + exit_status="stopped", + run_uids=plans_run_uids[0], + scan_ids=plans_scan_ids[0], + err_msg="Test", + err_tb="", + ) + assert plan == {} + assert pq.plan_queue_uid == pq_uid + assert pq.plan_history_uid == ph_uid + + # Execute the first plan + await pq.set_next_item_as_running() + running_uid1 = (await pq.get_running_item_info())["item_uid"] + pq_uid = pq.plan_queue_uid + plan = await pq.set_processed_item_as_stopped( + exit_status="failed", + run_uids=plans_run_uids[0], + scan_ids=plans_scan_ids[0], + err_msg="Plan failed", + err_tb="Traceback", + ) + assert pq.plan_queue_uid != pq_uid + assert pq.plan_history_uid != ph_uid + + assert await pq.get_queue_size() == 3 + assert await pq.get_history_size() == 1 + assert plan["name"] == plans[0]["name"] + assert plan["result"]["exit_status"] == "failed" + assert plan["result"]["run_uids"] == plans_run_uids[0] + assert plan["result"]["scan_ids"] == plans_scan_ids[0] + assert plan["result"]["msg"] == "Plan failed" + assert plan["result"]["traceback"] == "Traceback" + assert plan["result"]["time_stop"] > plan["result"]["time_start"] + assert plan["item_uid"] == plans[0]["item_uid"] + + # New plan UID is generated when the plan is pushed back into the queue + queue, _ = await pq.get_queue() + plan_modified_uid = queue[0]["item_uid"] + # Make sure that the item UID was changed + assert plan_modified_uid != plans[0]["item_uid"] + + plan_history, _ = await pq.get_history() + plan_history_expected = add_status_to_plans( + [plans[0]], [plans_run_uids[0]], [plans_scan_ids[0]], "failed" + ) + plan_history_expected = add_msg_to_plan_history( + plan_history_expected, [running_uid1], "Plan failed", "Traceback" + ) + check_plan_history(plan_history, plan_history_expected) + + # Execute the second plan + await pq.set_next_item_as_running() + running_uid2 = (await pq.get_running_item_info())["item_uid"] + plan = await pq.set_processed_item_as_stopped( + exit_status="stopped", + run_uids=plans_run_uids[1], + scan_ids=plans_scan_ids[1], + err_msg="Plan stopped", + err_tb="Traceback 2", + ) + + assert await pq.get_queue_size() == 2 + assert await pq.get_history_size() == 2 + assert plan["name"] == plans[0]["name"] + assert plan["result"]["exit_status"] == "stopped" + assert plan["result"]["run_uids"] == plans_run_uids[1] + assert plan["result"]["scan_ids"] == plans_scan_ids[1] + assert plan["result"]["msg"] == "Plan stopped" + assert plan["result"]["traceback"] == "Traceback 2" + assert plan["result"]["time_stop"] > plan["result"]["time_start"] + + plan_history, _ = await pq.get_history() + plan_history_expected = add_status_to_plans( + [plans[0].copy(), plans[0].copy()], + [plans_run_uids[0], plans_run_uids[1]], + [plans_scan_ids[0], plans_scan_ids[1]], + ["failed", "stopped"], + ) + # Plan 0 has different UID after it was inserted in the queue during the 1st attempt + plan_history_expected[1]["item_uid"] = plan_modified_uid + plan_history_expected = add_msg_to_plan_history( + plan_history_expected, [running_uid1], "Plan failed", "Traceback" + ) + plan_history_expected = add_msg_to_plan_history( + plan_history_expected, [running_uid2], "Plan stopped", "Traceback 2" + ) + check_plan_history(plan_history, plan_history_expected) + + # Verify that `_uid_dict` still has correct size. `_uid_dict` should never be accessed directly. + assert len(pq._uid_dict) == 2 + # Also it should not contain UIDs of already executed plans. + for plan in plan_history: + assert plan["item_uid"] not in pq._uid_dict + + asyncio.run(testing()) + + +# fmt: off +@pytest.mark.parametrize("func", ["completed", "unknown", "failed", "stopped", "aborted", "halted"]) +@pytest.mark.parametrize("loop_mode", [False, True]) +@pytest.mark.parametrize("immediate_execution", [False, True]) +# fmt: on +def test_set_processed_item_as_stopped_2(loop_mode, func, immediate_execution): + """ + ``set_processed_item_as_completed`` and ``set_processed_item_as_stopped`` processing of an item set + for normal and immediate execution with/without LOOP mode. + """ + plans = [ + {"item_type": "plan", "item_uid": 1, "name": "a"}, + {"item_type": "plan", "item_uid": 2, "name": "b"}, + {"item_type": "plan", "item_uid": 3, "name": "c"}, + ] + plan4 = {"item_type": "plan", "item_uid": 3, "name": "d"} + plan4_run_uids = ["abc2", "abc3"] + plan4_scan_ids = [100, 101] + + async def testing(): + async with PQ() as pq: + for plan in plans: + await pq.add_item_to_queue(plan) + + # Enable LOOP mode: the completed item will be pushed to the back of the queue + await pq.set_plan_queue_mode({"loop": loop_mode}) + + pq_uid = pq.plan_queue_uid + ph_uid = pq.plan_history_uid + queue_1, _ = await pq.get_queue() + + if immediate_execution: + await pq.process_next_item(item=plan4) + else: + await pq.process_next_item() + + assert pq.plan_queue_uid != pq_uid + pq_uid2 = pq.plan_queue_uid + assert pq.plan_history_uid == ph_uid + + err_msg, err_tb = f"Plan {func}", f"Traceback {func}" + if func in ("completed", "unknown"): + plan = await pq.set_processed_item_as_completed( + exit_status=func, + run_uids=plan4_run_uids, + scan_ids=plan4_scan_ids, + err_msg=err_msg, + err_tb=err_tb, + ) + elif func in ("failed", "stopped", "aborted", "halted"): + plan = await pq.set_processed_item_as_stopped( + exit_status=func, + run_uids=plan4_run_uids, + scan_ids=plan4_scan_ids, + err_msg=err_msg, + err_tb=err_tb, + ) + else: + assert False, f"Unknown value of parameter 'func': '{func}'" + + queue_2, _ = await pq.get_queue() + + if immediate_execution: + assert queue_1 == queue_2 + assert await pq.get_queue_size() == 3 + else: + assert queue_1 != queue_2 + qsize = 2 if func in ("completed", "unknown", "stopped") and not loop_mode else 3 + assert await pq.get_queue_size() == qsize + assert await pq.get_history_size() == 1 + assert pq.plan_queue_uid != pq_uid2 + assert pq.plan_history_uid != ph_uid + + def check_plan(p): + assert p["name"] == plan4["name"] if immediate_execution else plans[0]["name"] + assert p["result"]["exit_status"] == func + assert p["result"]["run_uids"] == plan4_run_uids + assert p["result"]["scan_ids"] == plan4_scan_ids + assert p["result"]["msg"] == err_msg + assert p["result"]["traceback"] == err_tb + assert plan["result"]["time_stop"] > plan["result"]["time_start"] + + check_plan(plan) + + history, _ = await pq.get_history() + check_plan(history[0]) + + assert history[0]["item_uid"] == plan["item_uid"] + + asyncio.run(testing()) + + +@pytest.mark.parametrize("loop_mode", [False, True]) +@pytest.mark.parametrize("func", ["completed", "stopped"]) +def test_set_processed_item_as_stopped_3(loop_mode, func): + """ + Test for ``PlanQueueOperations.set_processed_item_as_completed()`` and + ``PlanQueueOperations.set_processed_item_as_stopped()`` function. + Make sure that ``item["properties"]["immediate_execution"] is removed before the plan + is added to history and the plan is not pushed back into the queue if the plan is + stopped. The behavior of both functions is identical for the plans executed + in 'immediate execution' mode. + """ + + plan = {"item_type": "plan", "item_uid": 1, "name": "a", "properties": {"immediate_execution": True}} + + async def testing(): + async with PQ() as pq: + await pq.add_item_to_queue(plan) + await pq.set_plan_queue_mode({"loop": loop_mode}) + + assert await pq.get_queue_size() == 1 + assert await pq.get_history_size() == 0 + + await pq.set_next_item_as_running() + + assert await pq.get_queue_size() == 0 + assert await pq.get_history_size() == 0 + + if func == "completed": + plan1 = await pq.set_processed_item_as_completed( + exit_status="completed", run_uids=["abc"], scan_ids=[1], err_msg="", err_tb="" + ) + elif func == "stopped": + plan1 = await pq.set_processed_item_as_stopped( + exit_status="stopped", run_uids=["abc"], scan_ids=[1], err_msg="", err_tb="" + ) + else: + raise Exception(f"Unsupported parameter value func={func!r}") + + # assert False, plan1 + assert plan1["name"] == plan["name"] + assert "properties" not in plan1 + + assert await pq.get_queue_size() == 0 + assert await pq.get_history_size() == 1 + + history, _ = await pq.get_history() + assert history[0] == plan1 + assert "properties" not in history[0] + + asyncio.run(testing()) + + +# ============================================================================================== +# Saving and retrieving user group permissions +def test_user_group_permissions_1(): + """ + Test for saving and retrieving and clearing user group permissions, which are backed up in Redis. + """ + ug_permissions_1 = {"some_key_1": "some_value_1"} + ug_permissions_2 = {"some_key_2": "some_value_2"} + + async def testing(): + async with PQ() as pq: + assert await pq.user_group_permissions_retrieve() is None + + await pq.user_group_permissions_save(ug_permissions_1) + assert await pq.user_group_permissions_retrieve() == ug_permissions_1 + + await pq.user_group_permissions_save(ug_permissions_2) + assert await pq.user_group_permissions_retrieve() == ug_permissions_2 + + await pq.user_group_permissions_clear() + assert await pq.user_group_permissions_retrieve() is None + + asyncio.run(testing()) + + +# ============================================================================================== +# Saving and retrieving lock info +def test_lock_info_1(): + """ + Test for saving and retrieving and clearing user group permissions, which are backed up in Redis. + """ + lock_info_1 = {"environment": False, "queue": True, "lock_key": "abcde"} + lock_info_2 = {"environment": True, "queue": False, "lock_key": "fghijk"} + + async def testing(): + async with PQ() as pq: + assert await pq.lock_info_retrieve() is None + + await pq.lock_info_save(lock_info_1) + assert await pq.lock_info_retrieve() == lock_info_1 + + await pq.lock_info_save(lock_info_2) + assert await pq.lock_info_retrieve() == lock_info_2 + + await pq.lock_info_clear() + assert await pq.lock_info_retrieve() is None + + asyncio.run(testing()) + + +# ============================================================================================== +# Saving and retrieving 'stop pending' info +def test_stop_pending_info_1(): + """ + Test for saving and retrieving and clearing user group permissions, which are backed up in Redis. + """ + stop_pending_info_1 = {"enabled": False} + stop_pending_info_2 = {"enabled": True} + + async def testing(): + async with PQ() as pq: + assert await pq.stop_pending_retrieve() is None + + await pq.stop_pending_save(stop_pending_info_1) + assert await pq.stop_pending_retrieve() == stop_pending_info_1 + + await pq.stop_pending_save(stop_pending_info_2) + assert await pq.stop_pending_retrieve() == stop_pending_info_2 + + await pq.stop_pending_clear() + assert await pq.stop_pending_retrieve() is None + + asyncio.run(testing()) + + +# ============================================================================================== +# Saving and retrieving autostart info +def test_autostart_mode_info_1(): + """ + Test for saving and retrieving and clearing user group permissions, which are backed up in Redis. + """ + autostart_info_1 = {"enabled": False} + autostart_info_2 = {"enabled": True} + + async def testing(): + async with PQ() as pq: + assert await pq.autostart_mode_retrieve() is None + + await pq.autostart_mode_save(autostart_info_1) + assert await pq.autostart_mode_retrieve() == autostart_info_1 + + await pq.autostart_mode_save(autostart_info_2) + assert await pq.autostart_mode_retrieve() == autostart_info_2 + + await pq.autostart_mode_clear() + assert await pq.autostart_mode_retrieve() is None + + asyncio.run(testing()) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index dce64e56..060709ca 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +aiosqlite bluesky>=1.7.0 ipykernel jsonschema