Add check method to all storage classes

limits/storage.py

@@ -81,7 +81,12 @@ def get_expiry(self, key):
         """
         raise NotImplementedError
 
-
+    @abstractmethod
+    def check(self):
+        """
+        check if storage is healthy
+        """
+        raise NotImplementedError
 
 
 class LockableEntry(threading._RLock):
@@ -204,6 +209,13 @@ def get_moving_window(self, key, limit, expiry):
                 return int(item.atime), acquired
         return int(timestamp), acquired
 
+    def check(self):
+        """
+        check if storage is healthy
+        """
+        return True
+
+
 class RedisStorage(Storage):
     """
     rate limit storage with redis as backend
@@ -308,6 +320,15 @@ def get_expiry(self, key):
         """
         return int((self.storage.ttl(key) or 0) + time.time())
 
+    def check(self):
+        """
+        check if storage is healthy
+        """
+        try:
+            return self.storage.ping()
+        except: # noqa
+            return False
+
 class MemcachedStorage(Storage):
     """
     rate limit storage with memcached as backend
@@ -398,3 +419,12 @@ def get_expiry(self, key):
         """
         return int(float(self.storage.get(key + "/expires") or time.time()))
 
+    def check(self):
+        """
+        check if storage is healthy
+        """
+        try:
+            self.call_memcached_func(self.storage.stats)
+            return True
+        except: # noqa
+            return False

tests/test_storage.py

@@ -26,6 +26,11 @@ def test_storage_string(self):
         self.assertTrue(isinstance(storage_from_string("memcached://localhost:11211"), MemcachedStorage))
         self.assertRaises(ConfigurationError, storage_from_string, "blah://")
 
+    def test_storage_check(self):
+        self.assertTrue(storage_from_string("memory://").check())
+        self.assertTrue(storage_from_string("redis://localhost:6379").check())
+        self.assertTrue(storage_from_string("memcached://localhost:11211").check())
+
     def test_in_memory(self):
         with hiro.Timeline().freeze() as timeline:
             storage = MemoryStorage()