From 60db800d9c7bea7e6ece41681334059d4ab58b53 Mon Sep 17 00:00:00 2001
From: "Joshua M. Boniface" <joshua@boniface.me>
Date: Sun, 30 May 2021 15:59:37 -0400
Subject: [PATCH] Use full ZKHandler in provisioner

Required due to references to self from Celery that are replaced by the
ZKConnection self instance.
---
 api-daemon/pvcapid/provisioner.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/api-daemon/pvcapid/provisioner.py b/api-daemon/pvcapid/provisioner.py
index dcebdbfb..6e7cf159 100755
--- a/api-daemon/pvcapid/provisioner.py
+++ b/api-daemon/pvcapid/provisioner.py
@@ -26,7 +26,7 @@ import re
 
 from pvcapid.Daemon import config, strtobool
 
-from daemon_lib.zkhandler import ZKConnection
+from daemon_lib.zkhandler import ZKHandler
 
 import daemon_lib.common as pvc_common
 import daemon_lib.node as pvc_node
@@ -1051,8 +1051,7 @@ def delete_profile(name):
 #
 # Main VM provisioning function - executed by the Celery worker
 #
-@ZKConnection(config)
-def create_vm(self, zkhandler, vm_name, vm_profile, define_vm=True, start_vm=True, script_run_args=[]):
+def create_vm(self, vm_name, vm_profile, define_vm=True, start_vm=True, script_run_args=[]):
     # Runtime imports
     import time
     import importlib
@@ -1068,8 +1067,13 @@ def create_vm(self, zkhandler, vm_name, vm_profile, define_vm=True, start_vm=Tru
     try:
         db_conn, db_cur = open_database(config)
     except Exception:
-        print('FATAL - failed to connect to Postgres')
-        raise Exception
+        raise ClusterError('Failed to connect to Postgres')
+
+    try:
+        zkhandler = ZKHandler(config)
+        zkhandler.connect()
+    except Exception:
+        raise ClusterError('Failed to connect to Zookeeper')
 
     # Phase 1 - setup
     #  * Get the profile elements
@@ -1646,4 +1650,7 @@ def create_vm(self, zkhandler, vm_name, vm_profile, define_vm=True, start_vm=Tru
             retcode, retmsg = pvc_vm.start_vm(zkhandler, vm_name)
             print(retmsg)
 
+    zkhandler.disconnect()
+    del zkhandler
+
     return {'status': 'VM "{}" with profile "{}" has been provisioned and started successfully'.format(vm_name, vm_profile), 'current': 10, 'total': 10}