Repository: WordPress/openverse-api

Issue:
Audio waveform cache-warming Django command
## Description
<!-- Describe the feature and how it solves the problem. -->
Related to WordPress/openverse-catalog#510

We now have a table where audio waveforms are cached, but this table is not yet fully populated. We have plans to move this processing upstream into the catalog (WordPress/openverse#731). 

For the time being, we should create a Django command that can be run with `python manage.py <command>` which will iterate through all audio records and generate waveforms for them. Here is an outline of the necessary code:

```python
audios = Audio.objects.all()
with tqdm(total=audios.count()) as progress:
    for audio in audios:
        audio.get_or_create_waveform()
        progress.update(1)
```

The above code example uses `tqdm`, which I don't believe is a dependency we have right now. We could add it for this use and remove it after, since the progress bar and estimated time to completion are very useful. 

We have ~14k records, and a test run on staging showed that this could take anywhere from 12-48 hours.
![image (3)](https://user-images.githubusercontent.com/10214785/155229823-8f7a846f-11ff-4c32-8898-b6407f5aa0f6.png)



## Alternatives
<!-- Describe any alternative solutions or features you have considered. How is this feature better? -->

## Additional context
<!-- Add any other context about the feature here; or delete the section entirely. -->

## Implementation
<!-- Replace the [ ] with [x] to check the box. -->
- [ ] 🙋 I would be interested in implementing this feature.



Discussion:
It might be good to order the audio objects by `id`, so that the results are reproducible on multiple runs.

Resolving patch:
diff --git a/api/catalog/management/commands/generatewaveforms.py b/api/catalog/management/commands/generatewaveforms.py
new file mode 100644
--- /dev/null
+++ b/api/catalog/management/commands/generatewaveforms.py
@@ -0,0 +1,127 @@
+import logging
+import subprocess
+
+from catalog.api.models.audio import Audio, AudioAddOn
+from django_tqdm import BaseCommand
+from limit import limit
+
+
+def paginate_reducing_query(get_query_set, page_size=10):
+    """
+    We can't use `Paginator` because it can't handle the situation
+    where the query result changes each time a page is accessed.
+    Because the `audios` QuerySet result is naturally getting smaller
+    each time we successfully process waveforms, we can just take
+    the first ten for each "page" until the page comes back empty.
+    This should theoretically be faster/less DB latency inducing
+    anyway as we're never going to have huge OFFSET values to
+    access deep pages.
+    """
+    page = list(get_query_set()[0:page_size])
+    while len(page):
+        yield page
+        page = list(get_query_set()[0:page_size])
+
+
+class Command(BaseCommand):
+    help = "Generates waveforms for all audio records to populate the cache."
+    """
+    Note: We rely on the file download and waveform generation times
+    taking long enough to prevent us from either making too many requests
+    to the upstream provider or inserting into our database too quickly and
+    causing a slow down. In local tests and in tests run on the staging server
+    it appeared to take on average around 6 to 8 seconds for each audio file.
+    That should be enough latency to not cause any problems.
+    """
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--no_rate_limit", help="Remove self impose rate limits for testing."
+        )
+        parser.add_argument(
+            "--max_records", help="Limit the number of waveforms to create.", type=int
+        )
+
+    def get_audio_handler(self, options):
+        if options["no_rate_limit"]:
+            return lambda audio: audio.get_or_create_waveform()
+
+        @limit(limit=1, every=2)  # Call once per two seconds maximum
+        def limited(audio):
+            audio.get_or_create_waveform()
+
+        return limited
+
+    def _process_wavelengths(self, audios, audio_handler, count_to_process):
+        errored_identifiers = []
+        processed = 0
+        with self.tqdm(total=count_to_process) as progress:
+            paginator = paginate_reducing_query(
+                get_query_set=lambda: audios.exclude(identifier__in=errored_identifiers)
+            )
+            for page in paginator:
+                for audio in page:
+                    if processed > count_to_process:
+                        return errored_identifiers
+                    try:
+                        processed += 1
+                        audio_handler(audio)
+                    except subprocess.CalledProcessError as err:
+                        errored_identifiers.append(audio.identifier)
+                        self.error(
+                            f"Unable to process {audio.identifier}: "
+                            f"{err.stderr.decode().strip()}"
+                        )
+                    except KeyboardInterrupt:
+                        errored_identifiers.append(audio.identifier)
+                        return errored_identifiers
+                    except BaseException as err:
+                        errored_identifiers.append(audio.identifier)
+                        self.error(f"Unable to process {audio.identifier}: " f"{err}")
+                    progress.update(1)
+
+        return errored_identifiers
+
+    def handle(self, *args, **options):
+        # These logs really muck up the tqdm output and don't give us much helpful
+        # information, so they get silenced
+        logging.getLogger("catalog.api.utils.waveform").setLevel(logging.WARNING)
+
+        existing_waveform_audio_identifiers_query = AudioAddOn.objects.filter(
+            waveform_peaks__isnull=False
+        ).values_list("audio_identifier", flat=True)
+        audios = Audio.objects.exclude(
+            identifier__in=existing_waveform_audio_identifiers_query
+        ).order_by("id")
+
+        max_records = options["max_records"]
+        count = audios.count()
+
+        count_to_process = count
+
+        if max_records is not None:
+            count_to_process = max_records if max_records < count else count
+
+        self.info(
+            self.style.NOTICE(f"Generating waveforms for {count_to_process:,} records")
+        )
+
+        audio_handler = self.get_audio_handler(options)
+
+        errored_identifiers = self._process_wavelengths(
+            audios, audio_handler, count_to_process
+        )
+
+        self.info(self.style.SUCCESS("Finished generating waveforms!"))
+
+        if errored_identifiers:
+            errored_identifiers_joined = "\n".join(
+                str(identifier) for identifier in errored_identifiers
+            )
+
+            self.info(
+                self.style.WARNING(
+                    f"The following Audio identifiers were unable "
+                    f"to be processed\n\n{errored_identifiers_joined}"
+                )
+            )


Test patch:
diff --git a/api/test/factory/__init__.py b/api/test/factory/__init__.py
new file mode 100644
diff --git a/api/test/factory/faker.py b/api/test/factory/faker.py
new file mode 100644
--- /dev/null
+++ b/api/test/factory/faker.py
@@ -0,0 +1,17 @@
+from factory import Faker
+from faker.providers import BaseProvider
+from faker.utils.distribution import choices_distribution
+
+
+class WaveformProvider(BaseProvider):
+    _float_space = [x / 100.0 for x in range(101)] * 20
+
+    @classmethod
+    def generate_waveform(cls) -> list[float]:
+        return choices_distribution(cls._float_space, p=None, length=1000)
+
+    def waveform(self) -> list[float]:
+        return WaveformProvider.generate_waveform()
+
+
+Faker.add_provider(WaveformProvider)
diff --git a/api/test/factory/models/__init__.py b/api/test/factory/models/__init__.py
new file mode 100644
diff --git a/api/test/factory/models/audio.py b/api/test/factory/models/audio.py
new file mode 100644
--- /dev/null
+++ b/api/test/factory/models/audio.py
@@ -0,0 +1,19 @@
+from test.factory.faker import Faker
+from test.factory.models.media import IdentifierFactory, MediaFactory
+
+from catalog.api.models.audio import Audio, AudioAddOn
+from factory.django import DjangoModelFactory
+
+
+class AudioFactory(MediaFactory):
+    class Meta:
+        model = Audio
+
+
+class AudioAddOnFactory(DjangoModelFactory):
+    class Meta:
+        model = AudioAddOn
+
+    audio_identifier = IdentifierFactory(AudioFactory)
+
+    waveform_peaks = Faker("waveform")
diff --git a/api/test/factory/models/media.py b/api/test/factory/models/media.py
new file mode 100644
--- /dev/null
+++ b/api/test/factory/models/media.py
@@ -0,0 +1,38 @@
+from test.factory.faker import Faker
+from uuid import uuid4
+
+import factory
+from catalog.api.licenses import LICENSES
+from factory.django import DjangoModelFactory
+
+
+class MediaFactory(DjangoModelFactory):
+    """Base factory for models that extend from the AbstractMedia class."""
+
+    class Meta:
+        abstract = True
+
+    identifier = factory.sequence(lambda _: uuid4())
+
+    foreign_identifier = factory.sequence(lambda _: uuid4())
+    """The foreign identifier isn't necessarily a UUID but for test purposes it's fine if it looks like one"""
+
+    license = Faker(
+        "random_element", elements=[the_license[0] for the_license in LICENSES]
+    )
+
+    foreign_landing_url = Faker("url")
+
+
+class IdentifierFactory(factory.SubFactory):
+    """
+    A factory for creating a related model and returning the UUID.
+
+    Distinct from the `SubFactory` in that this creates the related model but
+    uses a specific attribute from it for the resulting value instead of the
+    related model itself.
+    """
+
+    def evaluate(self, instance, step, extra):
+        model = super().evaluate(instance, step, extra)
+        return model.identifier
diff --git a/api/test/run_test.sh b/api/test/run_test.sh
--- a/api/test/run_test.sh
+++ b/api/test/run_test.sh
@@ -18,8 +18,8 @@ succeeded=$?
 if [[ $succeeded -eq 0 ]]; then
 	printf "${green}:-) All tests passed${endcol}\n"
 else
-  printf "Full system logs:\n"
-  docker-compose logs
+    printf "Full system logs:\n"
+    docker-compose logs
 	printf "${red}:'( Some tests did not pass${endcol}\n"
 fi
 exit $succeeded
diff --git a/api/test/unit/management/__init__.py b/api/test/unit/management/__init__.py
new file mode 100644
diff --git a/api/test/unit/management/commands/__init__.py b/api/test/unit/management/commands/__init__.py
new file mode 100644
diff --git a/api/test/unit/management/commands/generatewaveforms_test.py b/api/test/unit/management/commands/generatewaveforms_test.py
new file mode 100644
--- /dev/null
+++ b/api/test/unit/management/commands/generatewaveforms_test.py
@@ -0,0 +1,180 @@
+import subprocess
+from io import StringIO
+from test.factory.faker import WaveformProvider
+from test.factory.models.audio import AudioAddOnFactory, AudioFactory
+from unittest import mock
+
+import psycopg2
+import pytest
+from catalog.api.models.audio import Audio, AudioAddOn
+from django.core.management import call_command
+from django.db import connections
+from django.test.utils import CaptureQueriesContext
+
+
+@mock.patch("catalog.api.models.audio.generate_peaks")
+def call_generatewaveforms(mock_generate_peaks: mock.MagicMock) -> tuple[str, str]:
+    mock_generate_peaks.side_effect = lambda _: WaveformProvider.generate_waveform()
+    out = StringIO()
+    err = StringIO()
+    call_command("generatewaveforms", no_rate_limit=True, stdout=out, stderr=err)
+
+    return out.getvalue(), err.getvalue()
+
+
+def assert_all_audio_have_waveforms():
+    assert (
+        list(
+            AudioAddOn.objects.filter(waveform_peaks__isnull=False).values_list(
+                "audio_identifier"
+            )
+        ).sort()
+        == list(Audio.objects.all().values_list("identifier")).sort()
+    )
+
+
+@pytest.mark.django_db
+def test_creates_waveforms_for_audio():
+    AudioFactory.create_batch(153)
+
+    assert AudioAddOn.objects.count() == 0
+
+    call_generatewaveforms()
+
+    assert_all_audio_have_waveforms()
+
+
+@pytest.mark.django_db
+def test_does_not_reprocess_existing_waveforms():
+    waveformless_audio = AudioFactory.create_batch(3)
+
+    # AudioAddOnFactory will create associated Audio objects as well
+    # so those three will serve as the audio that should _not_ get processed
+    AudioAddOnFactory.create_batch(3)
+
+    # Create an add on that doesn't have a waveform, this one should get processed as well
+    null_waveform_addon = AudioAddOnFactory.create(waveform_peaks=None)
+    waveformless_audio.append(
+        Audio.objects.get(identifier=null_waveform_addon.audio_identifier)
+    )
+
+    out, err = call_generatewaveforms()
+
+    assert f"Generating waveforms for {len(waveformless_audio)} records" in out
+    assert_all_audio_have_waveforms()
+
+
+@pytest.mark.django_db
+@mock.patch("catalog.api.models.audio.generate_peaks")
+def test_paginates_audio_waveforms_to_generate(
+    mock_generate_peaks, django_assert_num_queries
+):
+    mock_generate_peaks.return_value = WaveformProvider.generate_waveform()
+
+    audio_count = 53  # 6 pages
+    pages = 6
+    AudioFactory.create_batch(audio_count)
+
+    test_audio = AudioFactory.create()
+    with CaptureQueriesContext(connections["default"]) as capture:
+        test_audio.get_or_create_waveform()
+    test_audio.delete()
+
+    queries_per_iteration = len(capture.captured_queries)
+
+    # 1 per page + the final empty page's query
+    pagination_queries = pages + 1
+
+    # initializes the count for tqdm
+    count_queries = 1
+
+    # queries inside get_or_create_waveform
+    interation_queries = queries_per_iteration * audio_count
+
+    expected_queries = interation_queries + pagination_queries + count_queries
+
+    with django_assert_num_queries(expected_queries):
+        call_generatewaveforms()
+
+    assert_all_audio_have_waveforms()
+
+
+@pytest.mark.django_db
+@pytest.mark.parametrize(
+    ("exception_class", "exception_args", "exception_kwargs"),
+    (
+        (
+            subprocess.CalledProcessError,
+            (1, "audiowaveform"),
+            {"stderr": b"This is an error string"},
+        ),
+        (
+            psycopg2.errors.lookup(psycopg2.errorcodes.NOT_NULL_VIOLATION),
+            tuple(),
+            dict(),
+        ),
+    ),
+)
+@mock.patch("catalog.api.models.audio.generate_peaks")
+def test_logs_and_continues_if_waveform_generation_fails(
+    mock_generate_peaks, exception_class, exception_args, exception_kwargs
+):
+    audio_count = 23
+    return_values = [
+        exception_class(*exception_args, **exception_kwargs)
+        if i == 9
+        else WaveformProvider.generate_waveform()
+        for i in range(audio_count)
+    ]
+    mock_generate_peaks.side_effect = return_values
+    AudioFactory.create_batch(audio_count)
+
+    out = StringIO()
+    err = StringIO()
+    call_command("generatewaveforms", no_rate_limit=True, stdout=out, stderr=err)
+
+    failed_audio = Audio.objects.exclude(
+        identifier__in=AudioAddOn.objects.filter(
+            waveform_peaks__isnull=False
+        ).values_list("audio_identifier", flat=True)
+    )
+
+    assert failed_audio.count() == 1
+    assert f"Unable to process {failed_audio.first().identifier}" in err.getvalue()
+
+    assert (
+        AudioAddOn.objects.filter(waveform_peaks__isnull=False).count()
+        == audio_count - 1
+    )
+
+
+@pytest.mark.django_db
+@mock.patch("catalog.api.models.audio.generate_peaks")
+def test_keyboard_interrupt_should_halt_processing(mock_generate_peaks):
+    audio_count = 23
+    interrupt_at = 9
+    return_values = [
+        KeyboardInterrupt()
+        if i == interrupt_at
+        else WaveformProvider.generate_waveform()
+        for i in range(audio_count)
+    ]
+
+    mock_generate_peaks.side_effect = return_values
+    AudioFactory.create_batch(audio_count)
+
+    out = StringIO()
+    err = StringIO()
+    call_command("generatewaveforms", no_rate_limit=True, stdout=out, stderr=err)
+
+    failed_audio = Audio.objects.exclude(
+        identifier__in=AudioAddOn.objects.filter(
+            waveform_peaks__isnull=False
+        ).values_list("audio_identifier", flat=True)
+    )
+
+    assert failed_audio.count() == audio_count - interrupt_at
+
+    assert (
+        AudioAddOn.objects.filter(waveform_peaks__isnull=False).count() == interrupt_at
+    )
diff --git a/api/test/unit/models/audio_test.py b/api/test/unit/models/audio_test.py
--- a/api/test/unit/models/audio_test.py
+++ b/api/test/unit/models/audio_test.py
@@ -1,4 +1,5 @@
 import uuid
+from test.factory.faker import WaveformProvider
 from unittest import mock
 
 import pytest
@@ -20,7 +21,7 @@ def audio_fixture():
 @pytest.mark.django_db
 @mock.patch("catalog.api.models.audio.generate_peaks")
 def test_audio_waveform_caches(generate_peaks_mock, audio_fixture):
-    mock_waveform = [0.4, 0.3, 0.1, 0, 1, 0.6]
+    mock_waveform = WaveformProvider.generate_waveform()
     generate_peaks_mock.return_value = mock_waveform
 
     assert AudioAddOn.objects.count() == 0