def _set_engine(ctx, param, value):
if not value:
return
try:
from splitgraph.engine import get_engine, set_engine
engine = get_engine(value)
set_engine(engine)
except KeyError:
raise click.BadParameter("Unknown remote %s!" % value)
def wrapped(*args, **kwargs):
from splitgraph.engine import get_engine, set_engine
engine = get_engine()
try:
f(*args, **kwargs)
engine.commit()
except Exception:
engine.rollback()
raise
finally:
engine.close()
# In the context of a test run, we need to switch the global engine
# back to LOCAL (since the engine-switching decorator doesn't
# get control, so we can't do it there).
set_engine(get_engine("LOCAL"))