diff --git a/nomenklatura/db.py b/nomenklatura/db.py index 05d4e967..c586338e 100644 --- a/nomenklatura/db.py +++ b/nomenklatura/db.py @@ -38,16 +38,12 @@ def get_statement_table() -> Table: @contextmanager def ensure_tx(conn: Connish = None) -> Generator[Connection, None, None]: - try: - if conn is not None: - yield conn - else: - engine = get_engine() - with engine.begin() as conn: - yield conn - finally: - if conn is not None: - conn.commit() + if conn is not None: + yield conn + return + engine = get_engine() + with engine.begin() as conn: + yield conn @cache diff --git a/nomenklatura/store/sql.py b/nomenklatura/store/sql.py index 0e4e951a..222f2280 100644 --- a/nomenklatura/store/sql.py +++ b/nomenklatura/store/sql.py @@ -106,6 +106,7 @@ def flush(self) -> None: ) with ensure_tx(self.store.engine.connect()) as conn: conn.execute(stmt) + conn.commit() self.batch = set() self.batch_size = 0 @@ -129,6 +130,7 @@ def pop(self, entity_id: str) -> List[Statement]: statements: List[Statement] = list(self.store._iterate_stmts(q)) with ensure_tx(self.store.engine.connect()) as conn: conn.execute(q_delete) + conn.commit() return statements