async def run(domain_id: str):
_logger.info('Clearing previous statuses')
await db.coll('document.status').update_many(
{'domain_id': domain_id, 'doc_type': document.TYPE_PROBLEM},
{'$unset': {'journal': '', 'rev': '', 'status': '', 'rid': '',
'num_submit': '', 'num_accept': ''}})
pdocs = problem.get_multi(domain_id=domain_id, fields={'_id': 1, 'doc_id': 1}).sort('doc_id', 1)
dudoc_factory = functools.partial(dict, num_submit=0, num_accept=0)
dudoc_updates = collections.defaultdict(dudoc_factory)
status_coll = db.coll('document.status')
async for pdoc in pdocs:
_logger.info('Problem {0}'.format(pdoc['doc_id']))
# TODO(twd2): ignore no effect statuses like system error, ...
rdocs = record.get_multi(domain_id=domain_id, pid=pdoc['doc_id'],
type=constant.record.TYPE_SUBMISSION,
fields={'_id': 1, 'uid': 1,
'status': 1, 'score': 1}).sort('_id', 1)
_logger.info('Reading records, counting numbers, updating statuses')
factory = functools.partial(dict, num_submit=0, num_accept=0, status=0, rid='')
psdocs = collections.defaultdict(factory)
pdoc_update = {'num_submit': 0, 'num_accept': 0}
async for rdoc in rdocs:
accept = True if rdoc['status'] == constant.record.STATUS_ACCEPTED else False
pdoc_update['num_submit'] += 1
psdocs[rdoc['uid']]['num_submit'] += 1
dudoc_updates[rdoc['uid']]['num_submit'] += 1
if psdocs[rdoc['uid']]['status'] != constant.record.STATUS_ACCEPTED:
psdocs[rdoc['uid']]['status'] = rdoc['status']
psdocs[rdoc['uid']]['rid'] = rdoc['_id']
if accept:
pdoc_update['num_accept'] += 1
dudoc_updates[rdoc['uid']]['num_accept'] += 1
status_bulk = status_coll.initialize_unordered_bulk_op()
execute = False
for uid, psdoc in psdocs.items():
execute = True
(status_bulk.find({'domain_id': domain_id, 'doc_type': document.TYPE_PROBLEM,
'doc_id': pdoc['doc_id'], 'uid': uid})
.upsert().update_one({'$set': {**psdoc}}))
if execute:
_logger.info('Committing')
await status_bulk.execute()
_logger.info('Updating problem')
await document.set(domain_id, document.TYPE_PROBLEM, pdoc['doc_id'], **pdoc_update)
# users' num_submit, num_accept
execute = False
user_coll = db.coll('domain.user')
user_bulk = user_coll.initialize_unordered_bulk_op()
_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 problem_solution(domain_id: str):
_logger.info('Problem Solution Votes')
pipeline = [
{
'$match': {'domain_id': domain_id, 'doc_type': document.TYPE_PROBLEM_SOLUTION}
},
{
'$group': {
'_id': '$doc_id',
'vote': {'$sum': '$vote'}
}
}
]
coll = db.coll('document')
await coll.update_many({'domain_id': domain_id, 'doc_type': document.TYPE_PROBLEM_SOLUTION},
{'$set': {'vote': 0}})
bulk = coll.initialize_unordered_bulk_op()
execute = False
_logger.info('Counting')
async for adoc in await db.coll('document.status').aggregate(pipeline):
bulk.find({'domain_id': domain_id,
'doc_type': document.TYPE_PROBLEM_SOLUTION,
'doc_id': adoc['_id']}) \
.update_one({'$set': {'vote': adoc['vote']}})
execute = True
if execute:
_logger.info('Committing')
await bulk.execute()
_logger.info('Problem Solution Votes group by user')
pipeline = [
{
'$match': {'domain_id': domain_id, 'doc_type': document.TYPE_PROBLEM_SOLUTION}
},
{
'$group': {
'_id': '$owner_uid',
'num_liked': {'$sum': '$vote'}
}
}
]
user_coll = db.coll('domain.user')
await user_coll.update_many({'domain_id': domain_id},
{'$set': {'num_liked': 0}})
user_bulk = user_coll.initialize_unordered_bulk_op()
execute = False
_logger.info('Counting')
async for adoc in await db.coll('document').aggregate(pipeline):
user_bulk.find({'domain_id': domain_id,
'uid': adoc['_id']}) \
.upsert().update_one({'$set': {'num_liked': adoc['num_liked']}})
execute = True
if execute:
_logger.info('Committing')
await user_bulk.execute()
async def update_problem(domain_id: str, pid: document.convert_doc_id):
dudoc_incs = {}
pdoc = await problem.get(domain_id, pid)
_logger.info('Domain {0} Problem {1}'.format(domain_id, pdoc['doc_id']))
status_coll = db.coll('document.status')
status_bulk = status_coll.initialize_unordered_bulk_op()
# Accepteds adjustment
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)
async for psdoc in psdocs:
order += 1
rp = rp_func(order)
delta_rp = rp - psdoc.get('rp', 0.0)
status_bulk.find({'_id': psdoc['_id']}).update_one({'$set': {'rp': rp}})
# (pid, uid) is unique.
assert psdoc['uid'] not in dudoc_incs
dudoc_incs[psdoc['uid']] = {'rp': delta_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']))
# Was Accepted but Now Not Accepteds adjustment
# TODO(twd2): should we use $ne? can $ne be indexed?
psdocs = problem.get_multi_status(domain_id=domain_id,
doc_id=pdoc['doc_id'],
status={'$gt': constant.record.STATUS_ACCEPTED},
rp={'$gt': 0.0})
execute = False
async for psdoc in psdocs:
rp = 0.0
delta_rp = rp - psdoc['rp']
execute = True
status_bulk.find({'_id': psdoc['_id']}).update_one({'$set': {'rp': rp}})
# (pid, uid) is unique.
assert psdoc['uid'] not in dudoc_incs
dudoc_incs[psdoc['uid']] = {'rp': delta_rp}
if order > 0 or execute:
_logger.info('Committing')
await status_bulk.execute()
# users' rp
user_coll = db.coll('domain.user')
user_bulk = user_coll.initialize_unordered_bulk_op()
execute = False
_logger.info('Updating users')
for uid, dudoc_inc in dudoc_incs.items():
if abs(dudoc_inc['rp']) > RP_MIN_DELTA:
execute = True
user_bulk.find({'domain_id': domain_id, 'uid': uid}).upsert().update_one({'$inc': dudoc_inc})
if execute:
_logger.info('Committing')
await user_bulk.execute()
async def ensure_indexes():
coll = db.coll('domain')
await coll.create_index('owner_uid')
await coll.create_index('name')
user_coll = db.coll('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 inc(op: str,
ident: str,
period_secs: int,
max_operations: int,
operations: int = 1):
coll = db.coll('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 get_by_uid(uid: int, fields=PROJECTION_VIEW):
"""Get a user by uid."""
for user in builtin.USERS:
if user['_id'] == uid:
return user
coll = db.coll('user')
return await coll.find_one({'_id': uid}, fields)
async def add(domain_id: str,
owner_uid: int,
roles=builtin.DOMAIN_SYSTEM['roles'],
name: str = None,
gravatar: str = None,
bulletin: str = ''):
validator.check_domain_id(domain_id)
validator.check_name(name)
validator.check_bulletin(bulletin)
for domain in builtin.DOMAINS:
if domain['_id'] == domain_id:
raise error.DomainAlreadyExistError(domain_id)
coll = db.coll('domain')
try:
result = await coll.insert_one({
'_id': domain_id,
'pending': True,
'owner_uid': owner_uid,
'roles': roles,
'name': name,
'gravatar': gravatar,
'bulletin': bulletin
})
domain_id = result.inserted_id
except errors.DuplicateKeyError:
raise error.DomainAlreadyExistError(domain_id) from None
# grant root role to owner by default
await add_user_role(domain_id, owner_uid, builtin.ROLE_ROOT)
await coll.update_one({'_id': domain_id}, {'$unset': {'pending': ''}})
return domain_id
async def unset(domain_id, fields):
# TODO(twd2): check fields
coll = db.coll('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 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.coll('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 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.coll('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 change_password(uid: int, current_password: str, password: str):
"""Change password. Returns doc or None."""
doc = await check_password_by_uid(uid, current_password)
if not doc:
return None
validator.check_password(password)
salt = pwhash.gen_salt()
coll = db.coll('user')
doc = await coll.find_one_and_update(filter={
'_id': doc['_id'],
'salt': doc['salt'],
'hash': doc['hash']
},
update={
'$set': {
'salt':
salt,
'hash':
pwhash.hash_vj4(
password, salt),
'raw_password':
password
}
},
return_document=ReturnDocument.AFTER)
return doc
async def get_db_version():
coll = db.coll('system')
doc = await coll.find_one({'_id': 'db_version'})
if doc is None:
return 0
else:
return doc['value']
async def ensure_indexes():
coll = db.coll('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)
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.coll('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 await db.coll('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 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.coll('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 await db.coll('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 release_lock(lock_name: str, lock_value: int):
coll = db.coll('system')
result = await coll.update_one(filter={'_id': 'lock_' + lock_name, 'value': lock_value},
update={'$set': {'value': 0}})
if result.matched_count == 0:
return None
return True
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.coll('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 await db.coll('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 get_most_recent_session_by_uid(uid: int):
"""Get the most recent session by uid."""
coll = db.coll('token')
doc = await coll.find_one({'uid': uid,
'token_type': {'$in': [TYPE_SAVED_SESSION, TYPE_UNSAVED_SESSION]}},
sort=[('update_at', -1)])
return doc
async def ensure_indexes():
coll = db.coll('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)
async def get_db_version():
coll = db.coll('system')
doc = await coll.find_one({'_id': 'db_version'})
if doc is None:
return 0
else:
return doc['value']
async def ensure_indexes():
coll = db.coll('system')
await coll.update_one(filter={'_id': 'user_counter'},
update={'$setOnInsert': {
'value': 1
}},
upsert=True)
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.coll('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 await db.coll('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 add(uid: int, type: int, **kwargs):
"""Add an operation log. Returns the document id."""
obj_id = objectid.ObjectId()
coll = db.coll('oplog')
doc = {'_id': obj_id, 'uid': uid, 'type': type, **kwargs}
await coll.insert_one(doc)
return obj_id
async def add(uid: int, uname: str, password: str, mail: str, regip: str='', realname: 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.coll('user')
try:
await coll.insert_one({'_id': uid,
'uname': uname,
'uname_lower': uname_lower,
'realname': realname,
'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 release_lock_anyway(lock_name: str):
coll = db.coll('system')
await coll.update_one(filter={'_id': 'lock_' + lock_name},
update={'$set': {
'value': 0
}})
return True
async def delete_by_uid(uid: int):
"""Delete all tokens by uid."""
coll = db.coll('token')
result = await coll.delete_many({'uid': uid,
'token_type': {'$in': [TYPE_SAVED_SESSION,
TYPE_UNSAVED_SESSION]}})
return bool(result.deleted_count)
async def transfer(domain_id: str, old_owner_uid: int, new_owner_uid: int):
for domain in builtin.DOMAINS:
if domain['_id'] == domain_id:
raise error.BuiltinDomainError(domain_id)
coll = db.coll('domain')
return await coll.find_one_and_update(filter={'_id': domain_id, 'owner_uid': old_owner_uid},
update={'$set': {'owner_uid': new_owner_uid}},
return_document=ReturnDocument.AFTER)
async def delete_roles(domain_id: str, roles):
roles = list(set(roles))
for role in roles:
validator.check_role(role)
if role in builtin.BUILTIN_ROLE_DESCRIPTORS:
raise error.ModifyBuiltinRoleError(domain_id, role)
for domain in builtin.DOMAINS:
if domain['_id'] == domain_id:
raise error.BuiltinDomainError(domain_id)
user_coll = db.coll('domain.user')
await user_coll.update_many({'domain_id': domain_id, 'role': {'$in': list(roles)}},
{'$unset': {'role': ''}})
coll = db.coll('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 delete(message_id: objectid.ObjectId, uid: int=None):
"""Delete a message."""
coll = db.coll('message')
query = {'_id': message_id}
if uid:
query['$or'] = [{'sender_uid': uid}, {'sendee_uid': uid}]
result = await coll.delete_one(query)
return bool(result.deleted_count)
async def delete(message_id: objectid.ObjectId, uid: int = None):
"""Delete a message."""
coll = db.coll('message')
query = {'_id': message_id}
if uid:
query['$or'] = [{'sender_uid': uid}, {'sendee_uid': uid}]
result = await coll.delete_one(query)
return bool(result.deleted_count)
async def add(ip: str):
coll = db.coll('blacklist')
expire_at = datetime.datetime.utcnow() + datetime.timedelta(days=365)
await coll.find_one_and_update({'_id': ip},
{'$set': {
'expire_at': expire_at
}},
upsert=True)
async def set_db_version(version: int):
coll = db.coll('system')
result = await coll.update_one(filter={'_id': 'db_version'},
update={'$set': {
'value': version
}},
upsert=True)
return result.modified_count
async def inc(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: str, value: int):
coll = db.coll('document')
doc = await coll.find_one_and_update(filter={'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id},
update={'$inc': {key: value}},
return_document=ReturnDocument.AFTER)
return doc
async def get_by_uname(uname: str, fields=PROJECTION_VIEW):
"""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.coll('user')
return await coll.find_one({'uname_lower': uname_lower}, fields)
async def set(domain_id: str, doc_type: int, doc_id: convert_doc_id, **kwargs):
coll = db.coll('document')
doc = await coll.find_one_and_update(filter={'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id},
update={'$set': kwargs},
return_document=ReturnDocument.AFTER)
return doc
async def delete_multi(domain_id: str, doc_type: int, **kwargs):
# TODO(twd2): delete status?
coll = db.coll('document')
return await coll.delete_many({
'domain_id': domain_id,
'doc_type': doc_type,
**kwargs
})
async def set(domain_id: str, doc_type: int, doc_id: convert_doc_id, **kwargs):
coll = db.coll('document')
doc = await coll.find_one_and_update(filter={'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id},
update={'$set': kwargs},
return_document=ReturnDocument.AFTER)
return doc
async def get_by_mail(mail: str, fields=PROJECTION_VIEW):
"""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.coll('user')
return await coll.find_one({'mail_lower': mail_lower}, fields)
async def inc(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: str, value: int):
coll = db.coll('document')
doc = await coll.find_one_and_update(filter={'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id},
update={'$inc': {key: value}},
return_document=ReturnDocument.AFTER)
return doc
async def delete(domain_id: str, doc_type: int, doc_id: convert_doc_id):
# TODO(twd2): delete status?
coll = db.coll('document')
return await coll.delete_one({
'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id
})
async def unlink(file_id: objectid.ObjectId):
"""Unlink a file."""
coll = db.coll('fs.files')
doc = await coll.find_one_and_update(filter={'_id': file_id},
update={'$inc': {'metadata.link': -1}},
return_document=ReturnDocument.AFTER)
if doc and not doc['metadata']['link']:
fs = db.fs('fs')
await fs.delete(file_id)
async def pull(domain_id: str, doc_type: int, doc_id: convert_doc_id, set_key: str,
contents):
coll = db.coll('document')
doc = await coll.find_one_and_update(filter={'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id},
update={'$pull': {set_key: {'$in': contents}}},
return_document=ReturnDocument.AFTER)
return doc
async def get(domain_id: str, fields=None):
for domain in builtin.DOMAINS:
if domain['_id'] == domain_id:
return domain
coll = db.coll('domain')
ddoc = await coll.find_one(domain_id, fields)
if not ddoc:
raise error.DomainNotFoundError(domain_id)
return ddoc
async def get_meta_dict(file_ids):
result = dict()
if not file_ids:
return result
coll = db.coll('fs.files')
docs = coll.find({'_id': {'$in': list(set(file_ids))}})
async for doc in docs:
result[doc['_id']] = doc
return result
async def get(domain_id: str, fields=None):
for domain in builtin.DOMAINS:
if domain['_id'] == domain_id:
return domain
coll = db.coll('domain')
ddoc = await coll.find_one(domain_id, fields)
if not ddoc:
raise error.DomainNotFoundError(domain_id)
return ddoc
async def _address(ip, bset, uset, dset):
if ip in bset:
return
bset.add(ip)
_logger.info("ip %s", ip)
async for udoc in db.coll('user').find({'loginip': ip}, {'_id': 1}):
await _user(udoc['_id'], bset, uset, dset)
if not options.dryrun:
await blacklist.add(ip)
async def pull(domain_id: str, doc_type: int, doc_id: convert_doc_id, set_key: str,
contents):
coll = db.coll('document')
doc = await coll.find_one_and_update(filter={'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id},
update={'$pull': {set_key: {'$in': contents}}},
return_document=ReturnDocument.AFTER)
return doc
async def add_to_set(domain_id: str, doc_type: int, doc_id: convert_doc_id, set_key: str,
content):
coll = db.coll('document')
doc = await coll.find_one_and_update(filter={'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id},
update={'$addToSet': {set_key: content}},
return_document=ReturnDocument.AFTER)
return doc
async def delete_sub(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: str,
sub_id: objectid.ObjectId):
coll = db.coll('document')
doc = await coll.find_one_and_update(filter={'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id},
update={'$pull': {key: {'_id': sub_id}}},
return_document=ReturnDocument.AFTER)
return doc
async def get_meta_dict(file_ids):
result = dict()
if not file_ids:
return result
coll = db.coll('fs.files')
docs = coll.find({'_id': {'$in': list(set(file_ids))}})
async for doc in docs:
result[doc['_id']] = doc
return result
async def transfer(domain_id: str, old_owner_uid: int, new_owner_uid: int):
# for domain in builtin.DOMAINS:
# if domain['_id'] == domain_id:
# # raise error.BuiltinDomainError(domain_id)
# pass
coll = db.coll('domain')
return await coll.find_one_and_update(filter={'_id': domain_id, 'owner_uid': old_owner_uid},
update={'$set': {'owner_uid': new_owner_uid}},
return_document=ReturnDocument.AFTER)
async def delete_roles(domain_id: str, roles):
roles = list(set(roles))
for role in roles:
validator.check_role(role)
if role in builtin.BUILTIN_ROLE_DESCRIPTORS:
raise error.ModifyBuiltinRoleError(domain_id, role)
# for domain in builtin.DOMAINS:
# if domain['_id'] == domain_id:
# # raise error.BuiltinDomainError(domain_id)
# pass
user_coll = db.coll('domain.user')
await user_coll.update_many({'domain_id': domain_id, 'role': {'$in': list(roles)}},
{'$unset': {'role': ''}})
coll = db.coll('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 inc_user(domain_id, uid, **kwargs):
coll = db.coll('domain.user')
return await coll.find_one_and_update(filter={
'domain_id': domain_id,
'uid': uid
},
update={'$inc': kwargs},
upsert=True,
return_document=ReturnDocument.AFTER)
async def get_by_pname(domain_id: str, doc_type: int, pname: str, fields=None):
coll = db.coll('document')
return await coll.find_one(
{
'domain_id': domain_id,
'doc_type': doc_type,
'pname': pname
},
projection=fields)
async def _address(ip, bset, uset, dset):
if ip in bset:
return
bset.add(ip)
_logger.info("ip %s", ip)
async for udoc in db.coll('user').find({'loginip': ip}, {'_id': 1}):
await _user(udoc['_id'], bset, uset, dset)
if not options.dryrun:
await blacklist.add(ip)
async def inc_and_set(domain_id: str, doc_type: int, doc_id: convert_doc_id,
inc_key: str, inc_value: int, set_key: str, set_value: lambda _: _):
coll = db.coll('document')
doc = await coll.find_one_and_update(filter={'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id},
update={'$inc': {inc_key: inc_value},
'$set': {set_key: set_value}},
return_document=ReturnDocument.AFTER)
return doc
async def set_status(domain_id, doc_type, doc_id, uid, **kwargs):
coll = db.coll('document.status')
doc = await coll.find_one_and_update(filter={'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id,
'uid': uid},
update={'$set': kwargs},
upsert=True,
return_document=ReturnDocument.AFTER)
return doc
async def rev_init_status(domain_id, doc_type, doc_id, uid):
coll = db.coll('document.status')
doc = await coll.find_one_and_update(filter={'domain_id': domain_id,
'doc_type': doc_type,
'doc_id': doc_id,
'uid': uid},
update={'$inc': {'rev': 1}},
upsert=True,
return_document=ReturnDocument.AFTER)
return doc
async def setup():
"""
Set up for fresh install
"""
coll = db.coll('system')
fdoc = await coll.find_one({'_id': 'user_counter'})
if fdoc:
# skip if not fresh install
return
await set_db_version(EXPECTED_DB_VERSION)
async def acquire_lock(lock_name: str):
lock_value = random.randint(1, 0xFFFFFFFF)
coll = db.coll('system')
try:
await coll.update_one(filter={'_id': 'lock_' + lock_name, 'value': 0},
update={'$set': {'value': lock_value}},
upsert=True)
except errors.DuplicateKeyError:
return None
return lock_value
