mirror of
https://github.com/prowler-cloud/prowler.git
synced 2026-10-26 05:14:27 +00:00
48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
from celery import Celery, Task
|
|
|
|
celery_app = Celery("tasks")
|
|
|
|
celery_app.config_from_object("django.conf:settings", namespace="CELERY")
|
|
celery_app.conf.update(result_extended=True, result_expires=None)
|
|
|
|
celery_app.autodiscover_tasks(["api"])
|
|
|
|
|
|
class RLSTask(Task):
|
|
def apply_async(
|
|
self,
|
|
args=None,
|
|
kwargs=None,
|
|
task_id=None,
|
|
producer=None,
|
|
link=None,
|
|
link_error=None,
|
|
shadow=None,
|
|
**options,
|
|
):
|
|
from django_celery_results.models import TaskResult
|
|
|
|
from api.models import Task as APITask
|
|
|
|
result = super().apply_async(
|
|
args=args,
|
|
kwargs=kwargs,
|
|
task_id=task_id,
|
|
producer=producer,
|
|
link=link,
|
|
link_error=link_error,
|
|
shadow=shadow,
|
|
**options,
|
|
)
|
|
task_result_instance = TaskResult.objects.get(task_id=result.task_id)
|
|
from api.db_utils import rls_transaction
|
|
|
|
tenant_id = kwargs.get("tenant_id")
|
|
with rls_transaction(tenant_id):
|
|
APITask.objects.create(
|
|
id=task_result_instance.task_id,
|
|
tenant_id=tenant_id,
|
|
task_runner_task=task_result_instance,
|
|
)
|
|
return result
|