mirror of https://github.com/jumpserver/jumpserver
perf: Change the secret and retry in batches
parent
d1d6f3fe9c
commit
b72e8eba7c
|
@ -97,12 +97,13 @@ class ChangeSecretRecordViewSet(mixins.ListModelMixin, OrgGenericViewSet):
|
||||||
def execute(self, request, *args, **kwargs):
|
def execute(self, request, *args, **kwargs):
|
||||||
record_ids = request.data.get('record_ids')
|
record_ids = request.data.get('record_ids')
|
||||||
records = self.get_queryset().filter(id__in=record_ids)
|
records = self.get_queryset().filter(id__in=record_ids)
|
||||||
execution_count = records.values_list('execution_id', flat=True).distinct().count()
|
if not records.exists():
|
||||||
if execution_count != 1:
|
|
||||||
return Response(
|
return Response(
|
||||||
{'detail': 'Only one execution is allowed to execute'},
|
{'detail': 'No valid records found'},
|
||||||
status=status.HTTP_400_BAD_REQUEST
|
status=status.HTTP_400_BAD_REQUEST
|
||||||
)
|
)
|
||||||
|
|
||||||
|
record_ids = [str(_id) for _id in records.values_list('id', flat=True)]
|
||||||
task = execute_automation_record_task.delay(record_ids, self.tp)
|
task = execute_automation_record_task.delay(record_ids, self.tp)
|
||||||
return Response({'task': task.id}, status=status.HTTP_200_OK)
|
return Response({'task': task.id}, status=status.HTTP_200_OK)
|
||||||
|
|
||||||
|
|
|
@ -1,4 +1,5 @@
|
||||||
import datetime
|
import datetime
|
||||||
|
from collections import defaultdict
|
||||||
|
|
||||||
from celery import shared_task
|
from celery import shared_task
|
||||||
from django.db.models import Q
|
from django.db.models import Q
|
||||||
|
@ -72,24 +73,43 @@ def execute_automation_record_task(record_ids, tp):
|
||||||
task_name = gettext_noop('Execute automation record')
|
task_name = gettext_noop('Execute automation record')
|
||||||
|
|
||||||
with tmp_to_root_org():
|
with tmp_to_root_org():
|
||||||
records = ChangeSecretRecord.objects.filter(id__in=record_ids)
|
records = ChangeSecretRecord.objects.filter(id__in=record_ids).order_by('-date_updated')
|
||||||
|
|
||||||
if not records:
|
if not records:
|
||||||
logger.error('No automation record found: {}'.format(record_ids))
|
logger.error(f'No automation record found: {record_ids}')
|
||||||
return
|
return
|
||||||
|
|
||||||
record = records[0]
|
seen_accounts = set()
|
||||||
record_map = {f'{record.asset_id}-{record.account_id}': str(record.id) for record in records}
|
unique_records = []
|
||||||
task_snapshot = {
|
for rec in records:
|
||||||
'params': {},
|
acct = str(rec.account_id)
|
||||||
'record_map': record_map,
|
if acct not in seen_accounts:
|
||||||
'secret': record.new_secret,
|
seen_accounts.add(acct)
|
||||||
'secret_type': record.execution.snapshot.get('secret_type'),
|
unique_records.append(rec)
|
||||||
'assets': [str(instance.asset_id) for instance in records],
|
|
||||||
'accounts': [str(instance.account_id) for instance in records],
|
exec_groups = defaultdict(list)
|
||||||
}
|
for rec in unique_records:
|
||||||
with tmp_to_org(record.execution.org_id):
|
exec_groups[rec.execution_id].append(rec)
|
||||||
quickstart_automation_by_snapshot(task_name, tp, task_snapshot)
|
|
||||||
|
for __, group in exec_groups.items():
|
||||||
|
latest_rec = group[0]
|
||||||
|
snapshot = getattr(latest_rec.execution, 'snapshot', {}) or {}
|
||||||
|
|
||||||
|
record_map = {f"{r.asset_id}-{r.account_id}": str(r.id) for r in group}
|
||||||
|
assets = [str(r.asset_id) for r in group]
|
||||||
|
accounts = [str(r.account_id) for r in group]
|
||||||
|
|
||||||
|
task_snapshot = {
|
||||||
|
'params': {},
|
||||||
|
'record_map': record_map,
|
||||||
|
'secret': latest_rec.new_secret,
|
||||||
|
'secret_type': snapshot.get('secret_type'),
|
||||||
|
'assets': assets,
|
||||||
|
'accounts': accounts,
|
||||||
|
}
|
||||||
|
|
||||||
|
with tmp_to_org(latest_rec.execution.org_id):
|
||||||
|
quickstart_automation_by_snapshot(task_name, tp, task_snapshot)
|
||||||
|
|
||||||
|
|
||||||
@shared_task(
|
@shared_task(
|
||||||
|
|
Loading…
Reference in New Issue