async def training(domain_id: str): _logger.info('Training') pipeline = [{ '$match': { 'domain_id': domain_id, 'doc_type': document.TYPE_TRAINING } }, { '$group': { '_id': '$doc_id', 'enroll': { '$sum': '$enroll' } } }] coll = db.Collection('document') await coll.update_many( { 'domain_id': domain_id, 'doc_type': document.TYPE_TRAINING }, {'$set': { 'enroll': 0 }}) bulk = coll.initialize_unordered_bulk_op() execute = False _logger.info('Counting') async for adoc in db.Collection('document.status').aggregate(pipeline): bulk.find({'domain_id': domain_id, 'doc_type': document.TYPE_TRAINING, 'doc_id': adoc['_id']}) \ .update_one({'$set': {'enroll': adoc['enroll']}}) execute = True if execute: _logger.info('Committing') await bulk.execute()
async def discussion(domain_id: str): _logger.info('Discussion') pipeline = [{ '$match': { 'domain_id': domain_id, 'doc_type': document.TYPE_DISCUSSION_REPLY } }, { '$group': { '_id': '$parent_doc_id', 'num_replies': { '$sum': 1 } } }] coll = db.Collection('document') await coll.update_many( { 'domain_id': domain_id, 'doc_type': document.TYPE_DISCUSSION }, {'$set': { 'num_replies': 0 }}) bulk = coll.initialize_unordered_bulk_op() execute = False _logger.info('Counting') async for adoc in db.Collection('document').aggregate(pipeline): bulk.find({'domain_id': domain_id, 'doc_type': document.TYPE_DISCUSSION, 'doc_id': adoc['_id']}) \ .update_one({'$set': {'num_replies': adoc['num_replies']}}) execute = True if execute: _logger.info('Committing') await bulk.execute()
async def sync_usage(): _logger.info('Userfile length group by user') pipeline = [ { '$match': {'domain_id': userfile.STORE_DOMAIN_ID, 'doc_type': document.TYPE_USERFILE} }, { '$group': { '_id': '$owner_uid', 'usage_userfile': {'$sum': '$length'} } } ] coll = db.Collection('domain.user') await coll.update_many({'domain_id': userfile.STORE_DOMAIN_ID}, {'$set': {'usage_userfile': 0}}) bulk = coll.initialize_unordered_bulk_op() execute = False _logger.info('Counting') async for adoc in db.Collection('document').aggregate(pipeline): bulk.find({'domain_id': userfile.STORE_DOMAIN_ID, 'uid': adoc['_id']}) \ .update_one({'$set': {'usage_userfile': adoc['usage_userfile']}}) execute = True if execute: _logger.info('Committing') await bulk.execute()
async def delete_roles(domain_id: str, roles): roles = list(set(roles)) for role in roles: validator.check_role(role) for domain in builtin.DOMAINS: if domain['_id'] == domain_id: raise error.BuiltinDomainError(domain_id) user_coll = db.Collection('domain.user') await user_coll.update_many( { 'domain_id': domain_id, 'role': { '$in': list(roles) } }, {'$unset': { 'role': '' }}) coll = db.Collection('domain') return await coll.find_one_and_update(filter={'_id': domain_id}, update={ '$unset': dict(('roles.{0}'.format(role), '') for role in roles) }, return_document=ReturnDocument.AFTER)
async def problem(domain_id: str): _logger.info('Problem') pipeline = [{ '$match': { 'domain_id': domain_id, 'doc_type': document.TYPE_PROBLEM } }, { '$group': { '_id': '$owner_uid', 'num_problems': { '$sum': 1 } } }] user_coll = db.Collection('domain.user') await user_coll.update_many({'domain_id': domain_id}, {'$set': { 'num_problems': 0 }}) user_coll = user_coll.initialize_unordered_bulk_op() execute = False _logger.info('Counting') async for adoc in db.Collection('document').aggregate(pipeline): user_coll.find({'domain_id': domain_id, 'uid': adoc['_id']}) \ .upsert().update_one({'$set': {'num_problems': adoc['num_problems']}}) execute = True if execute: _logger.info('Committing') await user_coll.execute()
async def ensure_indexes(): coll = db.Collection('domain') await coll.create_index('owner_uid') user_coll = db.Collection('domain.user') await user_coll.create_index('uid') await user_coll.create_index([('domain_id', 1), ('uid', 1)], unique=True) await user_coll.create_index([('domain_id', 1), ('role', 1)], sparse=True) await user_coll.create_index([('domain_id', 1), ('rp', -1)]) await user_coll.create_index([('domain_id', 1), ('rank', 1)])
async def ensure_indexes(): coll = db.Collection('document') await coll.create_index([('domain_id', 1), ('doc_type', 1), ('doc_id', 1)], unique=True) await coll.create_index([('domain_id', 1), ('doc_type', 1), ('owner_uid', 1), ('doc_id', -1)]) # for problem solution await coll.create_index([('domain_id', 1), ('doc_type', 1), ('parent_doc_type', 1), ('parent_doc_id', 1), ('vote', -1), ('doc_id', -1)], sparse=True) # for discussion await coll.create_index([('domain_id', 1), ('doc_type', 1), ('parent_doc_type', 1), ('parent_doc_id', 1), ('update_at', -1), ('doc_id', -1)], sparse=True) # hidden doc await coll.create_index([('domain_id', 1), ('doc_type', 1), ('hidden', 1), ('doc_id', -1)], sparse=True) # for contest await coll.create_index([('domain_id', 1), ('doc_type', 1), ('pids', 1)], sparse=True) await coll.create_index([('domain_id', 1), ('doc_type', 1), ('rule', 1), ('doc_id', -1)], sparse=True) await coll.create_index([('domain_id', 1), ('doc_type', 1), ('pids', 1)], sparse=True) # for training await coll.create_index([('domain_id', 1), ('doc_type', 1), ('dag.pids', 1)], sparse=True) status_coll = db.Collection('document.status') await status_coll.create_index([('domain_id', 1), ('doc_type', 1), ('uid', 1), ('doc_id', 1)], unique=True) # for rp system await status_coll.create_index([('domain_id', 1), ('doc_type', 1), ('doc_id', 1), ('status', 1), ('rid', 1), ('rp', 1)], sparse=True) # for contest rule OI await status_coll.create_index([('domain_id', 1), ('doc_type', 1), ('doc_id', 1), ('score', -1)], sparse=True) # for contest rule ACM await status_coll.create_index([('domain_id', 1), ('doc_type', 1), ('doc_id', 1), ('accept', -1), ('time', 1)], sparse=True) # for training await status_coll.create_index([('domain_id', 1), ('doc_type', 1), ('uid', 1), ('enroll', 1), ('doc_id', 1)], sparse=True)
async def recalc(domain_id: str): user_coll = db.Collection('domain.user') await user_coll.update_many({'domain_id': domain_id}, {'$set': { 'rp': 0.0 }}) pdocs = problem.get_multi(domain_id=domain_id, fields={ '_id': 1, 'doc_id': 1, 'num_accept': 1 }).sort('doc_id', 1) dudoc_updates = {} status_coll = db.Collection('document.status') async for pdoc in pdocs: _logger.info('Problem {0}'.format(pdoc['doc_id'])) psdocs = problem.get_multi_status( domain_id=domain_id, doc_id=pdoc['doc_id'], status=constant.record.STATUS_ACCEPTED).sort('rid', 1) order = 0 rp_func = get_rp_func(pdoc) status_bulk = status_coll.initialize_unordered_bulk_op() async for psdoc in psdocs: order += 1 rp = rp_func(order) status_bulk.find({ '_id': psdoc['_id'] }).update_one({'$set': { 'rp': rp }}) if psdoc['uid'] not in dudoc_updates: dudoc_updates[psdoc['uid']] = {'rp': rp} else: dudoc_updates[psdoc['uid']]['rp'] += rp if order != pdoc['num_accept']: _logger.warning('{0} != {1}'.format(order, pdoc['num_accept'])) _logger.warning( 'Problem {0} num_accept may be inconsistent.'.format( pdoc['doc_id'])) if order > 0: _logger.info('Committing') await status_bulk.execute() # users' rp user_bulk = user_coll.initialize_unordered_bulk_op() execute = False _logger.info('Updating users') for uid, dudoc_update in dudoc_updates.items(): execute = True user_bulk.find({'domain_id': domain_id, 'uid': uid}) \ .upsert().update_one({'$set': dudoc_update}) if execute: _logger.info('Committing') await user_bulk.execute()
async def main(): pipeline = [ { '$match': { 'hidden': False, 'type': constant.record.TYPE_SUBMISSION } }, { '$group': { '_id': { 'domain_id': '$domain_id', 'pid': '$pid', 'uid': '$uid' }, 'num_submit': { '$sum': 1 }, 'num_accept': { '$sum': { '$cond': [{ '$eq': ['$status', constant.record.STATUS_ACCEPTED] }, 1, 0] } } } }, { '$group': { '_id': { 'domain_id': '$_id.domain_id', 'pid': '$_id.pid' }, 'num_submit': { '$sum': '$num_submit' }, 'num_accept': { '$sum': { '$min': ['$num_accept', 1] } } } }, ] bulk = db.Collection('document').initialize_unordered_bulk_op() async for adoc in db.Collection('record').aggregate(pipeline): bulk.find({'domain_id': adoc['_id']['domain_id'], 'doc_type': document.TYPE_PROBLEM, 'doc_id': adoc['_id']['pid']}) \ .update_one({'$set': {'num_submit': adoc['num_submit'], 'num_accept': adoc['num_accept']}}) await bulk.execute()
async def ensure_indexes(): coll = db.Collection('document') await coll.ensure_index([('domain_id', 1), ('doc_type', 1), ('doc_id', 1)], unique=True) await coll.ensure_index([('domain_id', 1), ('doc_type', 1), ('parent_doc_type', 1), ('parent_doc_id', 1), ('vote', -1), ('doc_id', -1)], sparse=True) status_coll = db.Collection('document.status') await status_coll.ensure_index([('domain_id', 1), ('doc_type', 1), ('uid', 1), ('doc_id', 1)], unique=True)
async def ensure_indexes(): coll = db.Collection('record') await coll.create_index([('hidden', 1), ('_id', -1)]) await coll.create_index([('hidden', 1), ('uid', 1), ('_id', -1)]) await coll.create_index([('hidden', 1), ('domain_id', 1), ('pid', 1), ('uid', 1), ('_id', -1)]) await coll.create_index([('hidden', 1), ('domain_id', 1), ('tid', 1), ('pid', 1), ('uid', 1), ('_id', -1)], sparse=True) # for job record await coll.create_index([('domain_id', 1), ('pid', 1), ('type', 1), ('_id', 1)]) await coll.create_index([('domain_id', 1), ('pid', 1), ('uid', 1), ('type', 1), ('_id', 1)])
async def capped_inc_status(domain_id: str, doc_type: int, doc_id: convert_doc_id, uid: int, key: str, value: int, min_value: int = -1, max_value: int = 1): assert value != 0 if value > 0: not_expr = {'$gte': max_value} else: not_expr = {'$lte': min_value} coll = db.Collection('document.status') doc = await coll.find_one_and_update(filter={ 'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id, 'uid': uid, key: { '$not': not_expr } }, update={'$inc': { key: value }}, upsert=True, return_document=ReturnDocument.AFTER) return doc
async def inc(op: str, ident: str, period_secs: int, max_operations: int, operations: int = 1): coll = db.Collection('opcount') cur_time = int(time.time()) begin_at = datetime.datetime.utcfromtimestamp(cur_time - cur_time % period_secs) expire_at = begin_at + datetime.timedelta(seconds=period_secs) try: doc = await coll.find_one_and_update( filter={ 'ident': ident, 'begin_at': begin_at, 'expire_at': expire_at, op: { '$not': { '$gte': max_operations } } }, update={'$inc': { op: operations }}, upsert=True, return_document=ReturnDocument.AFTER) return doc except errors.DuplicateKeyError: raise error.OpcountExceededError(op, period_secs, max_operations)
async def rejudge(record_id: objectid.ObjectId, enqueue: bool = True): coll = db.Collection('record') doc = await coll.find_one_and_update( filter={'_id': record_id}, update={ '$unset': { 'judge_uid': '', 'judge_token': '', 'judge_at': '', 'compiler_texts': '', 'judge_texts': '', 'cases': '' }, '$set': { 'status': constant.record.STATUS_WAITING, 'score': 0, 'time_ms': 0, 'memory_kb': 0, 'rejudged': True } }, return_document=ReturnDocument.BEFORE) post_coros = [bus.publish('record_change', doc['_id'])] if enqueue: post_coros.append(queue.publish('judge', rid=doc['_id'])) await asyncio.gather(*post_coros)
async def add(uid: int, type: int, **kwargs): """Add an operation log. Returns the document id.""" obj_id = objectid.ObjectId() coll = db.Collection('oplog') doc = {'_id': obj_id, 'uid': uid, 'type': type, **kwargs} await coll.insert_one(doc) return obj_id
async def unset(domain_id, fields): # TODO(twd2): check fields coll = db.Collection('domain') return await coll.find_one_and_update( filter={'_id': domain_id}, update={'$unset': dict((f, '') for f in set(fields))}, return_document=ReturnDocument.AFTER)
async def get(domain_id: str, doc_type: int, doc_id: convert_doc_id): coll = db.Collection('document') return await coll.find_one({ 'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id })
async def add(domain_id: str, content: str, owner_uid: int, doc_type: int, doc_id: convert_doc_id = None, parent_doc_type: int = None, parent_doc_id: convert_doc_id = None, **kwargs): """Add a document. Returns the document id.""" obj_id = objectid.ObjectId() coll = db.Collection('document') doc = { '_id': obj_id, 'content': content, 'owner_uid': owner_uid, 'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id or obj_id, **kwargs } if parent_doc_type or parent_doc_id: assert parent_doc_type and parent_doc_id doc['parent_doc_type'], doc[ 'parent_doc_id'] = parent_doc_type, parent_doc_id await coll.insert_one(doc) return doc['doc_id']
async def add(domain_id: str, pid: document.convert_doc_id, type: int, uid: int, lang: str, code: str, data_id: objectid.ObjectId=None, tid: objectid.ObjectId=None, hidden=False): validator.check_lang(lang) coll = db.Collection('record') rid = (await coll.insert_one({'hidden': hidden, 'status': constant.record.STATUS_WAITING, 'score': 0, 'time_ms': 0, 'memory_kb': 0, 'domain_id': domain_id, 'pid': pid, 'uid': uid, 'lang': lang, 'code': code, 'tid': tid, 'data_id': data_id, 'type': type})).inserted_id post_coros = [queue.publish('judge', rid=rid), bus.publish('record_change', rid)] if type == constant.record.TYPE_SUBMISSION: post_coros.extend([problem.inc_status(domain_id, pid, uid, 'num_submit', 1), problem.inc(domain_id, pid, 'num_submit', 1), domain.inc_user(domain_id, uid, num_submit=1)]) await asyncio.gather(*post_coros) return rid
async def ensure_indexes(): coll = db.Collection('system') await coll.find_one_and_update(filter={'_id': 'user_counter'}, update={'$setOnInsert': { 'value': 1 }}, upsert=True)
async def update(token_id: str, token_type: int, expire_seconds: int, **kwargs): """Update a token. Args: token_id: token ID. token_type: type of the token. expire_seconds: expire time, in seconds. **kwargs: extra data. Returns: The token document, or None. """ id_binary = binascii.unhexlify(token_id) coll = db.Collection('token') assert 'token_type' not in kwargs now = datetime.datetime.utcnow() doc = await coll.find_and_modify( query={ '_id': _get_id(id_binary), 'token_type': token_type }, update={ '$set': { **kwargs, 'update_at': now, 'expire_at': now + datetime.timedelta(seconds=expire_seconds) } }, new=True) return doc
async def get_by_uid(uid: int): """Get a user by uid.""" for user in builtin.USERS: if user['_id'] == uid: return user coll = db.Collection('user') return await coll.find_one({'_id': uid})
async def add(uid: int, uname: str, password: str, mail: str, regip: str = ''): """Add a user.""" validator.check_uname(uname) # TODO(iceboy): Filter uname by keywords. validator.check_password(password) validator.check_mail(mail) uname_lower = uname.strip().lower() mail_lower = mail.strip().lower() for user in builtin.USERS: if user['_id'] == uid or user['uname_lower'] == uname_lower or user[ 'mail_lower'] == mail_lower: raise error.UserAlreadyExistError(uname) salt = pwhash.gen_salt() coll = db.Collection('user') try: await coll.insert_one({ '_id': uid, 'uname': uname, 'uname_lower': uname_lower, 'mail': mail, 'mail_lower': mail_lower, 'salt': salt, 'hash': pwhash.hash_vj4(password, salt), 'regat': datetime.datetime.utcnow(), 'regip': regip, 'priv': builtin.DEFAULT_PRIV, 'loginat': datetime.datetime.utcnow(), 'loginip': regip, 'gravatar': mail }) except errors.DuplicateKeyError: raise error.UserAlreadyExistError(uid, uname, mail) from None
async def ensure_indexes(): coll = db.Collection('oplog') await coll.create_index('uid') # type delete document await coll.create_index([('doc.domain_id', 1), ('doc.doc_type', 1), ('doc.doc_id', 1)], sparse=True)
def get_all_multi(end_id: objectid.ObjectId=None, get_hidden: bool=False, *, fields=None, **kwargs): coll = db.Collection('record') query = {**kwargs, 'hidden': False if not get_hidden else {'$gte': False}} if end_id: query['_id'] = {'$lt': end_id} return coll.find(query, projection=fields)
async def next_judge(record_id, judge_uid, judge_token, **kwargs): coll = db.Collection('record') doc = await coll.find_one_and_update(filter={'_id': record_id, 'judge_uid': judge_uid, 'judge_token': judge_token}, update=kwargs, return_document=ReturnDocument.AFTER) return doc
async def get_by_mail(mail: str): """Get a user by mail.""" mail_lower = mail.strip().lower() for user in builtin.USERS: if user['mail_lower'] == mail_lower: return user coll = db.Collection('user') return await coll.find_one({'mail_lower': mail_lower})
async def get_by_uname(uname: str): """Get a user by uname.""" uname_lower = uname.strip().lower() for user in builtin.USERS: if user['uname_lower'] == uname_lower: return user coll = db.Collection('user') return await coll.find_one({'uname_lower': uname_lower})
async def delete(domain_id: str, doc_type: int, doc_id: convert_doc_id): # TODO(twd2): delete status? coll = db.Collection('document') return await coll.delete_one({ 'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id })
async def delete_multi(domain_id: str, doc_type: int, **kwargs): # TODO(twd2): delete status? coll = db.Collection('document') return await coll.delete_many({ 'domain_id': domain_id, 'doc_type': doc_type, **kwargs })