raise TokensExhausted when maximum ShortUrl objects are reached

1 files changed, 5 insertions, 1 deletions

diff --git a/brevisurl/backends/local.py b/brevisurl/backends/local.py
index 1b72e64..8622f5a 100644
--- a/brevisurl/backends/local.py
+++ b/brevisurl/backends/local.py
@@ -12,6 +12,8 @@ from brevisurl.models import ShortUrl
 
 log = logging.getLogger(__name__)
 
+class TokensExhausted(Exception):
+    pass
 
 class BrevisUrlBackend(BaseBrevisUrlBackend):
 
@@ -53,9 +55,11 @@ class BrevisUrlBackend(BaseBrevisUrlBackend):
 
     def __generate_token(self, size=5):
         chars = list(string.ascii_letters + string.digits)
+        if ShortUrl.objects.count >= chars ** size:
+            raise TokensExhausted('Consider incrementing the token length or chars list')
         random.shuffle(chars)
         while True:
             token = ''.join([random.choice(chars) for x in range(size)])
             if not ShortUrl.objects.filter(backend=self.class_path, shortened_url__endswith=token).count():
                 break
-        return token
\ No newline at end of file
+        return token