<s> from __future__ import absolute_import , unicode_literals <EOL> import logging <EOL> import simple_audit <EOL> from django . db import models <EOL> from django . utils . translation import ugettext_lazy as _ <EOL> from dateutil import tz <EOL> from datetime import datetime <EOL> from physical . models import Host <EOL> from util . models import BaseModel <EOL> from django . db . models . signals import post_save <EOL> from django . db . models . signals import pre_delete <EOL> from django . dispatch import receiver <EOL> from celery . task import control <EOL> from maintenance . tasks import execute_scheduled_maintenance <EOL> from . registered_functions . functools import _get_registered_functions <EOL> LOG = logging . getLogger ( __name__ ) <EOL> class Maintenance ( BaseModel ) : <EOL> WAITING = <NUM_LIT:0> <EOL> RUNNING = <NUM_LIT:1> <EOL> FINISHED = <NUM_LIT:2> <EOL> REJECTED = <NUM_LIT:3> <EOL> REVOKED = <NUM_LIT:4> <EOL> MAINTENANCE_STATUS = ( <EOL> ( FINISHED , '<STR_LIT>' ) , <EOL> ( RUNNING , '<STR_LIT>' ) , <EOL> ( WAITING , '<STR_LIT>' ) , <EOL> ( REJECTED , '<STR_LIT>' ) , <EOL> ( REVOKED , '<STR_LIT>' ) <EOL> ) <EOL> description = models . CharField ( verbose_name = _ ( "<STR_LIT>" ) , <EOL> null = False , blank = False , max_length = <NUM_LIT> , ) <EOL> scheduled_for = models . DateTimeField ( verbose_name = _ ( "<STR_LIT>" ) , <EOL> unique = True , ) <EOL> main_script = models . TextField ( verbose_name = _ ( "<STR_LIT>" ) , <EOL> null = False , blank = False ) <EOL> rollback_script = models . TextField ( verbose_name = _ ( "<STR_LIT>" ) , <EOL> null = True , blank = True ) <EOL> maximum_workers = models . PositiveSmallIntegerField ( verbose_name = _ ( "<STR_LIT>" ) , <EOL> null = False , default = <NUM_LIT:1> ) <EOL> celery_task_id = models . CharField ( verbose_name = _ ( "<STR_LIT>" ) , <EOL> null = True , blank = True , max_length = <NUM_LIT:50> , ) <EOL> status = models . IntegerField ( choices = MAINTENANCE_STATUS , default = WAITING ) <EOL> affected_hosts = models . IntegerField ( <EOL> verbose_name = _ ( "<STR_LIT>" ) , default = <NUM_LIT:0> ) <EOL> started_at = models . DateTimeField ( <EOL> verbose_name = _ ( "<STR_LIT>" ) , null = True , blank = True ) <EOL> finished_at = models . DateTimeField ( <EOL> verbose_name = _ ( "<STR_LIT>" ) , null = True , blank = True ) <EOL> hostsid = models . CommaSeparatedIntegerField ( verbose_name = _ ( "<STR_LIT>" ) , <EOL> null = False , blank = False , max_length = <NUM_LIT> ) <EOL> created_by = models . CharField ( <EOL> verbose_name = _ ( "<STR_LIT>" ) , max_length = <NUM_LIT:255> , null = True , blank = True ) <EOL> revoked_by = models . CharField ( <EOL> verbose_name = _ ( "<STR_LIT>" ) , max_length = <NUM_LIT:255> , null = True , blank = True ) <EOL> def __unicode__ ( self ) : <EOL> return "<STR_LIT:%s>" % self . description <EOL> class Meta : <EOL> permissions = ( <EOL> ( "<STR_LIT>" , "<STR_LIT>" ) , <EOL> ) <EOL> def save_host_maintenance ( self , ) : <EOL> save_host_ok = True <EOL> total_hosts = <NUM_LIT:0> <EOL> try : <EOL> hostsid_list = self . hostsid . split ( '<STR_LIT:U+002C>' ) <EOL> hosts = Host . objects . filter ( pk__in = hostsid_list ) <EOL> for host in hosts : <EOL> hm = HostMaintenance ( ) <EOL> hm . host = host <EOL> hm . maintenance = self <EOL> hm . hostname = host . hostname <EOL> hm . save ( ) <EOL> total_hosts += <NUM_LIT:1> <EOL> except Exception , e : <EOL> error = e . args [ <NUM_LIT:1> ] <EOL> LOG . warn ( "<STR_LIT>" . format ( error ) ) <EOL> self . status = self . REJECTED <EOL> save_host_ok = False <EOL> else : <EOL> self . affected_hosts = total_hosts <EOL> self . status = self . WAITING <EOL> finally : <EOL> post_save . disconnect ( maintenance_post_save , sender = Maintenance ) <EOL> self . save ( ) <EOL> post_save . connect ( maintenance_post_save , sender = Maintenance ) <EOL> return save_host_ok <EOL> def revoke_maintenance ( self , request ) : <EOL> if self . is_waiting_to_run : <EOL> control . revoke ( self . celery_task_id , ) <EOL> self . status = self . REVOKED <EOL> self . revoked_by = request . user . username <EOL> self . save ( ) <EOL> HostMaintenance . objects . filter ( maintenance = self , <EOL> ) . update ( status = HostMaintenance . REVOKED ) <EOL> return True <EOL> return False <EOL> def is_waiting_to_run ( self ) : <EOL> if self . status == self . FINISHED : <EOL> LOG . info ( "<STR_LIT>" . format ( self ) ) <EOL> return False <EOL> elif self . status == self . REJECTED : <EOL> LOG . info ( "<STR_LIT>" . format ( self ) ) <EOL> return False <EOL> inspect = control . inspect ( ) <EOL> scheduled_tasks = inspect . scheduled ( ) <EOL> try : <EOL> hosts = scheduled_tasks . keys ( ) <EOL> except Exception , e : <EOL> LOG . info ( "<STR_LIT>" . format ( e ) ) <EOL> return False <EOL> for host in hosts : <EOL> try : <EOL> scheduled_tasks = scheduled_tasks [ host ] <EOL> except TypeError : <EOL> LOG . warn ( "<STR_LIT>" ) <EOL> LOG . info ( scheduled_tasks ) <EOL> continue <EOL> for task in scheduled_tasks : <EOL> if task [ '<STR_LIT>' ] [ '<STR_LIT:id>' ] == self . celery_task_id : <EOL> return True <EOL> return False <EOL> class HostMaintenance ( BaseModel ) : <EOL> ERROR = <NUM_LIT:0> <EOL> SUCCESS = <NUM_LIT:1> <EOL> RUNNING = <NUM_LIT:2> <EOL> ROLLBACK = <NUM_LIT:3> <EOL> WAITING = <NUM_LIT:4> <EOL> ROLLBACK_ERROR = <NUM_LIT:5> <EOL> ROLLBACK_SUCCESS = <NUM_LIT:6> <EOL> REVOKED = <NUM_LIT:7> <EOL> UNAVAILABLEHOST = <NUM_LIT:8> <EOL> UNAVAILABLECSHOSTATTR = <NUM_LIT:9> <EOL> MAINTENANCE_STATUS = ( <EOL> ( ERROR , '<STR_LIT>' ) , <EOL> ( SUCCESS , '<STR_LIT>' ) , <EOL> ( RUNNING , '<STR_LIT>' ) , <EOL> ( ROLLBACK , '<STR_LIT>' ) , <EOL> ( WAITING , '<STR_LIT>' ) , <EOL> ( ROLLBACK_ERROR , '<STR_LIT>' ) , <EOL> ( ROLLBACK_SUCCESS , '<STR_LIT>' ) , <EOL> ( REVOKED , '<STR_LIT>' ) , <EOL> ( UNAVAILABLEHOST , '<STR_LIT>' ) , <EOL> ( UNAVAILABLECSHOSTATTR , '<STR_LIT>' ) , <EOL> ) <EOL> started_at = models . DateTimeField ( verbose_name = _ ( "<STR_LIT>" ) , null = True ) <EOL> finished_at = models . DateTimeField ( <EOL> verbose_name = _ ( "<STR_LIT>" ) , null = True ) <EOL> main_log = models . TextField ( verbose_name = _ ( "<STR_LIT>" ) , <EOL> null = True , blank = True ) <EOL> rollback_log = models . TextField ( verbose_name = _ ( "<STR_LIT>" ) , <EOL> null = True , blank = True ) <EOL> status = models . IntegerField ( choices = MAINTENANCE_STATUS , default = WAITING ) <EOL> host = models . ForeignKey ( Host , related_name = "<STR_LIT>" , <EOL> on_delete = models . SET_NULL , null = True ) <EOL> hostname = models . CharField ( verbose_name = _ ( "<STR_LIT>" ) , max_length = <NUM_LIT:255> , <EOL> default = '<STR_LIT>' ) <EOL> maintenance = models . ForeignKey ( Maintenance , related_name = "<STR_LIT>" , ) <EOL> class Meta : <EOL> unique_together = ( ( "<STR_LIT:host>" , "<STR_LIT>" ) , ) <EOL> index_together = [ [ "<STR_LIT:host>" , "<STR_LIT>" ] , ] <EOL> permissions = ( <EOL> ( "<STR_LIT>" , "<STR_LIT>" ) , <EOL> ) <EOL> def __unicode__ ( self ) : <EOL> return "<STR_LIT>" % ( self . host , self . maintenance ) <EOL> class MaintenanceParameters ( BaseModel ) : <EOL> parameter_name = models . CharField ( verbose_name = _ ( "<STR_LIT>" ) , <EOL> null = False , blank = False , max_length = <NUM_LIT:100> , ) <EOL> function_name = models . CharField ( verbose_name = _ ( "<STR_LIT>" ) , <EOL> null = False , blank = False , max_length = <NUM_LIT:100> , choices = ( _get_registered_functions ( ) ) ) <EOL> maintenance = models . ForeignKey ( Maintenance , <EOL> related_name = "<STR_LIT>" , ) <EOL> def __unicode__ ( self ) : <EOL> return "<STR_LIT>" . format ( self . parameter_name , self . function_name ) <EOL> class Meta : <EOL> permissions = ( <EOL> ( "<STR_LIT>" , "<STR_LIT>" ) , <EOL> ) <EOL> simple_audit . register ( Maintenance ) <EOL> simple_audit . register ( HostMaintenance ) <EOL> simple_audit . register ( MaintenanceParameters ) <EOL> @ receiver ( pre_delete , sender = Maintenance ) <EOL> def maintenance_pre_delete ( sender , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> maintenance = kwargs . get ( "<STR_LIT>" ) <EOL> LOG . debug ( "<STR_LIT>" ) <EOL> HostMaintenance . objects . filter ( ) . delete ( ) <EOL> control . revoke ( task_id = maintenance . celery_task_id ) <EOL> @ receiver ( post_save , sender = Maintenance ) <EOL> def maintenance_post_save ( sender , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> maintenance = kwargs . get ( "<STR_LIT>" ) <EOL> is_new = kwargs . get ( "<STR_LIT>" ) <EOL> LOG . debug ( "<STR_LIT>" ) <EOL> if is_new or not maintenance . celery_task_id : <EOL> LOG . info ( "<STR_LIT>" ) <EOL> if maintenance . save_host_maintenance ( ) : <EOL> if maintenance . scheduled_for > datetime . now ( ) : <EOL> task = execute_scheduled_maintenance . apply_async ( args = [ maintenance . id ] , <EOL> eta = maintenance . scheduled_for . replace ( tzinfo = tz . tzlocal ( ) ) . astimezone ( tz . tzutc ( ) ) ) <EOL> else : <EOL> task = execute_scheduled_maintenance . apply_async ( args = [ maintenance . id ] , <EOL> countdown = <NUM_LIT:5> ) <EOL> maintenance . celery_task_id = task . task_id <EOL> maintenance . save ( ) </s>
