diff options
| -rw-r--r-- | brevisurl/backends/local.py | 6 |
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 |