Changeset View
Changeset View
Standalone View
Standalone View
swh/core/db/common.py
# Copyright (C) 2015-2019 The Software Heritage developers | # Copyright (C) 2015-2019 The Software Heritage developers | ||||
# See the AUTHORS file at the top-level directory of this distribution | # See the AUTHORS file at the top-level directory of this distribution | ||||
# License: GNU General Public License version 3, or any later version | # License: GNU General Public License version 3, or any later version | ||||
# See top-level LICENSE file for more information | # See top-level LICENSE file for more information | ||||
import inspect | import inspect | ||||
import functools | import functools | ||||
def remove_kwargs(names): | def remove_kwargs(names): | ||||
def decorator(f): | def decorator(f): | ||||
sig = inspect.signature(f) | sig = inspect.signature(f) | ||||
params = sig.parameters | params = sig.parameters | ||||
params = [param for param in params.values() | params = [param for param in params.values() if param.name not in names] | ||||
if param.name not in names] | |||||
sig = sig.replace(parameters=params) | sig = sig.replace(parameters=params) | ||||
f.__signature__ = sig | f.__signature__ = sig | ||||
return f | return f | ||||
return decorator | return decorator | ||||
def apply_options(cursor, options): | def apply_options(cursor, options): | ||||
Show All 12 Lines | |||||
def db_transaction(**client_options): | def db_transaction(**client_options): | ||||
"""decorator to execute Backend methods within DB transactions | """decorator to execute Backend methods within DB transactions | ||||
The decorated method must accept a `cur` and `db` keyword argument | The decorated method must accept a `cur` and `db` keyword argument | ||||
Client options are passed as `set` options to the postgresql server | Client options are passed as `set` options to the postgresql server | ||||
""" | """ | ||||
def decorator(meth, __client_options=client_options): | def decorator(meth, __client_options=client_options): | ||||
if inspect.isgeneratorfunction(meth): | if inspect.isgeneratorfunction(meth): | ||||
raise ValueError( | raise ValueError('Use db_transaction_generator for generator functions.') | ||||
'Use db_transaction_generator for generator functions.') | |||||
@remove_kwargs(['cur', 'db']) | @remove_kwargs(['cur', 'db']) | ||||
@functools.wraps(meth) | @functools.wraps(meth) | ||||
def _meth(self, *args, **kwargs): | def _meth(self, *args, **kwargs): | ||||
if 'cur' in kwargs and kwargs['cur']: | if 'cur' in kwargs and kwargs['cur']: | ||||
cur = kwargs['cur'] | cur = kwargs['cur'] | ||||
old_options = apply_options(cur, __client_options) | old_options = apply_options(cur, __client_options) | ||||
ret = meth(self, *args, **kwargs) | ret = meth(self, *args, **kwargs) | ||||
apply_options(cur, old_options) | apply_options(cur, old_options) | ||||
return ret | return ret | ||||
else: | else: | ||||
db = self.get_db() | db = self.get_db() | ||||
try: | try: | ||||
with db.transaction() as cur: | with db.transaction() as cur: | ||||
apply_options(cur, __client_options) | apply_options(cur, __client_options) | ||||
return meth(self, *args, db=db, cur=cur, **kwargs) | return meth(self, *args, db=db, cur=cur, **kwargs) | ||||
finally: | finally: | ||||
self.put_db(db) | self.put_db(db) | ||||
return _meth | return _meth | ||||
return decorator | return decorator | ||||
def db_transaction_generator(**client_options): | def db_transaction_generator(**client_options): | ||||
"""decorator to execute Backend methods within DB transactions, while | """decorator to execute Backend methods within DB transactions, while | ||||
returning a generator | returning a generator | ||||
The decorated method must accept a `cur` and `db` keyword argument | The decorated method must accept a `cur` and `db` keyword argument | ||||
Client options are passed as `set` options to the postgresql server | Client options are passed as `set` options to the postgresql server | ||||
""" | """ | ||||
def decorator(meth, __client_options=client_options): | def decorator(meth, __client_options=client_options): | ||||
if not inspect.isgeneratorfunction(meth): | if not inspect.isgeneratorfunction(meth): | ||||
raise ValueError( | raise ValueError('Use db_transaction for non-generator functions.') | ||||
'Use db_transaction for non-generator functions.') | |||||
@remove_kwargs(['cur', 'db']) | @remove_kwargs(['cur', 'db']) | ||||
@functools.wraps(meth) | @functools.wraps(meth) | ||||
def _meth(self, *args, **kwargs): | def _meth(self, *args, **kwargs): | ||||
if 'cur' in kwargs and kwargs['cur']: | if 'cur' in kwargs and kwargs['cur']: | ||||
cur = kwargs['cur'] | cur = kwargs['cur'] | ||||
old_options = apply_options(cur, __client_options) | old_options = apply_options(cur, __client_options) | ||||
yield from meth(self, *args, **kwargs) | yield from meth(self, *args, **kwargs) | ||||
apply_options(cur, old_options) | apply_options(cur, old_options) | ||||
else: | else: | ||||
db = self.get_db() | db = self.get_db() | ||||
try: | try: | ||||
with db.transaction() as cur: | with db.transaction() as cur: | ||||
apply_options(cur, __client_options) | apply_options(cur, __client_options) | ||||
yield from meth(self, *args, db=db, cur=cur, **kwargs) | yield from meth(self, *args, db=db, cur=cur, **kwargs) | ||||
finally: | finally: | ||||
self.put_db(db) | self.put_db(db) | ||||
return _meth | return _meth | ||||
return decorator | return decorator |