feat: LDAP HA

apps/settings/tasks/ldap.py

@@ -1,5 +1,4 @@
 # coding: utf-8
-#
 import time
 from celery import shared_task
 from django.conf import settings
@@ -16,45 +15,44 @@ from settings.notifications import LDAPImportMessage
 from users.models import User
 from ..utils import LDAPSyncUtil, LDAPServerUtil, LDAPImportUtil
 
-__all__ = ['sync_ldap_user', 'import_ldap_user_periodic', 'import_ldap_user']
+__all__ = [
+    'sync_ldap_user', 'import_ldap_user_periodic', 'import_ldap_ha_user_periodic',
+    'import_ldap_user', 'import_ldap_ha_user'
+]
 
 logger = get_logger(__file__)
 
 
-def sync_ldap_user():
-    LDAPSyncUtil().perform_sync()
+def sync_ldap_user(category='ldap'):
+    LDAPSyncUtil(category=category).perform_sync()
 
 
-@shared_task(
-    verbose_name=_('Periodic import ldap user'),
-    description=_(
-        """
-        When LDAP auto-sync is configured, this task will be invoked to synchronize users
-        """
-    )
-)
-def import_ldap_user():
+def perform_import(category, util_server):
     start_time = time.time()
     time_start_display = local_now_display()
-    logger.info("Start import ldap user task")
-    util_server = LDAPServerUtil()
+    logger.info(f"Start import {category} ldap user task")
+
     util_import = LDAPImportUtil()
     users = util_server.search()
+
     if settings.XPACK_ENABLED:
-        org_ids = settings.AUTH_LDAP_SYNC_ORG_IDS
+        org_ids = getattr(settings, f"AUTH_{category.upper()}_SYNC_ORG_IDS")
         default_org = None
     else:
-        # 社区版默认导入Default组织
         org_ids = [Organization.DEFAULT_ID]
         default_org = Organization.default()
+
     orgs = list(set([Organization.get_instance(org_id, default=default_org) for org_id in org_ids]))
     new_users, errors = util_import.perform_import(users, orgs)
+
     if errors:
-        logger.error("Imported LDAP users errors: {}".format(errors))
+        logger.error(f"Imported {category} LDAP users errors: {errors}")
     else:
-        logger.info('Imported {} users successfully'.format(len(users)))
-    if settings.AUTH_LDAP_SYNC_RECEIVERS:
-        user_ids = settings.AUTH_LDAP_SYNC_RECEIVERS
+        logger.info(f"Imported {len(users)} {category} users successfully")
+
+    receivers_setting = f"AUTH_{category.upper()}_SYNC_RECEIVERS"
+    if getattr(settings, receivers_setting, None):
+        user_ids = getattr(settings, receivers_setting)
         recipient_list = User.objects.filter(id__in=list(user_ids))
         end_time = time.time()
         extra_kwargs = {
@@ -70,6 +68,54 @@ def import_ldap_user():
             LDAPImportMessage(user, extra_kwargs).publish()
 
 
+@shared_task(
+    verbose_name=_('Periodic import ldap user'),
+    description=_(
+        """
+        When LDAP auto-sync is configured, this task will be invoked to synchronize users
+        """
+    )
+)
+def import_ldap_user():
+    perform_import('ldap', LDAPServerUtil())
+
+
+@shared_task(
+    verbose_name=_('Periodic import ldap ha user'),
+    description=_(
+        """
+        When LDAP auto-sync is configured, this task will be invoked to synchronize users
+        """
+    )
+)
+def import_ldap_ha_user():
+    perform_import('ldap_ha', LDAPServerUtil(category='ldap_ha'))
+
+
+def register_periodic_task(task_name, task_func, interval_key, enabled_key, crontab_key, **kwargs):
+    interval = kwargs.get(interval_key, settings.AUTH_LDAP_SYNC_INTERVAL)
+    enabled = kwargs.get(enabled_key, settings.AUTH_LDAP_SYNC_IS_PERIODIC)
+    crontab = kwargs.get(crontab_key, settings.AUTH_LDAP_SYNC_CRONTAB)
+
+    if isinstance(interval, int):
+        interval = interval * 3600
+    else:
+        interval = None
+
+    if crontab:
+        interval = None  # 优先使用 crontab
+
+    tasks = {
+        task_name: {
+            'task': task_func.name,
+            'interval': interval,
+            'crontab': crontab,
+            'enabled': enabled
+        }
+    }
+    create_or_update_celery_periodic_tasks(tasks)
+
+
 @shared_task(
     verbose_name=_('Registration periodic import ldap user task'),
     description=_(
@@ -81,23 +127,26 @@ def import_ldap_user():
 )
 @after_app_ready_start
 def import_ldap_user_periodic(**kwargs):
-    task_name = 'import_ldap_user_periodic'
-    interval = kwargs.get('AUTH_LDAP_SYNC_INTERVAL', settings.AUTH_LDAP_SYNC_INTERVAL)
-    enabled = kwargs.get('AUTH_LDAP_SYNC_IS_PERIODIC', settings.AUTH_LDAP_SYNC_IS_PERIODIC)
-    crontab = kwargs.get('AUTH_LDAP_SYNC_CRONTAB', settings.AUTH_LDAP_SYNC_CRONTAB)
-    if isinstance(interval, int):
-        interval = interval * 3600
-    else:
-        interval = None
-    if crontab:
-        # 优先使用 crontab
-        interval = None
-    tasks = {
-        task_name: {
-            'task': import_ldap_user.name,
-            'interval': interval,
-            'crontab': crontab,
-            'enabled': enabled
-        }
-    }
-    create_or_update_celery_periodic_tasks(tasks)
+    register_periodic_task(
+        'import_ldap_user_periodic', import_ldap_user,
+        'AUTH_LDAP_SYNC_INTERVAL', 'AUTH_LDAP_SYNC_IS_PERIODIC',
+        'AUTH_LDAP_SYNC_CRONTAB', **kwargs
+    )
+
+
+@shared_task(
+    verbose_name=_('Registration periodic import ldap ha user task'),
+    description=_(
+        """
+        When LDAP HA auto-sync parameters change, such as Crontab parameters, the LDAP HA sync task 
+        will be re-registered or updated, and this task will be invoked
+        """
+    )
+)
+@after_app_ready_start
+def import_ldap_ha_user_periodic(**kwargs):
+    register_periodic_task(
+        'import_ldap_ha_user_periodic', import_ldap_ha_user,
+        'AUTH_LDAP_HA_SYNC_INTERVAL', 'AUTH_LDAP_HA_SYNC_IS_PERIODIC',
+        'AUTH_LDAP_HA_SYNC_CRONTAB', **kwargs
+    )