Spaces:
Running
Running
| """Account for every accepted Celery delivery, including retries and chains.""" | |
| import uuid | |
| from celery import Task, current_task | |
| from . import maintenance as gate | |
| from .state import enabled | |
| from .git import BackupError | |
| HEADER = 'brainapi-backup-work' | |
| class ContentTask(Task): | |
| abstract = True | |
| def apply_async(self, args=None, kwargs=None, **options): | |
| if not enabled(): | |
| return super().apply_async(args=args, kwargs=kwargs, **options) | |
| parent = getattr(getattr(current_task, 'request', None), 'headers', None) or {} | |
| work_id = uuid.uuid4().hex | |
| parent_id = getattr(getattr(current_task, 'request', None), 'backup_execution', None) or parent.get(HEADER) | |
| gate.register_task(work_id, parent_id) | |
| headers = dict(options.pop('headers', None) or {}) | |
| headers[HEADER] = work_id | |
| # An ambiguous publication failure deliberately stays pending: fail closed | |
| # rather than claim a durable barrier while a message might be in flight. | |
| return super().apply_async(args=args, kwargs=kwargs, headers=headers, **options) | |
| def __call__(self, *args, **kwargs): | |
| if not enabled(): | |
| return super().__call__(*args, **kwargs) | |
| work_id = (self.request.headers or {}).get(HEADER) | |
| self.request.backup_execution = gate.begin_task(work_id) | |
| token = gate._admitted.set(True) | |
| try: | |
| return super().__call__(*args, **kwargs) | |
| finally: | |
| gate._admitted.reset(token) | |
| def after_return(self, status, retval, task_id, args, kwargs, einfo): | |
| gate.finish_task((self.request.headers or {}).get(HEADER), getattr(self.request, 'backup_execution', None)) | |
| return super().after_return(status, retval, task_id, args, kwargs, einfo) | |
| def on_retry(self, exc, task_id, args, kwargs, einfo): | |
| gate.finish_task((self.request.headers or {}).get(HEADER), getattr(self.request, 'backup_execution', None)) | |
| return super().on_retry(exc, task_id, args, kwargs, einfo) | |