Merge pull request #20 from Vladg90/master

URL shortening speed improvement

4 files changed, 108 insertions, 33 deletions

diff --git a/brevisurl/backends/local.py b/brevisurl/backends/local.py
index 77e0a7e..a00b0fe 100644
--- a/brevisurl/backends/local.py
+++ b/brevisurl/backends/local.py
@@ -3,7 +3,9 @@ import random
 import logging
 
 from django.contrib.sites.models import Site
+from django.core.exceptions import ValidationError
 from django.core.urlresolvers import reverse
+from django.db import IntegrityError, transaction
 
 import brevisurl.settings
 from brevisurl import Error
@@ -26,32 +28,48 @@ class BrevisUrlBackend(BaseBrevisUrlBackend):
         :raises: ImproperlyConfigured, django.core.exceptions.ValidationError
         :raises: brevisurl.backends.local.TokensExhaustedError
         """
-        try:
-            short_url = ShortUrl.objects.get(backend=self.class_path, original_url=original_url)
-            log.info('Url "%s" already shortened to "%s"', original_url, short_url.shortened_url)
-            return short_url
-        except ShortUrl.DoesNotExist:
-            pass
+
+        if self.kwargs.get('domain') is not None:
+            # Domain is present in keyword arguments supplied by constructor.
+            domain = self.kwargs.get('domain')
+        elif brevisurl.settings.LOCAL_BACKEND_DOMAIN is not None:
+            # Domain is defined in settings.
+            domain = brevisurl.settings.LOCAL_BACKEND_DOMAIN
+        else:
+            # Domain is taken from django site framework.
+            domain = Site.objects.get_current().domain
 
         try:
-            short_url = ShortUrl()
-            if self.kwargs.get('domain') is not None:
-                # Domain is present in keyword arguments supplied by constructor.
-                domain = self.kwargs.get('domain')
-            elif brevisurl.settings.LOCAL_BACKEND_DOMAIN is not None:
-                # Domain is defined in settings.
-                domain = brevisurl.settings.LOCAL_BACKEND_DOMAIN
-            else:
-                # Domain is taken from django site framework.
-                domain = Site.objects.get_current().domain
-            # Saving newly generated shortened url.
-            short_url.original_url = original_url
-            short_url.shortened_url = absurl(domain=domain, path=reverse('brevisurl_redirect',
-                                                                         kwargs={'token': self.__generate_token()}))
-            short_url.backend = self.class_path
-            short_url.save()
-            log.info('Url "%s" shortened to "%s"', original_url, short_url.shortened_url)
-            return short_url
+            shortened_url = self.__generate_shortened_url(domain)
+            try:
+                short_url, created = ShortUrl.objects.get_or_create(backend=self.class_path,
+                                                                    original_url=original_url,
+                                                                    defaults={'shortened_url': shortened_url})
+                if created:
+                     log.info('Url "%s" shortened to "%s"', original_url, shortened_url)
+                else:
+                     log.info('Url "%s" already shortened to "%s"', original_url, short_url.shortened_url)
+                return short_url
+            except (IntegrityError, ValidationError) as e:
+                # Check if the error is an URL validation error.
+                if e.message_dict.has_key('original_url'):
+                    raise
+
+                # Generate another token.
+                self.__check_tokens_exhausted()
+                while True:
+                    shortened_url = self.__generate_shortened_url(domain)
+                    sid = transaction.savepoint()
+                    try:
+                        short_url = ShortUrl.objects.create(backend=self.class_path,
+                                                            original_url=original_url,
+                                                            shortened_url=shortened_url)
+                        log.info('Url "%s" shortened to "%s"', original_url, shortened_url)
+                        transaction.savepoint_commit(sid)
+                        return short_url
+                    except (IntegrityError, ValidationError) as e:
+                        transaction.savepoint_rollback(sid)
+                        self.__check_tokens_exhausted()
         except Exception:
             if self.fail_silently:
                 return None
@@ -59,14 +77,18 @@ class BrevisUrlBackend(BaseBrevisUrlBackend):
                 log.exception('Unknown exception raised while shortening url "%s"', original_url)
                 raise
 
-    def __generate_token(self):
+    def __check_tokens_exhausted(self):
         chars = brevisurl.settings.LOCAL_BACKEND_TOKEN_CHARS
         size = brevisurl.settings.LOCAL_BACKEND_TOKEN_LENGTH
         if ShortUrl.objects.count() >= math.pow(len(chars), size):
             raise TokensExhaustedError('Consider incrementing the token length or change the char list')
+
+    def __generate_shortened_url(self, domain):
+        chars = brevisurl.settings.LOCAL_BACKEND_TOKEN_CHARS
+        size = brevisurl.settings.LOCAL_BACKEND_TOKEN_LENGTH
         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
+        token = ''.join([random.choice(chars) for x in range(size)])
+        shortened_url = absurl(domain=domain,
+                               path=reverse('brevisurl_redirect',
+                                            kwargs={'token': token}))
+        return shortened_url
diff --git a/brevisurl/migrations/0002_auto__add_index_shorturl_original_url.py b/brevisurl/migrations/0002_auto__add_index_shorturl_original_url.py
new file mode 100644
index 0000000..5e5811d
--- /dev/null
+++ b/brevisurl/migrations/0002_auto__add_index_shorturl_original_url.py
@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+from south.utils import datetime_utils as datetime
+from south.db import db
+from south.v2 import SchemaMigration
+from django.db import models
+
+
+class Migration(SchemaMigration):
+
+    def forwards(self, orm):
+        # Adding index on 'ShortUrl', fields ['original_url']
+        db.create_index('brevisurl_shorturl', ['original_url'])
+
+
+    def backwards(self, orm):
+        # Removing index on 'ShortUrl', fields ['original_url']
+        db.delete_index('brevisurl_shorturl', ['original_url'])
+
+
+    models = {
+        'brevisurl.shorturl': {
+            'Meta': {'ordering': "['-created']", 'unique_together': "(('original_url_hash', 'backend'),)", 'object_name': 'ShortUrl'},
+            'backend': ('django.db.models.fields.CharField', [], {'max_length': '200'}),
+            'created': ('django.db.models.fields.DateTimeField', [], {'auto_now_add': 'True', 'db_index': 'True', 'blank': 'True'}),
+            'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
+            'original_url': ('django.db.models.fields.URLField', [], {'max_length': '200', 'db_index': 'True'}),
+            'original_url_hash': ('django.db.models.fields.CharField', [], {'max_length': '64'}),
+            'shortened_url': ('django.db.models.fields.URLField', [], {'unique': 'True', 'max_length': '200'})
+        }
+    }
+
+    complete_apps = ['brevisurl']
\ No newline at end of file
diff --git a/brevisurl/models.py b/brevisurl/models.py
index cb5575b..4b46bdd 100644
--- a/brevisurl/models.py
+++ b/brevisurl/models.py
@@ -15,7 +15,7 @@ log = logging.getLogger(__name__)
 class ShortUrl(models.Model):
     """Model that represents shortened url."""
     original_url = models.URLField(max_length=brevisurl.settings.LOCAL_BACKEND_ORIGINAL_URL_MAX_LENGTH,
-                                   null=False, blank=False)
+                                   null=False, blank=False, db_index=True)
     original_url_hash = models.CharField(max_length=64, null=False, blank=False)
     shortened_url = models.URLField(max_length=200, null=False, blank=False, unique=True)
     backend = models.CharField(max_length=200, null=False, blank=False)
diff --git a/brevisurl/tests/backends/test_local.py b/brevisurl/tests/backends/test_local.py
index 4577376..8ddc7ed 100644
--- a/brevisurl/tests/backends/test_local.py
+++ b/brevisurl/tests/backends/test_local.py
@@ -1,5 +1,6 @@
+from django.db import transaction
 from django.core.exceptions import ValidationError
-from django.test import TestCase
+from django.test import TestCase, TransactionTestCase
 from django.core.validators import URLValidator
 
 import brevisurl.settings
@@ -7,6 +8,9 @@ from brevisurl import get_connection
 from brevisurl.models import ShortUrl
 from brevisurl.backends.local import TokensExhaustedError
 
+import random
+
+
 
 class TestLocalBrevisUrlBackend(TestCase):
 
@@ -128,4 +132,21 @@ class TestLocalBrevisUrlBackend(TestCase):
         original_url = 'http://www.codescale.net/'
         connection = get_connection('brevisurl.backends.local.BrevisUrlBackend', domain='http://test.com/d/')
         short_url = connection.shorten_url(original_url)
-        self.assertRegexpMatches(short_url.shortened_url, r'^http://test\.com/d/[^/]{5}$')
\ No newline at end of file
+        self.assertRegexpMatches(short_url.shortened_url, r'^http://test\.com/d/[^/]{5}$')
+
+class TestDuration(TransactionTestCase):
+    
+    def setUp(self):
+        self.connection = get_connection('brevisurl.backends.local.BrevisUrlBackend')
+
+    def test_shorten_urls_duration(self):
+        for i in range(0, 10000):
+            url = 'http://www.codescale.net/%s' % random.getrandbits(30)
+            self.connection.shorten_url(url)
+
+    @transaction.commit_on_success
+    def test_shorten_urls_duration_commit_on_success(self):
+        for i in range(0, 10000):
+            url = 'http://www.codescale.net/%s' % random.getrandbits(30)
+            self.connection.shorten_url(url)  
+