update mal, anilist, kitsu import

Author: FuzzyGrim

src/app/database.py

@@ -2,6 +2,7 @@
 from datetime import timedelta
 
 from django.apps import apps
+from django.db import models
 from django.db.models import Count, F
 from django.db.models.functions import TruncDate
 
@@ -45,6 +46,14 @@ def get_fields(model):
     return [f.name for f in model._meta.fields]  # noqa: SLF001
 
 
+def get_unique_constraint_fields(model):
+    """Get fields that make up the unique constraint for the model."""
+    for constraint in model._meta.constraints:  # noqa: SLF001
+        if isinstance(constraint, models.UniqueConstraint):
+            return constraint.fields
+    return None
+
+
 def get_properties(model):
     """Get properties of a model."""
     return [name for name in dir(model) if isinstance(getattr(model, name), property)]

src/integrations/helpers.py

@@ -1,17 +1,107 @@
-from simple_history.utils import bulk_create_with_history
+from simple_history.utils import bulk_create_with_history, bulk_update_with_history
 
+import app
 
-def bulk_chunk_import(media_list, model, user):
-    """Bulk import media in chunks.
 
-    Fixes bulk_create_with_history limit
-    https://github.com/jazzband/django-simple-history/issues/1216#issuecomment-1903240831
-    """
-    chunk_size = 500
-    for i in range(0, len(media_list), chunk_size):
-        bulk_create_with_history(
-            media_list[i : i + chunk_size],
+def bulk_chunk_import(bulk_media, model, user, mode):
+    """Bulk import media in chunks."""
+    if mode == "new":
+        num_imported = bulk_create_new_with_history(bulk_media, model, user)
+
+    elif mode == "overwrite":
+        num_imported = bulk_create_update_with_history(
+            bulk_media,
             model,
-            ignore_conflicts=True,
+            user,
+        )
+
+    return num_imported
+
+
+def bulk_create_new_with_history(bulk_media, model, user):
+    """Filter out existing records and bulk create only new ones."""
+    # Get existing records' unique IDs since bulk_create_with_history
+    # returns all objects even if they weren't created due to conflicts
+    unique_fields = app.database.get_unique_constraint_fields(model)
+    existing_combos = set(
+        model.objects.values_list(*unique_fields),
+    )
+
+    new_records = [
+        record
+        for record in bulk_media
+        if tuple(getattr(record, field + "_id") for field in unique_fields)
+        not in existing_combos
+    ]
+
+    bulk_create_with_history(
+        new_records,
+        model,
+        batch_size=500,
+        default_user=user,
+    )
+
+    return len(new_records)
+
+
+def bulk_create_update_with_history(
+    bulk_media,
+    model,
+    user,
+):
+    """Bulk create new records and update existing ones with history tracking."""
+    unique_fields = app.database.get_unique_constraint_fields(model)
+    model_fields = app.database.get_fields(model)
+    update_fields = [
+        field for field in model_fields if field not in unique_fields and field != "id"
+    ]
+
+    # Get existing objects with their unique fields and id
+    existing_objs = model.objects.filter(
+        **{
+            f"{field}__in": [getattr(obj, field + "_id") for obj in bulk_media]
+            for field in unique_fields
+        },
+    ).values(*unique_fields, "id")
+
+    # Create lookup dictionary using unique field combinations
+    existing_lookup = {
+        tuple(obj[field] for field in unique_fields): obj["id"] for obj in existing_objs
+    }
+
+    # Split records into new and existing based on unique constraints
+    create_objs = []
+    update_objs = []
+
+    for record in bulk_media:
+        record_key = tuple(getattr(record, field + "_id") for field in unique_fields)
+        if record_key in existing_lookup:
+            # Set the primary key for update
+            record.id = existing_lookup[record_key]
+            update_objs.append(record)
+        else:
+            create_objs.append(record)
+
+    # Bulk create new records
+    num_created = 0
+    if create_objs:
+        created_objs = bulk_create_with_history(
+            create_objs,
+            model,
+            batch_size=500,
             default_user=user,
         )
+        num_created = len(created_objs)
+
+    # Bulk update existing records
+    num_updated = 0
+    if update_objs and update_fields:
+        num_updated = bulk_update_with_history(
+            update_objs,
+            model,
+            fields=update_fields,
+            batch_size=500,
+            default_user=user,
+        )
+
+    return num_created + num_updated

src/integrations/imports/anilist.py

@@ -10,7 +10,7 @@
 logger = logging.getLogger(__name__)
 
 
-def importer(username, user):
+def importer(username, user, mode):
     """Import anime and manga ratings from Anilist."""
     query = """
     query ($userName: String){
@@ -94,19 +94,21 @@ def importer(username, user):
         response["data"]["anime"],
         "anime",
         user,
+        mode,
     )
 
     manga_imported, manga_warnings = import_media(
         response["data"]["manga"],
         "manga",
         user,
+        mode,
     )
 
     warning_messages = anime_warnings + manga_warnings
     return anime_imported, manga_imported, "\n".join(warning_messages)
 
 
-def import_media(media_data, media_type, user):
+def import_media(media_data, media_type, user, mode):
     """Import media of a specific type from Anilist."""
     logger.info("Importing %s from Anilist", media_type)
 
@@ -123,10 +125,7 @@ def import_media(media_data, media_type, user):
             )
 
     model = apps.get_model(app_label="app", model_name=media_type)
-    num_before = model.objects.filter(user=user).count()
-    helpers.bulk_chunk_import(bulk_media, model, user)
-    num_after = model.objects.filter(user=user).count()
-    num_imported = num_after - num_before
+    num_imported = helpers.bulk_chunk_import(bulk_media, model, user, mode)
 
     logger.info("Imported %s %s", num_imported, media_type)
 

src/integrations/imports/kitsu.py

@@ -15,13 +15,50 @@
 KITSU_PAGE_LIMIT = 500
 
 
-def import_by_user_id(kitsu_id, user):
+def get_kitsu_id(username):
+    """Get the user ID from Kitsu."""
+    url = f"{KITSU_API_BASE_URL}/users"
+    response = app.providers.services.api_request(
+        "KITSU",
+        "GET",
+        url,
+        params={"filter[name]": username},
+    )
+
+    if not response["data"]:
+        msg = f"User {username} not found."
+        raise ValueError(msg)
+    if len(response["data"]) > 1:
+        msg = (
+            f"Multiple users found for {username}, please use your user ID. "
+            "User IDs can be found in the URL when viewing your Kitsu profile."
+        )
+        raise ValueError(msg)
+
+    return response["data"][0]["id"]
+
+
+def importer(kitsu_id, user, mode):
     """Import anime and manga ratings from Kitsu by user ID."""
+    # Check if given ID is a username
+    if not kitsu_id.isdigit():
+        kitsu_id = get_kitsu_id(kitsu_id)
+
     anime_response = get_media_response(kitsu_id, "anime")
-    num_anime_imported, anime_warnings = importer(anime_response, "anime", user)
+    num_anime_imported, anime_warnings = import_media(
+        anime_response,
+        "anime",
+        user,
+        mode,
+    )
 
     manga_response = get_media_response(kitsu_id, "manga")
-    num_manga_imported, manga_warning = importer(manga_response, "manga", user)
+    num_manga_imported, manga_warning = import_media(
+        manga_response,
+        "manga",
+        user,
+        mode,
+    )
 
     warning_messages = anime_warnings + manga_warning
     return num_anime_imported, num_manga_imported, "\n".join(warning_messages)
@@ -51,7 +88,7 @@ def get_media_response(kitsu_id, media_type):
     return all_data
 
 
-def importer(response, media_type, user):
+def import_media(response, media_type, user, mode):
     """Import media from Kitsu and return the number of items imported."""
     logger.info("Importing %s from Kitsu", media_type)
 
@@ -85,10 +122,7 @@ def importer(response, media_type, user):
             else:
                 bulk_data.append(instance)
 
-    num_before = model.objects.filter(user=user).count()
-    helpers.bulk_chunk_import(bulk_data, model, user)
-    num_after = model.objects.filter(user=user).count()
-    num_imported = num_after - num_before
+    num_imported = helpers.bulk_chunk_import(bulk_data, model, user, mode)
 
     logger.info("Imported %s %s", num_imported, media_type)
 

src/integrations/imports/mal.py

@@ -12,14 +12,14 @@
 base_url = "https://api.myanimelist.net/v2/users"
 
 
-def importer(username, user):
+def importer(username, user, mode):
     """Import anime and manga from MyAnimeList."""
-    anime_imported = import_media(username, user, "anime")
-    manga_imported = import_media(username, user, "manga")
+    anime_imported = import_media(username, user, "anime", mode)
+    manga_imported = import_media(username, user, "manga", mode)
     return anime_imported, manga_imported
 
 
-def import_media(username, user, media_type):
+def import_media(username, user, media_type, mode):
     """Import media of a specific type from MyAnimeList."""
     logger.info("Fetching %s from MyAnimeList", media_type)
     params = {
@@ -32,11 +32,7 @@ def import_media(username, user, media_type):
     bulk_media = add_media_list(media_data, media_type, user)
 
     model = apps.get_model(app_label="app", model_name=media_type)
-    num_before = model.objects.filter(user=user).count()
-    helpers.bulk_chunk_import(bulk_media, model, user)
-    num_after = model.objects.filter(user=user).count()
-
-    num_imported = num_after - num_before
+    num_imported = helpers.bulk_chunk_import(bulk_media, model, user, mode)
     logger.info("Imported %s %s", num_imported, media_type)
 
     return num_imported