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): | ||||
"""Applies the given postgresql client options to the given cursor. | """Applies the given postgresql client options to the given cursor. | ||||
Returns a dictionary with the old values if they changed.""" | Returns a dictionary with the old values if they changed.""" | ||||
old_options = {} | old_options = {} | ||||
for option, value in options.items(): | for option, value in options.items(): | ||||
cursor.execute('SHOW %s' % option) | cursor.execute("SHOW %s" % option) | ||||
old_value = cursor.fetchall()[0][0] | old_value = cursor.fetchall()[0][0] | ||||
if old_value != value: | if old_value != value: | ||||
cursor.execute('SET LOCAL %s TO %%s' % option, (value,)) | cursor.execute("SET LOCAL %s TO %%s" % option, (value,)) | ||||
old_options[option] = old_value | old_options[option] = old_value | ||||
return old_options | return old_options | ||||
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 |