def check_file(file_id: str, upsert: bool = False) -> File: """Checks that the file with file_id exists in the DB Args: file_id: The id for the requested file. upsert: If the file doesn't exist create a placeholder file Returns: The file object Raises: NotFoundError: File with the requested ID doesn't exist and is expected to ModelValidationError: Incorrectly formatted ID is given """ try: ObjectId(file_id) except (InvalidId, TypeError): raise ModelValidationError( f"Cannot create a file id with the string {file_id}. " "Requires 24-character hex string.") res = db.query_unique(File, id=file_id) if res is None: if upsert: create_file("BG_placeholder", 0, 0, file_id) res = db.query_unique(File, id=file_id) else: raise NotFoundError(f"Tried to fetch an unsaved file {file_id}") db.modify(res, updated_at=datetime.utcnow()) return res
def _from_kwargs( system: System = None, instance: Instance = None, system_id: str = None, instance_name: str = None, instance_id: str = None, **_, ) -> Tuple[System, Instance]: if system and instance: return system, instance if not system: if system_id: system = db.query_unique(System, raise_missing=True, id=system_id) elif instance: system = db.query_unique(System, raise_missing=True, instances__contains=instance) elif instance_id: system = db.query_unique(System, raise_missing=True, instances__id=instance_id) else: raise NotFoundException("Unable to find System") if not instance: if instance_name: instance = system.get_instance_by_name(instance_name) elif instance_id: instance = system.get_instance_by_id(instance_id) else: raise NotFoundException("Unable to find Instance") return system, instance
def upsert(system: System) -> System: """Helper to create or update a system Args: system: The system to create or update Returns: The created / updated system """ try: return create_system(system, _publish_error=False) except NotUniqueException: existing = db.query_unique( System, namespace=system.namespace, name=system.name, version=system.version ) return update_system( system=existing, new_commands=system.commands, add_instances=system.instances, description=system.description, display_name=system.display_name, icon_name=system.icon_name, metadata=system.metadata, )
def _determine_target_garden(operation: Operation) -> str: """Determine the system the operation is targeting""" # Certain operations are ASSUMED to be targeted at the local garden if ("READ" in operation.operation_type or "JOB" in operation.operation_type or "FILE" in operation.operation_type or operation.operation_type in ("PLUGIN_LOG_RELOAD", "SYSTEM_CREATE", "SYSTEM_RESCAN") or "PUBLISH_EVENT" in operation.operation_type or "RUNNER" in operation.operation_type or operation.operation_type in ("PLUGIN_LOG_RELOAD", "SYSTEM_CREATE")): return config.get("garden.name") # Otherwise, each operation needs to be "parsed" if operation.operation_type in ("SYSTEM_RELOAD", "SYSTEM_UPDATE"): return _system_id_lookup(operation.args[0]) if operation.operation_type == "SYSTEM_DELETE": # Force deletes get routed to local garden if operation.kwargs.get("force"): return config.get("garden.name") return _system_id_lookup(operation.args[0]) if "INSTANCE" in operation.operation_type: if "system_id" in operation.kwargs and "instance_name" in operation.kwargs: return _system_id_lookup(operation.kwargs["system_id"]) else: return _instance_id_lookup(operation.args[0]) if operation.operation_type == "REQUEST_CREATE": target_system = System( namespace=operation.model.namespace, name=operation.model.system, version=operation.model.system_version, ) return _system_name_lookup(target_system) if operation.operation_type.startswith("REQUEST"): request = db.query_unique(Request, id=operation.args[0]) operation.kwargs["request"] = request return config.get("garden.name") if "GARDEN" in operation.operation_type: if operation.operation_type == "GARDEN_SYNC": sync_target = operation.kwargs.get("sync_target") if sync_target: return sync_target return config.get("garden.name") if operation.operation_type == "QUEUE_DELETE": # Need to deconstruct the queue name parts = operation.args[0].split(".") version = parts[2].replace("-", ".") return _system_name_lookup( System(namespace=parts[0], name=parts[1], version=version)) raise Exception(f"Bad operation type {operation.operation_type}")
def local_garden(all_systems: bool = False) -> Garden: """Get the local garden definition Args: all_systems: If False, only include "local" systems in the garden systems list Returns: The local Garden """ # This function is still necessary because there are various things that expect # the system information to be embedded in the garden document itself (as opposed # Systems just having a reference to their garden). There is nothing that would # keep a LOCAL garden's embedded list of systems up to date currently, so we instead # build the list of systems (and namespaces) at call time. Once the System # relationship has been refactored, the need for this function should go away. garden: Garden = db.query_unique(Garden, connection_type="LOCAL") filter_params = {} if not all_systems: filter_params["local"] = True garden.systems = get_systems(filter_params=filter_params) garden.namespaces = get_namespaces() return garden
def _fetch_file(file_id: str) -> FileStatus: """Fetches the entire requested file Args: file_id: This should be a valid file id. Returns: The file data if the file is valid; None otherwise. Raises: NotFoundError: File with the requested ID doesn't exist and is expected to. ModelValidationError: File with the requested ID has no associated chunks. """ # This is going to get big, try our best to be efficient check = _verify_chunks(file_id) if check.valid: file = check_chunks(file_id) all_data = [ db.query_unique(FileChunk, id=file.chunks[str(x)]).data for x in range(len(file.chunks)) ] return _safe_build_object( FileStatus, file, operation_complete=True, # Each chunk should be base64 encoded, and # we can't just concat those strings. data=b64encode(b"".join(map(b64decode, all_data))).decode("utf-8"), ) else: return check
def process_wait(request: Request, timeout: float) -> Request: """Helper to process a request and wait for completion using a threading.Event Args: request: Request to create timeout: Timeout used for wait Returns: The completed request """ # We need a better solution for this. Because the Request library is imported # everywhere it causes issues when importing the router at the top because all of # the functions are not initialized. So we either leave this as is, or move the # requests import to the end of all of the files. import beer_garden.router as router req_complete = threading.Event() # Send the request through the router to allow for commands to work across Gardens created_request = router.route( Operation( operation_type="REQUEST_CREATE", model=request, model_type="Request", kwargs={"wait_event": req_complete}, )) if not req_complete.wait(timeout): raise TimeoutError( "Request did not complete before the specified timeout") return db.query_unique(Request, id=created_request.id)
def get_and_validate_system(self, request): """Ensure there is a system in the DB that corresponds to this Request. :param request: The request to validate :return: The system corresponding to this Request :raises ModelValidationError: There is no system that corresponds to this Request """ system = db.query_unique( System, namespace=request.namespace, name=request.system, version=request.system_version, ) if system is None: raise ModelValidationError( "Could not find System named '%s' matching version '%s'" % (request.system, request.system_version)) if request.instance_name not in system.instance_names: raise ModelValidationError( "Could not find instance with name '%s' in system '%s'" % (request.instance_name, system.name)) self.logger.debug("Found System %s-%s" % (request.system, request.instance_name)) return system
def test_existing_multiple(self, tmp_path, loader, registry, plugin_1, bg_instance): """This is mainly to test that Instance IDs are correct We save a system with 2 instances: - instance1, 58542eb571afd47ead90beef - instance2, 58542eb571afd47ead90beee Then we load a plugin that defines instances [instance2, instance3]. Correct behavior is: - instance1 removed from the database - instance3 created in the database - instance2 remains in the database, and the ID remains the same """ instance1 = Instance(name="instance1", id="58542eb571afd47ead90beef") instance2 = Instance(name="instance2", id="58542eb571afd47ead90beee") create_system( System(name="foo", version="1.0", instances=[instance1, instance2]) ) plugin = tmp_path / "plugin" plugin.mkdir() write_file( plugin, textwrap.dedent( """ NAME='foo' VERSION='1.0' PLUGIN_ENTRY='entry.py' INSTANCES=["instance2", "instance3"] """ ), ) plugin_runners = loader.load_plugin(plugin) assert len(plugin_runners) == 2 assert db.query_unique(Instance, name="instance1") is None assert db.query_unique(Instance, name="instance3") is not None instance2_db = db.query_unique(Instance, name="instance2") assert instance2_db is not None assert instance2_db.id == instance2.id
def set_owner(file_id: str, owner_id: str = None, owner_type: str = None) -> FileStatus: """Sets the owner field of the file. This is used for DB pruning. Args: file_id: The id of the file. owner_id: The id of the owner. owner_type: The type of the owner (job/request). """ if (owner_id is not None) and (owner_type is not None): file = check_file(file_id) old_owner_priority = OWNERSHIP_PRIORITY.get(file.owner_type, 1_000_000) new_owner_priority = OWNERSHIP_PRIORITY.get(owner_type, 1_000_000) # Case 1 : New owner has equal or higher priority # Case 2 : The old owner is a higher priority than the new one, but it was deleted. if new_owner_priority <= old_owner_priority or ( file.owner_type in OWNERSHIP_PRIORITY and (file.job is None and file.request is None)): if owner_type in OWNERSHIP_MAP: owner = db.query_unique(OWNERSHIP_MAP[owner_type], id=owner_id) file = db.modify( file, owner_id=owner_id, owner_type=owner_type, job=owner.id if owner is not None and owner_type == "JOB" else None, request=owner.id if owner is not None and owner_type == "REQUEST" else None, ) else: file = db.modify(file, owner_id=owner_id, owner_type=owner_type) return _safe_build_object(FileStatus, file, operation_complete=True) return _safe_build_object( FileStatus, operation_complete=False, message= (f"Owner type {owner_type} has lower priority than {file.owner_type}" ), ) return _safe_build_object( FileStatus, operation_complete=False, message=("Operation FILE_OWN requires an owner type " f"and id. Got {owner_type} and {owner_id}"), )
def get_job(self, job_id): """Looks up a job Args: job_id: The job id """ if job_id in self._async_jobs: return db.query_unique(Job, id=job_id) else: return self._sync_scheduler.get_job(job_id)
def get_instance( instance_id: str = None, system_id: str = None, instance_name: str = None, instance: Instance = None, **_, ) -> Instance: """Retrieve an individual Instance Args: instance_id: The Instance ID system_id: The System ID instance_name: The Instance name instance: The Instance Returns: The Instance """ if instance: return instance if system_id and instance_name: system = db.query_unique(System, raise_missing=True, id=system_id) try: return system.get_instance_by_name(instance_name, raise_missing=True) except BrewtilsException: raise NotFoundException( f"System {system} does not have an instance with name '{instance_name}'" ) from None elif instance_id: system = db.query_unique(System, raise_missing=True, instances__id=instance_id) try: return system.get_instance_by_id(instance_id, raise_missing=True) except BrewtilsException: raise NotFoundException( f"System {system} does not have an instance with id '{instance_id}'" ) from None raise NotFoundException()
def remove_job(job_id: str) -> None: """Remove a Job Args: job_id: The Job ID Returns: The Job ID """ # The scheduler takes care of removing the Job from the database return db.query_unique(Job, id=job_id)
def get_system(system_id: str) -> System: """Retrieve an individual System Args: system_id: The System ID Returns: The System """ return db.query_unique(System, id=system_id)
def get_garden(garden_name: str) -> Garden: """Retrieve an individual Garden Args: garden_name: The name of Garden Returns: The Garden """ if garden_name == config.get("garden.name"): return local_garden() return db.query_unique(Garden, name=garden_name, raise_missing=True)
def handle_event(event): # Whenever a request is completed check to see if this process is waiting for it if event.name == Events.REQUEST_COMPLETED.name: completion_event = request_map.pop(event.payload.id, None) if completion_event: completion_event.set() # Only care about local garden if event.garden == config.get("garden.name"): if event.name == Events.GARDEN_STOPPED.name: # When shutting down we need to close all handing connections/threads # waiting for a response. This will invoke each connection/thread to be # returned the current status of the Request. for request_event in request_map: request_map[request_event].set() # Only care about downstream garden elif event.garden != config.get("garden.name"): if event.name in ( Events.REQUEST_CREATED.name, Events.REQUEST_STARTED.name, Events.REQUEST_COMPLETED.name, ): # When we send child requests to child gardens where the parent was on # the local garden we remove the parent before sending them. Only setting # the subset of fields that change "corrects" the parent existing_request = db.query_unique(Request, id=event.payload.id) if existing_request: for field in ("status", "output", "error_class"): setattr(existing_request, field, getattr(event.payload, field)) try: db.update(existing_request) except RequestStatusTransitionError: pass else: # Attempt to create the request, if it already exists then continue on try: db.create(event.payload) except NotUniqueException: pass # Required if the main process spawns a wait Request if event.name == Events.REQUEST_COMPLETED.name: if str(event.payload.id) in request_map: request_map[str(event.payload.id)].set()
def resume_job(job_id: str) -> Job: """Resume a Job Args: job_id: The Job ID Returns: The Job definition """ job = db.query_unique(Job, id=job_id) job.status = "RUNNING" job = db.update(job) return job
def pause_job(job_id: str) -> Job: """Pause a Job Args: job_id: The Job ID Returns: The Job definition """ job = db.query_unique(Job, id=job_id) job.status = "PAUSED" job = db.update(job) return job
def reload_system(system_id: str = None, system: System = None) -> None: """Reload a local plugin System Args: system_id: The System ID system: The System Returns: None """ system = system or db.query_unique(System, id=system_id) lpm.reload(system=system) return system
def remove_job(self, job_id, **kwargs): """Removes the job from the corresponding scheduler Args: job_id: The job id to lookup kwargs: Any other scheduler-specific arguments """ if job_id in self._async_jobs: self._async_jobs.pop(job_id) # Clean up the if job_id in self._async_paused_jobs: self._async_paused_jobs.remove(job_id) db.delete(db.query_unique(Job, id=job_id)) else: self._sync_scheduler.remove_job(job_id, **kwargs)
def update_garden_config(garden: Garden) -> Garden: """Update Garden configuration parameters Args: garden: The Garden to Update Returns: The Garden updated """ db_garden = db.query_unique(Garden, id=garden.id) db_garden.connection_params = garden.connection_params db_garden.connection_type = garden.connection_type db_garden.status = "INITIALIZING" return update_garden(db_garden)
def get_request(request_id: str = None, request: Request = None) -> Request: """Retrieve an individual Request Args: request_id: The Request ID request: The Request Returns: The Request """ request = request or db.query_unique( Request, id=request_id, raise_missing=True) request.children = db.query(Request, filter_params={"parent": request}) return request
def update_garden_status(garden_name: str, new_status: str) -> Garden: """Update an Garden status. Will also update the status_info heartbeat. Args: garden_name: The Garden Name new_status: The new status Returns: The updated Garden """ garden = db.query_unique(Garden, name=garden_name) garden.status = new_status garden.status_info["heartbeat"] = datetime.utcnow() return update_garden(garden)
def remove_garden(garden_name: str) -> None: """Remove a garden Args: garden_name: The Garden name Returns: None """ garden = db.query_unique(Garden, name=garden_name) for system in garden.systems: remove_system(system.id) db.delete(garden) return garden
def get_command(system_id: str, command_name: str) -> Command: """Retrieve an individual Command Args: system_id: The System ID command_name: The Command name Returns: The Command """ system = db.query_unique( System, raise_missing=True, id=system_id, commands__name=command_name ) for command in system.commands: if command.name == command_name: return command
def execute_job(job_id: str, reset_interval=False) -> Job: """Execute a Job ad-hoc Creates a new job with a trigger for now. Args: job_id: The Job ID Returns: The spawned Request """ job = db.query_unique(Job, id=job_id, raise_missing=True) job.reset_interval = reset_interval if reset_interval and job.trigger_type != "interval": raise ModelValidationError( "reset_interval can only be used with trigger type of interval") return job
def remove_system(system_id: str = None, system: System = None) -> System: """Remove a system Args: system_id: The System ID system: The System Returns: The removed System """ system = system or db.query_unique(System, id=system_id) db.delete(system) # Also need to let the routing module know from beer_garden.router import remove_routing_system remove_routing_system(system=system) return system
def check_chunk(chunk_id: str): """Checks that the file with file_id exists in the DB Args: chunk_id: The id for the requested chunk. Returns: The file object. Raises: NotFoundError: Chunk with the requested ID doesn't exist. ModelValidationError: Incorrectly formatted ID is given """ try: ObjectId(chunk_id) except (InvalidId, TypeError): raise ModelValidationError( f"Cannot create a chunk id with the string {chunk_id}. " "Requires 24-character hex string.") return db.query_unique(FileChunk, id=chunk_id, raise_missing=True)
def cancel_request(request_id: str = None, request: Request = None) -> Request: """Mark a Request as CANCELED Args: request_id: The Request ID to cancel request: The Request to cancel Returns: The modified Request Raises: ModelValidationError: The Request is already completed """ request = request or db.query_unique( Request, raise_missing=True, id=request_id) request.status = "CANCELED" request = db.update(request) # TODO - Metrics here? return request
def handle_event(event): # Only care about downstream garden if event.garden != config.get("garden.name") and not event.error: if event.name in ( Events.REQUEST_CREATED.name, Events.REQUEST_STARTED.name, Events.REQUEST_COMPLETED.name, Events.REQUEST_UPDATED.name, Events.REQUEST_CANCELED.name, ): # When we send child requests to child gardens where the parent was on # the local garden we remove the parent before sending them. Only setting # the subset of fields that change "corrects" the parent existing_request = db.query_unique(Request, id=event.payload.id) if existing_request is None: # Attempt to create the request, if it already exists then continue on try: db.create(event.payload) except NotUniqueException: pass elif event.name != Events.REQUEST_CREATED.name: request_changed = False for field in ("status", "output", "error_class", "status_updated_at"): new_value = getattr(event.payload, field) if getattr(existing_request, field) != new_value: request_changed = True setattr(existing_request, field, new_value) if request_changed: try: update_request(existing_request, _publish_error=False) except RequestStatusTransitionError: pass