Skip to content

Commit

Permalink
Add Startup Annotation Support
Browse files Browse the repository at this point in the history
  • Loading branch information
TheBurchLog committed Oct 31, 2024
1 parent 22d6552 commit 4b89b59
Show file tree
Hide file tree
Showing 7 changed files with 175 additions and 40 deletions.
3 changes: 3 additions & 0 deletions CHANGELOG.rst
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@ TBD

- Added new annotation/configuration support for shutdown functions. These functions will be executed at the start
of the shutdown process.
- Added new annotation/configuration support for startup functions. These functions will be executed after `Plugin().run()`
has completed startup processes


3.28.0
------
Expand Down
17 changes: 13 additions & 4 deletions brewtils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,21 @@
from brewtils.__version__ import __version__
from brewtils.auto_decorator import AutoDecorator
from brewtils.config import get_argument_parser, get_connection_info, load_config
from brewtils.decorators import client, command, parameter, shutdown, subscribe, system
from brewtils.decorators import (
client,
command,
parameter,
shutdown,
startup,
subscribe,
system,
)
from brewtils.log import configure_logging
from brewtils.plugin import (
get_current_request_read_only,
from brewtils.plugin import ( # noqa F401
Plugin,
RemotePlugin,
) # noqa F401
get_current_request_read_only,
)
from brewtils.rest import normalize_url_prefix
from brewtils.rest.easy_client import EasyClient, get_easy_client
from brewtils.rest.publish_client import PublishClient
Expand All @@ -20,6 +28,7 @@
"command",
"parameter",
"shutdown",
"startup",
"system",
"subscribe",
"Plugin",
Expand Down
36 changes: 36 additions & 0 deletions brewtils/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -469,6 +469,24 @@ def pre_shutdown(self):
return _wrapped


def startup(_wrapped=None):
"""Decorator for specifying a function to run before a plugin is running.
for example::
@startup
def pre_running(self):
# Run pre-running processing
return
Args:
_wrapped: The function to decorate. This is handled as a positional argument and
shouldn't be explicitly set.
"""
_wrapped._startup = True
return _wrapped


def subscribe(_wrapped=None, topic: str = None, topics=[]):
"""Decorator for specifiying topic to listen to.
Expand Down Expand Up @@ -529,6 +547,24 @@ def _parse_shutdown_functions(client):
return shutdown_functions


def _parse_startup_functions(client):
# type: (object) -> List[Callable]
"""Get a list of callable fields labeled with the startup annotation
This will iterate over everything returned from dir, looking for metadata added
by the startup decorator.
"""

startup_functions = []

for attr in dir(client):
method = getattr(client, attr)
if callable(method) and getattr(method, "_startup", False):
startup_functions.append(method)

return startup_functions


def _parse_client(client):
# type: (object) -> List[Command]
"""Get a list of Beergarden Commands from a client object
Expand Down
118 changes: 86 additions & 32 deletions brewtils/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,11 @@

import brewtils
from brewtils.config import load_config
from brewtils.decorators import _parse_client, _parse_shutdown_functions
from brewtils.decorators import (
_parse_client,
_parse_shutdown_functions,
_parse_startup_functions,
)
from brewtils.display import resolve_template
from brewtils.errors import (
ConflictError,
Expand Down Expand Up @@ -183,9 +187,22 @@ class Plugin(object):
group (str): Grouping label applied to plugin
groups (list): Grouping labels applied to plugin
client_shutdown_function (str): Function to be executed at start of shutdown
within client code
client_shutdown_functions (list): Functions to be executed at start of shutdown
within client code
shutdown_function (func): Function to be executed at start of shutdown
shutdown_functions (list): Functions to be executed at start of shutdown
client_startup_function (str): Function to be executed at start of running plugin
within client code
client_startup_functions (list): Functions to be executed at start of running plugin
within client code
startup_function (func): Function to be executed at start of running plugin
startup_functions (list): Functions to be executed at start of running plugin
prefix_topic (str): Prefix for Generated Command Topics
logger (:py:class:`logging.Logger`): Logger that will be used by the Plugin.
Expand Down Expand Up @@ -223,9 +240,42 @@ def __init__(self, client=None, system=None, logger=None, **kwargs):

# Need to pop out shutdown functions because these are not processed
# until shutdown
self._arg_shutdown_functions = kwargs.pop("shutdown_functions", [])
self.shutdown_functions = []
for shutdown_function in kwargs.pop("shutdown_functions", []):
if callable(shutdown_function):
self.shutdown_functions.append(shutdown_function)
else:
raise PluginValidationError(
f"Provided un-callable shutdown function {shutdown_function}"
)

if "shutdown_function" in kwargs:
self._arg_shutdown_functions.append(kwargs.pop("shutdown_function"))
shutdown_function = kwargs.pop("shutdown_function")
if callable(shutdown_function):
self.shutdown_functions.append(shutdown_function)
else:
raise PluginValidationError(
f"Provided un-callable shutdown function {shutdown_function}"
)

self.startup_functions = []

for startup_function in kwargs.pop("startup_functions", []):
if callable(startup_function):
self.startup_functions.append(startup_function)
else:
raise PluginValidationError(
f"Provided un-callable startup function {shutdown_function}"
)

if "startup_function" in kwargs:
startup_function = kwargs.pop("startup_function")
if callable(startup_function):
self.startup_functions.append(startup_function)
else:
raise PluginValidationError(
f"Provided un-callable startup function {shutdown_function}"
)

# Now that logging is configured we can load the real config
self._config = load_config(**kwargs)
Expand Down Expand Up @@ -275,6 +325,18 @@ def run(self):

try:
self._startup()

# Run provided startup functions
self._logger.debug("About to run annotated startup functions")
startup_functions = _parse_startup_functions(self._client)
startup_functions.extend(self.startup_functions)
startup_functions.extend(self._config.client_startup_functions)

if getattr(self._config, "client_startup_function"):
startup_functions.append(self._config.client_startup_function)

self._run_configured_functions(startup_functions)

self._logger.info("Plugin %s has started", self.unique_name)

try:
Expand Down Expand Up @@ -309,38 +371,35 @@ def client(self, new_client):

self._set_client(new_client)

def _run_shutdown_functions(
self, shutdown_functions, executed_shutdown_functions=None
):
if executed_shutdown_functions is None:
executed_shutdown_functions = []
def _run_configured_functions(self, functions):
executed_functions = []

for shutdown_function in shutdown_functions:
if callable(shutdown_function):
if shutdown_function not in executed_shutdown_functions:
shutdown_function()
executed_shutdown_functions.append(shutdown_function)
for function in functions:
if callable(function):
if function not in executed_functions:
function()
executed_functions.append(function)

elif self._client and hasattr(self._client, shutdown_function):
client_function = getattr(self._client, shutdown_function)
elif self._client and hasattr(self._client, function):
client_function = getattr(self._client, function)
if callable(client_function):
if client_function not in executed_shutdown_functions:
if client_function not in executed_functions:
client_function()
executed_shutdown_functions.append(client_function)
executed_functions.append(client_function)
else:
self._logger.error(
f"Provided non callable function for shutdown function: {shutdown_function}"
f"Provided non callable function for function: {function}"
)
elif self._client:
self._logger.error(
(
"Provided function not existing on client "
f"for shutdown function: {shutdown_function}"
f"for function: {function}"
)
)
else:
self._logger.error(
f"No client provided to check for shutdown function: {shutdown_function}"
f"No client provided to check for function: {function}"
)

def _set_client(self, new_client):
Expand Down Expand Up @@ -539,20 +598,15 @@ def _shutdown(self, status="STOPPED"):
# Run shutdown functions prior to setting shutdown event to allow for
# any functions that might generate Requests

self._logger.debug("About to run annotated shutdown functions")
executed_shutdown_functions = self._run_shutdown_functions(
_parse_shutdown_functions(self._client)
)
self._logger.debug("About to run shutdown functions")
shutdown_functions = _parse_shutdown_functions(self._client)
shutdown_functions.extend(self.shutdown_functions)
shutdown_functions.extend(self._config.client_shutdown_functions)

self._logger.debug("About to run plugin shutdown functions")
executed_shutdown_functions = self._run_shutdown_functions(
self._arg_shutdown_functions, executed_shutdown_functions
)
if getattr(self._config, "client_shutdown_function"):
shutdown_functions.append(self._config.client_shutdown_function)

self._logger.debug("About to run config shutdown functions")
executed_shutdown_functions = self._run_shutdown_functions(
self._config.shutdown_functions, executed_shutdown_functions
)
self._run_configured_functions(shutdown_functions)

self._shutdown_event.set()

Expand Down
16 changes: 14 additions & 2 deletions brewtils/specification.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,18 +188,30 @@ def _is_json_dict(s):
"description": "The dependency timeout to use",
"default": 300,
},
"shutdown_function": {
"client_shutdown_function": {
"type": "str",
"description": "The function in client to be executed at shutdown",
"required": False,
},
"shutdown_functions": {
"client_shutdown_functions": {
"type": "list",
"description": "The functions in client to be executed at shutdown",
"items": {"name": {"type": "str"}},
"required": False,
"default": [],
},
"client_startup_function": {
"type": "str",
"description": "The function in client to be executed at run",
"required": False,
},
"client_startup_functions": {
"type": "list",
"description": "The functions in client to be executed at run",
"items": {"name": {"type": "str"}},
"required": False,
"default": [],
},
}

_PLUGIN_SPEC = {
Expand Down
21 changes: 21 additions & 0 deletions test/decorators_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
plugin_param,
register,
shutdown,
startup,
system,
)
from brewtils.errors import PluginParamError
Expand Down Expand Up @@ -1459,3 +1460,23 @@ def cmd():

assert not hasattr(cmd, "_shutdown")
assert not getattr(cmd, "_shutdown", False)


class TestStartup(object):
"""Test shutdown decorator"""

def test_startup(self):
@startup
def cmd():
return True

assert hasattr(cmd, "_startup")
assert cmd._startup

def test_missing_startup(self):
@command
def cmd():
return True

assert not hasattr(cmd, "_startup")
assert not getattr(cmd, "_startup", False)
4 changes: 2 additions & 2 deletions test/plugin_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -434,7 +434,7 @@ def test_shutdown_plugin_shutdown_function_executed(self, plugin):
plugin._admin_processor = Mock()

mock_shutdown_function = Mock()
plugin._arg_shutdown_functions = [mock_shutdown_function]
plugin.shutdown_functions = [mock_shutdown_function]

plugin._shutdown()
assert mock_shutdown_function.called is True
Expand All @@ -444,7 +444,7 @@ def test_shutdown_config_shutdown_function_executed(self, plugin):
plugin._admin_processor = Mock()

mock_shutdown_function = Mock()
plugin._config.shutdown_functions = [mock_shutdown_function]
plugin._config.client_shutdown_functions = [mock_shutdown_function]

plugin._shutdown()
assert mock_shutdown_function.called is True
Expand Down

0 comments on commit 4b89b59

Please sign in to comment.