|
| 1 | +# |
| 2 | +# Copyright (c) nexB Inc. and others. All rights reserved. |
| 3 | +# VulnerableCode is a trademark of nexB Inc. |
| 4 | +# SPDX-License-Identifier: Apache-2.0 |
| 5 | +# See http://www.apache.org/licenses/LICENSE-2.0 for the license text. |
| 6 | +# See https://github.com/aboutcode-org/vulnerablecode for support or download. |
| 7 | +# See https://aboutcode.org for more information about nexB OSS projects. |
| 8 | +# |
| 9 | + |
| 10 | +from aboutcode.pipeline import LoopProgress |
| 11 | +from django.db import migrations |
| 12 | +from django.db import models |
| 13 | + |
| 14 | +""" |
| 15 | +Model and data migration for converting the Advisory aliases |
| 16 | +JSON field to a concrete M2M Advisory Alias relationship. |
| 17 | +""" |
| 18 | + |
| 19 | +def bulk_update(model, items, fields, logger): |
| 20 | + item_count = 0 |
| 21 | + if items: |
| 22 | + try: |
| 23 | + model.objects.bulk_update(objs=items, fields=fields) |
| 24 | + item_count += len(items) |
| 25 | + except Exception as e: |
| 26 | + logger(f"Error updating Advisory: {e}") |
| 27 | + items.clear() |
| 28 | + return item_count |
| 29 | + |
| 30 | + |
| 31 | +class Migration(migrations.Migration): |
| 32 | + |
| 33 | + dependencies = [ |
| 34 | + ("vulnerabilities", "0088_fix_alpine_purl_type"), |
| 35 | + ] |
| 36 | + |
| 37 | + def populate_new_advisory_aliases_field(apps, schema_editor): |
| 38 | + Advisory = apps.get_model("vulnerabilities", "Advisory") |
| 39 | + Alias = apps.get_model("vulnerabilities", "Alias") |
| 40 | + advisories = Advisory.objects.all() |
| 41 | + |
| 42 | + chunk_size = 10000 |
| 43 | + advisories_count = advisories.count() |
| 44 | + print(f"\nPopulate new advisory aliases relationship.") |
| 45 | + progress = LoopProgress( |
| 46 | + total_iterations=advisories_count, |
| 47 | + logger=print, |
| 48 | + progress_step=1, |
| 49 | + ) |
| 50 | + for advisory in progress.iter(advisories.iterator(chunk_size=chunk_size)): |
| 51 | + aliases = Alias.objects.filter(alias__in=advisory.old_aliases) |
| 52 | + advisory.aliases.set(aliases) |
| 53 | + |
| 54 | + def reverse_populate_new_advisory_aliases_field(apps, schema_editor): |
| 55 | + Advisory = apps.get_model("vulnerabilities", "Advisory") |
| 56 | + advisories = Advisory.objects.all() |
| 57 | + |
| 58 | + updated_advisory_count = 0 |
| 59 | + batch_size = 10000 |
| 60 | + chunk_size = 10000 |
| 61 | + updated_advisory = [] |
| 62 | + progress = LoopProgress( |
| 63 | + total_iterations=advisories.count(), |
| 64 | + logger=print, |
| 65 | + progress_step=1, |
| 66 | + ) |
| 67 | + for advisory in progress.iter(advisories.iterator(chunk_size=chunk_size)): |
| 68 | + aliases = advisory.aliases.all() |
| 69 | + advisory.old_aliases = [alias.alias for alias in aliases] |
| 70 | + updated_advisory.append(advisory) |
| 71 | + |
| 72 | + if len(updated_advisory) > batch_size: |
| 73 | + updated_advisory_count += bulk_update( |
| 74 | + model=Advisory, |
| 75 | + items=updated_advisory, |
| 76 | + fields=["old_aliases"], |
| 77 | + logger=print, |
| 78 | + ) |
| 79 | + |
| 80 | + updated_advisory_count += bulk_update( |
| 81 | + model=Advisory, |
| 82 | + items=updated_advisory, |
| 83 | + fields=["old_aliases"], |
| 84 | + logger=print, |
| 85 | + ) |
| 86 | + |
| 87 | + operations = [ |
| 88 | + # Rename aliases field to old_aliases |
| 89 | + migrations.AlterModelOptions( |
| 90 | + name="advisory", |
| 91 | + options={"ordering": ["date_published", "unique_content_id"]}, |
| 92 | + ), |
| 93 | + migrations.AlterUniqueTogether( |
| 94 | + name="advisory", |
| 95 | + unique_together={("unique_content_id", "date_published", "url")}, |
| 96 | + ), |
| 97 | + migrations.RenameField( |
| 98 | + model_name="advisory", |
| 99 | + old_name="aliases", |
| 100 | + new_name="old_aliases", |
| 101 | + ), |
| 102 | + migrations.AddField( |
| 103 | + model_name="advisory", |
| 104 | + name="aliases", |
| 105 | + field=models.ManyToManyField(related_name="advisories", to="vulnerabilities.alias"), |
| 106 | + ), |
| 107 | + # Populate the new M2M aliases relation |
| 108 | + migrations.RunPython( |
| 109 | + code=populate_new_advisory_aliases_field, |
| 110 | + reverse_code=reverse_populate_new_advisory_aliases_field, |
| 111 | + ), |
| 112 | + # Delete old_alias field |
| 113 | + migrations.RemoveField( |
| 114 | + model_name="advisory", |
| 115 | + name="old_aliases", |
| 116 | + ), |
| 117 | + ] |
0 commit comments