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
