Expand exceptions

src/epplibwrapper/__init__.py

@@ -45,9 +45,6 @@ except NameError:
 # Attn: these imports should NOT be at the top of the file
 try:
     from .client import CLIENT, commands
-except ImportError:
-    pass
-try:
     from .errors import RegistryError, ErrorCode
     from epplib.models import common, info
     from epplib.responses import extensions

src/epplibwrapper/client.py

@@ -3,6 +3,8 @@
 import logging
 from time import sleep
 
+from epplibwrapper.utility.pool_status import PoolStatus
+
 try:
     from epplib.client import Client
     from epplib import commands
@@ -16,7 +18,7 @@ from django.conf import settings
 from .cert import Cert, Key
 from .errors import LoginError, RegistryError
 from .socket import Socket
-from .utility.pool import EppConnectionPool
+from .utility.pool import EPPConnectionPool
 
 logger = logging.getLogger(__name__)
 
@@ -32,7 +34,6 @@ except Exception:
         exc_info=True,
     )
 
-
 class EPPLibWrapper:
     """
     A wrapper over epplib's client.
@@ -61,35 +62,33 @@ class EPPLibWrapper:
                 password=settings.SECRET_REGISTRY_KEY_PASSPHRASE,
             )
         )
-        options = {
+
+        self.pool_options = {
             # Pool size
             "size": settings.EPP_CONNECTION_POOL_SIZE,
             # Which errors the pool should look out for
             "exc_classes": (
-                LoginError,
-                RegistryError,
+                TransportError,
             ),
-            # Occasionally pings the registry to keep the connection alive
+            # Occasionally pings the registry to keep the connection alive.
+            # Value in seconds => (keepalive / size)
             "keepalive": settings.POOL_KEEP_ALIVE,
         }
 
         self._pool = None
-        # Since we reuse the same creds for each pool, we can test on
-        # one socket, and if successful, then we know we can connect.
-        if not settings.DEBUG or self._test_registry_connection_success():
-            self._pool = EppConnectionPool(
-                client=self._client, login=self._login, options=options
-            )
-        else:
-            logger.warning("Cannot contact the Registry")
-            # TODO - signal that the app may need to restart?
+        # Tracks the status of the pool
+        self.pool_status = PoolStatus()
+
+        self.start_connection_pool()
 
     def _send(self, command):
         """Helper function used by `send`."""
         cmd_type = command.__class__.__name__
         try:
-            if self._pool is None:
-                raise LoginError
+            if not self.pool_status.connection_success:
+                raise LoginError(
+                    "Couldn't connect to the registry after three attempts"
+                )
             # TODO - add a timeout
             with self._pool.get() as connection:
                 response = connection.send(command)
@@ -123,6 +122,21 @@ class EPPLibWrapper:
         if not cleaned:
             raise ValueError("Please sanitize user input before sending it.")
         
+        # Reopen the pool if its closed
+        if not self.pool_status.pool_running:
+            # We want to reopen the connection pool,
+            # but we don't want the end user to wait while it opens.
+            # Raise syntax doesn't allow this, so we use a try/catch
+            # block.
+            try:
+                raise RegistryError(
+                    "Can't contact the Registry. Please try again later"
+                )
+            except RegistryError as err:
+                raise err
+            finally:
+                self.start_connection_pool()
+
         counter = 0  # we'll try 3 times
         while True:
             try:
@@ -134,6 +148,53 @@ class EPPLibWrapper:
                 else:  # don't try again
                     raise err
 
+    def start_connection_pool(self, restart_pool_if_exists = True):
+        """Starts a connection pool for the registry. 
+
+        restart_pool_if_exists -> bool: 
+        If an instance of the pool already exists,
+        then then that instance will be killed first. 
+        It is generally recommended to keep this enabled."""
+        
+        # Since we reuse the same creds for each pool, we can test on
+        # one socket, and if successful, then we know we can connect.
+        if settings.DEBUG or self._test_registry_connection_success():
+            logger.warning("Cannot contact the Registry")
+            self.pool_status.connection_success = False
+            # Q: Should err be raised instead?
+            # Q2: Since we try to connect 3 times,
+            # this indicates that the Registry isn't responsive.
+            # What should we do in this case?
+            return
+        else:
+            self.pool_status.connection_success = True
+        
+        # If this function is reinvoked, then ensure
+        # that we don't have duplicate data sitting around.
+        if self._pool is not None and restart_pool_if_exists:
+            logger.info("Connection pool restarting...")
+            self.kill_pool()
+        
+        self._pool = EPPConnectionPool(
+            client=self._client, login=self._login, options=self.pool_options
+        )
+        self.pool_status.pool_running = True
+
+        logger.info("Connection pool started")
+    
+    def kill_pool(self):
+        """Kills the existing pool. Use this instead
+        of self._pool = None, as that doesn't clear
+        gevent instances."""
+        if self._pool is not None:
+            self._pool.kill_all_connections()
+            self._pool = None
+            self.pool_status.pool_running = False
+            return
+        logger.info(
+            "kill_pool() was invoked but there was no pool to delete"
+        )
+
     def _test_registry_connection_success(self):
         """Check that determines if our login
         credentials are valid, and/or if the Registrar

src/epplibwrapper/errors.py

@@ -80,6 +80,9 @@ class RegistryError(Exception):
     def is_client_error(self):
         return self.code is not None and (self.code >= 2000 and self.code <= 2308)
     
+    def is_not_retryable(self):
+        pass
+
 
 class LoginError(RegistryError):
     pass

src/epplibwrapper/utility/pool.py

@@ -1,4 +1,6 @@
+from collections import deque
 import logging
+import gevent
 from geventconnpool import ConnectionPool
 from epplibwrapper.errors import RegistryError, LoginError
 from epplibwrapper.socket import Socket
@@ -11,15 +13,23 @@ except ImportError:
 logger = logging.getLogger(__name__)
 
 
-class EppConnectionPool(ConnectionPool):
-    def __init__(self, client, login, options):
+class EPPConnectionPool(ConnectionPool):
+    """A connection pool for EPPLib.
+
+    Args:
+        client (Client): The client
+        login (commands.Login): Login creds
+        options (dict): Options for the ConnectionPool
+        base class
+    """
+    def __init__(self, client, login, options: dict):
         # For storing shared credentials
         self._client = client
         self._login = login
         super().__init__(**options)
 
     def _new_connection(self):
-        socket = self.create_socket(self._client, self._login)
+        socket = self._create_socket(self._client, self._login)
         try:
             connection = socket.connect()
             return connection
@@ -37,7 +47,39 @@ class EppConnectionPool(ConnectionPool):
             logger.error("Failed to keep the connection alive.", exc_info=True)
             raise RegistryError("Failed to keep the connection alive.") from err
 
-    def create_socket(self, client, login) -> Socket:
+    def _create_socket(self, client, login) -> Socket:
         """Creates and returns a socket instance"""
         socket = Socket(client, login)
         return socket
+    
+    def get_connections(self):
+        """Returns the connection queue"""
+        return self.conn
+    
+    def kill_all_connections(self):
+        """Kills all active connections in the pool."""
+        try:
+            gevent.killall(self.conn)
+            self.conn.clear()
+            # Clear the semaphore
+            for i in range(self.lock.counter):
+                self.lock.release()
+        # TODO - connection pool err
+        except Exception as err:
+            logger.error(
+                "Could not kill all connections."
+            )
+            raise err
+    
+    def repopulate_all_connections(self):
+        """Regenerates the connection pool.
+        If any connections exist, kill them first.
+        """
+        if len(self.conn) > 0:
+            self.kill_all_connections()
+        for i in range(self.size):
+            self.lock.acquire()
+        for i in range(self.size):
+            gevent.spawn_later(self.SPAWN_FREQUENCY*i, self._addOne)
+
+

src/epplibwrapper/utility/pool_status.py

@@ -0,0 +1,6 @@
+class PoolStatus:
+    """A list of Booleans to keep track of Pool Status"""
+    def __init__(self):
+        self.pool_running = False
+        self.connection_success = False
+        self.pool_hanging = False