From 8ea47aafa4eb0dee924878e16aa764a794a4c3a6 Mon Sep 17 00:00:00 2001 From: Jan Losinski Date: Wed, 21 Mar 2018 23:50:01 +0100 Subject: [PATCH] Add context based session. Signed-off-by: Jan Losinski --- wallabag_kindle_consumer/models.py | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/wallabag_kindle_consumer/models.py b/wallabag_kindle_consumer/models.py index d5fffce..eb7889f 100644 --- a/wallabag_kindle_consumer/models.py +++ b/wallabag_kindle_consumer/models.py @@ -31,14 +31,30 @@ class Job(Base): format = Column(Enum('pdf', 'mobi')) -def session_maker(uri): +class ContextSession: + def __init__(self, session_maker): + self.session_maker = session_maker + + def __enter__(self): + self.session = self.session_maker() + return self.session + + def __exit__(self, exc_type, exc_val, exc_tb): + self.session.close() + + +def context_session(config): + return ContextSession(session_maker(config)) + + +def session_maker(config): Session = sessionmaker(autocommit=False, autoflush=False, - bind=create_engine(uri)) + bind=create_engine(config.db_uri)) return Session -def re_create_db(uri): - engine = create_engine(uri) +def re_create_db(config): + engine = create_engine(config.db_uri) Base.metadata.drop_all(engine) Base.metadata.create_all(engine)