From c25856746865a0ea1e548807b7904c083e7c662a Mon Sep 17 00:00:00 2001 From: rolin999 Date: Tue, 3 Dec 2024 14:45:39 +0800 Subject: [PATCH] feat: update tenant user status in periodic task (#1993) --- src/bk-user/bkuser/apps/tenant/tasks.py | 33 +++++++++++- src/bk-user/bkuser/settings.py | 4 ++ src/bk-user/tests/apps/tenant/conftest.py | 60 +++++++++++++++++++++ src/bk-user/tests/apps/tenant/test_tasks.py | 31 +++++++++++ 4 files changed, 127 insertions(+), 1 deletion(-) create mode 100644 src/bk-user/tests/apps/tenant/conftest.py create mode 100644 src/bk-user/tests/apps/tenant/test_tasks.py diff --git a/src/bk-user/bkuser/apps/tenant/tasks.py b/src/bk-user/bkuser/apps/tenant/tasks.py index a934f19a0..b1f0cbeeb 100644 --- a/src/bk-user/bkuser/apps/tenant/tasks.py +++ b/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) diff --git a/src/bk-user/bkuser/settings.py b/src/bk-user/bkuser/settings.py index b3c96fbf6..16908d333 100644 --- a/src/bk-user/bkuser/settings.py +++ b/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="") diff --git a/src/bk-user/tests/apps/tenant/conftest.py b/src/bk-user/tests/apps/tenant/conftest.py new file mode 100644 index 000000000..84947a988 --- /dev/null +++ b/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), + ) diff --git a/src/bk-user/tests/apps/tenant/test_tasks.py b/src/bk-user/tests/apps/tenant/test_tasks.py new file mode 100644 index 000000000..aadcc59c7 --- /dev/null +++ b/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