from datetime import timedelta, timezone from collections import defaultdict import requests from celery.utils.log import get_task_logger from django.conf import settings from django.utils import timezone as django_timezone from pingcap_sso.celery import app as celery_app from pingcap_sso.models import User, UserProfileCollections, UserAwardedPoint from pingcap_sso.lib.points import ( AWARD_PE_PASS_PCTP, AWARD_PE_PASS_PCTA, AWARD_PE_PASS_PCSD, AWARD_PE_PASS_PRACZH, AWARD_PE_COMPLETE_101, AWARD_PE_COMPLETE_201, AWARD_PE_COMPLETE_301, AWARD_PE_COMPLETE_302, AWARD_PE_COMPLETE_304, AWARD_PE_COMPLETE_305, AWARD_PE_COMPLETE_307, AWARD_PE_COMPLETE_V2_960001, AWARD_PE_COMPLETE_V2_1050001, AWARD_PE_COMPLETE_V2_750002, AWARD_PE_COMPLETE_V2_840002, AWARD_PE_COMPLETE_V2_780002, AWARD_PE_COMPLETE_V2_1110001, award_points_by_type, ) from pingcap_sso.discourse import api as discourse_api logger = get_task_logger(__name__) def get_gender_by_chinese_idcard(number): if not number: return None sex_flag = int(number[-2]) if sex_flag % 2 == 0: return '女' else: return '男' def _iter_pe_sidecar_users(): """Paginate through all users from the PE Sidecar API.""" url = f'{settings.PE_SIDECAR_API_URL}/api/v1/users' headers = {'Authorization': f'Bearer {settings.PE_SIDECAR_API_TOKEN}'} page = 1 page_size = 500 while True: resp = requests.get(url, params={'page': page, 'page_size': page_size}, headers=headers, timeout=30) resp.raise_for_status() data = resp.json() items = data['items'] if not items: break yield from items if page * page_size >= data['total']: break page += 1 @celery_app.task def sync_profile_from_pe_to_sso(): for item in _iter_pe_sidecar_users(): idcard_number = item.get('idcard_number') full_name = item.get('full_name') attributes = item.get('attributes') or {} openid = attributes.get('openid') if not openid: continue if idcard_number is None and full_name is None: continue try: user = User.objects.get(openid=openid) except User.DoesNotExist: logger.error('User %s does not exist', openid) continue upc, _ = UserProfileCollections.objects.get_or_create( user=user, source=UserProfileCollections.SOURCE_PE_PROFILE ) upc.name = full_name or attributes.get('name') or None upc.gender = get_gender_by_chinese_idcard(idcard_number) company_changed = upc.company != attributes.get('company') upc.company = attributes.get('company') upc.email = attributes.get('mail-address') upc.tidb_use_experience = attributes.get('use-experience') upc.tidb_use_stage = attributes.get('tidb-stage') upc.job_title = attributes.get('position') upc.address = attributes.get('city') update_fields=[ "name", "gender", "company", "email", "tidb_use_experience", "tidb_use_stage", "job_title", "address", ] if company_changed: update_fields.append("updated_at") upc.save(update_fields=update_fields) def _iter_pe_sidecar_certificates(created_date_from, created_date_to): """Paginate through certificates from the PE Sidecar API within a date range.""" url = f'{settings.PE_SIDECAR_API_URL}/api/v1/certificates' headers = {'Authorization': f'Bearer {settings.PE_SIDECAR_API_TOKEN}'} page = 1 page_size = 500 while True: resp = requests.get(url, params={ 'page': page, 'page_size': page_size, 'created_date_from': created_date_from.isoformat(), 'created_date_to': created_date_to.isoformat(), }, headers=headers, timeout=30) resp.raise_for_status() data = resp.json() items = data['items'] if not items: break yield from items if page * page_size >= data['total']: break page += 1 @celery_app.task def sync_certificate_from_pe_to_sso(start_at=None, end_at=None): if not (start_at and end_at): now = django_timezone.now().astimezone( timezone(timedelta(hours=8)) ) start_at = ( now - timedelta(days=1) ).replace( hour=0, minute=0, second=0, microsecond=0 ) end_at = start_at + timedelta(days=1) grant_user_badges = defaultdict(list) for item in _iter_pe_sidecar_certificates(start_at, end_at): cert_code = (item.get('cert_code') or '').upper() if cert_code.startswith('PCTA'): badge_id = 115 award_type = AWARD_PE_PASS_PCTA elif cert_code.startswith('PCTP'): badge_id = 116 award_type = AWARD_PE_PASS_PCTP elif cert_code.startswith('PCSD'): badge_id = 155 award_type = AWARD_PE_PASS_PCSD elif cert_code.startswith('PRACZH'): badge_id = 176 award_type = AWARD_PE_PASS_PRACZH else: continue openid = item.get('openid') if not openid: logger.error('PE certificate %s has no openid', cert_code) continue try: sso_user = User.objects.get(openid=openid) except User.DoesNotExist: logger.error('PE certificate %s, user with openid %s not found in sso', cert_code, openid) continue award_points_by_type(sso_user, award_type) grant_user_badges[badge_id].append(sso_user.username) if award_type == AWARD_PE_PASS_PCTA: User.objects.filter(id=sso_user.id).update(pass_pcta=True) elif award_type == AWARD_PE_PASS_PCTP: User.objects.filter(id=sso_user.id).update(pass_pctp=True) elif award_type == AWARD_PE_PASS_PCSD: User.objects.filter(id=sso_user.id).update(pass_pcsd=True) elif award_type == AWARD_PE_PASS_PRACZH: User.objects.filter(id=sso_user.id).update(pass_praczh=True) for badge_id, usernames in grant_user_badges.items(): discourse_api.grant_badge_to_users(usernames, badge_id) PE_COURSES_WITH_REWARDS = { '101': {'offering_course_id': 6, 'award_type': AWARD_PE_COMPLETE_101}, '201': {'offering_course_id': 600003, 'award_type': AWARD_PE_COMPLETE_201}, '301': {'offering_course_id': 30002, 'award_type': AWARD_PE_COMPLETE_301}, '302': {'offering_course_id': 120005, 'award_type': AWARD_PE_COMPLETE_302}, '304': {'offering_course_id': 570012, 'award_type': AWARD_PE_COMPLETE_304}, '305': {'offering_course_id': 690007, 'award_type': AWARD_PE_COMPLETE_305}, '307': {'offering_course_id': 540005, 'award_type': AWARD_PE_COMPLETE_307}, 'v2_960001': {'offering_course_id': 960001, 'award_type': AWARD_PE_COMPLETE_V2_960001}, 'v2_1050001': {'offering_course_id': 1050001, 'award_type': AWARD_PE_COMPLETE_V2_1050001}, 'v2_750002': {'offering_course_id': 750002, 'award_type': AWARD_PE_COMPLETE_V2_750002}, 'v2_840002': {'offering_course_id': 840002, 'award_type': AWARD_PE_COMPLETE_V2_840002}, 'v2_780002': {'offering_course_id': 780002, 'award_type': AWARD_PE_COMPLETE_V2_780002}, 'v2_1110001': {'offering_course_id': 1110001, 'award_type': AWARD_PE_COMPLETE_V2_1110001}, } @celery_app.task def award_completed_pe_course_users(): url = f'{settings.PE_SIDECAR_API_URL}/api/v1/course-completions' headers = {'Authorization': f'Bearer {settings.PE_SIDECAR_API_TOKEN}'} for name, c in PE_COURSES_WITH_REWARDS.items(): page = 1 while True: resp = requests.get(url, params={ 'offering_course_id': c['offering_course_id'], 'page': page, 'page_size': 100, }, headers=headers, timeout=30) resp.raise_for_status() data = resp.json() if not data['items']: break openids = [item['openid'] for item in data['items']] users = User.objects.filter(openid__in=openids).exclude( id__in=UserAwardedPoint.objects.filter( award_type=c['award_type'] ).values('target_user_id') ) for u in users: award_points_by_type(u, c['award_type']) logger.info("User %s awarded %s for course %s", u.username, c['award_type'], name) if page * 100 >= data['total']: break page += 1