diff --git a/src/khoj/database/migrations/0090_alter_khojuser_uuid.py b/src/khoj/database/migrations/0090_alter_khojuser_uuid.py new file mode 100644 index 00000000..72aeea1b --- /dev/null +++ b/src/khoj/database/migrations/0090_alter_khojuser_uuid.py @@ -0,0 +1,106 @@ +# Generated by Django 5.1.9 on 2025-06-04 01:11 + +import uuid + +from django.db import migrations, models + + +def fix_malformed_uuids(apps, schema_editor): + KhojUser = apps.get_model("database", "KhojUser") + + # Track UUID changes for automation cleanup + uuid_mappings = {} + + # Handle null or empty user UUIDs + for user in KhojUser.objects.filter(uuid__isnull=True): + old_uuid = str(user.uuid) if user.uuid else "None" + user.uuid = uuid.uuid4() + user.save() + uuid_mappings[old_uuid] = str(user.uuid) + + # Handle malformed user UUIDs + for user in KhojUser.objects.all(): + current_uuid_val = user.uuid + try: + if not isinstance(current_uuid_val, uuid.UUID): + # Attempt to parse it as UUID. This will catch "None", "null" strings or other malformed hex. + uuid.UUID(str(current_uuid_val)) + except (ValueError, TypeError, AttributeError): + old_uuid_str = str(current_uuid_val) + new_uuid_obj = uuid.uuid4() + user.uuid = new_uuid_obj + user.save( + update_fields=["uuid"] + ) # Important to use update_fields to avoid triggering full save logic if not needed + uuid_mappings[old_uuid_str] = str(new_uuid_obj) + print(f"Fixed malformed UUID for user (old: '{old_uuid_str}', new: {str(new_uuid_obj)})") + + # Clean up orphaned automations + cleanup_orphaned_automations(uuid_mappings) + + +def cleanup_orphaned_automations(uuid_mappings): + """Remove automations with malformed UUIDs in job_ids""" + from apscheduler.jobstores.base import JobLookupError + + from khoj.utils import state + + if not state.scheduler: + return + + all_jobs = state.scheduler.get_jobs() + removed_orphaned_count = 0 + removed_malformed_count = 0 + + for job in all_jobs: + if job.id.startswith("automation_"): + # Extract UUID from job_id: "automation_{uuid}_{query_id}" + job_parts = job.id.split("_", 2) + if len(job_parts) >= 2: + job_uuid = job_parts[1] + + # Check if this UUID was malformed + if job_uuid in uuid_mappings: + # Remove orphaned automation + try: + state.scheduler.remove_job(job.id) + removed_orphaned_count += 1 + print(f"Removed orphaned automation: {job.id}") + except JobLookupError: + pass # Job already removed + + # Also remove jobs with clearly malformed UUIDs + elif job_uuid in ["None", "null"] or not is_valid_uuid(job_uuid): + try: + state.scheduler.remove_job(job.id) + removed_malformed_count += 1 + print(f"Removed automation with malformed UUID: {job.id}") + except JobLookupError: + pass + + if removed_orphaned_count > 0 or removed_malformed_count > 0: + print(f"Removed {removed_orphaned_count} orphaned and {removed_malformed_count} malformed automations.") + + +def is_valid_uuid(uuid_string): + """Check if string is a valid UUID""" + try: + uuid.UUID(str(uuid_string)) + return True + except (ValueError, TypeError, AttributeError): + return False + + +class Migration(migrations.Migration): + dependencies = [ + ("database", "0089_chatmodel_price_tier_and_more"), + ] + + operations = [ + migrations.RunPython(fix_malformed_uuids, reverse_code=migrations.RunPython.noop), + migrations.AlterField( + model_name="khojuser", + name="uuid", + field=models.UUIDField(default=uuid.uuid4, editable=False, unique=True), + ), + ] diff --git a/src/khoj/database/models/__init__.py b/src/khoj/database/models/__init__.py index a854dc91..84948015 100644 --- a/src/khoj/database/models/__init__.py +++ b/src/khoj/database/models/__init__.py @@ -137,7 +137,7 @@ class ClientApplication(DbBaseModel): class KhojUser(AbstractUser): - uuid = models.UUIDField(models.UUIDField(default=uuid.uuid4, editable=False)) + uuid = models.UUIDField(default=uuid.uuid4, editable=False, unique=True) phone_number = PhoneNumberField(null=True, default=None, blank=True) verified_phone_number = models.BooleanField(default=False) verified_email = models.BooleanField(default=False)