Initial connection hook and worker integration

Author: mktrevor

.gitignore

@@ -159,4 +159,4 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+.idea/

hyrex/__init__.py

@@ -6,6 +6,7 @@
 from .hyrex_queue import HyrexQueue
 from .hyrex_registry import HyrexRegistry
 from .worker.worker import HyrexWorker
+from .connection_hook import ConnectionHook
 
 # Set up null handler at library root level
 logging.getLogger("hyrex").addHandler(logging.NullHandler())

hyrex/connection_hook.py

@@ -0,0 +1,40 @@
+from typing import Dict, Any
+from abc import ABC, abstractmethod
+
+
+class ConnectionHook(ABC):
+    """Base class for connection hooks"""
+
+    def __init__(self):
+        self._connection = None
+
+    @property
+    def name(self) -> str:
+        """Unique identifier for this connection type"""
+        return self.__class__.__name__
+
+    @abstractmethod
+    def create_connection(self) -> Any:
+        """Create a new connection"""
+        pass
+
+    @abstractmethod
+    def close_connection(self, connection: Any) -> None:
+        """Close the given connection"""
+        pass
+
+    def get_connection(self) -> Any:
+        """Get the cached connection, creating it if necessary"""
+        if self._connection is None:
+            self._connection = self.create_connection()
+        return self._connection
+
+    def cleanup(self) -> None:
+        """Internal method to clean up the current connection"""
+        if self._connection is not None:
+            self.close_connection(self._connection)
+            self._connection = None
+
+    def health_check(self) -> bool:
+        """Verify the connection is still viable"""
+        return True

hyrex/worker/worker.py

@@ -1,6 +1,7 @@
-from typing import Callable
+from typing import Any, Callable
 
 from hyrex import constants
+from hyrex.connection_hook import ConnectionHook
 from hyrex.hyrex_registry import HyrexRegistry
 
 
@@ -14,6 +15,26 @@ def __init__(
     ):
         self.queue = queue
         self.task_registry: HyrexRegistry = HyrexRegistry()
+        self._connection_hooks: dict[str, ConnectionHook] = {}
 
     def add_registry(self, registry: HyrexRegistry):
         self.task_registry.add_registry(registry)
+
+    def register_connection(self, hook: ConnectionHook) -> None:
+        """Register a connection hook with this worker"""
+        if hook.name in self._connection_hooks:
+            raise ValueError(f"A connection hook is already registered with name: {hook.name}")
+        self._connection_hooks[hook.name] = hook
+        
+    def get_connection(self, name: str) -> Any:
+        """Get a connection by name"""
+        hook = self._connection_hooks.get(name)
+        if not hook:
+            raise KeyError(f"No connection hook registered for: {name}")
+        return hook.get_connection()
+    
+    def cleanup_connections(self) -> None:
+        """Cleanup all connections"""
+        for hook in self._connection_hooks.values():
+            hook.cleanup()
+        self._connection_hooks.clear()