refactor(accounts): simplify application credential switch workflow

2026-08-08 23:35:34 +00:00 · 2026-07-31 17:55:05 +08:00
parent 803f2e559b
commit 5018b975f0
9 changed files with 146 additions and 139 deletions
--- a/apps/accounts/api/account/application.py
+++ b/apps/accounts/api/account/application.py
@@ -1,7 +1,6 @@
 import os

 from django.conf import settings
-from django.db import transaction
 from django.utils.translation import gettext_lazy as _, get_language
 from rest_framework.decorators import action
 from rest_framework.response import Response
@@ -80,12 +79,7 @@ class IntegrationApplicationViewSet(OrgBulkModelViewSet):
    )
    def reset_agent(self, request, *args, **kwargs):
        instance = self.get_object()
-        with transaction.atomic():
-            instance = IntegrationApplication.objects.select_for_update().get(pk=instance.pk)
-            agent = getattr(instance, 'agent', None)
-            if agent:
-                agent.delete()
-            instance.refresh_secret()
+        instance.reset_agent()
        return Response(data={
            'id': instance.id,
            'msg': _('Agent reset and App Secret refreshed successfully'),
--- a/apps/accounts/api/account/application_agent.py
+++ b/apps/accounts/api/account/application_agent.py
@@ -1,34 +1,32 @@
 import json
 import time
+from textwrap import dedent

-from django.db import transaction
 from django.http import StreamingHttpResponse
 from django.utils.translation import gettext_lazy as _
 from rest_framework import mixins, permissions, status
 from rest_framework.decorators import action
-from rest_framework.renderers import JSONRenderer
 from rest_framework.response import Response

 from accounts import serializers
 from accounts.const import ApplicationAgentEventStatus
 from accounts.models import (
-    ApplicationAccountBinding, ApplicationAccountSwitch, IntegrationApplication,
-    IntegrationApplicationAgentEvent,
+    ApplicationAccountSwitch, IntegrationApplication, IntegrationApplicationAgentEvent,
 )
 from authentication.permissions import IsValidUser
 from common.api import JMSGenericViewSet
 from common.db.utils import close_old_connections
+from common.drf.renders import EventStreamRenderer
 from orgs.mixins.api import OrgGenericViewSet


-class EventStreamRenderer(JSONRenderer):
-    media_type = 'text/event-stream'
-    format = 'event-stream'
+EVENT_STREAM_TIMEOUT_SECONDS = 55
+EVENT_STREAM_POLL_INTERVAL_SECONDS = 5


 def stream_agent_events(application_id):
    sent = set()
-    deadline = time.monotonic() + 55
+    deadline = time.monotonic() + EVENT_STREAM_TIMEOUT_SECONDS
    while time.monotonic() < deadline:
        close_old_connections()
        events = IntegrationApplicationAgentEvent.objects.filter(
@@ -39,10 +37,15 @@ def stream_agent_events(application_id):
            if event.id in sent:
                continue
            data = json.dumps(event.as_payload())
-            yield f'id: {event.id}\nevent: credential.change\ndata: {data}\n\n'
+            yield dedent(f"""\
+                id: {event.id}
+                event: credential.change
+                data: {data}
+
+                """)
            sent.add(event.id)
        yield ': heartbeat\n\n'
-        time.sleep(5)
+        time.sleep(EVENT_STREAM_POLL_INTERVAL_SECONDS)


 class ApplicationAccountSwitchViewSet(
@@ -55,6 +58,7 @@ class ApplicationAccountSwitchViewSet(
        'default': serializers.ApplicationAccountSwitchSerializer,
        'create': serializers.ApplicationAccountSwitchCreateSerializer,
        'confirm': serializers.ApplicationAccountSwitchConfirmSerializer,
+        'credentials': serializers.ApplicationAccountCredentialSerializer,
    }
    filterset_fields = ['source_account_id', 'target_account_id', 'status']
    search_fields = [
@@ -71,37 +75,9 @@ class ApplicationAccountSwitchViewSet(

    @action(['GET'], detail=False)
    def credentials(self, request, *args, **kwargs):
-        bindings = ApplicationAccountBinding.objects.filter(
-            application__agent__isnull=False,
-            application__is_active=True,
-            current_account__is_active=True,
-        ).select_related('application', 'current_account__asset').order_by(
-            'current_account__name', 'application__name'
-        )
-        credentials = {}
-        for binding in bindings:
-            account = binding.current_account
-            credential = credentials.setdefault(str(account.id), {
-                'account': {
-                    'id': str(account.id),
-                    'name': account.name,
-                    'username': account.username,
-                },
-                'asset': {
-                    'id': str(account.asset_id),
-                    'name': account.asset.name,
-                    'address': account.asset.address,
-                },
-                'bindings': [],
-            })
-            credential['bindings'].append({
-                'credential_id': str(binding.id),
-                'application': {
-                    'id': str(binding.application_id),
-                    'name': binding.application.name,
-                },
-            })
-        return Response(list(credentials.values()))
+        queryset = self.get_serializer_class().get_queryset()
+        serializer = self.get_serializer(queryset, many=True)
+        return Response(serializer.data)

    def create(self, request, *args, **kwargs):
        serializer = self.get_serializer(data=request.data)
@@ -148,9 +124,8 @@ class IntegrationApplicationAgentViewSet(JMSGenericViewSet):
    def register(self, request):
        serializer = self.get_serializer(data=request.data)
        serializer.is_valid(raise_exception=True)
-        with transaction.atomic():
-            agent = serializer.save()
-            data = serializers.AgentRegisterResultSerializer(agent).data
+        agent = serializer.save()
+        data = serializers.AgentRegisterResultSerializer(agent).data
        return Response(data)

    @action(['POST'], detail=False)
@@ -158,7 +133,7 @@ class IntegrationApplicationAgentViewSet(JMSGenericViewSet):
        serializer = self.get_serializer(data=request.data)
        serializer.is_valid(raise_exception=True)
        agent = serializer.save()
-        return Response({'server_time': agent.last_seen})
+        return Response({'server_time': agent.date_last_used})

    @action(['GET'], detail=False)
    def credentials(self, request):
@@ -170,21 +145,18 @@ class IntegrationApplicationAgentViewSet(JMSGenericViewSet):
    def reports(self, request):
        serializer = self.get_serializer(data=request.data)
        serializer.is_valid(raise_exception=True)
-        event = serializer.save()
-        return Response({
-            'event_id': event.id,
-            'status': event.status,
-            'switch_status': event.item.switch.status,
-            'duplicate': serializer.duplicate,
-        })
+        serializer.save()
+        return Response(serializer.data)

    @action(['GET'], detail=False, renderer_classes=(EventStreamRenderer,))
    def events(self, request):
        serializer = self.get_serializer(data=request.query_params)
        serializer.is_valid(raise_exception=True)
-        response = StreamingHttpResponse(
-            stream_agent_events(request.user.id), content_type='text/event-stream'
+        return StreamingHttpResponse(
+            stream_agent_events(request.user.id),
+            content_type='text/event-stream',
+            headers={
+                'Cache-Control': 'no-cache',
+                'X-Accel-Buffering': 'no',
+            },
        )
-        response['Cache-Control'] = 'no-cache'
-        response['X-Accel-Buffering'] = 'no'
-        return response
--- a/apps/accounts/migrations/0009_integrationapplication_agent_error_and_more.py
+++ b/apps/accounts/migrations/0009_integrationapplication_agent_error_and_more.py
@@ -32,7 +32,7 @@ class Migration(migrations.Migration):
                ('hostname', models.CharField(blank=True, max_length=255, verbose_name='Hostname')),
                ('platform', models.CharField(blank=True, max_length=64, verbose_name='Platform')),
                ('version', models.CharField(blank=True, max_length=64, verbose_name='Version')),
-                ('last_seen', models.DateTimeField(blank=True, null=True, verbose_name='Last seen')),
+                ('date_last_used', models.DateTimeField(blank=True, null=True, verbose_name='Date last used')),
                ('error', models.TextField(blank=True, verbose_name='Error')),
                ('application', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, related_name='agent', to='accounts.integrationapplication', verbose_name='Application')),
            ],
--- a/apps/accounts/models/application.py
+++ b/apps/accounts/models/application.py
@@ -36,27 +36,6 @@ class IntegrationApplication(JMSOrgBaseModel):
        query = RelatedManager.get_to_filter_qs(self.accounts.value, Account)
        return qs.filter(*query)

-    def get_managed_accounts(self):
-        if (self.accounts.value or {}).get('type') != 'ids':
-            return Account.objects.none()
-        return self.get_accounts()
-
-    def sync_account_bindings(self):
-        from accounts.models import ApplicationAccountBinding
-
-        account_ids = set(self.get_managed_accounts().values_list('id', flat=True))
-        bindings = self.account_bindings.all()
-        bindings.exclude(current_account_id__in=account_ids).delete()
-        existing_ids = set(bindings.values_list('current_account_id', flat=True))
-        ApplicationAccountBinding.objects.bulk_create([
-            ApplicationAccountBinding(
-                org_id=self.org_id,
-                application=self,
-                current_account_id=account_id,
-            )
-            for account_id in account_ids - existing_ids
-        ])
-
    @property
    def accounts_amount(self) -> int:
        return self.get_accounts().count()
@@ -79,6 +58,12 @@ class IntegrationApplication(JMSOrgBaseModel):
        self.save(update_fields=['secret'])
        return self.secret

+    def reset_agent(self):
+        agent = getattr(self, 'agent', None)
+        if agent:
+            agent.delete()
+        self.refresh_secret()
+
    def get_account(self, asset='', asset_id='', account='', account_id=''):
        qs = Account.objects.all()
        if account_id:
--- a/apps/accounts/models/application_agent.py
+++ b/apps/accounts/models/application_agent.py
@@ -1,6 +1,6 @@
 from datetime import timedelta

-from django.db import models, transaction
+from django.db import models
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 from rest_framework.exceptions import ValidationError
@@ -9,9 +9,11 @@ from accounts.const import (
    ApplicationAgentEventStatus, ApplicationAgentEventType, ApplicationAgentStatus,
    ApplicationSwitchItemStatus, ApplicationSwitchStatus,
 )
-from common.const.signals import OP_LOG_SKIP_SIGNAL
 from orgs.mixins.models import JMSOrgBaseModel

+from .account import Account
+from .application import IntegrationApplication
+

 ACTIVE_SWITCH_STATUSES = (
    ApplicationSwitchStatus.RUNNING,
@@ -28,7 +30,9 @@ class IntegrationApplicationAgent(JMSOrgBaseModel):
    hostname = models.CharField(max_length=255, blank=True, verbose_name=_('Hostname'))
    platform = models.CharField(max_length=64, blank=True, verbose_name=_('Platform'))
    version = models.CharField(max_length=64, blank=True, verbose_name=_('Version'))
-    last_seen = models.DateTimeField(null=True, blank=True, verbose_name=_('Last seen'))
+    date_last_used = models.DateTimeField(
+        null=True, blank=True, verbose_name=_('Date last used')
+    )
    error = models.TextField(blank=True, verbose_name=_('Error'))

    class Meta:
@@ -38,27 +42,26 @@ class IntegrationApplicationAgent(JMSOrgBaseModel):
    def status(self):
        if self.error:
            return ApplicationAgentStatus.ERROR
-        if not self.last_seen or self.last_seen < timezone.now() - timedelta(seconds=90):
+        if (
+            not self.date_last_used or
+            self.date_last_used < timezone.now() - timedelta(seconds=90)
+        ):
            return ApplicationAgentStatus.OFFLINE
        return ApplicationAgentStatus.ONLINE

-    @transaction.atomic
    def touch(self, error=None):
-        application = self.application.__class__.objects.select_for_update().get(
-            pk=self.application_id
-        )
-        agent = self.__class__.objects.select_for_update().get(pk=self.pk)
-        agent.last_seen = timezone.now()
-        fields = ['last_seen', 'date_updated']
+        now = timezone.now()
+        updates = {'date_last_used': now, 'date_updated': now}
        if error is not None:
-            agent.error = error
-            fields.append('error')
-        application.date_last_used = agent.last_seen
-        setattr(application, OP_LOG_SKIP_SIGNAL, True)
-        application.save(update_fields=['date_last_used', 'date_updated'])
-        setattr(agent, OP_LOG_SKIP_SIGNAL, True)
-        agent.save(update_fields=fields)
-        return agent
+            updates['error'] = error
+            self.error = error
+        IntegrationApplicationAgent.objects.filter(pk=self.pk).update(**updates)
+        IntegrationApplication.objects.filter(pk=self.application_id).update(
+            date_last_used=now, date_updated=now
+        )
+        self.date_last_used = now
+        self.date_updated = now
+        return self


 class ApplicationAccountBinding(JMSOrgBaseModel):
@@ -66,6 +69,7 @@ class ApplicationAccountBinding(JMSOrgBaseModel):
        'accounts.IntegrationApplication', on_delete=models.CASCADE,
        related_name='account_bindings', verbose_name=_('Application')
    )
+    # Do not silently remove a credential still used by an application.
    current_account = models.ForeignKey(
        'accounts.Account', on_delete=models.PROTECT,
        related_name='application_bindings', verbose_name=_('Current account')
@@ -76,22 +80,38 @@ class ApplicationAccountBinding(JMSOrgBaseModel):
        ordering = ['application__name', 'current_account__name']
        verbose_name = _('Application account binding')

-    @transaction.atomic
+    @classmethod
+    def sync_application(cls, application):
+        accounts = application.accounts.value or {}
+        account_ids = accounts.get('ids', []) if accounts.get('type') == 'ids' else []
+        account_ids = set(Account.objects.filter(
+            id__in=account_ids
+        ).values_list('id', flat=True))
+        application.account_bindings.exclude(
+            current_account_id__in=account_ids
+        ).delete()
+        cls.objects.bulk_create([
+            cls(
+                org_id=application.org_id,
+                application=application,
+                current_account_id=account_id,
+            )
+            for account_id in account_ids
+        ], ignore_conflicts=True)
+
    def move_to(self, account):
-        binding = self.__class__.objects.select_for_update().get(pk=self.pk)
+        application = IntegrationApplication.objects.select_for_update().get(
+            pk=self.application_id
+        )
+        binding = ApplicationAccountBinding.objects.get(pk=self.pk)
        if binding.current_account_id == account.pk:
            return binding
-        if self.__class__.objects.filter(
+        if ApplicationAccountBinding.objects.filter(
            application_id=binding.application_id, current_account=account
        ).exclude(pk=binding.pk).exists():
            raise ValidationError(_(
                'The target account is already bound to this application.'
            ))
-
-        application_model = self._meta.get_field('application').remote_field.model
-        application = application_model.objects.select_for_update().get(
-            pk=binding.application_id
-        )
        accounts = dict(application.accounts.value or {})
        account_ids = list(accounts.get('ids') or [])
        source_id = str(binding.current_account_id)
@@ -104,7 +124,7 @@ class ApplicationAccountBinding(JMSOrgBaseModel):
            target_id if account_id == source_id else account_id
            for account_id in account_ids
        ]
-        application.accounts.set(accounts)
+        application.accounts = accounts
        application.save(update_fields=['accounts', 'date_updated'])
        binding.current_account = account
        binding.save(update_fields=['current_account', 'date_updated'])
@@ -131,11 +151,7 @@ class ApplicationAccountSwitch(JMSOrgBaseModel):
        verbose_name = _('Application account switch')

    @classmethod
-    @transaction.atomic
    def start(cls, source_account, target_account, user, comment=''):
-        source_account = source_account.__class__.objects.select_for_update().get(
-            pk=source_account.pk
-        )
        cls._validate_accounts(source_account, target_account)
        bindings = list(
            cls.get_affected_bindings(source_account).select_for_update().select_related(
@@ -236,7 +252,6 @@ class ApplicationAccountSwitch(JMSOrgBaseModel):
            'status', 'updated_by', 'date_finished', 'date_updated'
        ])

-    @transaction.atomic
    def rollback(self, user):
        switch = self._lock_active(_('Only an active switch task can be rolled back.'))
        switch.status = ApplicationSwitchStatus.ROLLING_BACK
@@ -247,7 +262,6 @@ class ApplicationAccountSwitch(JMSOrgBaseModel):
            item.prepare_rollback()
        return switch

-    @transaction.atomic
    def end(self, user):
        switch = self._lock_active(_('The switch task is already finished.'))
        switch.updated_by = str(user)
@@ -256,7 +270,7 @@ class ApplicationAccountSwitch(JMSOrgBaseModel):
        return switch

    def _lock_active(self, message):
-        switch = self.__class__.objects.select_for_update().get(pk=self.pk)
+        switch = ApplicationAccountSwitch.objects.select_for_update().get(pk=self.pk)
        if switch.status not in ACTIVE_SWITCH_STATUSES:
            raise ValidationError(message)
        return switch
@@ -296,10 +310,9 @@ class ApplicationAccountSwitchItem(JMSOrgBaseModel):
    def application(self):
        return self.binding.application

-    @transaction.atomic
    def confirm(self, user):
        switch = ApplicationAccountSwitch.objects.select_for_update().get(pk=self.switch_id)
-        item = self.__class__.objects.select_for_update().get(pk=self.pk)
+        item = ApplicationAccountSwitchItem.objects.get(pk=self.pk)
        item.switch = switch
        status, account = item._confirmation_result()
        item.binding.move_to(account)
@@ -352,12 +365,11 @@ class IntegrationApplicationAgentEvent(JMSOrgBaseModel):
        ordering = ['date_created']
        verbose_name = _('Integration application Agent event')

-    @transaction.atomic
    def report(self, success, error=''):
        switch = ApplicationAccountSwitch.objects.select_for_update().get(
            pk=self.item.switch_id
        )
-        event = self.__class__.objects.select_for_update().select_related(
+        event = IntegrationApplicationAgentEvent.objects.select_related(
            'item__binding__application', 'item'
        ).get(pk=self.pk)
        event.item.switch = switch
--- a/apps/accounts/serializers/account/application_agent.py
+++ b/apps/accounts/serializers/account/application_agent.py
@@ -1,5 +1,6 @@
 from django.conf import settings
 from django.db import IntegrityError, transaction
+from django.db.models import Prefetch
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 from rest_framework import serializers
@@ -49,6 +50,7 @@ def account_credential(account, credential_id):

 class IntegrationApplicationAgentSerializer(CommonModelSerializer):
    status = serializers.SerializerMethodField(label=_('Status'))
+    last_seen = serializers.DateTimeField(source='date_last_used', read_only=True)

    class Meta:
        model = IntegrationApplicationAgent
@@ -77,6 +79,41 @@ class ApplicationAccountBindingSerializer(CommonModelSerializer):
        fields = ['id', 'account', 'asset']


+class ApplicationAccountCredentialBindingSerializer(serializers.Serializer):
+    credential_id = serializers.UUIDField(source='id', read_only=True)
+    application = ObjectRelatedField(
+        read_only=True, attrs=('id', 'name')
+    )
+
+
+class ApplicationAccountCredentialSerializer(serializers.Serializer):
+    account = ObjectRelatedField(
+        source='*', read_only=True, attrs=('id', 'name', 'username')
+    )
+    asset = ObjectRelatedField(
+        read_only=True, attrs=('id', 'name', 'address')
+    )
+    bindings = ApplicationAccountCredentialBindingSerializer(
+        source='agent_application_bindings', many=True, read_only=True
+    )
+
+    @staticmethod
+    def get_queryset():
+        bindings = ApplicationAccountBinding.objects.filter(
+            application__agent__isnull=False,
+            application__is_active=True,
+            current_account__is_active=True,
+        ).select_related('application').order_by('application__name')
+        return Account.objects.filter(
+            id__in=bindings.values('current_account_id')
+        ).select_related('asset').prefetch_related(
+            Prefetch(
+                'application_bindings', queryset=bindings,
+                to_attr='agent_application_bindings',
+            )
+        ).order_by('name')
+
+
 class ApplicationAccountSwitchItemSerializer(CommonModelSerializer):
    credential_id = serializers.UUIDField(source='binding_id', read_only=True)
    application = ObjectRelatedField(
@@ -186,7 +223,6 @@ class AgentRegisterSerializer(serializers.Serializer):
    platform = serializers.CharField(max_length=64, allow_blank=True, required=False)
    version = serializers.CharField(max_length=64, allow_blank=True, required=False)

-    @transaction.atomic
    def create(self, validated_data):
        authenticated_application = self.context['request'].user
        application = IntegrationApplication.objects.select_for_update().get(
@@ -214,7 +250,7 @@ class AgentRegisterSerializer(serializers.Serializer):
            )
        for field in ('hostname', 'platform', 'version'):
            setattr(agent, field, validated_data.get(field, ''))
-        agent.last_seen = timezone.now()
+        agent.date_last_used = timezone.now()
        agent.error = ''
        try:
            with transaction.atomic():
@@ -223,7 +259,7 @@ class AgentRegisterSerializer(serializers.Serializer):
            raise serializers.ValidationError({
                'agent_id': _('This Agent ID is already registered.')
            })
-        application.date_last_used = agent.last_seen
+        application.date_last_used = agent.date_last_used
        application.save(update_fields=['date_last_used', 'date_updated'])
        return agent

@@ -321,6 +357,14 @@ class AgentEventReportSerializer(AgentEventSerializerMixin, AgentIdentitySeriali
            self.notify_owner(event, validated_data['success'])
        return event

+    def to_representation(self, event):
+        return {
+            'event_id': event.id,
+            'status': event.status,
+            'switch_status': event.item.switch.status,
+            'duplicate': self.duplicate,
+        }
+
    @staticmethod
    def notify_owner(event, success):
        result = _('succeeded') if success else _('failed')
--- a/apps/accounts/serializers/account/service.py
+++ b/apps/accounts/serializers/account/service.py
@@ -1,17 +1,18 @@
-from django.db import transaction
 from django.templatetags.static import static
 from django.utils.translation import gettext_lazy as _
 from rest_framework import serializers

 from accounts.const import ApplicationAgentStatus, ApplicationSwitchStatus
-from accounts.models import Account, ApplicationAccountSwitch, IntegrationApplication
+from accounts.models import (
+    Account, ApplicationAccountBinding, ApplicationAccountSwitch,
+    IntegrationApplication,
+)
 from acls.serializers.rules import ip_group_child_validator, ip_group_help_text
 from common.db.fields import RelatedManager
 from common.serializers.fields import JSONManyToManyField
 from common.serializers.fields import ObjectRelatedField
 from common.utils import random_string
 from orgs.mixins.serializers import BulkOrgResourceModelSerializer
-from users import utils as user_utils
 from users.models import User

 from .application_agent import (
@@ -54,10 +55,6 @@ class IntegrationApplicationSerializer(BulkOrgResourceModelSerializer):
            data['logo'] = static('img/logo.png')
        return data

-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.fields['owner'].queryset = user_utils.get_current_org_members()
-
    @classmethod
    def setup_eager_loading(cls, queryset):
        return queryset.select_related('owner', 'agent').prefetch_related(
@@ -111,17 +108,15 @@ class IntegrationApplicationSerializer(BulkOrgResourceModelSerializer):
            })
        return attrs

-    @transaction.atomic
    def create(self, validated_data):
        instance = super().create(validated_data)
-        instance.sync_account_bindings()
+        ApplicationAccountBinding.sync_application(instance)
        instance.refresh_secret()
        return instance

-    @transaction.atomic
    def update(self, instance, validated_data):
        instance = super().update(instance, validated_data)
-        instance.sync_account_bindings()
+        ApplicationAccountBinding.sync_application(instance)
        return instance


--- a/apps/accounts/tests/init.py
+++ b/apps/accounts/tests/init.py
--- a/apps/common/drf/renders/init.py
+++ b/apps/common/drf/renders/init.py
@@ -13,3 +13,8 @@ class PassthroughRenderer(renderers.BaseRenderer):

    def render(self, data, accepted_media_type=None, renderer_context=None):
        return data
+
+
+class EventStreamRenderer(renderers.JSONRenderer):
+    media_type = 'text/event-stream'
+    format = 'event-stream'