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('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 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 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 _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 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 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 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