feat: celery tasks to replace ietf/bin scripts (#6971)

* refactor: Change import style for clarity * feat: Add iana_changes_updates_task() * chore: Squelch lint warning My linter does not like variables defined outside of __init__() * feat: Add PeriodicTask for iana_changes_updates_task * refactor: tasks instead of scripts on sync.views.notify() * test: Test iana_changes_updates_task * refactor: rename task for consistency * feat: Add iana_protocols_update_task * feat: Add PeriodicTask for iana protocols sync * refactor: Use protocol sync task instead of script in view * refactor: itertools.batched() not available until py312 * test: test iana_protocols_update_task * feat: Add idindex_update_task() Calls idindex generation functions and does the file update dance to put them in place. * chore: Add comments to bin/hourly * fix: annotate types and fix bug * feat: Create PeriodicTask for idindex_update_task * refactor: Move helpers into a class More testable this way * refactor: Make TempFileManager a context mgr * test: Test idindex_update_task * test: Test TempFileManager * fix: Fix bug in TestFileManager yay testing * feat: Add expire_ids_task() * feat: Create PeriodicTask for expire_ids_task * test: Test expire_ids_task * test: Test request timeout in iana_protocols_update_task * refactor: do not re-raise timeout exception Not sure this is the right thing to do, but it's the same as rfc_editor_index_update_task * feat: Add notify_expirations_task * feat: Add "weekly" celery beat crontab * refactor: Reorder crontab fields This matches the crontab file field order * feat: Add PeriodicTask for notify_expirations * test: Test notify_expirations_task * test: Add annotation to satisfy mypy
ietf-tools · Jan 31, 2024 · b4cf04a · b4cf04a
1 parent 118b00d
commit b4cf04a
Show file tree

Hide file tree

Showing 9 changed files with 526 additions and 23 deletions.
diff --git a/bin/hourly b/bin/hourly
@@ -45,6 +45,7 @@ ID=/a/ietfdata/doc/draft/repository
 DERIVED=/a/ietfdata/derived
 DOWNLOAD=/a/www/www6s/download
 
+## Start of script refactored into idindex_update_task() ===
 export TMPDIR=/a/tmp
 
 TMPFILE1=`mktemp` || exit 1
@@ -85,6 +86,8 @@ mv $TMPFILE9 $DERIVED/1id-index.txt
 mv $TMPFILEA $DERIVED/1id-abstracts.txt
 mv $TMPFILEB $DERIVED/all_id2.txt
 
+## End of script refactored into idindex_update_task() ===
+
 $DTDIR/ietf/manage.py generate_idnits2_rfc_status
 $DTDIR/ietf/manage.py generate_idnits2_rfcs_obsoleted
 

diff --git a/ietf/doc/tasks.py b/ietf/doc/tasks.py
@@ -0,0 +1,56 @@
+# Copyright The IETF Trust 2024, All Rights Reserved
+#
+# Celery task definitions
+#
+import datetime
+import debug  # pyflakes:ignore
+
+from celery import shared_task
+
+from ietf.utils import log
+from ietf.utils.timezone import datetime_today
+
+from .expire import (
+    in_draft_expire_freeze,
+    get_expired_drafts,
+    expirable_drafts,
+    send_expire_notice_for_draft,
+    expire_draft,
+    clean_up_draft_files,
+    get_soon_to_expire_drafts,
+    send_expire_warning_for_draft,
+)
+from .models import Document
+
+
+@shared_task
+def expire_ids_task():
+    try:
+        if not in_draft_expire_freeze():
+            log.log("Expiring drafts ...")
+            for doc in get_expired_drafts():
+                # verify expirability -- it might have changed after get_expired_drafts() was run
+                # (this whole loop took about 2 minutes on 04 Jan 2018)
+                # N.B., re-running expirable_drafts() repeatedly is fairly expensive. Where possible,
+                # it's much faster to run it once on a superset query of the objects you are going
+                # to test and keep its results. That's not desirable here because it would defeat
+                # the purpose of double-checking that a document is still expirable when it is actually
+                # being marked as expired.
+                if expirable_drafts(
+                    Document.objects.filter(pk=doc.pk)
+                ).exists() and doc.expires < datetime_today() + datetime.timedelta(1):
+                    send_expire_notice_for_draft(doc)
+                    expire_draft(doc)
+                    log.log(f"  Expired draft {doc.name}-{doc.rev}")
+
+        log.log("Cleaning up draft files")
+        clean_up_draft_files()
+    except Exception as e:
+        log.log("Exception in expire-ids: %s" % e)
+        raise
+
+
+@shared_task
+def notify_expirations_task(notify_days=14):
+    for doc in get_soon_to_expire_drafts(notify_days):
+        send_expire_warning_for_draft(doc)
diff --git a/ietf/doc/tests_tasks.py b/ietf/doc/tests_tasks.py
@@ -0,0 +1,63 @@
+# Copyright The IETF Trust 2024, All Rights Reserved
+import mock
+
+from ietf.utils.test_utils import TestCase
+from ietf.utils.timezone import datetime_today
+
+from .factories import DocumentFactory
+from .models import Document
+from .tasks import expire_ids_task, notify_expirations_task
+
+
+class TaskTests(TestCase):
+
+    @mock.patch("ietf.doc.tasks.in_draft_expire_freeze")
+    @mock.patch("ietf.doc.tasks.get_expired_drafts")
+    @mock.patch("ietf.doc.tasks.expirable_drafts")
+    @mock.patch("ietf.doc.tasks.send_expire_notice_for_draft")
+    @mock.patch("ietf.doc.tasks.expire_draft")
+    @mock.patch("ietf.doc.tasks.clean_up_draft_files")
+    def test_expire_ids_task(
+        self,
+        clean_up_draft_files_mock,
+        expire_draft_mock,
+        send_expire_notice_for_draft_mock,
+        expirable_drafts_mock,
+        get_expired_drafts_mock,
+        in_draft_expire_freeze_mock,
+    ):
+        # set up mocks
+        in_draft_expire_freeze_mock.return_value = False
+        doc, other_doc = DocumentFactory.create_batch(2)
+        doc.expires = datetime_today()
+        get_expired_drafts_mock.return_value = [doc, other_doc]
+        expirable_drafts_mock.side_effect = [
+            Document.objects.filter(pk=doc.pk),
+            Document.objects.filter(pk=other_doc.pk),
+        ]
+
+        # call task
+        expire_ids_task()
+
+        # check results
+        self.assertTrue(in_draft_expire_freeze_mock.called)
+        self.assertEqual(expirable_drafts_mock.call_count, 2)
+        self.assertEqual(send_expire_notice_for_draft_mock.call_count, 1)
+        self.assertEqual(send_expire_notice_for_draft_mock.call_args[0], (doc,))
+        self.assertEqual(expire_draft_mock.call_count, 1)
+        self.assertEqual(expire_draft_mock.call_args[0], (doc,))
+        self.assertTrue(clean_up_draft_files_mock.called)
+
+        # test that an exception is raised
+        in_draft_expire_freeze_mock.side_effect = RuntimeError
+        with self.assertRaises(RuntimeError):(
+            expire_ids_task())
+
+    @mock.patch("ietf.doc.tasks.send_expire_warning_for_draft")
+    @mock.patch("ietf.doc.tasks.get_soon_to_expire_drafts")
+    def test_notify_expirations_task(self, get_drafts_mock, send_warning_mock):
+        # Set up mocks
+        get_drafts_mock.return_value = ["sentinel"]
+        notify_expirations_task()
+        self.assertEqual(send_warning_mock.call_count, 1)
+        self.assertEqual(send_warning_mock.call_args[0], ("sentinel",))
diff --git a/ietf/idindex/tasks.py b/ietf/idindex/tasks.py
@@ -0,0 +1,85 @@
+# Copyright The IETF Trust 2024, All Rights Reserved
+#
+# Celery task definitions
+#
+import shutil
+
+import debug    # pyflakes:ignore
+
+from celery import shared_task
+from contextlib import AbstractContextManager
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+
+from .index import all_id_txt, all_id2_txt, id_index_txt
+
+
+class TempFileManager(AbstractContextManager):
+    def __init__(self, tmpdir=None) -> None:
+        self.cleanup_list: set[Path] = set()
+        self.dir = tmpdir
+
+    def make_temp_file(self, content):
+        with NamedTemporaryFile(mode="wt", delete=False, dir=self.dir) as tf:
+            tf_path = Path(tf.name)
+            self.cleanup_list.add(tf_path)
+            tf.write(content)
+        return tf_path
+
+    def move_into_place(self, src_path: Path, dest_path: Path):
+        shutil.move(src_path, dest_path)
+        dest_path.chmod(0o644)
+        self.cleanup_list.remove(src_path)
+
+    def cleanup(self):
+        for tf_path in self.cleanup_list:
+            tf_path.unlink(missing_ok=True)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.cleanup()
+        return False  # False: do not suppress the exception
+
+
+@shared_task
+def idindex_update_task():
+    """Update I-D indexes"""
+    id_path = Path("/a/ietfdata/doc/draft/repository")
+    derived_path = Path("/a/ietfdata/derived")
+    download_path = Path("/a/www/www6s/download")
+
+    with TempFileManager("/a/tmp") as tmp_mgr:
+        # Generate copies of new contents
+        all_id_content = all_id_txt()
+        all_id_tmpfile = tmp_mgr.make_temp_file(all_id_content)
+        derived_all_id_tmpfile = tmp_mgr.make_temp_file(all_id_content)
+        download_all_id_tmpfile = tmp_mgr.make_temp_file(all_id_content)
+
+        id_index_content = id_index_txt()
+        id_index_tmpfile = tmp_mgr.make_temp_file(id_index_content)
+        derived_id_index_tmpfile = tmp_mgr.make_temp_file(id_index_content)
+        download_id_index_tmpfile = tmp_mgr.make_temp_file(id_index_content)
+
+        id_abstracts_content = id_index_txt(with_abstracts=True)
+        id_abstracts_tmpfile = tmp_mgr.make_temp_file(id_abstracts_content)
+        derived_id_abstracts_tmpfile = tmp_mgr.make_temp_file(id_abstracts_content)
+        download_id_abstracts_tmpfile = tmp_mgr.make_temp_file(id_abstracts_content)
+
+        all_id2_content = all_id2_txt()
+        all_id2_tmpfile = tmp_mgr.make_temp_file(all_id2_content)
+        derived_all_id2_tmpfile = tmp_mgr.make_temp_file(all_id2_content)
+
+        # Move temp files as-atomically-as-possible into place
+        tmp_mgr.move_into_place(all_id_tmpfile, id_path / "all_id.txt")
+        tmp_mgr.move_into_place(derived_all_id_tmpfile, derived_path / "all_id.txt")
+        tmp_mgr.move_into_place(download_all_id_tmpfile, download_path / "id-all.txt")
+
+        tmp_mgr.move_into_place(id_index_tmpfile, id_path / "1id-index.txt")
+        tmp_mgr.move_into_place(derived_id_index_tmpfile, derived_path / "1id-index.txt")
+        tmp_mgr.move_into_place(download_id_index_tmpfile, download_path / "id-index.txt")
+
+        tmp_mgr.move_into_place(id_abstracts_tmpfile, id_path / "1id-abstracts.txt")
+        tmp_mgr.move_into_place(derived_id_abstracts_tmpfile, derived_path / "1id-abstracts.txt")
+        tmp_mgr.move_into_place(download_id_abstracts_tmpfile, download_path / "id-abstract.txt")
+
+        tmp_mgr.move_into_place(all_id2_tmpfile, id_path / "all_id2.txt")
+        tmp_mgr.move_into_place(derived_all_id2_tmpfile, derived_path / "all_id2.txt")
diff --git a/ietf/idindex/tests.py b/ietf/idindex/tests.py
@@ -3,8 +3,10 @@
 
 
 import datetime
+import mock
 
 from pathlib import Path
+from tempfile import TemporaryDirectory
 
 from django.conf import settings
 from django.utils import timezone
@@ -16,6 +18,7 @@
 from ietf.group.factories import GroupFactory
 from ietf.name.models import DocRelationshipName
 from ietf.idindex.index import all_id_txt, all_id2_txt, id_index_txt
+from ietf.idindex.tasks import idindex_update_task, TempFileManager
 from ietf.person.factories import PersonFactory, EmailFactory
 from ietf.utils.test_utils import TestCase
 
@@ -151,3 +154,51 @@ def test_id_index_txt(self):
         txt = id_index_txt(with_abstracts=True)
 
         self.assertTrue(draft.abstract[:20] in txt)
+
+
+class TaskTests(TestCase):
+    @mock.patch("ietf.idindex.tasks.all_id_txt")
+    @mock.patch("ietf.idindex.tasks.all_id2_txt")
+    @mock.patch("ietf.idindex.tasks.id_index_txt")
+    @mock.patch.object(TempFileManager, "__enter__")
+    def test_idindex_update_task(
+        self,
+        temp_file_mgr_enter_mock,
+        id_index_mock,
+        all_id2_mock,
+        all_id_mock,
+    ):
+        # Replace TempFileManager's __enter__() method with one that returns a mock.
+        # Pass a spec to the mock so we validate that only actual methods are called.
+        mgr_mock = mock.Mock(spec=TempFileManager)
+        temp_file_mgr_enter_mock.return_value = mgr_mock
+
+        idindex_update_task()
+
+        self.assertEqual(all_id_mock.call_count, 1)
+        self.assertEqual(all_id2_mock.call_count, 1)
+        self.assertEqual(id_index_mock.call_count, 2)
+        self.assertEqual(id_index_mock.call_args_list[0], (tuple(), dict()))
+        self.assertEqual(
+            id_index_mock.call_args_list[1], 
+            (tuple(), {"with_abstracts": True}),
+        )
+        self.assertEqual(mgr_mock.make_temp_file.call_count, 11)
+        self.assertEqual(mgr_mock.move_into_place.call_count, 11)
+
+    def test_temp_file_manager(self):
+        with TemporaryDirectory() as temp_dir:
+            temp_path = Path(temp_dir)
+            with TempFileManager(temp_path) as tfm:
+                path1 = tfm.make_temp_file("yay")
+                path2 = tfm.make_temp_file("boo")  # do not keep this one
+                self.assertTrue(path1.exists())
+                self.assertTrue(path2.exists())
+                dest = temp_path / "yay.txt"
+                tfm.move_into_place(path1, dest)
+            # make sure things were cleaned up...
+            self.assertFalse(path1.exists())  # moved to dest
+            self.assertFalse(path2.exists())  # left behind
+            # check destination contents and permissions
+            self.assertEqual(dest.read_text(), "yay")
+            self.assertEqual(dest.stat().st_mode & 0o777, 0o644)