diff --git a/requirements.txt b/requirements.txt index 9d70fab62..cc8e70eeb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,6 @@ aiohttp>=2.0.7 jinja2>=2.9.0 sockjs>=0.6.0 -motor hoedown accept aioamqp @@ -12,3 +11,4 @@ httpagentparser geoip2 GitPython PyYAML +git+https://github.com/iceb0y/aiomongo diff --git a/vj4/app.py b/vj4/app.py index 6908567c7..54635fb57 100644 --- a/vj4/app.py +++ b/vj4/app.py @@ -5,6 +5,7 @@ import sockjs from aiohttp import web +from vj4 import db from vj4 import error from vj4.service import bus from vj4.service import smallcache @@ -46,8 +47,9 @@ def __init__(self): # Initialize components. staticmanifest.init(static_path) locale.load_translations(translation_path) - asyncio.get_event_loop().run_until_complete( - asyncio.gather(tools.ensure_all_indexes(), bus.init())) + loop = asyncio.get_event_loop() + loop.run_until_complete(db.init()) + loop.run_until_complete(asyncio.gather(tools.ensure_all_indexes(), bus.init())) smallcache.init() # Load views. diff --git a/vj4/db.py b/vj4/db.py index 9b0a0fa5a..b8758aea1 100644 --- a/vj4/db.py +++ b/vj4/db.py @@ -1,4 +1,5 @@ -from motor import motor_asyncio +import aiomongo +import functools from vj4.util import options @@ -6,29 +7,17 @@ options.define('db_name', default='test', help='Database name.') -class Database(object): - _instance = None +async def init(): + client = await aiomongo.create_client('mongodb://' + options.db_host) + global _db + _db = client.get_database(options.db_name) - def __new__(cls): - if not cls._instance: - client = motor_asyncio.AsyncIOMotorClient(options.db_host) - cls._instance = motor_asyncio.AsyncIOMotorDatabase(client, options.db_name) - return cls._instance +@functools.lru_cache() +def coll(name): + return aiomongo.Collection(_db, name) -class Collection(object): - _instances = {} - def __new__(cls, name): - if name not in cls._instances: - cls._instances[name] = motor_asyncio.AsyncIOMotorCollection(Database(), name) - return cls._instances[name] - - -class GridFS(object): - _instances = {} - - def __new__(cls, name): - if name not in cls._instances: - cls._instances[name] = motor_asyncio.AsyncIOMotorGridFS(Database(), name) - return cls._instances[name] +@functools.lru_cache() +def fs(name): + return aiomongo.GridFS(_db, name) diff --git a/vj4/handler/contest.py b/vj4/handler/contest.py index a7eb71873..a2f944306 100644 --- a/vj4/handler/contest.py +++ b/vj4/handler/contest.py @@ -219,7 +219,7 @@ async def get(self, *, tid: objectid.ObjectId, pid: document.convert_doc_id): rdocs = await record.get_user_in_problem_multi(uid, self.domain_id, pdoc['doc_id']) \ .sort([('_id', -1)]) \ .limit(10) \ - .to_list(None) + .to_list() else: rdocs = [] if not self.prefer_json: @@ -326,7 +326,7 @@ async def post(self, *, title: str, content: str, rule: int, pdocs = await problem.get_multi(domain_id=self.domain_id, doc_id={'$in': pids}, fields={'doc_id': 1}) \ .sort('doc_id', 1) \ - .to_list(None) + .to_list() exist_pids = [pdoc['doc_id'] for pdoc in pdocs] if len(pids) != len(exist_pids): for pid in pids: diff --git a/vj4/handler/domain.py b/vj4/handler/domain.py index 2eaf6114c..44f41967b 100644 --- a/vj4/handler/domain.py +++ b/vj4/handler/domain.py @@ -22,7 +22,7 @@ async def prepare_contest(self): if self.has_perm(builtin.PERM_VIEW_CONTEST): tdocs = await contest.get_multi(self.domain_id) \ .limit(self.CONTESTS_ON_MAIN) \ - .to_list(None) + .to_list() tsdict = await contest.get_dict_status(self.domain_id, self.user['_id'], (tdoc['doc_id'] for tdoc in tdocs)) else: @@ -35,7 +35,7 @@ async def prepare_training(self): tdocs = await training.get_multi(self.domain_id) \ .sort('doc_id', 1) \ .limit(self.TRAININGS_ON_MAIN) \ - .to_list(None) + .to_list() tsdict = await training.get_dict_status(self.domain_id, self.user['_id'], (tdoc['doc_id'] for tdoc in tdocs)) else: @@ -47,7 +47,7 @@ async def prepare_discussion(self): if self.has_perm(builtin.PERM_VIEW_DISCUSSION): ddocs = await discussion.get_multi(self.domain_id) \ .limit(self.DISCUSSIONS_ON_MAIN) \ - .to_list(None) + .to_list() vndict = await discussion.get_dict_vnodes(self.domain_id, map(discussion.node_id, ddocs)) else: ddocs = [] diff --git a/vj4/handler/home.py b/vj4/handler/home.py index 6f98355f9..054365933 100644 --- a/vj4/handler/home.py +++ b/vj4/handler/home.py @@ -165,7 +165,7 @@ def modify_udoc(self, udict, key): @base.require_priv(builtin.PRIV_USER_PROFILE) async def get(self): # TODO(iceboy): projection, pagination. - mdocs = await message.get_multi(self.user['_id']).sort([('_id', -1)]).limit(50).to_list(None) + mdocs = await message.get_multi(self.user['_id']).sort([('_id', -1)]).limit(50).to_list() udict = await user.get_dict( itertools.chain.from_iterable((mdoc['sender_uid'], mdoc['sendee_uid']) for mdoc in mdocs), fields=user.PROJECTION_PUBLIC) @@ -240,7 +240,7 @@ async def get(self): dids = list(dudict.keys()) ddocs = await domain.get_multi(**{'$or': [{'_id': {'$in': dids}}, {'owner_uid': self.user['_id']}]}) \ - .to_list(None) + .to_list() can_manage = {} for ddoc in builtin.DOMAINS + ddocs: role = dudict.get(ddoc['_id'], {}).get('role', builtin.ROLE_DEFAULT) @@ -276,7 +276,7 @@ def file_url(self, fdoc): @base.require_priv(builtin.PRIV_USER_PROFILE) async def get(self): - ufdocs = await userfile.get_multi(owner_uid=self.user['_id']).to_list(None) + ufdocs = await userfile.get_multi(owner_uid=self.user['_id']).to_list() fdict = await fs.get_meta_dict(ufdoc.get('file_id') for ufdoc in ufdocs) self.render('home_file.html', ufdocs=ufdocs, fdict=fdict) diff --git a/vj4/handler/problem.py b/vj4/handler/problem.py index 3fb300757..873ce07df 100644 --- a/vj4/handler/problem.py +++ b/vj4/handler/problem.py @@ -196,9 +196,9 @@ async def get(self, *, pid: document.convert_doc_id): if pdoc.get('hidden', False): self.check_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN) udoc = await user.get_by_uid(pdoc['owner_uid']) - tdocs = await training.get_multi(self.domain_id, **{'dag.pids': pid}).to_list(None) \ + tdocs = await training.get_multi(self.domain_id, **{'dag.pids': pid}).to_list() \ if self.has_perm(builtin.PERM_VIEW_TRAINING) else None - ctdocs = await contest.get_multi(self.domain_id, pids=pid).to_list(None) \ + ctdocs = await contest.get_multi(self.domain_id, pids=pid).to_list() \ if self.has_perm(builtin.PERM_VIEW_CONTEST) else None path_components = self.build_path( (self.translate('problem_main'), self.reverse_url('problem_main')), @@ -227,7 +227,7 @@ async def get(self, *, pid: document.convert_doc_id): .get_user_in_problem_multi(uid, self.domain_id, pdoc['doc_id']) \ .sort([('_id', -1)]) \ .limit(10) \ - .to_list(None) + .to_list() if not self.prefer_json: path_components = self.build_path( (self.translate('problem_main'), self.reverse_url('problem_main')), diff --git a/vj4/handler/record.py b/vj4/handler/record.py index eedbd9e78..bfae4917d 100644 --- a/vj4/handler/record.py +++ b/vj4/handler/record.py @@ -57,7 +57,7 @@ async def get(self, *, uid_or_name: str='', pid: str='', tid: str=''): query = await self.get_filter_query(uid_or_name, pid, tid) # TODO(iceboy): projection, pagination. rdocs = await record.get_all_multi(**query, - get_hidden=self.has_priv(builtin.PRIV_VIEW_HIDDEN_RECORD)).sort([('_id', -1)]).limit(50).to_list(None) + get_hidden=self.has_priv(builtin.PRIV_VIEW_HIDDEN_RECORD)).sort([('_id', -1)]).limit(50).to_list() # TODO(iceboy): projection. udict, pdict = await asyncio.gather( user.get_dict(rdoc['uid'] for rdoc in rdocs), diff --git a/vj4/handler/training.py b/vj4/handler/training.py index fae755ffc..4b9605161 100644 --- a/vj4/handler/training.py +++ b/vj4/handler/training.py @@ -181,7 +181,7 @@ async def post(self, *, title: str, content: str, dag: str): pdocs = await problem.get_multi(domain_id=self.domain_id, doc_id={'$in': pids}, fields={'doc_id': 1, 'hidden': 1}) \ .sort('doc_id', 1) \ - .to_list(None) + .to_list() exist_pids = [pdoc['doc_id'] for pdoc in pdocs] if len(pids) != len(exist_pids): for pid in pids: @@ -229,7 +229,7 @@ async def post(self, *, tid: objectid.ObjectId, title: str, content: str, dag: s pdocs = await problem.get_multi(domain_id=self.domain_id, doc_id={'$in': pids}, fields={'doc_id': 1, 'hidden': 1}) \ .sort('doc_id', 1) \ - .to_list(None) + .to_list() exist_pids = [pdoc['doc_id'] for pdoc in pdocs] if len(pids) != len(exist_pids): for pid in pids: diff --git a/vj4/handler/user.py b/vj4/handler/user.py index e92eb3367..e167ba3c4 100644 --- a/vj4/handler/user.py +++ b/vj4/handler/user.py @@ -194,17 +194,17 @@ async def get(self, *, uid: int): bg = random.randint(1, 21) rdocs = record.get_multi(get_hidden=self.has_priv(builtin.PRIV_VIEW_HIDDEN_RECORD), uid=uid).sort([('_id', -1)]) - rdocs = await rdocs.limit(10).to_list(None) + rdocs = await rdocs.limit(10).to_list() # TODO(twd2): check status, eg. test, hidden problem, ... pdocs = problem.get_multi(domain_id=self.domain_id, owner_uid=uid).sort([('_id', -1)]) pcount = await pdocs.count() - pdocs = await pdocs.limit(10).to_list(None) + pdocs = await pdocs.limit(10).to_list() psdocs = problem.get_multi_solution_by_uid(self.domain_id, uid) pscount = await psdocs.count() - psdocs = await psdocs.limit(10).to_list(None) + psdocs = await psdocs.limit(10).to_list() ddocs = discussion.get_multi(self.domain_id, owner_uid=uid) dcount = await ddocs.count() - ddocs = await ddocs.limit(10).to_list(None) + ddocs = await ddocs.limit(10).to_list() self.render('user_detail.html', is_self_profile=is_self_profile, udoc=udoc, dudoc=dudoc, sdoc=sdoc, email=email, bg=bg, rdocs=rdocs, pdocs=pdocs, pcount=pcount, psdocs=psdocs, pscount=pscount, diff --git a/vj4/job/difficulty.py b/vj4/job/difficulty.py index ec670449c..d33a1ccac 100644 --- a/vj4/job/difficulty.py +++ b/vj4/job/difficulty.py @@ -110,7 +110,7 @@ async def update_problem(domain_id: str, pid: document.convert_doc_id): @domainjob.wrap async def recalc(domain_id: str): pdocs = problem.get_multi(domain_id=domain_id) - coll = db.Collection('document') + coll = db.coll('document') bulk = coll.initialize_unordered_bulk_op() execute = False _logger.info('Calculating') diff --git a/vj4/job/fs.py b/vj4/job/fs.py index 6341d7327..9bbbfc6bb 100644 --- a/vj4/job/fs.py +++ b/vj4/job/fs.py @@ -17,7 +17,7 @@ @argmethod.wrap async def sync_length(): _logger.info('Userfile length') - coll = db.Collection('document') + coll = db.coll('document') ufdocs = userfile.get_multi() bulk = coll.initialize_unordered_bulk_op() execute = False @@ -51,13 +51,13 @@ async def sync_usage(): } } ] - coll = db.Collection('domain.user') + 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 db.Collection('document').aggregate(pipeline): + async for adoc in db.coll('document').aggregate(pipeline): bulk.find({'domain_id': userfile.STORE_DOMAIN_ID, 'uid': adoc['_id']}) \ .update_one({'$set': {'usage_userfile': adoc['usage_userfile']}}) diff --git a/vj4/job/num.py b/vj4/job/num.py index e2aa4c6a5..bf7ac42c8 100644 --- a/vj4/job/num.py +++ b/vj4/job/num.py @@ -27,13 +27,13 @@ async def discussion(domain_id: str): } } ] - coll = db.Collection('document') + 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 db.Collection('document').aggregate(pipeline): + async for adoc in db.coll('document').aggregate(pipeline): bulk.find({'domain_id': domain_id, 'doc_type': document.TYPE_DISCUSSION, 'doc_id': adoc['_id']}) \ @@ -58,13 +58,13 @@ async def contest(domain_id: str): } } ] - coll = db.Collection('document') + coll = db.coll('document') await coll.update_many({'domain_id': domain_id, 'doc_type': document.TYPE_CONTEST}, {'$set': {'attend': 0}}) bulk = coll.initialize_unordered_bulk_op() execute = False _logger.info('Counting') - async for adoc in db.Collection('document.status').aggregate(pipeline): + async for adoc in db.coll('document.status').aggregate(pipeline): bulk.find({'domain_id': domain_id, 'doc_type': document.TYPE_CONTEST, 'doc_id': adoc['_id']}) \ @@ -89,13 +89,13 @@ async def training(domain_id: str): } } ] - coll = db.Collection('document') + 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 db.Collection('document.status').aggregate(pipeline): + async for adoc in db.coll('document.status').aggregate(pipeline): bulk.find({'domain_id': domain_id, 'doc_type': document.TYPE_TRAINING, 'doc_id': adoc['_id']}) \ @@ -120,13 +120,13 @@ async def problem(domain_id: str): } } ] - user_coll = db.Collection('domain.user') + 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 db.Collection('document').aggregate(pipeline): + async for adoc in db.coll('document').aggregate(pipeline): user_coll.find({'domain_id': domain_id, 'uid': adoc['_id']}) \ .upsert().update_one({'$set': {'num_problems': adoc['num_problems']}}) @@ -150,13 +150,13 @@ async def problem_solution(domain_id: str): } } ] - coll = db.Collection('document') + 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 db.Collection('document.status').aggregate(pipeline): + async for adoc in db.coll('document.status').aggregate(pipeline): bulk.find({'domain_id': domain_id, 'doc_type': document.TYPE_PROBLEM_SOLUTION, 'doc_id': adoc['_id']}) \ @@ -178,13 +178,13 @@ async def problem_solution(domain_id: str): } } ] - user_coll = db.Collection('domain.user') + 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 db.Collection('document').aggregate(pipeline): + async for adoc in db.coll('document').aggregate(pipeline): user_bulk.find({'domain_id': domain_id, 'uid': adoc['_id']}) \ .upsert().update_one({'$set': {'num_liked': adoc['num_liked']}}) diff --git a/vj4/job/rank.py b/vj4/job/rank.py index 2369ebe3b..683e9d9ec 100644 --- a/vj4/job/rank.py +++ b/vj4/job/rank.py @@ -18,7 +18,7 @@ async def run(domain_id: str, keyword: str='rp', rank_field: str='rank', level_f last_dudoc = {keyword: None} rank = 0 count = 0 - user_coll = db.Collection('domain.user') + user_coll = db.coll('domain.user') user_bulk = user_coll.initialize_unordered_bulk_op() async for dudoc in dudocs: count += 1 diff --git a/vj4/job/record.py b/vj4/job/record.py index d494c4231..dbc9257ee 100644 --- a/vj4/job/record.py +++ b/vj4/job/record.py @@ -57,14 +57,14 @@ async def user_in_problem(uid: int, domain_id: str, pid: document.convert_doc_id @domainjob.wrap async def run(domain_id: str): _logger.info('Clearing previous statuses') - await db.Collection('document.status').update_many( + 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.Collection('document.status') + 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, ... @@ -101,7 +101,7 @@ async def run(domain_id: str): await document.set(domain_id, document.TYPE_PROBLEM, pdoc['doc_id'], **pdoc_update) # users' num_submit, num_accept execute = False - user_coll = db.Collection('domain.user') + 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(): diff --git a/vj4/job/rp.py b/vj4/job/rp.py index f580fb062..5e21aac9c 100644 --- a/vj4/job/rp.py +++ b/vj4/job/rp.py @@ -47,7 +47,7 @@ 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.Collection('document.status') + 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, @@ -85,7 +85,7 @@ async def update_problem(domain_id: str, pid: document.convert_doc_id): _logger.info('Committing') await status_bulk.execute() # users' rp - user_coll = db.Collection('domain.user') + user_coll = db.coll('domain.user') user_bulk = user_coll.initialize_unordered_bulk_op() execute = False _logger.info('Updating users') @@ -99,12 +99,12 @@ async def update_problem(domain_id: str, pid: document.convert_doc_id): @domainjob.wrap async def recalc(domain_id: str): - user_coll = db.Collection('domain.user') + user_coll = db.coll('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') + status_coll = db.coll('document.status') async for pdoc in pdocs: _logger.info('Problem {0}'.format(pdoc['doc_id'])) psdocs = problem.get_multi_status(domain_id=domain_id, diff --git a/vj4/model/adaptor/contest.py b/vj4/model/adaptor/contest.py index 5b3b34119..5e49ac07a 100644 --- a/vj4/model/adaptor/contest.py +++ b/vj4/model/adaptor/contest.py @@ -128,7 +128,7 @@ async def get_and_list_status(domain_id: str, tid: objectid.ObjectId, fields=Non doc_id=tdoc['doc_id'], fields=fields) \ .sort(RULES[tdoc['rule']].status_sort) \ - .to_list(None) + .to_list() return tdoc, tsdocs diff --git a/vj4/model/adaptor/problem.py b/vj4/model/adaptor/problem.py index 5a0e09659..b42a306b7 100644 --- a/vj4/model/adaptor/problem.py +++ b/vj4/model/adaptor/problem.py @@ -205,7 +205,7 @@ async def get_list_solution(domain_id: str, pid: document.convert_doc_id, .sort([('vote', -1), ('doc_id', -1)]) \ .skip(skip) \ .limit(limit) \ - .to_list(None) + .to_list() @argmethod.wrap @@ -290,7 +290,7 @@ async def set_hidden(domain_id: str, pid: document.convert_doc_id, hidden: bool) async def get_data_list(last: int): last_datetime = datetime.datetime.fromtimestamp(last) # TODO(twd2): performance improve, more elegant - coll = db.Collection('document') + coll = db.coll('document') pdocs = coll.find({'doc_type': document.TYPE_PROBLEM}) pids = [] # with domain_id async for pdoc in pdocs: diff --git a/vj4/model/document.py b/vj4/model/document.py index 0c5996334..e25e50e72 100644 --- a/vj4/model/document.py +++ b/vj4/model/document.py @@ -44,7 +44,7 @@ async def add(domain_id: str, content: str, owner_uid: int, 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') + coll = db.coll('document') doc = {'_id': obj_id, 'content': content, 'owner_uid': owner_uid, @@ -61,14 +61,14 @@ async def add(domain_id: str, content: str, owner_uid: int, @argmethod.wrap async def get(domain_id: str, doc_type: int, doc_id: convert_doc_id, fields=None): - coll = db.Collection('document') + coll = db.coll('document') return await coll.find_one({'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id}, projection=fields) async def set(domain_id: str, doc_type: int, doc_id: convert_doc_id, **kwargs): - coll = db.Collection('document') + coll = db.coll('document') doc = await coll.find_one_and_update(filter={'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id}, @@ -79,7 +79,7 @@ async def set(domain_id: str, doc_type: int, doc_id: convert_doc_id, **kwargs): async def delete(domain_id: str, doc_type: int, doc_id: convert_doc_id): # TODO(twd2): delete status? - coll = db.Collection('document') + coll = db.coll('document') return await coll.delete_one({'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id}) @@ -87,14 +87,14 @@ async def delete(domain_id: str, doc_type: int, doc_id: convert_doc_id): async def delete_multi(domain_id: str, doc_type: int, **kwargs): # TODO(twd2): delete status? - coll = db.Collection('document') + coll = db.coll('document') return await coll.delete_many({'domain_id': domain_id, 'doc_type': doc_type, **kwargs}) def get_multi(*, fields=None, **kwargs): - coll = db.Collection('document') + coll = db.coll('document') return coll.find(kwargs, projection=fields) @@ -115,7 +115,7 @@ async def get_dict(domain_id: str, dtuples, *, fields=None): @argmethod.wrap async def inc(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: str, value: int): - coll = db.Collection('document') + coll = db.coll('document') doc = await coll.find_one_and_update(filter={'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id}, @@ -127,7 +127,7 @@ async def inc(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: str, v @argmethod.wrap 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.Collection('document') + coll = db.coll('document') doc = await coll.find_one_and_update(filter={'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id}, @@ -140,7 +140,7 @@ async def inc_and_set(domain_id: str, doc_type: int, doc_id: convert_doc_id, @argmethod.wrap async def push(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: str, content: str, owner_uid: int, **kwargs): - coll = db.Collection('document') + coll = db.coll('document') obj_id = objectid.ObjectId() doc = await coll.find_one_and_update(filter={'domain_id': domain_id, 'doc_type': doc_type, @@ -156,7 +156,7 @@ async def push(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: str, @argmethod.wrap async def delete_sub(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: str, sub_id: objectid.ObjectId): - coll = db.Collection('document') + coll = db.coll('document') doc = await coll.find_one_and_update(filter={'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id}, @@ -168,7 +168,7 @@ async def delete_sub(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: @argmethod.wrap async def get_sub(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: str, sub_id: objectid.ObjectId): - coll = db.Collection('document') + coll = db.coll('document') doc = await coll.find_one({'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id, @@ -184,7 +184,7 @@ async def get_sub(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: st @argmethod.wrap async def set_sub(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: str, sub_id: objectid.ObjectId, **kwargs): - coll = db.Collection('document') + coll = db.coll('document') mod = dict(('{0}.$.{1}'.format(key, k), v) for k, v in kwargs.items()) doc = await coll.find_one_and_update(filter={'domain_id': domain_id, 'doc_type': doc_type, @@ -198,7 +198,7 @@ async def set_sub(domain_id: str, doc_type: int, doc_id: convert_doc_id, key: st @argmethod.wrap async def add_to_set(domain_id: str, doc_type: int, doc_id: convert_doc_id, set_key: str, content): - coll = db.Collection('document') + coll = db.coll('document') doc = await coll.find_one_and_update(filter={'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id}, @@ -210,7 +210,7 @@ async def add_to_set(domain_id: str, doc_type: int, doc_id: convert_doc_id, set_ @argmethod.wrap async def pull(domain_id: str, doc_type: int, doc_id: convert_doc_id, set_key: str, contents): - coll = db.Collection('document') + coll = db.coll('document') doc = await coll.find_one_and_update(filter={'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id}, @@ -222,19 +222,19 @@ async def pull(domain_id: str, doc_type: int, doc_id: convert_doc_id, set_key: s @argmethod.wrap async def get_status(domain_id: str, doc_type: int, doc_id: convert_doc_id, uid: int, *, fields=None): - coll = db.Collection('document.status') + coll = db.coll('document.status') return await coll.find_one({'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id, 'uid': uid}, projection=fields) def get_multi_status(*, fields=None, **kwargs): - coll = db.Collection('document.status') + coll = db.coll('document.status') return coll.find(kwargs, projection=fields) async def set_status(domain_id, doc_type, doc_id, uid, **kwargs): - coll = db.Collection('document.status') + 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, @@ -248,7 +248,7 @@ async def set_status(domain_id, doc_type, doc_id, uid, **kwargs): @argmethod.wrap async def set_if_not_status(domain_id: str, doc_type: int, doc_id: convert_doc_id, uid: int, key: str, value: int, if_not: int, **kwargs): - coll = db.Collection('document.status') + coll = db.coll('document.status') return await coll.find_one_and_update(filter={'domain_id': domain_id, 'doc_type': doc_type, 'doc_id': doc_id, @@ -273,7 +273,7 @@ async def capped_inc_status(domain_id: str, not_expr = {'$gte': max_value} else: not_expr = {'$lte': min_value} - coll = db.Collection('document.status') + 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, @@ -288,7 +288,7 @@ async def capped_inc_status(domain_id: str, @argmethod.wrap async def inc_status(domain_id: str, doc_type: int, doc_id: convert_doc_id, uid: int, key: str, value: int): - coll = db.Collection('document.status') + 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, @@ -300,7 +300,7 @@ async def inc_status(domain_id: str, doc_type: int, doc_id: convert_doc_id, uid: async def rev_push_status(domain_id, doc_type, doc_id, uid, key, value): - coll = db.Collection('document.status') + 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, @@ -313,7 +313,7 @@ async def rev_push_status(domain_id, doc_type, doc_id, uid, key, value): async def rev_init_status(domain_id, doc_type, doc_id, uid): - coll = db.Collection('document.status') + 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, @@ -325,7 +325,7 @@ async def rev_init_status(domain_id, doc_type, doc_id, uid): async def rev_set_status(domain_id, doc_type, doc_id, uid, rev, **kwargs): - coll = db.Collection('document.status') + 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, @@ -339,7 +339,7 @@ async def rev_set_status(domain_id, doc_type, doc_id, uid, rev, **kwargs): @argmethod.wrap async def ensure_indexes(): - coll = db.Collection('document') + coll = db.coll('document') await coll.create_index([('domain_id', 1), ('doc_type', 1), ('doc_id', 1)], unique=True) @@ -404,7 +404,7 @@ async def ensure_indexes(): await coll.create_index([('domain_id', 1), ('doc_type', 1), ('dag.pids', 1)], sparse=True) - status_coll = db.Collection('document.status') + status_coll = db.coll('document.status') await status_coll.create_index([('domain_id', 1), ('doc_type', 1), ('uid', 1), diff --git a/vj4/model/domain.py b/vj4/model/domain.py index 7552e095c..ac1580a0c 100644 --- a/vj4/model/domain.py +++ b/vj4/model/domain.py @@ -19,7 +19,7 @@ async def add(domain_id: str, owner_uid: int, for domain in builtin.DOMAINS: if domain['_id'] == domain_id: raise error.DomainAlreadyExistError(domain_id) - coll = db.Collection('domain') + coll = db.coll('domain') try: return (await coll.insert_one({'_id': domain_id, 'owner_uid': owner_uid, 'roles': roles, 'name': name, @@ -33,18 +33,18 @@ async def get(domain_id: str, fields=None): for domain in builtin.DOMAINS: if domain['_id'] == domain_id: return domain - coll = db.Collection('domain') + coll = db.coll('domain') return await coll.find_one(domain_id, fields) def get_multi(*, fields=None, **kwargs): - coll = db.Collection('domain') + coll = db.coll('domain') return coll.find(kwargs, fields) @argmethod.wrap async def get_list(*, fields=None, limit: int=None, **kwargs): - coll = db.Collection('domain') + coll = db.coll('domain') return await coll.find(kwargs, fields).limit(limit).to_list(None) @@ -53,7 +53,7 @@ async def edit(domain_id: str, **kwargs): for domain in builtin.DOMAINS: if domain['_id'] == domain_id: raise error.BuiltinDomainError(domain_id) - coll = db.Collection('domain') + coll = db.coll('domain') if 'owner_uid' in kwargs: del kwargs['owner_uid'] if 'name' in kwargs: @@ -66,7 +66,7 @@ async def edit(domain_id: str, **kwargs): async def unset(domain_id, fields): # TODO(twd2): check fields - coll = db.Collection('domain') + 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) @@ -78,7 +78,7 @@ async def set_role(domain_id: str, role: str, perm: int): for domain in builtin.DOMAINS: if domain['_id'] == domain_id: raise error.BuiltinDomainError(domain_id) - coll = db.Collection('domain') + coll = db.coll('domain') return await coll.find_one_and_update(filter={'_id': domain_id}, update={'$set': {'roles.{0}'.format(role): perm}}, return_document=ReturnDocument.AFTER) @@ -96,10 +96,10 @@ async def delete_roles(domain_id: str, roles): for domain in builtin.DOMAINS: if domain['_id'] == domain_id: raise error.BuiltinDomainError(domain_id) - user_coll = db.Collection('domain.user') + user_coll = db.coll('domain.user') await user_coll.update_many({'domain_id': domain_id, 'role': {'$in': list(roles)}}, {'$unset': {'role': ''}}) - coll = db.Collection('domain') + 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)}, @@ -111,7 +111,7 @@ 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.Collection('domain') + 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) @@ -119,12 +119,12 @@ async def transfer(domain_id: str, old_owner_uid: int, new_owner_uid: int): @argmethod.wrap async def get_user(domain_id: str, uid: int, fields=None): - coll = db.Collection('domain.user') + coll = db.coll('domain.user') return await coll.find_one({'domain_id': domain_id, 'uid': uid}, fields) async def set_user(domain_id, uid, **kwargs): - coll = db.Collection('domain.user') + coll = db.coll('domain.user') return await coll.find_one_and_update(filter={'domain_id': domain_id, 'uid': uid}, update={'$set': kwargs}, upsert=True, @@ -132,7 +132,7 @@ async def set_user(domain_id, uid, **kwargs): async def unset_user(domain_id, uid, fields): - coll = db.Collection('domain.user') + coll = db.coll('domain.user') return await coll.find_one_and_update(filter={'domain_id': domain_id, 'uid': uid}, update={'$unset': dict((f, '') for f in set(fields))}, upsert=True, @@ -140,14 +140,14 @@ async def unset_user(domain_id, uid, fields): async def set_users(domain_id, uids, **kwargs): - coll = db.Collection('domain.user') + coll = db.coll('domain.user') await coll.update_many({'domain_id': domain_id, 'uid': {'$in': list(set(uids))}}, {'$set': kwargs}, upsert=False) async def unset_users(domain_id, uids, fields): - coll = db.Collection('domain.user') + coll = db.coll('domain.user') await coll.update_many({'domain_id': domain_id, 'uid': {'$in': list(set(uids))}}, {'$unset': dict((f, '') for f in set(fields))}, upsert=True) @@ -174,7 +174,7 @@ async def unset_users_role(domain_id: str, uids): async def inc_user(domain_id, uid, **kwargs): - coll = db.Collection('domain.user') + coll = db.coll('domain.user') return await coll.find_one_and_update(filter={'domain_id': domain_id, 'uid': uid}, update={'$inc': kwargs}, upsert=True, @@ -182,7 +182,7 @@ async def inc_user(domain_id, uid, **kwargs): async def inc_user_usage(domain_id: str, uid: int, usage_field: str, usage: int, quota: int): - coll = db.Collection('domain.user') + coll = db.coll('domain.user') try: return await coll.find_one_and_update(filter={'domain_id': domain_id, 'uid': uid, usage_field: {'$not': {'$gte': quota - usage}}}, @@ -194,7 +194,7 @@ async def inc_user_usage(domain_id: str, uid: int, usage_field: str, usage: int, def get_multi_user(*, fields=None, **kwargs): - coll = db.Collection('domain.user') + coll = db.coll('domain.user') return coll.find(kwargs, fields) @@ -215,9 +215,9 @@ async def get_dict_user_by_domain_id(uid, *, fields=None): @argmethod.wrap async def ensure_indexes(): - coll = db.Collection('domain') + coll = db.coll('domain') await coll.create_index('owner_uid') - user_coll = db.Collection('domain.user') + user_coll = db.coll('domain.user') await user_coll.create_index('uid') await user_coll.create_index([('domain_id', 1), ('uid', 1)], unique=True) diff --git a/vj4/model/fs.py b/vj4/model/fs.py index 13cbe3436..5ecd72b3f 100644 --- a/vj4/model/fs.py +++ b/vj4/model/fs.py @@ -12,7 +12,7 @@ async def add(content_type): """Add a file. Returns MotorGridIn.""" - fs = db.GridFS('fs') + fs = db.fs('fs') secret = pwhash.gen_secret() return await fs.new_file(content_type=content_type, metadata={'link': 1, 'secret': secret}) @@ -44,7 +44,7 @@ async def add_file_object(content_type, file_object): async def get(file_id): """Get a file. Returns MotorGridOut.""" - fs = db.GridFS('fs') + fs = db.fs('fs') return await fs.get(file_id) @@ -60,7 +60,7 @@ async def get_by_secret(secret): @argmethod.wrap async def get_file_id(secret: str): """Get the _id of a file by secret.""" - coll = db.Collection('fs.files') + coll = db.coll('fs.files') doc = await coll.find_one({'metadata.secret': secret}) if doc: return doc['_id'] @@ -69,7 +69,7 @@ async def get_file_id(secret: str): @argmethod.wrap async def get_md5(file_id: objectid.ObjectId): """Get the MD5 checksum of a file.""" - coll = db.Collection('fs.files') + coll = db.coll('fs.files') doc = await coll.find_one(file_id) if doc: return doc['md5'] @@ -78,7 +78,7 @@ async def get_md5(file_id: objectid.ObjectId): @argmethod.wrap async def get_datetime(file_id: objectid.ObjectId): """Get the upload date and time of a file.""" - coll = db.Collection('fs.files') + coll = db.coll('fs.files') doc = await coll.find_one(file_id) if doc: return doc['uploadDate'] @@ -87,7 +87,7 @@ async def get_datetime(file_id: objectid.ObjectId): @argmethod.wrap async def get_secret(file_id: objectid.ObjectId): """Get the secret of a file.""" - coll = db.Collection('fs.files') + coll = db.coll('fs.files') doc = await coll.find_one(file_id) if doc: return doc['metadata']['secret'] @@ -96,7 +96,7 @@ async def get_secret(file_id: objectid.ObjectId): @argmethod.wrap async def get_meta(file_id: objectid.ObjectId): """Get all metadata of a file.""" - coll = db.Collection('fs.files') + coll = db.coll('fs.files') doc = await coll.find_one(file_id) return doc @@ -105,7 +105,7 @@ async def get_meta_dict(file_ids): result = dict() if not file_ids: return result - coll = db.Collection('fs.files') + coll = db.coll('fs.files') docs = coll.find({'_id': {'$in': list(set(file_ids))}}) async for doc in docs: result[doc['_id']] = doc @@ -128,7 +128,7 @@ async def link_by_md5(file_md5: str, except_id: objectid.ObjectId=None): query = {} if except_id: query['_id'] = {'$ne': except_id} - coll = db.Collection('fs.files') + coll = db.coll('fs.files') doc = await coll.find_one_and_update(filter={'md5': file_md5, **query}, update={'$inc': {'metadata.link': 1}}) if doc: @@ -138,18 +138,18 @@ async def link_by_md5(file_md5: str, except_id: objectid.ObjectId=None): @argmethod.wrap async def unlink(file_id: objectid.ObjectId): """Unlink a file.""" - coll = db.Collection('fs.files') + 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.GridFS('fs') + fs = db.fs('fs') await fs.delete(file_id) @argmethod.wrap async def ensure_indexes(): - coll = db.Collection('fs.files') + coll = db.coll('fs.files') await coll.create_index('metadata.secret', unique=True) await coll.create_index('md5') diff --git a/vj4/model/message.py b/vj4/model/message.py index 4ef87c573..bda57c145 100644 --- a/vj4/model/message.py +++ b/vj4/model/message.py @@ -15,7 +15,7 @@ async def add(sender_uid: int, sendee_uid: int, content: str): """Send a message from sender to sendee with specified content.""" validator.check_content(content) - coll = db.Collection('message') + coll = db.coll('message') mdoc = {'sender_uid': sender_uid, 'sendee_uid': sendee_uid, 'status': STATUS_UNREAD, @@ -30,7 +30,7 @@ async def add(sender_uid: int, sendee_uid: int, content: str): @argmethod.wrap def get_multi(uid: int, *, fields=None): """Get messages related to a specified user.""" - coll = db.Collection('message') + coll = db.coll('message') return coll.find({'$or': [{'sender_uid': uid}, {'sendee_uid': uid}]}, projection=fields) @@ -38,7 +38,7 @@ def get_multi(uid: int, *, fields=None): async def add_reply(message_id: objectid.ObjectId, sender_uid: int, content: str): """Reply a message with specified content.""" validator.check_content(content) - coll = db.Collection('message') + coll = db.coll('message') reply = {'sender_uid': sender_uid, 'content': content, 'status': STATUS_UNREAD, @@ -52,7 +52,7 @@ async def add_reply(message_id: objectid.ObjectId, sender_uid: int, content: str @argmethod.wrap async def delete(message_id: objectid.ObjectId, uid: int=None): """Delete a message.""" - coll = db.Collection('message') + coll = db.coll('message') query = {'_id': message_id} if uid: query['$or'] = [{'sender_uid': uid}, {'sendee_uid': uid}] @@ -62,7 +62,7 @@ async def delete(message_id: objectid.ObjectId, uid: int=None): @argmethod.wrap async def ensure_indexes(): - coll = db.Collection('message') + coll = db.coll('message') await coll.create_index([('sender_uid', 1), ('_id', -1)]) await coll.create_index([('sendee_uid', 1), ('_id', -1)]) diff --git a/vj4/model/opcount.py b/vj4/model/opcount.py index b2557f939..0acbfdc56 100644 --- a/vj4/model/opcount.py +++ b/vj4/model/opcount.py @@ -35,7 +35,7 @@ @argmethod.wrap async def inc(op: str, ident: str, period_secs: int, max_operations: int, operations: int=1): - coll = db.Collection('opcount') + 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) @@ -54,7 +54,7 @@ async def inc(op: str, ident: str, period_secs: int, max_operations: int, operat @argmethod.wrap async def force_inc(op: str, ident: str, period_secs: int, max_operations: int, operations: int=1): - coll = db.Collection('opcount') + 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) @@ -69,7 +69,7 @@ async def force_inc(op: str, ident: str, period_secs: int, max_operations: int, @argmethod.wrap async def get(op: str, ident: str, period_secs: int, max_operations: int): - coll = db.Collection('opcount') + 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) @@ -84,7 +84,7 @@ async def get(op: str, ident: str, period_secs: int, max_operations: int): @argmethod.wrap async def ensure_indexes(): - coll = db.Collection('opcount') + coll = db.coll('opcount') await coll.create_index([('ident', 1), ('begin_at', 1), ('expire_at', 1)], unique=True) diff --git a/vj4/model/oplog.py b/vj4/model/oplog.py index 0778ccf83..68b4e0352 100644 --- a/vj4/model/oplog.py +++ b/vj4/model/oplog.py @@ -12,7 +12,7 @@ async def add(uid: int, type: int, **kwargs): """Add an operation log. Returns the document id.""" obj_id = objectid.ObjectId() - coll = db.Collection('oplog') + coll = db.coll('oplog') doc = {'_id': obj_id, 'uid': uid, 'type': type, @@ -23,7 +23,7 @@ async def add(uid: int, type: int, **kwargs): @argmethod.wrap async def ensure_indexes(): - coll = db.Collection('oplog') + coll = db.coll('oplog') await coll.create_index('uid') # type delete document await coll.create_index([('doc.domain_id', 1), diff --git a/vj4/model/record.py b/vj4/model/record.py index dadc67d63..04b2c9264 100644 --- a/vj4/model/record.py +++ b/vj4/model/record.py @@ -22,7 +22,7 @@ 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') + coll = db.coll('record') doc = {'hidden': hidden, 'status': constant.record.STATUS_WAITING, 'score': 0, @@ -49,13 +49,13 @@ async def add(domain_id: str, pid: document.convert_doc_id, type: int, uid: int, @argmethod.wrap async def get(record_id: objectid.ObjectId, fields=PROJECTION_ALL): - coll = db.Collection('record') + coll = db.coll('record') return await coll.find_one(record_id, fields) @argmethod.wrap async def rejudge(record_id: objectid.ObjectId, enqueue: bool=True): - coll = db.Collection('record') + coll = db.coll('record') doc = await coll.find_one_and_update(filter={'_id': record_id}, update={'$unset': {'judge_uid': '', 'judge_token': '', @@ -77,7 +77,7 @@ async def rejudge(record_id: objectid.ObjectId, enqueue: bool=True): @argmethod.wrap def get_all_multi(end_id: objectid.ObjectId=None, get_hidden: bool=False, *, fields=None, **kwargs): - coll = db.Collection('record') + coll = db.coll('record') query = {**kwargs, 'hidden': False if not get_hidden else {'$gte': False}} if end_id: query['_id'] = {'$lt': end_id} @@ -86,14 +86,14 @@ def get_all_multi(end_id: objectid.ObjectId=None, get_hidden: bool=False, *, fie @argmethod.wrap def get_multi(get_hidden: bool=False, fields=None, **kwargs): - coll = db.Collection('record') + coll = db.coll('record') kwargs['hidden'] = False if not get_hidden else {'$gte': False} return coll.find(kwargs, projection=fields) @argmethod.wrap async def get_count(begin_id: objectid.ObjectId=None): - coll = db.Collection('record') + coll = db.coll('record') query = {} if begin_id: query['_id'] = {'$gte': begin_id} @@ -103,7 +103,7 @@ async def get_count(begin_id: objectid.ObjectId=None): @argmethod.wrap def get_problem_multi(domain_id: str, pid: document.convert_doc_id, get_hidden: bool=False, type: int=None, *, fields=None): - coll = db.Collection('record') + coll = db.coll('record') query = {'hidden': False if not get_hidden else {'$gte': False}, 'domain_id': domain_id, 'pid': pid} if type != None: @@ -114,7 +114,7 @@ def get_problem_multi(domain_id: str, pid: document.convert_doc_id, @argmethod.wrap def get_user_in_problem_multi(uid: int, domain_id: str, pid: document.convert_doc_id, get_hidden: bool=False, type: int=None, *, fields=None): - coll = db.Collection('record') + coll = db.coll('record') query = {'hidden': False if not get_hidden else {'$gte': False}, 'domain_id': domain_id, 'pid': pid, 'uid': uid} if type != None: @@ -133,7 +133,7 @@ async def get_dict(rids, *, get_hidden=False, fields=None): @argmethod.wrap async def begin_judge(record_id: objectid.ObjectId, judge_uid: int, judge_token: str, status: int): - coll = db.Collection('record') + coll = db.coll('record') doc = await coll.find_one_and_update(filter={'_id': record_id}, update={'$set': {'status': status, 'judge_uid': judge_uid, @@ -148,7 +148,7 @@ async def begin_judge(record_id: objectid.ObjectId, async def next_judge(record_id, judge_uid, judge_token, **kwargs): - coll = db.Collection('record') + coll = db.coll('record') doc = await coll.find_one_and_update(filter={'_id': record_id, 'judge_uid': judge_uid, 'judge_token': judge_token}, @@ -160,7 +160,7 @@ async def next_judge(record_id, judge_uid, judge_token, **kwargs): @argmethod.wrap async def end_judge(record_id: objectid.ObjectId, judge_uid: int, judge_token: str, status: int, score: int, time_ms: int, memory_kb: int): - coll = db.Collection('record') + coll = db.coll('record') doc = await coll.find_one_and_update(filter={'_id': record_id, 'judge_uid': judge_uid, 'judge_token': judge_token}, @@ -176,7 +176,7 @@ async def end_judge(record_id: objectid.ObjectId, judge_uid: int, judge_token: s @argmethod.wrap async def ensure_indexes(): - coll = db.Collection('record') + coll = db.coll('record') await coll.create_index([('hidden', 1), ('_id', -1)]) await coll.create_index([('hidden', 1), diff --git a/vj4/model/system.py b/vj4/model/system.py index 95b0f2c17..fcb1e9d42 100644 --- a/vj4/model/system.py +++ b/vj4/model/system.py @@ -11,7 +11,7 @@ async def inc_user_counter(): Returns: Integer value after increment. """ - coll = db.Collection('system') + coll = db.coll('system') doc = await coll.find_one_and_update(filter={'_id': 'user_counter'}, update={'$inc': {'value': 1}}, upsert=True, @@ -21,7 +21,7 @@ async def inc_user_counter(): @argmethod.wrap async def ensure_indexes(): - coll = db.Collection('system') + coll = db.coll('system') await coll.find_one_and_update(filter={'_id': 'user_counter'}, update={'$setOnInsert': {'value': 1}}, upsert=True) diff --git a/vj4/model/token.py b/vj4/model/token.py index 458f0db84..8d5fb7ff0 100644 --- a/vj4/model/token.py +++ b/vj4/model/token.py @@ -39,7 +39,7 @@ async def add(token_type: int, expire_seconds: int, **kwargs): 'create_at': now, 'update_at': now, 'expire_at': now + datetime.timedelta(seconds=expire_seconds)} - coll = db.Collection('token') + coll = db.coll('token') await coll.insert_one(doc) return binascii.hexlify(id_binary).decode(), doc @@ -56,7 +56,7 @@ async def get(token_id: str, token_type: int): The token document, or None. """ id_binary = binascii.unhexlify(token_id) - coll = db.Collection('token') + coll = db.coll('token') doc = await coll.find_one({'_id': _get_id(id_binary), 'token_type': token_type}) return doc @@ -64,7 +64,7 @@ async def get(token_id: str, token_type: int): @argmethod.wrap async def get_most_recent_session_by_uid(uid: int): """Get the most recent session by uid.""" - coll = db.Collection('token') + coll = db.coll('token') doc = await coll.find_one({'uid': uid, 'token_type': {'$in': [TYPE_SAVED_SESSION, TYPE_UNSAVED_SESSION]}}, sort=[('update_at', -1)]) @@ -74,10 +74,10 @@ async def get_most_recent_session_by_uid(uid: int): @argmethod.wrap async def get_session_list_by_uid(uid: int): """Get the session list by uid.""" - coll = db.Collection('token') + coll = db.coll('token') return await coll.find({'uid': uid, 'token_type': {'$in': [TYPE_SAVED_SESSION, TYPE_UNSAVED_SESSION]}}, - sort=[('create_at', 1)]).to_list(None) + sort=[('create_at', 1)]).to_list() @argmethod.wrap @@ -94,7 +94,7 @@ async def update(token_id: str, token_type: int, expire_seconds: int, **kwargs): The token document, or None. """ id_binary = binascii.unhexlify(token_id) - coll = db.Collection('token') + coll = db.coll('token') assert 'token_type' not in kwargs now = datetime.datetime.utcnow() doc = await coll.find_one_and_update( @@ -123,7 +123,7 @@ async def delete(token_id: str, token_type: int): @argmethod.wrap async def delete_by_hashed_id(hashed_id: str, token_type: int): """Delete a token by the hashed ID.""" - coll = db.Collection('token') + coll = db.coll('token') result = await coll.delete_one({'_id': hashed_id, 'token_type': token_type}) return bool(result.deleted_count) @@ -131,7 +131,7 @@ async def delete_by_hashed_id(hashed_id: str, token_type: int): @argmethod.wrap async def delete_by_uid(uid: int): """Delete all tokens by uid.""" - coll = db.Collection('token') + coll = db.coll('token') result = await coll.delete_many({'uid': uid, 'token_type': {'$in': [TYPE_SAVED_SESSION, TYPE_UNSAVED_SESSION]}}) @@ -140,7 +140,7 @@ async def delete_by_uid(uid: int): @argmethod.wrap async def ensure_indexes(): - coll = db.Collection('token') + coll = db.coll('token') await coll.create_index([('uid', 1), ('token_type', 1), ('update_at', -1)], sparse=True) await coll.create_index('expire_at', expireAfterSeconds=0) diff --git a/vj4/model/user.py b/vj4/model/user.py index 40a71e99f..9b8e4f9d9 100644 --- a/vj4/model/user.py +++ b/vj4/model/user.py @@ -34,7 +34,7 @@ async def add(uid: int, uname: str, password: str, mail: str, regip: str=''): raise error.UserAlreadyExistError(uname) salt = pwhash.gen_salt() - coll = db.Collection('user') + coll = db.coll('user') try: await coll.insert_one({'_id': uid, 'uname': uname, @@ -59,7 +59,7 @@ async def get_by_uid(uid: int, fields=PROJECTION_VIEW): for user in builtin.USERS: if user['_id'] == uid: return user - coll = db.Collection('user') + coll = db.coll('user') return await coll.find_one({'_id': uid}, fields) @@ -70,7 +70,7 @@ async def get_by_uname(uname: str, fields=PROJECTION_VIEW): for user in builtin.USERS: if user['uname_lower'] == uname_lower: return user - coll = db.Collection('user') + coll = db.coll('user') return await coll.find_one({'uname_lower': uname_lower}, fields) @@ -81,13 +81,13 @@ async def get_by_mail(mail: str, fields=PROJECTION_VIEW): for user in builtin.USERS: if user['mail_lower'] == mail_lower: return user - coll = db.Collection('user') + coll = db.coll('user') return await coll.find_one({'mail_lower': mail_lower}, fields) def get_multi(*, fields=PROJECTION_VIEW, **kwargs): """Get multiple users.""" - coll = db.Collection('user') + coll = db.coll('user') return coll.find(kwargs, fields) @@ -124,7 +124,7 @@ async def set_password(uid: int, password: str): """Set password. Returns doc or None.""" validator.check_password(password) salt = pwhash.gen_salt() - coll = db.Collection('user') + coll = db.coll('user') doc = await coll.find_one_and_update(filter={'_id': uid}, update={'$set': {'salt': salt, 'hash': pwhash.hash_vj4(password, salt)}}, @@ -147,7 +147,7 @@ async def change_password(uid: int, current_password: str, password: str): return None validator.check_password(password) salt = pwhash.gen_salt() - coll = db.Collection('user') + coll = db.coll('user') doc = await coll.find_one_and_update(filter={'_id': doc['_id'], 'salt': doc['salt'], 'hash': doc['hash']}, @@ -158,7 +158,7 @@ async def change_password(uid: int, current_password: str, password: str): async def set_by_uid(uid, **kwargs): - coll = db.Collection('user') + coll = db.coll('user') doc = await coll.find_one_and_update(filter={'_id': uid}, update={'$set': kwargs}, return_document=ReturnDocument.AFTER) return doc @@ -191,10 +191,10 @@ async def set_default(uid: int): async def get_prefix_list(prefix: str, fields=PROJECTION_VIEW, limit: int=50): prefix = prefix.lower() regex = '\\A\\Q{0}\\E'.format(prefix.replace('\\E', '\\E\\\\E\\Q')) - coll = db.Collection('user') + coll = db.coll('user') udocs = await coll.find({'uname_lower': {'$regex': regex}}, projection=fields) \ .limit(limit) \ - .to_list(None) + .to_list() for udoc in builtin.USERS: if udoc['uname_lower'].startswith(prefix): udocs.append(udoc) @@ -203,13 +203,13 @@ async def get_prefix_list(prefix: str, fields=PROJECTION_VIEW, limit: int=50): @argmethod.wrap async def count(**kwargs): - coll = db.Collection('user') + coll = db.coll('user') return coll.find({**kwargs}).count() @argmethod.wrap async def ensure_indexes(): - coll = db.Collection('user') + coll = db.coll('user') await coll.create_index('uname_lower', unique=True) await coll.create_index('mail_lower', sparse=True) diff --git a/vj4/pipeline/problem_stat.py b/vj4/pipeline/problem_stat.py index a852ff7d7..d66300ca5 100644 --- a/vj4/pipeline/problem_stat.py +++ b/vj4/pipeline/problem_stat.py @@ -31,8 +31,8 @@ async def main(): }, ] - bulk = db.Collection('document').initialize_unordered_bulk_op() - async for adoc in db.Collection('record').aggregate(pipeline): + bulk = db.coll('document').initialize_unordered_bulk_op() + async for adoc in db.coll('record').aggregate(pipeline): bulk.find({'domain_id': adoc['_id']['domain_id'], 'doc_type': document.TYPE_PROBLEM, 'doc_id': adoc['_id']['pid']}) \ diff --git a/vj4/test/base.py b/vj4/test/base.py index ac54db06b..663edb727 100644 --- a/vj4/test/base.py +++ b/vj4/test/base.py @@ -18,8 +18,11 @@ class DatabaseTestCase(unittest.TestCase): def setUp(self): - db.Database._instance, db.Collection._instances, db.GridFS._instances = None, {}, {} + db._db = None + db.coll.cache_clear() + db.fs.cache_clear() options.db_name = 'unittest_' + str(os.getpid()) + wait(db.init()) wait(tools.ensure_all_indexes()) def tearDown(self): diff --git a/vj4/test/test_contest.py b/vj4/test/test_contest.py index 1f5dedc13..9727433ca 100644 --- a/vj4/test/test_contest.py +++ b/vj4/test/test_contest.py @@ -111,7 +111,7 @@ async def test_add_get(self): self.assertEqual(tdoc['owner_uid'], OWNER_UID) self.assertEqual(tdoc['title'], TITLE) self.assertEqual(tdoc['content'], CONTENT) - tdocs = await contest.get_multi(DOMAIN_ID_DUMMY, fields=['title']).to_list(None) + tdocs = await contest.get_multi(DOMAIN_ID_DUMMY, fields=['title']).to_list() self.assertEqual(len(tdocs), 1) self.assertEqual(tdocs[0]['title'], TITLE) self.assertFalse('content' in tdocs[0]) diff --git a/vj4/test/test_discussion.py b/vj4/test/test_discussion.py index f6b34ef00..d4dcb74cb 100644 --- a/vj4/test/test_discussion.py +++ b/vj4/test/test_discussion.py @@ -95,7 +95,7 @@ async def test_add_get(self): vnode = await discussion.get_vnode(DOMAIN_ID_DUMMY, 'meow') ddocs = await discussion.get_multi(DOMAIN_ID_DUMMY, parent_doc_type=vnode['doc_type'], - parent_doc_id=vnode['doc_id']).to_list(None) + parent_doc_id=vnode['doc_id']).to_list() self.assertEqual(len(ddocs), 0) did = await discussion.add(DOMAIN_ID_DUMMY, 'meow', OWNER_UID, TITLE, CONTENT) ddoc = await discussion.get(DOMAIN_ID_DUMMY, did) @@ -112,7 +112,7 @@ async def test_add_get(self): parent_doc_id=vnode['doc_id'], fields=['title', 'owner_uid', - 'parent_doc_id']).to_list(None) + 'parent_doc_id']).to_list() self.assertEqual(len(ddocs), 1) self.assertEqual(ddocs[0]['title'], TITLE) self.assertFalse('content' in ddocs[0]) @@ -124,7 +124,7 @@ async def test_reply_add_get_del(self): vnode = await discussion.get_vnode(DOMAIN_ID_DUMMY, 'meow') ddocs = await discussion.get_multi(DOMAIN_ID_DUMMY, parent_doc_type=vnode['doc_type'], - parent_doc_id=vnode['doc_id']).to_list(None) + parent_doc_id=vnode['doc_id']).to_list() self.assertEqual(len(ddocs), 0) did = await discussion.add(DOMAIN_ID_DUMMY, 'meow', OWNER_UID, TITLE, CONTENT) drid = await discussion.add_reply(DOMAIN_ID_DUMMY, did, OWNER_UID, REPLY_CONTENT) diff --git a/vj4/test/test_problem.py b/vj4/test/test_problem.py index c095b9db4..2b4ee902a 100644 --- a/vj4/test/test_problem.py +++ b/vj4/test/test_problem.py @@ -24,7 +24,7 @@ async def test_add_get(self): self.assertEqual(pdoc['content'], CONTENT) self.assertEqual(pdoc['owner_uid'], UID) self.assertEqual(pdoc['doc_id'], pid) - pdocs = await problem.get_multi(domain_id=DOMAIN_ID, fields=['doc_id', 'title']).to_list(None) + pdocs = await problem.get_multi(domain_id=DOMAIN_ID, fields=['doc_id', 'title']).to_list() self.assertEqual(len(pdocs), 1) self.assertEqual(pdocs[0]['doc_id'], pid) self.assertEqual(pdocs[0]['title'], TITLE) diff --git a/vj4/util/argmethod.py b/vj4/util/argmethod.py index 2986b1b15..9b8b1c455 100644 --- a/vj4/util/argmethod.py +++ b/vj4/util/argmethod.py @@ -11,6 +11,7 @@ import logging import logging.config +from vj4 import db from vj4.util import options options.define('pretty', default=False, help='Pretty print the result.') @@ -73,6 +74,7 @@ def invoke_by_args(): parser.print_help() else: loop = asyncio.get_event_loop() + loop.run_until_complete(db.init()) try: result = _methods[name](**vars(args)) if inspect.iscoroutine(result): diff --git a/vj4/util/pagination.py b/vj4/util/pagination.py index 98513809c..e2d82d3d0 100644 --- a/vj4/util/pagination.py +++ b/vj4/util/pagination.py @@ -8,6 +8,6 @@ async def paginate(cursor, page: int, page_size: int): count, page_docs = await asyncio.gather(cursor.count(), cursor.skip((page - 1) * page_size) \ .limit(page_size) \ - .to_list(None)) + .to_list()) num_pages = (count + page_size - 1) // page_size return page_docs, num_pages, count