Error & exception handling for the annotation, reduced indentation

Author: Invictus17

leaderelection/leaderelection.py

@@ -20,6 +20,7 @@
 import logging
 logging.basicConfig(level=logging.INFO)
 
+
 class LeaderElection:
     def __init__(self, election_config):
         if election_config is None:
@@ -36,12 +37,12 @@ def __init__(self, election_config):
 
 
     # Annotation used in the lock object
-    def leaderelector_record(self, election_config):
+    def leaderelector_record(self, election_config, now):
         return {
                     "holderIdentity": str(election_config.lock.identity),
                     "leaseDurationSeconds": str(election_config.lease_duration),
-                    "acquireTime": str(datetime.datetime.now()),
-                    "renewTime": str(datetime.datetime.now())
+                    "acquireTime": str(now),
+                    "renewTime": str(now)
                 }
 
     # Point of entry to Leader election
@@ -97,51 +98,22 @@ def renew_loop(self):
             return
 
     def try_acquire_or_renew(self):
+        now = datetime.datetime.now()
+
         # Check if lock is created
-        lock_status, lock_response = self.election_config.lock.get(self.election_config.lock.name, self.election_config.lock.namespace)
+        lock_status, lock_response, lock_record = self.election_config.lock.get(self.election_config.lock.name,
+                                                                                self.election_config.lock.namespace)
 
         # create a default Election record for this candidate
-        leader_election_record = self.leaderelector_record(self.election_config)
-        
-        # If a lock is already created with that name
-        if lock_status:
-            old_election_record = ast.literal_eval(lock_response.metadata.annotations[self.election_config.lock.leader_electionrecord_annotationkey])
-
-            # Report transitions
-            if self.observed_record and self.observed_record['holderIdentity'] != old_election_record['holderIdentity']:
-                logging.info("Leader has switched to {}".format(old_election_record['holderIdentity']))
-
-
-            if old_election_record != self.observed_record:
-                self.observed_record = old_election_record
-                self.observed_time_milliseconds = int(time.time()*1000)
-
-
-            # If This candidate is not the leader and lease duration is yet to finish
-            if str(self.election_config.lock.identity) != self.observed_record['holderIdentity'] and self.observed_time_milliseconds + self.election_config.lease_duration*1000 > int(time.time()*1000):
-                logging.info("yet to finish lease_duration, lease held by {} and has not expired".format(old_election_record['holderIdentity']))
-                return False
+        leader_election_record = self.leaderelector_record(self.election_config, now)
 
-            # If this candidate is the Leader
-            if str(self.election_config.lock.identity) == self.observed_record['holderIdentity']:
-                # Leader sets acquireTime
-                leader_election_record['acquireTime'] = self.observed_record['acquireTime']
-
-            # Update object with latest election record
-            lock_response.metadata.annotations[self.election_config.lock.leader_electionrecord_annotationkey] = str(leader_election_record)
-            update_status, update_response = self.election_config.lock.update(self.election_config.lock.name, self.election_config.lock.namespace, lock_response)
-
-            if update_status is False:
-                logging.info("{} failed to acquire lease".format(leader_election_record['holderIdentity']))
+        # A lock is not created with that name, try to create one
+        if not lock_status:
+            if ast.literal_eval(lock_response.body)['code'] != 404:
+                logging.info("Error retrieving resource lock {} as {}".format(self.election_config.lock.name, lock_response.reason))
                 return False
 
-            self.observed_record = leader_election_record
-            self.observed_time_milliseconds = int(time.time()*1000)
-            logging.info("leader {} has successfully acquired lease".format(leader_election_record['holderIdentity']))
-            return True
 
-        # A lock is not created with that name, try to create one
-        else:
             logging.info("{} is trying to create a lock".format(leader_election_record['holderIdentity']))
             create_status, create_response = self.election_config.lock.create(name=self.election_config.lock.name,
                                                                               namespace=self.election_config.lock.namespace,
@@ -151,6 +123,60 @@ def try_acquire_or_renew(self):
                 return False
 
             self.observed_record = leader_election_record
-            self.observed_time_milliseconds = int(time.time()*1000)
+            self.observed_time_milliseconds = int(time.time() * 1000)
             return True
 
+        # A lock exists with that name
+        # Validate lock_record
+        if lock_record is None:
+            # try to update lock with proper annotation and election record
+            return self.update_lock(lock_response, leader_election_record)
+
+        # check for any key, value errors in the record
+        try:
+            old_election_record = ast.literal_eval(lock_record)
+            if (old_election_record['holderIdentity'] == '' or old_election_record['leaseDurationSeconds'] == ''
+                    or old_election_record['acquireTime'] == '' or old_election_record['renewTime'] == ''):
+                # try to update lock with proper annotation and election record
+                return self.update_lock(lock_response, leader_election_record)
+        except:
+            # try to update lock with proper annotation and election record
+            return self.update_lock(lock_response, leader_election_record)
+
+
+        # Report transitions
+        if self.observed_record and self.observed_record['holderIdentity'] != old_election_record['holderIdentity']:
+            logging.info("Leader has switched to {}".format(old_election_record['holderIdentity']))
+
+        if old_election_record != self.observed_record:
+            self.observed_record = old_election_record
+            self.observed_time_milliseconds = int(time.time() * 1000)
+
+        # If This candidate is not the leader and lease duration is yet to finish
+        if (str(self.election_config.lock.identity) != self.observed_record['holderIdentity']
+                and self.observed_time_milliseconds + self.election_config.lease_duration*1000 > int(now.timestamp()*1000)):
+            logging.info("yet to finish lease_duration, lease held by {} and has not expired".format(old_election_record['holderIdentity']))
+            return False
+
+        # If this candidate is the Leader
+        if str(self.election_config.lock.identity) == self.observed_record['holderIdentity']:
+            # Leader sets acquireTime
+            leader_election_record['acquireTime'] = self.observed_record['acquireTime']
+
+        return self.update_lock(lock_response, leader_election_record)
+
+    def update_lock(self, lock_response, leader_election_record):
+        # Update object with latest election record
+        update_status, update_response = self.election_config.lock.update(self.election_config.lock.name,
+                                                                          self.election_config.lock.namespace,
+                                                                          lock_response, leader_election_record)
+
+        if update_status is False:
+            logging.info("{} failed to acquire lease".format(leader_election_record['holderIdentity']))
+            return False
+
+        self.observed_record = leader_election_record
+        self.observed_time_milliseconds = int(time.time() * 1000)
+        logging.info("leader {} has successfully acquired lease".format(leader_election_record['holderIdentity']))
+        return True
+

leaderelection/resourcelock/configmaplock.py

@@ -16,7 +16,6 @@
 from kubernetes import client, config
 from kubernetes.client.api_client import ApiClient
 
-
 class ConfigMapLock:
     def __init__(self, name, namespace, identity):
         """
@@ -35,13 +34,24 @@ def get(self, name, namespace):
         """
         :param name: Name of the configmap object information to get
         :param namespace: Namespace in which the configmap object is to be searched
-        :return: True, response if object found else False, exception
+        :return: 'True, response, record' if object found else 'False, exception, None'
         """
         try:
             api_response = self.api_instance.read_namespaced_config_map(name, namespace)
-            return True, api_response
+
+            # If an annotation does not exist - add the leader_electionrecord_annotationkey
+            if api_response.metadata.annotations is None or api_response.metadata.annotations == '':
+                api_response.metadata.annotations = {self.leader_electionrecord_annotationkey: ''}
+                return True, api_response, None
+
+            # If an annotation exists but, the leader_electionrecord_annotationkey does not then add it as a key
+            if not api_response.metadata.annotations.get(self.leader_electionrecord_annotationkey):
+                api_response.metadata.annotations = {self.leader_electionrecord_annotationkey: ''}
+                return True, api_response, None
+
+            return True, api_response, api_response.metadata.annotations[self.leader_electionrecord_annotationkey]
         except ApiException as e:
-            return False, e
+            return False, e, None
 
     def create(self, name, namespace, election_record):
         """
@@ -59,15 +69,18 @@ def create(self, name, namespace, election_record):
         except ApiException as e:
             return False, e
 
-    def update(self, name, namespace, updated_record):
+    def update(self, name, namespace, update_object, updated_record):
         """
         :param name: name of the lock to be updated
         :param namespace: namespace the lock is in
-        :param ElectionRecord: the updated record
+        :param update_object: the object used to update the lock
+        :param updated_record: the updated election record
         :return: True if update is succesful False if it fails
         """
         try:
-            api_response = self.api_instance.replace_namespaced_config_map(name=name, namespace=namespace, body=updated_record)
+            # Set the updated record
+            update_object.metadata.annotations[self.leader_electionrecord_annotationkey] = str(updated_record)
+            api_response = self.api_instance.replace_namespaced_config_map(name=name, namespace=namespace, body=update_object)
             return True, api_response
         except ApiException as e:
             return False, e