77from datetime import datetime
88
99import pytz
10+ from celery import states as celery_states
1011from django .conf import settings
1112from django .contrib .auth .base_user import AbstractBaseUser
1213from django .contrib .auth .base_user import BaseUserManager
7475from contentcuration .db .models .manager import CustomManager
7576from contentcuration .statistics import record_channel_stats
7677from contentcuration .utils .cache import delete_public_channel_cache_keys
78+ from contentcuration .utils .celery .tasks import generate_task_signature
7779from contentcuration .utils .parser import load_json_string
7880from contentcuration .viewsets .sync .constants import ALL_CHANGES
7981from contentcuration .viewsets .sync .constants import ALL_TABLES
@@ -2436,13 +2438,20 @@ def serialize_to_change_dict(self):
24362438class TaskResultCustom (object ):
24372439 """
24382440 Custom fields to add to django_celery_results's TaskResult model
2441+
2442+ If adding fields to this class, run `makemigrations` then move the generated migration from the
2443+ `django_celery_results` app to the `contentcuration` app and override the constructor to change
2444+ the app_label. See `0141_add_task_signature` for an example
24392445 """
24402446 # user shouldn't be null, but in order to append the field, this needs to be allowed
24412447 user = models .ForeignKey (settings .AUTH_USER_MODEL , related_name = "tasks" , on_delete = models .CASCADE , null = True )
24422448 channel_id = DjangoUUIDField (db_index = True , null = True , blank = True )
24432449 progress = models .IntegerField (null = True , blank = True , validators = [MinValueValidator (0 ), MaxValueValidator (100 )])
2450+ # a hash of the task name and kwargs for identifying repeat tasks
2451+ signature = models .CharField (null = True , blank = False , max_length = 32 )
24442452
24452453 super_as_dict = TaskResult .as_dict
2454+ super_save = TaskResult .save
24462455
24472456 def as_dict (self ):
24482457 """
@@ -2456,16 +2465,45 @@ def as_dict(self):
24562465 )
24572466 return super_dict
24582467
2468+ def set_signature (self ):
2469+ """
2470+ Generates and sets the signature for the task if it isn't set
2471+ """
2472+ if self .signature is not None :
2473+ # nothing to do
2474+ return
2475+ self .signature = generate_task_signature (self .task_name , task_kwargs = self .task_kwargs , channel_id = self .channel_id )
2476+
2477+ def save (self , * args , ** kwargs ):
2478+ """
2479+ Override save to ensure signature is generated
2480+ """
2481+ self .set_signature ()
2482+ return self .super_save (* args , ** kwargs )
2483+
24592484 @classmethod
24602485 def contribute_to_class (cls , model_class = TaskResult ):
24612486 """
24622487 Adds fields to model, by default TaskResult
24632488 :param model_class: TaskResult model
24642489 """
24652490 for field in dir (cls ):
2466- if not field .startswith ("_" ):
2491+ if not field .startswith ("_" ) and field not in ( 'contribute_to_class' , 'Meta' ) :
24672492 model_class .add_to_class (field , getattr (cls , field ))
24682493
2494+ # manually add Meta afterwards
2495+ setattr (model_class ._meta , 'indexes' , getattr (model_class ._meta , 'indexes' , []) + cls .Meta .indexes )
2496+
2497+ class Meta :
2498+ indexes = [
2499+ # add index that matches query usage for signature
2500+ models .Index (
2501+ fields = ['signature' ],
2502+ name = 'task_result_signature_idx' ,
2503+ condition = Q (status__in = celery_states .UNREADY_STATES ),
2504+ ),
2505+ ]
2506+
24692507
24702508# trigger class contributions immediately
24712509TaskResultCustom .contribute_to_class ()
0 commit comments