[backport] Improving rebuild timeline command

taiga/timeline/management/commands/rebuild_timeline.py

@@ -14,7 +14,13 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+# Examples:
+# python manage.py rebuild_timeline --settings=settings.local_timeline --initial_date 2014-10-02 --final_date 2014-10-03
+# python manage.py rebuild_timeline --settings=settings.local_timeline --purge
+# python manage.py rebuild_timeline --settings=settings.local_timeline --initial_date 2014-10-02
+
 from django.conf import settings
+from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ObjectDoesNotExist
 from django.core.management.base import BaseCommand
 from django.db.models import Model
@@ -31,7 +37,7 @@ from taiga.timeline.signals import on_new_history_entry, _push_to_timelines
 from taiga.users.models import User
 
 from unittest.mock import patch
-from django.contrib.contenttypes.models import ContentType
+from optparse import make_option
 
 import gc
 
@@ -51,26 +57,6 @@ class BulkCreator(object):
 bulk_creator = BulkCreator()
 
 
-def queryset_iterator(queryset, chunksize=1000):
-    '''''
-    Iterate over a Django Queryset ordered by the primary key
-
-    This method loads a maximum of chunksize (default: 1000) rows in it's
-    memory at the same time while django normally would load all rows in it's
-    memory. Using the iterator() method only causes it to not preload all the
-    classes.
-
-    Note that the implementation of the iterator does not support ordered query sets.
-    '''
-    queryset = queryset.order_by('pk')
-    pk = queryset[0].pk
-    last_pk = queryset.order_by('-pk')[0].pk
-    while pk < last_pk:
-        for row in queryset.filter(pk__gt=pk)[:chunksize]:
-            pk = row.pk
-            yield row
-        gc.collect()
-
 def custom_add_to_object_timeline(obj:object, instance:object, event_type:str, namespace:str="default", extra_data:dict={}):
     assert isinstance(obj, Model), "obj must be a instance of Model"
     assert isinstance(instance, Model), "instance must be a instance of Model"
@@ -88,10 +74,30 @@ def custom_add_to_object_timeline(obj:object, instance:object, event_type:str, n
     ))
 
 
-def generate_timeline():
+def generate_timeline(initial_date, final_date):
+    if initial_date or final_date:
+        timelines = Timeline.objects.all()
+        if initial_date:
+            timelines = timelines.filter(created__gte=initial_date)
+        if final_date:
+            timelines = timelines.filter(created__lt=final_date)
+
+        timelines.delete()
+
     with patch('taiga.timeline.service._add_to_object_timeline', new=custom_add_to_object_timeline):
         # Projects api wasn't a HistoryResourceMixin so we can't interate on the HistoryEntries in this case
-        for project in queryset_iterator(Project.objects.order_by("created_date")):
+        projects = Project.objects.order_by("created_date")
+        history_entries = HistoryEntry.objects.order_by("created_at")
+
+        if initial_date:
+            projects = projects.filter(created_date__gte=initial_date)
+            history_entries = history_entries.filter(created_at__gte=initial_date)
+
+        if final_date:
+            projects = projects.filter(created_date__lt=final_date)
+            history_entries = history_entries.filter(created_at__lt=final_date)
+
+        for project in projects.iterator():
             bulk_creator.created = project.created_date
             print("Project:", bulk_creator.created)
             extra_data = {
@@ -101,7 +107,7 @@ def generate_timeline():
             _push_to_timelines(project, project.owner, project, "create", extra_data=extra_data)
             del extra_data
 
-        for historyEntry in queryset_iterator(HistoryEntry.objects.order_by("created_at")):
+        for historyEntry in history_entries.iterator():
             print("History entry:", historyEntry.created_at)
             try:
                 bulk_creator.created = historyEntry.created_at
@@ -111,11 +117,35 @@ def generate_timeline():
 
 
 class Command(BaseCommand):
+    help = 'Regenerate project timeline'
+    option_list = BaseCommand.option_list + (
+        make_option('--purge',
+                    action='store_true',
+                    dest='purge',
+                    default=False,
+                    help='Purge existing timelines'),
+        ) + (
+        make_option('--initial_date',
+                    action='store',
+                    dest='initial_date',
+                    default=None,
+                    help='Initial date for timeline generation'),
+        ) + (
+        make_option('--final_date',
+                    action='store',
+                    dest='final_date',
+                    default=None,
+                    help='Final date for timeline generation'),
+        )
+
+
     def handle(self, *args, **options):
         debug_enabled = settings.DEBUG
         if debug_enabled:
             print("Please, execute this script only with DEBUG mode disabled (DEBUG=False)")
             return
 
-        Timeline.objects.all().delete()
-        generate_timeline()
+        if options["purge"] == True:
+            Timeline.objects.all().delete()
+
+        generate_timeline(options["initial_date"], options["final_date"])