Skip to content

Commit b4cf04a

Browse files
jennifer-richardsjennifer-richards
authored
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
1 parent 118b00d commit b4cf04a

File tree

9 files changed

+526
-23
lines changed

9 files changed

+526
-23
lines changed

‎bin/hourly

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ ID=/a/ietfdata/doc/draft/repository
4545
DERIVED=/a/ietfdata/derived
4646
DOWNLOAD=/a/www/www6s/download
4747

48+
## Start of script refactored into idindex_update_task() ===
4849
export TMPDIR=/a/tmp
4950

5051
TMPFILE1=`mktemp` || exit 1
@@ -85,6 +86,8 @@ mv $TMPFILE9 $DERIVED/1id-index.txt
8586
mv $TMPFILEA $DERIVED/1id-abstracts.txt
8687
mv $TMPFILEB $DERIVED/all_id2.txt
8788

89+
## End of script refactored into idindex_update_task() ===
90+
8891
$DTDIR/ietf/manage.py generate_idnits2_rfc_status
8992
$DTDIR/ietf/manage.py generate_idnits2_rfcs_obsoleted
9093

‎ietf/doc/tasks.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
# Copyright The IETF Trust 2024, All Rights Reserved
2+
#
3+
# Celery task definitions
4+
#
5+
import datetime
6+
import debug # pyflakes:ignore
7+
8+
from celery import shared_task
9+
10+
from ietf.utils import log
11+
from ietf.utils.timezone import datetime_today
12+
13+
from .expire import (
14+
in_draft_expire_freeze,
15+
get_expired_drafts,
16+
expirable_drafts,
17+
send_expire_notice_for_draft,
18+
expire_draft,
19+
clean_up_draft_files,
20+
get_soon_to_expire_drafts,
21+
send_expire_warning_for_draft,
22+
)
23+
from .models import Document
24+
25+
26+
@shared_task
27+
def expire_ids_task():
28+
try:
29+
if not in_draft_expire_freeze():
30+
log.log("Expiring drafts ...")
31+
for doc in get_expired_drafts():
32+
# verify expirability -- it might have changed after get_expired_drafts() was run
33+
# (this whole loop took about 2 minutes on 04 Jan 2018)
34+
# N.B., re-running expirable_drafts() repeatedly is fairly expensive. Where possible,
35+
# it's much faster to run it once on a superset query of the objects you are going
36+
# to test and keep its results. That's not desirable here because it would defeat
37+
# the purpose of double-checking that a document is still expirable when it is actually
38+
# being marked as expired.
39+
if expirable_drafts(
40+
Document.objects.filter(pk=doc.pk)
41+
).exists() and doc.expires < datetime_today() + datetime.timedelta(1):
42+
send_expire_notice_for_draft(doc)
43+
expire_draft(doc)
44+
log.log(f" Expired draft {doc.name}-{doc.rev}")
45+
46+
log.log("Cleaning up draft files")
47+
clean_up_draft_files()
48+
except Exception as e:
49+
log.log("Exception in expire-ids: %s" % e)
50+
raise
51+
52+
53+
@shared_task
54+
def notify_expirations_task(notify_days=14):
55+
for doc in get_soon_to_expire_drafts(notify_days):
56+
send_expire_warning_for_draft(doc)

‎ietf/doc/tests_tasks.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
# Copyright The IETF Trust 2024, All Rights Reserved
2+
import mock
3+
4+
from ietf.utils.test_utils import TestCase
5+
from ietf.utils.timezone import datetime_today
6+
7+
from .factories import DocumentFactory
8+
from .models import Document
9+
from .tasks import expire_ids_task, notify_expirations_task
10+
11+
12+
class TaskTests(TestCase):
13+
14+
@mock.patch("ietf.doc.tasks.in_draft_expire_freeze")
15+
@mock.patch("ietf.doc.tasks.get_expired_drafts")
16+
@mock.patch("ietf.doc.tasks.expirable_drafts")
17+
@mock.patch("ietf.doc.tasks.send_expire_notice_for_draft")
18+
@mock.patch("ietf.doc.tasks.expire_draft")
19+
@mock.patch("ietf.doc.tasks.clean_up_draft_files")
20+
def test_expire_ids_task(
21+
self,
22+
clean_up_draft_files_mock,
23+
expire_draft_mock,
24+
send_expire_notice_for_draft_mock,
25+
expirable_drafts_mock,
26+
get_expired_drafts_mock,
27+
in_draft_expire_freeze_mock,
28+
):
29+
# set up mocks
30+
in_draft_expire_freeze_mock.return_value = False
31+
doc, other_doc = DocumentFactory.create_batch(2)
32+
doc.expires = datetime_today()
33+
get_expired_drafts_mock.return_value = [doc, other_doc]
34+
expirable_drafts_mock.side_effect = [
35+
Document.objects.filter(pk=doc.pk),
36+
Document.objects.filter(pk=other_doc.pk),
37+
]
38+
39+
# call task
40+
expire_ids_task()
41+
42+
# check results
43+
self.assertTrue(in_draft_expire_freeze_mock.called)
44+
self.assertEqual(expirable_drafts_mock.call_count, 2)
45+
self.assertEqual(send_expire_notice_for_draft_mock.call_count, 1)
46+
self.assertEqual(send_expire_notice_for_draft_mock.call_args[0], (doc,))
47+
self.assertEqual(expire_draft_mock.call_count, 1)
48+
self.assertEqual(expire_draft_mock.call_args[0], (doc,))
49+
self.assertTrue(clean_up_draft_files_mock.called)
50+
51+
# test that an exception is raised
52+
in_draft_expire_freeze_mock.side_effect = RuntimeError
53+
with self.assertRaises(RuntimeError):(
54+
expire_ids_task())
55+
56+
@mock.patch("ietf.doc.tasks.send_expire_warning_for_draft")
57+
@mock.patch("ietf.doc.tasks.get_soon_to_expire_drafts")
58+
def test_notify_expirations_task(self, get_drafts_mock, send_warning_mock):
59+
# Set up mocks
60+
get_drafts_mock.return_value = ["sentinel"]
61+
notify_expirations_task()
62+
self.assertEqual(send_warning_mock.call_count, 1)
63+
self.assertEqual(send_warning_mock.call_args[0], ("sentinel",))

‎ietf/idindex/tasks.py

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
# Copyright The IETF Trust 2024, All Rights Reserved
2+
#
3+
# Celery task definitions
4+
#
5+
import shutil
6+
7+
import debug # pyflakes:ignore
8+
9+
from celery import shared_task
10+
from contextlib import AbstractContextManager
11+
from pathlib import Path
12+
from tempfile import NamedTemporaryFile
13+
14+
from .index import all_id_txt, all_id2_txt, id_index_txt
15+
16+
17+
class TempFileManager(AbstractContextManager):
18+
def __init__(self, tmpdir=None) -> None:
19+
self.cleanup_list: set[Path] = set()
20+
self.dir = tmpdir
21+
22+
def make_temp_file(self, content):
23+
with NamedTemporaryFile(mode="wt", delete=False, dir=self.dir) as tf:
24+
tf_path = Path(tf.name)
25+
self.cleanup_list.add(tf_path)
26+
tf.write(content)
27+
return tf_path
28+
29+
def move_into_place(self, src_path: Path, dest_path: Path):
30+
shutil.move(src_path, dest_path)
31+
dest_path.chmod(0o644)
32+
self.cleanup_list.remove(src_path)
33+
34+
def cleanup(self):
35+
for tf_path in self.cleanup_list:
36+
tf_path.unlink(missing_ok=True)
37+
38+
def __exit__(self, exc_type, exc_val, exc_tb):
39+
self.cleanup()
40+
return False # False: do not suppress the exception
41+
42+
43+
@shared_task
44+
def idindex_update_task():
45+
"""Update I-D indexes"""
46+
id_path = Path("/a/ietfdata/doc/draft/repository")
47+
derived_path = Path("/a/ietfdata/derived")
48+
download_path = Path("/a/www/www6s/download")
49+
50+
with TempFileManager("/a/tmp") as tmp_mgr:
51+
# Generate copies of new contents
52+
all_id_content = all_id_txt()
53+
all_id_tmpfile = tmp_mgr.make_temp_file(all_id_content)
54+
derived_all_id_tmpfile = tmp_mgr.make_temp_file(all_id_content)
55+
download_all_id_tmpfile = tmp_mgr.make_temp_file(all_id_content)
56+
57+
id_index_content = id_index_txt()
58+
id_index_tmpfile = tmp_mgr.make_temp_file(id_index_content)
59+
derived_id_index_tmpfile = tmp_mgr.make_temp_file(id_index_content)
60+
download_id_index_tmpfile = tmp_mgr.make_temp_file(id_index_content)
61+
62+
id_abstracts_content = id_index_txt(with_abstracts=True)
63+
id_abstracts_tmpfile = tmp_mgr.make_temp_file(id_abstracts_content)
64+
derived_id_abstracts_tmpfile = tmp_mgr.make_temp_file(id_abstracts_content)
65+
download_id_abstracts_tmpfile = tmp_mgr.make_temp_file(id_abstracts_content)
66+
67+
all_id2_content = all_id2_txt()
68+
all_id2_tmpfile = tmp_mgr.make_temp_file(all_id2_content)
69+
derived_all_id2_tmpfile = tmp_mgr.make_temp_file(all_id2_content)
70+
71+
# Move temp files as-atomically-as-possible into place
72+
tmp_mgr.move_into_place(all_id_tmpfile, id_path / "all_id.txt")
73+
tmp_mgr.move_into_place(derived_all_id_tmpfile, derived_path / "all_id.txt")
74+
tmp_mgr.move_into_place(download_all_id_tmpfile, download_path / "id-all.txt")
75+
76+
tmp_mgr.move_into_place(id_index_tmpfile, id_path / "1id-index.txt")
77+
tmp_mgr.move_into_place(derived_id_index_tmpfile, derived_path / "1id-index.txt")
78+
tmp_mgr.move_into_place(download_id_index_tmpfile, download_path / "id-index.txt")
79+
80+
tmp_mgr.move_into_place(id_abstracts_tmpfile, id_path / "1id-abstracts.txt")
81+
tmp_mgr.move_into_place(derived_id_abstracts_tmpfile, derived_path / "1id-abstracts.txt")
82+
tmp_mgr.move_into_place(download_id_abstracts_tmpfile, download_path / "id-abstract.txt")
83+
84+
tmp_mgr.move_into_place(all_id2_tmpfile, id_path / "all_id2.txt")
85+
tmp_mgr.move_into_place(derived_all_id2_tmpfile, derived_path / "all_id2.txt")

‎ietf/idindex/tests.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,10 @@
33

44

55
import datetime
6+
import mock
67

78
from pathlib import Path
9+
from tempfile import TemporaryDirectory
810

911
from django.conf import settings
1012
from django.utils import timezone
@@ -16,6 +18,7 @@
1618
from ietf.group.factories import GroupFactory
1719
from ietf.name.models import DocRelationshipName
1820
from ietf.idindex.index import all_id_txt, all_id2_txt, id_index_txt
21+
from ietf.idindex.tasks import idindex_update_task, TempFileManager
1922
from ietf.person.factories import PersonFactory, EmailFactory
2023
from ietf.utils.test_utils import TestCase
2124

@@ -151,3 +154,51 @@ def test_id_index_txt(self):
151154
txt = id_index_txt(with_abstracts=True)
152155

153156
self.assertTrue(draft.abstract[:20] in txt)
157+
158+
159+
class TaskTests(TestCase):
160+
@mock.patch("ietf.idindex.tasks.all_id_txt")
161+
@mock.patch("ietf.idindex.tasks.all_id2_txt")
162+
@mock.patch("ietf.idindex.tasks.id_index_txt")
163+
@mock.patch.object(TempFileManager, "__enter__")
164+
def test_idindex_update_task(
165+
self,
166+
temp_file_mgr_enter_mock,
167+
id_index_mock,
168+
all_id2_mock,
169+
all_id_mock,
170+
):
171+
# Replace TempFileManager's __enter__() method with one that returns a mock.
172+
# Pass a spec to the mock so we validate that only actual methods are called.
173+
mgr_mock = mock.Mock(spec=TempFileManager)
174+
temp_file_mgr_enter_mock.return_value = mgr_mock
175+
176+
idindex_update_task()
177+
178+
self.assertEqual(all_id_mock.call_count, 1)
179+
self.assertEqual(all_id2_mock.call_count, 1)
180+
self.assertEqual(id_index_mock.call_count, 2)
181+
self.assertEqual(id_index_mock.call_args_list[0], (tuple(), dict()))
182+
self.assertEqual(
183+
id_index_mock.call_args_list[1],
184+
(tuple(), {"with_abstracts": True}),
185+
)
186+
self.assertEqual(mgr_mock.make_temp_file.call_count, 11)
187+
self.assertEqual(mgr_mock.move_into_place.call_count, 11)
188+
189+
def test_temp_file_manager(self):
190+
with TemporaryDirectory() as temp_dir:
191+
temp_path = Path(temp_dir)
192+
with TempFileManager(temp_path) as tfm:
193+
path1 = tfm.make_temp_file("yay")
194+
path2 = tfm.make_temp_file("boo") # do not keep this one
195+
self.assertTrue(path1.exists())
196+
self.assertTrue(path2.exists())
197+
dest = temp_path / "yay.txt"
198+
tfm.move_into_place(path1, dest)
199+
# make sure things were cleaned up...
200+
self.assertFalse(path1.exists()) # moved to dest
201+
self.assertFalse(path2.exists()) # left behind
202+
# check destination contents and permissions
203+
self.assertEqual(dest.read_text(), "yay")
204+
self.assertEqual(dest.stat().st_mode & 0o777, 0o644)

0 commit comments

Comments
 (0)