somakeit · sjefferson99 · May 29, 2024 · May 28, 2024 · May 29, 2024 · May 29, 2024
diff --git a/smib/logging.json b/smib/logging.json
@@ -51,6 +51,12 @@
     },
     "pymongo": {
       "level": "WARNING"
+    },
+    "uvicorn": {
+      "level": "WARNING"
+    },
+    "asyncio": {
+      "level": "WARNING"
     }
   }
 }
diff --git a/smib/slack/__main__.py b/smib/slack/__main__.py
@@ -38,7 +38,7 @@ def create_slack_bolt_app():
     logger.info(f"Created SlackApp: {APPLICATION_NAME}")
     app.error(handle_errors)
     app.middleware(inject_logger_to_slack_context)
-    logger.info(f"Registered SlackApp error handler: {handle_errors}")
+    logger.debug(f"Registered SlackApp error handler: {handle_errors}")
     return app
 
 

diff --git a/smib/slack/plugin/loaders/abstract_plugin_loader.py b/smib/slack/plugin/loaders/abstract_plugin_loader.py
@@ -50,10 +50,16 @@ def load_all(self) -> list[Plugin]:
 
             # If the plugin ID already exists, give it a new one
             if plugin.id in loaded_plugin_ids:
+                logger.debug(f"Plugin {plugin.id} already exists, giving it new id")
                 plugin.id = f"{plugin.id}_{id(plugin)}"
+                logger.debug(f"New plugin id: {plugin.id}")
 
             plugins.append(plugin)
 
+            logger.debug(f"Plugin {plugin.id} loaded. Enabled: {plugin.enabled}")
+
+        logger.info(f"Loaded {len(plugins)} {self.type} plugins")
+
         return plugins
 
     def plugin_path_to_id(self, plugin_path: Path) -> str:
@@ -76,6 +82,7 @@ def load_plugin(self, plugin_path: Path) -> Plugin:
         returned_plugin = self.register_plugin(plugin)
 
         if not plugin.enabled:
+            logger.debug(f"Plugin {plugin.id} is not enabled... Unloading...")
             self.unload_plugin(plugin)
 
         if plugin.error:
@@ -90,13 +97,18 @@ def unload_plugin(self, plugin: Plugin) -> None:
         self._remove_middlewares(plugin)
 
     def _remove_listeners(self, plugin: Plugin) -> None:
+        logger = inject('logger')
+        logger.info(f"Removing listeners for plugin {plugin.id}")
         listeners = self.app._listeners[::]
         for listener in listeners:
             listener_path = inspect.getfile(inspect.unwrap(listener.ack_function))
             if Path(listener_path).is_relative_to(plugin.directory):
+                logger.debug(f"Listener {listener.__name__} from {Path(listener_path).relative_to(plugin.directory).as_posix()} removed")
                 self.app._listeners.remove(listener)
 
     def _remove_scheduled_jobs(self, plugin: Plugin) -> None:
+        logger = inject('logger')
+        logger.info(f"Removing scheduled jobs for plugin {plugin.id}")
         listeners = self.app._listeners[::]
         for listener in listeners:
             raw_listener_ack = inspect.unwrap(listener.ack_function)
@@ -106,6 +118,7 @@ def _remove_scheduled_jobs(self, plugin: Plugin) -> None:
                 continue
 
             if job := self._find_job_from_plugin_function(raw_listener_ack):
+                logger.debug(f"Scheduled job {job.id} from {Path(listener_path).relative_to(plugin.directory).as_posix()} removed")
                 self.scheduler.remove_job(job.id)
 
     def _find_job_from_plugin_function(self, plugin_function: callable) -> Job:
@@ -115,6 +128,8 @@ def _find_job_from_plugin_function(self, plugin_function: callable) -> Job:
             )), None)
 
     def _remove_middlewares(self, plugin: Plugin) -> None:
+        logger = inject('logger')
+        logger.info(f"Removing middlewares for plugin {plugin.id}")
         middlewares = self.app._middleware_list[::]
         for middleware in middlewares:
             func = getattr(middleware, 'func', None)
@@ -123,11 +138,12 @@ def _remove_middlewares(self, plugin: Plugin) -> None:
 
             middleware_path = inspect.getfile(inspect.unwrap(func))
             if Path(middleware_path).is_relative_to(plugin.directory):
+                logger.debug(f"Middleware {middleware.__name__} from {Path(middleware_path).relative_to(plugin.directory).as_posix()} removed")
                 self.app._middleware_list.remove(middleware)
 
     def reload_plugin(self, plugin: Plugin) -> Plugin:
         logger: logging.Logger = inject("logger")
-        logger.debug(f"Reloading: {plugin}")
+        logger.debug(f"Reloading plugin {plugin.id}")
         self.unload_plugin(plugin)
         reloaded_plugin = self.load_plugin(plugin.directory)
         return reloaded_plugin

diff --git a/smib/slack/plugin/manager.py b/smib/slack/plugin/manager.py
@@ -2,6 +2,11 @@
 from pathlib import Path
 from typing import Iterator
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from smib.slack.plugin import PluginType
+    from smib.slack.plugin.loaders.abstract_plugin_loader import AbstractPluginLoader
 from smib.slack.plugin.plugin import Plugin
 
 from injectable import inject, inject_multiple, injectable, autowired, Autowired
@@ -29,8 +34,13 @@ def reload_all_plugins(self):
                 self.plugins.append(reloaded_plugin)
 
     def disable_plugin(self, plugin: Plugin):
-        loader = next((x for x in self.plugin_loaders if x.type == plugin.type), None)
+        logger = inject("logger")
+
+        logger.info(f"Disabling plugin {plugin.id}")
+
+        loader = self._get_plugin_loader_from_type(plugin.type)
         if not loader:
+            logger.warning(f"Unable to disable plugin {plugin.id} - no loader found")
             return
 
         open(plugin.directory / '.disable', "wb")
@@ -39,22 +49,37 @@ def disable_plugin(self, plugin: Plugin):
         reloaded_plugin = loader.reload_plugin(plugin)
         self.plugins.append(reloaded_plugin)
 
+        logger.info(f"Plugin {plugin.id} disabled")
+
     def enable_plugin(self, plugin: Plugin):
-        loader = next((x for x in self.plugin_loaders if x.type == plugin.type), None)
+        logger = inject("logger")
+
+        logger.info(f"Enabling plugin {plugin.id}")
+
+        loader = self._get_plugin_loader_from_type(plugin.type)
         if not loader:
+            logger.warning(f"Unable to enable plugin {plugin.id} - no loader found")
             return
 
         if (disable_file := plugin.directory / '.disable').exists():
             os.remove(disable_file)
+        else:
+            logger.debug(f"Plugin {plugin.id} already enabled! Reloading anyway...")
 
         self.plugins.remove(plugin)
 
         reloaded_plugin = loader.reload_plugin(plugin)
         self.plugins.append(reloaded_plugin)
 
+        logger.info(f"Plugin {plugin.id} enabled")
+
     def get_plugin_from_file(self, file: Path) -> Plugin:
         return next((plugin for plugin in self.plugins if file.is_relative_to(plugin.directory)), None)
 
     def __iter__(self) -> Iterator[Plugin]:
         yield from self.plugins
 
+    def _get_plugin_loader_from_type(self, plugin_type: "PluginType") -> "AbstractPluginLoader":
+        loader = next((x for x in self.plugin_loaders if x.type == plugin_type), None)
+        return loader
+
diff --git a/smib/slack/signal_handlers.py b/smib/slack/signal_handlers.py
@@ -6,6 +6,7 @@
 def sigterm_handler(signum, frame):
     logger: logging.Logger = inject("logger")
 
+    logger.info("Received shutdown signal")
     logger.debug(f'Signal handler called with signal, {signum}')
 
     raise SystemExit('Exiting...')
diff --git a/smib/slack/websocket/server.py b/smib/slack/websocket/server.py
@@ -29,17 +29,17 @@ def handle(self):
         slack_app: App
 
         event_type: str = bolt_request.body.get('event').get('type')
-        logger.debug(f"Received event: {event_type}")
+        logger.info(f"Received http event: {event_type}")
 
         bolt_response: BoltResponse = slack_app.dispatch(bolt_request)
         self.send_message(pickle.dumps(bolt_response))
 
         http_status: HTTPStatus = HTTPStatus(bolt_response.status)
-        logger.debug(f"Sent status: {bolt_response.status} - {http_status.name}: {http_status.description}")
+        logger.debug(f"Sent response status: {bolt_response.status} - {http_status.name}: {http_status.description}")
 
     def connected(self):
         logger: logging.Logger = inject("logger")
-        logger.info(f"{self.address} connected")
+        logger.info(f"New websocket connection from {self.address}")
 
         hostname = None
         hostname_ip = None
@@ -56,10 +56,12 @@ def connected(self):
         else:
             logger.debug(f'Address {address} resolved to hostname {hostname}')
 
+        logger.info(f"{hostname} connected")
+
         if {address, hostname, hostname_ip}.intersection(set(WEBSOCKET_ALLOWED_HOSTS)):
             return
 
-        logger.warning(f"Connection from {self.address} is {NOT_AUTHORIZED}")
+        logger.warning(f"Connection from {self.address} is {NOT_AUTHORIZED}. Closing connection.")
         self.close(reason=NOT_AUTHORIZED)
 
     def handle_close(self):

diff --git a/smib/webserver/__main__.py b/smib/webserver/__main__.py
@@ -1,5 +1,6 @@
 import logging
 import re
+from contextlib import asynccontextmanager
 from pathlib import Path
 
 from fastapi import FastAPI, Request, APIRouter
@@ -57,13 +58,30 @@ async def generate_bolt_request(fastapi_request: Request):
 
 
 def create_directories():
+    logger = inject("logger")
+    logger.debug(f"Resolved Webserver Template Directory to: {WEBSERVER_TEMPLATES_DIRECTORY}")
     if not WEBSERVER_TEMPLATES_DIRECTORY.exists():
+        logger.info(f"Creating webserver templates directory: {WEBSERVER_TEMPLATES_DIRECTORY}")
         WEBSERVER_TEMPLATES_DIRECTORY.mkdir()
 
+    logger.debug(f"Resolved Webserver Static Directory to: {WEBSERVER_STATIC_DIRECTORY}")
     if not WEBSERVER_STATIC_DIRECTORY.exists():
+        logger.info(f"Creating webserver static directory: {WEBSERVER_STATIC_DIRECTORY}")
         WEBSERVER_STATIC_DIRECTORY.mkdir()
 
 
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Load the ML model
+    logger = inject("logger")
+    logger.info(f"Webserver started")
+
+    yield
+
+    # Clean up the ML models and release the resources
+    logger.info(f"Webserver Stopping")
+
+
 def get_readme() -> str:
     description = None
     description_path = Path(__file__).parent / "README.md"
@@ -94,13 +112,13 @@ def get_title() -> str:
 
 ws_handler = WebSocketHandler()
 
-app = FastAPI(title=get_title(), version=get_version(), description=get_readme_without_title(), redoc_url=None)
+app = FastAPI(lifespan=lifespan, title=get_title(), version=get_version(), description=get_readme_without_title(), redoc_url=None)
 smib_router = APIRouter(prefix=WEBSERVER_PATH_PREFIX)
 event_router = APIRouter(prefix='/event', tags=['S.M.I.B. Events'], responses=event_responses)
 
 create_directories()
 
-smib_router.mount("/static", StaticFiles(directory=WEBSERVER_STATIC_DIRECTORY), name="static")
+app.mount("/static", StaticFiles(directory=WEBSERVER_STATIC_DIRECTORY), name="static")
 templates = Jinja2Templates(directory=str(WEBSERVER_TEMPLATES_DIRECTORY))
 
 
@@ -109,7 +127,7 @@ def get_title() -> str:
 @event_router.put('/{event}', name="S.M.I.B. PUT Event")
 async def smib_event_handler(request: Request, event: str):
     logger = inject("logger")
-    logger.debug(f"Received event {event}")
+    logger.info(f"Received event {event}")
     ws_handler.check_and_reconnect_websocket_conn()
     bolt_request: BoltRequest = await generate_bolt_request(request)
     logger.debug(f"Request: {request} -> Bolt Request: {bolt_request}")
@@ -133,9 +151,14 @@ def main(app: FastAPI, ws_handler: WebSocketHandler):
         import uvicorn
         logger.info(f"Starting WebServer v{get_version()}")
         uvicorn.run(app, host=WEBSERVER_HOST, port=WEBSERVER_PORT, log_config=read_logging_json(), headers=[("server", APPLICATION_NAME)])
+    except KeyboardInterrupt:
+        ...
     finally:
+        logger.info(f"Stopping WebsocketHandler")
         ws_handler.close_conn()
 
+        logger.info(f"Webserver Stopped")
+
 
 if __name__ == '__main__':
     main(app, ws_handler)