SQL: mess around with ensure_tx contextmanager

opensanctions · Aug 9, 2023 · 84600b6 · 84600b6
1 parent 33e5456
commit 84600b6
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/nomenklatura/db.py b/nomenklatura/db.py
@@ -25,9 +25,13 @@ def get_metadata() -> MetaData:
 
 @contextmanager
 def ensure_tx(conn: Connish = None) -> Generator[Connection, None, None]:
-    if conn is not None:
-        yield conn
-        return
-    engine = get_engine()
-    with engine.begin() as conn:
-        yield conn
+    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()