optimize export to csv

FuzzyGrim · FuzzyGrim · commit 02b9fb310a5f · 2025-02-13T16:51:49.000+01:00
diff --git a/src/integrations/exports.py b/src/integrations/exports.py
@@ -10,26 +10,56 @@
 logger = logging.getLogger(__name__)
 
 
-def db_to_csv(response, user):
-    """Export a CSV file of the user's media."""
+class Echo:
+    """An object that implements just the write method of the file-like interface."""
+
+    def write(self, value):
+        """Write the value by returning it, instead of storing in a buffer."""
+        return value
+
+
+def generate_rows(user):
+    """Generate CSV rows."""
+    pseudo_buffer = Echo()
+    writer = csv.writer(pseudo_buffer, quoting=csv.QUOTE_ALL)
+
+    # Get fields
     fields = {
         "item": get_model_fields(Item),
         "track": get_track_fields(),
     }
 
-    writer = csv.writer(response, quoting=csv.QUOTE_ALL)
-    writer.writerow(fields["item"] + fields["track"])
+    # Yield header row
+    yield writer.writerow(fields["item"] + fields["track"])
 
+    # Yield data rows
     media_types = Item.MediaTypes.values
     for media_type in media_types:
         model = apps.get_model("app", media_type)
         if media_type == "episode":
             queryset = model.objects.filter(related_season__user=user)
         else:
             queryset = model.objects.filter(user=user)
-        write_model_to_csv(writer, fields, queryset, media_type)
 
-    return response
+        logger.debug("Streaming %ss to CSV", media_type)
+
+        # Stream each row
+        for media in queryset.iterator():  # Using iterator() for memory efficiency
+            # Build row data
+            row = [getattr(media.item, field, "") for field in fields["item"]] + [
+                getattr(media, field, "") for field in fields["track"]
+            ]
+
+            if media_type == "game":
+                # calculate index of progress field
+                progress_index = fields["track"].index("progress")
+                row[progress_index + len(fields["item"])] = helpers.minutes_to_hhmm(
+                    media.progress,
+                )
+
+            yield writer.writerow(row)
+
+        logger.debug("Finished streaming %ss to CSV", media_type)
 
 
 def get_model_fields(model):
@@ -64,23 +94,3 @@ def get_track_fields():
         all_fields.insert(end_idx, "start_date")
 
     return list(all_fields)
-
-
-def write_model_to_csv(writer, fields, queryset, media_type):
-    """Export entries from a model to a CSV file."""
-    logger.info("Adding %ss to CSV", media_type)
-
-    for media in queryset:
-        # row with item and track fields
-        row = [getattr(media.item, field, "") for field in fields["item"]] + [
-            getattr(media, field, "") for field in fields["track"]
-        ]
-
-        if media_type == "game":
-            # calculate index of progress field
-            progress_index = fields["track"].index("progress")
-            row[progress_index + len(fields["item"])] = helpers.minutes_to_hhmm(
-                media.progress,
-            )
-
-        writer.writerow(row)
diff --git a/src/integrations/tests/test_exports.py b/src/integrations/tests/test_exports.py
@@ -3,6 +3,7 @@
 from io import StringIO
 
 from django.contrib.auth import get_user_model
+from django.db.models import Q
 from django.test import TestCase
 from django.urls import reverse
 
@@ -156,52 +157,26 @@ def test_export_csv(self):
         # Assert that the response content type is text/csv
         self.assertEqual(response["Content-Type"], "text/csv")
 
-        # Read the CSV content from the response
-        csv_content = response.content.decode("utf-8")
+        # Read the streaming content and decode it
+        content = b"".join(response.streaming_content).decode("utf-8")
 
         # Create a CSV reader from the CSV content
-        reader = csv.DictReader(StringIO(csv_content))
+        reader = csv.DictReader(StringIO(content))
 
-        # Get all media IDs from the database
         db_media_ids = set(
-            TV.objects.values_list("item__media_id", flat=True).filter(user=self.user),
-        )
-        db_media_ids.update(
-            Movie.objects.values_list("item__media_id", flat=True).filter(
-                user=self.user,
-            ),
-        )
-        db_media_ids.update(
-            Season.objects.values_list("item__media_id", flat=True).filter(
-                user=self.user,
-            ),
-        )
-        db_media_ids.update(
-            Episode.objects.values_list("item__media_id", flat=True).filter(
-                related_season__user=self.user,
-            ),
-        )
-        db_media_ids.update(
-            Anime.objects.values_list("item__media_id", flat=True).filter(
-                user=self.user,
-            ),
-        )
-        db_media_ids.update(
-            Manga.objects.values_list("item__media_id", flat=True).filter(
-                user=self.user,
-            ),
-        )
-        db_media_ids.update(
-            Game.objects.values_list("item__media_id", flat=True).filter(
-                user=self.user,
-            ),
-        )
-        db_media_ids.update(
-            Book.objects.values_list("item__media_id", flat=True).filter(
-                user=self.user,
-            ),
-        )
-
+            Item.objects.filter(
+                Q(tv__user=self.user)
+                | Q(movie__user=self.user)
+                | Q(season__user=self.user)
+                | Q(episode__related_season__user=self.user)
+                | Q(anime__user=self.user)
+                | Q(manga__user=self.user)
+                | Q(game__user=self.user)
+                | Q(book__user=self.user),
+            ).values_list("media_id", flat=True),
+        )
+
+        # Verify each row in the CSV exists in the database
         for row in reader:
             media_id = row["media_id"]
             self.assertIn(media_id, db_media_ids)
diff --git a/src/integrations/views.py b/src/integrations/views.py
@@ -7,7 +7,7 @@
 from django.contrib import messages
 from django.contrib.auth.decorators import login_not_required
 from django.core.exceptions import ObjectDoesNotExist
-from django.http import HttpResponse
+from django.http import HttpResponse, StreamingHttpResponse
 from django.shortcuts import redirect
 from django.urls import reverse
 from django.utils import timezone
@@ -184,16 +184,12 @@ def import_yamtrack(request):
 def export_csv(request):
     """View for exporting all media data to a CSV file."""
     today = timezone.now().strftime("%Y-%m-%d")
-    # Create the HttpResponse object with the appropriate CSV header.
-    response = HttpResponse(
+    response = StreamingHttpResponse(
+        streaming_content=exports.generate_rows(request.user),
         content_type="text/csv",
         headers={"Content-Disposition": f'attachment; filename="yamtrack_{today}.csv"'},
     )
-
-    response = exports.db_to_csv(response, request.user)
-
-    logger.info("User %s successfully exported their data", request.user.username)
-
+    logger.info("User %s started CSV export", request.user.username)
     return response
 
 
diff --git a/src/templates/users/profile.html b/src/templates/users/profile.html
@@ -443,8 +443,13 @@ <h5 class="card-title mb-1">YamTrack CSV</h5>
       <div class="card-body d-flex flex-column justify-content-between">
         <div class="d-flex flex-column flex-lg-row flex-wrap justify-content-between align-items-lg-center row-gap-4">
           <div class="d-flex gap-3 flex-wrap">
-            <svg xmlns="http://www.w3.org/2000/svg" width="58" height="58" fill="currentColor" class="bi bi-filetype-csv" viewBox="0 0 16 16">
-              <path fill-rule="evenodd" d="M14 4.5V14a2 2 0 0 1-2 2h-1v-1h1a1 1 0 0 0 1-1V4.5h-2A1.5 1.5 0 0 1 9.5 3V1H4a1 1 0 0 0-1 1v9H2V2a2 2 0 0 1 2-2h5.5zM3.517 14.841a1.13 1.13 0 0 0 .401.823q.195.162.478.252.284.091.665.091.507 0 .859-.158.354-.158.539-.44.187-.284.187-.656 0-.336-.134-.56a1 1 0 0 0-.375-.357 2 2 0 0 0-.566-.21l-.621-.144a1 1 0 0 1-.404-.176.37.37 0 0 1-.144-.299q0-.234.185-.384.188-.152.512-.152.214 0 .37.068a.6.6 0 0 1 .246.181.56.56 0 0 1 .12.258h.75a1.1 1.1 0 0 0-.2-.566 1.2 1.2 0 0 0-.5-.41 1.8 1.8 0 0 0-.78-.152q-.439 0-.776.15-.337.149-.527.421-.19.273-.19.639 0 .302.122.524.124.223.352.367.228.143.539.213l.618.144q.31.073.463.193a.39.39 0 0 1 .152.326.5.5 0 0 1-.085.29.56.56 0 0 1-.255.193q-.167.07-.413.07-.175 0-.32-.04a.8.8 0 0 1-.248-.115.58.58 0 0 1-.255-.384zM.806 13.693q0-.373.102-.633a.87.87 0 0 1 .302-.399.8.8 0 0 1 .475-.137q.225 0 .398.097a.7.7 0 0 1 .272.26.85.85 0 0 1 .12.381h.765v-.072a1.33 1.33 0 0 0-.466-.964 1.4 1.4 0 0 0-.489-.272 1.8 1.8 0 0 0-.606-.097q-.534 0-.911.223-.375.222-.572.632-.195.41-.196.979v.498q0 .568.193.976.197.407.572.626.375.217.914.217.439 0 .785-.164t.55-.454a1.27 1.27 0 0 0 .226-.674v-.076h-.764a.8.8 0 0 1-.118.363.7.7 0 0 1-.272.25.9.9 0 0 1-.401.087.85.85 0 0 1-.478-.132.83.83 0 0 1-.299-.392 1.7 1.7 0 0 1-.102-.627zm8.239 2.238h-.953l-1.338-3.999h.917l.896 3.138h.038l.888-3.138h.879z"/>
+            <svg xmlns="http://www.w3.org/2000/svg"
+                 width="58"
+                 height="58"
+                 fill="currentColor"
+                 class="bi bi-filetype-csv"
+                 viewBox="0 0 16 16">
+              <path fill-rule="evenodd" d="M14 4.5V14a2 2 0 0 1-2 2h-1v-1h1a1 1 0 0 0 1-1V4.5h-2A1.5 1.5 0 0 1 9.5 3V1H4a1 1 0 0 0-1 1v9H2V2a2 2 0 0 1 2-2h5.5zM3.517 14.841a1.13 1.13 0 0 0 .401.823q.195.162.478.252.284.091.665.091.507 0 .859-.158.354-.158.539-.44.187-.284.187-.656 0-.336-.134-.56a1 1 0 0 0-.375-.357 2 2 0 0 0-.566-.21l-.621-.144a1 1 0 0 1-.404-.176.37.37 0 0 1-.144-.299q0-.234.185-.384.188-.152.512-.152.214 0 .37.068a.6.6 0 0 1 .246.181.56.56 0 0 1 .12.258h.75a1.1 1.1 0 0 0-.2-.566 1.2 1.2 0 0 0-.5-.41 1.8 1.8 0 0 0-.78-.152q-.439 0-.776.15-.337.149-.527.421-.19.273-.19.639 0 .302.122.524.124.223.352.367.228.143.539.213l.618.144q.31.073.463.193a.39.39 0 0 1 .152.326.5.5 0 0 1-.085.29.56.56 0 0 1-.255.193q-.167.07-.413.07-.175 0-.32-.04a.8.8 0 0 1-.248-.115.58.58 0 0 1-.255-.384zM.806 13.693q0-.373.102-.633a.87.87 0 0 1 .302-.399.8.8 0 0 1 .475-.137q.225 0 .398.097a.7.7 0 0 1 .272.26.85.85 0 0 1 .12.381h.765v-.072a1.33 1.33 0 0 0-.466-.964 1.4 1.4 0 0 0-.489-.272 1.8 1.8 0 0 0-.606-.097q-.534 0-.911.223-.375.222-.572.632-.195.41-.196.979v.498q0 .568.193.976.197.407.572.626.375.217.914.217.439 0 .785-.164t.55-.454a1.27 1.27 0 0 0 .226-.674v-.076h-.764a.8.8 0 0 1-.118.363.7.7 0 0 1-.272.25.9.9 0 0 1-.401.087.85.85 0 0 1-.478-.132.83.83 0 0 1-.299-.392 1.7 1.7 0 0 1-.102-.627zm8.239 2.238h-.953l-1.338-3.999h.917l.896 3.138h.038l.888-3.138h.879z" />
             </svg>
             <div>
               <h5 class="card-title mb-1">Export Data</h5>
@@ -458,8 +463,6 @@ <h5 class="card-title mb-1">Export Data</h5>
             </button>
           </form>
         </div>
-
-        {% include "users/components/import_status.html" with import_task=import_tasks.kitsu source="kitsu" %}
       </div>
     </div>
   </div>
@@ -468,9 +471,7 @@ <h5 class="card-title mb-1">Export Data</h5>
     <legend class="border-bottom mb-4">Other</legend>
 
     <div class="input-group profile-grid">
-      <form class="p-2 grid-item"
-            method="post"
-            action="{% url 'reload_calendar' %}">
+      <form class="p-2" method="post" action="{% url 'reload_calendar' %}">
         {% csrf_token %}
         <button class="btn btn-secondary w-100" type="submit">Reload Calendar</button>
       </form>
