From 16915ed5075b6a93585bc8e8d2cab61f1b0bab77 Mon Sep 17 00:00:00 2001
From: "Joshua M. Boniface" <joshua@boniface.me>
Date: Wed, 5 Oct 2022 17:35:05 -0400
Subject: [PATCH] Add better exception handling with ctxtmgrs

---
 api-daemon/pvcapid/vmbuilder.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/api-daemon/pvcapid/vmbuilder.py b/api-daemon/pvcapid/vmbuilder.py
index 43c43f5d..6ee16ffd 100755
--- a/api-daemon/pvcapid/vmbuilder.py
+++ b/api-daemon/pvcapid/vmbuilder.py
@@ -159,9 +159,13 @@ def open_db(config):
             password=config["database_password"],
         )
         cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
-        yield cur
     except Exception:
         raise ClusterError("Failed to connect to Postgres")
+
+    try:
+        yield cur
+    except Exception:
+        raise
     finally:
         conn.commit()
         cur.close()
@@ -174,9 +178,13 @@ def open_zk(config):
     try:
         zkhandler = ZKHandler(config)
         zkhandler.connect()
-        yield zkhandler
     except Exception:
         raise ClusterError("Failed to connect to Zookeeper")
+
+    try:
+        yield zkhandler
+    except Exception:
+        raise
     finally:
         zkhandler.disconnect()
         del zkhandler