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 functools | import functools | ||||
import inspect | |||||
import sys | |||||
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() if param.name not in names] | params = [param for param in params.values() if param.name not in names] | ||||
sig = sig.replace(parameters=params) | sig = sig.replace(parameters=params) | ||||
Show All 12 Lines | 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 check_no_transaction(f_name): | |||||
"""Checks there is no open DB transaction in the stack | |||||
(db + cur variables). | |||||
While it is not necessarily broken code, it is very likely | |||||
a mistake.""" | |||||
frame = sys._getframe().f_back.f_back | |||||
while frame: | |||||
if "db" in frame.f_locals and "cur" in frame.f_locals: | |||||
raise AssertionError( | |||||
f'Calling function {f_name} without "db" and "cur" arguments ' | |||||
"from a function ({frame.f_code.co_name}) with these variables." | |||||
) | |||||
frame = frame.f_back | |||||
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("Use db_transaction_generator for generator functions.") | raise ValueError("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: | ||||
check_no_transaction(meth.__name__) | |||||
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) | ||||
Show All 19 Lines | def decorator(meth, __client_options=client_options): | ||||
@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: | ||||
check_no_transaction(meth.__name__) | |||||
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 |