feat: update tenant user status in periodic task (#1993)

src/bk-user/bkuser/apps/tenant/tasks.py

@@ -16,7 +16,10 @@
 # to the current version of the project delivered to anyone in the future.
 import logging
 
-from bkuser.apps.tenant.models import CollaborationStrategy
+from django.utils import timezone
+
+from bkuser.apps.tenant.constants import TenantUserStatus
+from bkuser.apps.tenant.models import CollaborationStrategy, TenantUser
 from bkuser.celery import app
 from bkuser.common.task import BaseTask
 
@@ -41,3 +44,31 @@ def remove_dropped_field_in_collaboration_strategy_field_mapping(tenant_id: str,
             mp for mp in strategy.target_config["field_mapping"] if mp["target_field"] != field_name
         ]
         strategy.save(update_fields=["target_config", "updated_at"])
+
+
+@app.task(base=BaseTask, ignore_result=True)
+def update_expired_tenant_user_status():
+    """定时任务：批量更新过期用户的状态"""
+    logger.info("[celery] receive task: update_expired_tenant_user_status")
+
+    now = timezone.now()
+
+    expired_users = TenantUser.objects.filter(
+        status=TenantUserStatus.ENABLED,
+        account_expired_at__lte=now,
+    )
+
+    expired_count = expired_users.count()
+
+    if expired_count == 0:
+        logger.info("No expired users found.")
+        return
+
+    # Q: 为什么不直接使用 expired_users.update(...)
+    # A: 为避免 update 数据量过大，这里直接使用 bulk_update 支持 batch_size 分批量处理
+    for user in expired_users:
+        user.status = TenantUserStatus.EXPIRED
+        user.updated_at = now
+
+    TenantUser.objects.bulk_update(expired_users, fields=["status", "updated_at"], batch_size=500)
+    logger.info("Updated %d expired users to EXPIRED status.", expired_count)

src/bk-user/bkuser/settings.py

@@ -289,6 +289,10 @@
         "task": "bkuser.apps.sync.periodic_tasks.mark_running_sync_task_as_failed_if_exceed_one_day",
         "schedule": crontab(minute="0", hour="9"),
     },
+    "periodic_update_tenant_user_status": {
+        "task": "bkuser.apps.tenant.tasks.update_expired_tenant_user_status",
+        "schedule": crontab(minute="0", hour="3"),
+    },
 }
 # Celery 消息队列配置
 CELERY_BROKER_URL = env.str("BK_BROKER_URL", default="")

src/bk-user/tests/apps/tenant/conftest.py

@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+# TencentBlueKing is pleased to support the open source community by making
+# 蓝鲸智云 - 用户管理 (bk-user) available.
+# Copyright (C) 2017 THL A29 Limited, a Tencent company. All rights reserved.
+# Licensed under the MIT License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+#     http://opensource.org/licenses/MIT
+#
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+# either express or implied. See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# We undertake not to change the open source license (MIT license) applicable
+# to the current version of the project delivered to anyone in the future.
+from datetime import timedelta
+
+import pytest
+from bkuser.apps.data_source.models import DataSourceUser
+from bkuser.apps.tenant.constants import TenantUserStatus
+from bkuser.apps.tenant.models import TenantUser
+from bkuser.utils.time import get_midnight
+from django.utils import timezone
+
+from tests.test_utils.helpers import generate_random_string
+
+
+@pytest.fixture
+def not_expired_tenant_user(bare_local_data_source, random_tenant):
+    data_source_user = DataSourceUser.objects.create(
+        username=generate_random_string(length=8),
+        full_name=generate_random_string(length=8),
+        data_source=bare_local_data_source,
+    )
+    return TenantUser.objects.create(
+        id=generate_random_string(),
+        tenant=random_tenant,
+        data_source=bare_local_data_source,
+        data_source_user=data_source_user,
+        status=TenantUserStatus.ENABLED,
+        account_expired_at=timezone.now() + timedelta(days=1),
+    )
+
+
+@pytest.fixture
+def expired_tenant_user(bare_local_data_source, random_tenant):
+    data_source_user = DataSourceUser.objects.create(
+        username=generate_random_string(length=8),
+        full_name=generate_random_string(length=8),
+        data_source=bare_local_data_source,
+    )
+    return TenantUser.objects.create(
+        id=generate_random_string(),
+        tenant=random_tenant,
+        data_source=bare_local_data_source,
+        data_source_user=data_source_user,
+        status=TenantUserStatus.ENABLED,
+        account_expired_at=get_midnight() - timedelta(days=1),
+    )

src/bk-user/tests/apps/tenant/test_tasks.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+# TencentBlueKing is pleased to support the open source community by making
+# 蓝鲸智云 - 用户管理 (bk-user) available.
+# Copyright (C) 2017 THL A29 Limited, a Tencent company. All rights reserved.
+# Licensed under the MIT License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+#     http://opensource.org/licenses/MIT
+#
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+# either express or implied. See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# We undertake not to change the open source license (MIT license) applicable
+# to the current version of the project delivered to anyone in the future.
+import pytest
+from bkuser.apps.tenant.constants import TenantUserStatus
+from bkuser.apps.tenant.tasks import update_expired_tenant_user_status
+
+pytestmark = pytest.mark.django_db
+
+
+class TestUpdateExpiredTenantUserStatus:
+    def test_success(self, not_expired_tenant_user, expired_tenant_user):
+        update_expired_tenant_user_status()
+        not_expired_tenant_user.refresh_from_db()
+        expired_tenant_user.refresh_from_db()
+
+        assert not_expired_tenant_user.status == TenantUserStatus.ENABLED
+        assert expired_tenant_user.status == TenantUserStatus.EXPIRED