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
})
