<s> import os <EOL> import Queue <EOL> import eventlet <EOL> from st2common import log as logging <EOL> from st2exporter . exporter . file_writer import TextFileWriter <EOL> from st2exporter . exporter . json_converter import JsonConverter <EOL> from st2common . models . db . marker import DumperMarkerDB <EOL> from st2common . persistence . marker import DumperMarker <EOL> from st2common . util import date as date_utils <EOL> from st2common . util import isotime <EOL> __all__ = [ <EOL> '<STR_LIT>' <EOL> ] <EOL> ALLOWED_EXTENSIONS = [ '<STR_LIT>' ] <EOL> CONVERTERS = { <EOL> '<STR_LIT>' : JsonConverter <EOL> } <EOL> LOG = logging . getLogger ( __name__ ) <EOL> class Dumper ( object ) : <EOL> def __init__ ( self , queue , export_dir , file_format = '<STR_LIT>' , <EOL> file_prefix = '<STR_LIT>' , <EOL> batch_size = <NUM_LIT:1000> , sleep_interval = <NUM_LIT> , <EOL> max_files_per_sleep = <NUM_LIT:5> , <EOL> file_writer = None ) : <EOL> if not queue : <EOL> raise Exception ( '<STR_LIT>' ) <EOL> if not export_dir : <EOL> raise Exception ( '<STR_LIT>' ) <EOL> self . _export_dir = export_dir <EOL> if not os . path . exists ( self . _export_dir ) : <EOL> raise Exception ( '<STR_LIT>' % <EOL> self . _export_dir ) <EOL> self . _file_format = file_format . lower ( ) <EOL> if self . _file_format not in ALLOWED_EXTENSIONS : <EOL> raise ValueError ( '<STR_LIT>' % file_format ) <EOL> self . _file_prefix = file_prefix <EOL> self . _batch_size = batch_size <EOL> self . _max_files_per_sleep = max_files_per_sleep <EOL> self . _queue = queue <EOL> self . _flush_thread = None <EOL> self . _sleep_interval = sleep_interval <EOL> self . _converter = CONVERTERS [ self . _file_format ] ( ) <EOL> self . _shutdown = False <EOL> self . _persisted_marker = None <EOL> if not file_writer : <EOL> self . _file_writer = TextFileWriter ( ) <EOL> def start ( self , wait = False ) : <EOL> self . _flush_thread = eventlet . spawn ( self . _flush ) <EOL> if wait : <EOL> self . wait ( ) <EOL> def wait ( self ) : <EOL> self . _flush_thread . wait ( ) <EOL> def stop ( self ) : <EOL> self . _shutdown = True <EOL> return eventlet . kill ( self . _flush_thread ) <EOL> def _get_batch ( self ) : <EOL> if self . _queue . empty ( ) : <EOL> return None <EOL> executions_to_write = [ ] <EOL> for _ in range ( self . _batch_size ) : <EOL> try : <EOL> item = self . _queue . get ( block = False ) <EOL> except Queue . Empty : <EOL> break <EOL> else : <EOL> executions_to_write . append ( item ) <EOL> LOG . debug ( '<STR_LIT>' , len ( executions_to_write ) ) <EOL> LOG . debug ( '<STR_LIT>' , self . _queue . qsize ( ) ) <EOL> return executions_to_write <EOL> def _flush ( self ) : <EOL> while not self . _shutdown : <EOL> while self . _queue . empty ( ) : <EOL> eventlet . sleep ( self . _sleep_interval ) <EOL> try : <EOL> self . _write_to_disk ( ) <EOL> except : <EOL> LOG . error ( '<STR_LIT>' ) <EOL> def _write_to_disk ( self ) : <EOL> count = <NUM_LIT:0> <EOL> self . _create_date_folder ( ) <EOL> for _ in range ( self . _max_files_per_sleep ) : <EOL> batch = self . _get_batch ( ) <EOL> if not batch : <EOL> return count <EOL> try : <EOL> self . _write_batch_to_disk ( batch ) <EOL> self . _update_marker ( batch ) <EOL> count += <NUM_LIT:1> <EOL> except : <EOL> LOG . exception ( '<STR_LIT>' ) <EOL> return count <EOL> def _create_date_folder ( self ) : <EOL> folder_name = self . _get_date_folder ( ) <EOL> folder_path = os . path . join ( self . _export_dir , folder_name ) <EOL> if not os . path . exists ( folder_path ) : <EOL> try : <EOL> os . makedirs ( folder_path ) <EOL> except : <EOL> LOG . exception ( '<STR_LIT>' , folder_name ) <EOL> raise <EOL> def _write_batch_to_disk ( self , batch ) : <EOL> doc_to_write = self . _converter . convert ( batch ) <EOL> self . _file_writer . write_text ( doc_to_write , self . _get_file_name ( ) ) <EOL> def _get_file_name ( self ) : <EOL> timestring = date_utils . get_datetime_utc_now ( ) . strftime ( '<STR_LIT>' ) <EOL> file_name = self . _file_prefix + timestring + '<STR_LIT:.>' + self . _file_format <EOL> file_name = os . path . join ( self . _export_dir , self . _get_date_folder ( ) , file_name ) <EOL> return file_name <EOL> def _get_date_folder ( self ) : <EOL> return date_utils . get_datetime_utc_now ( ) . strftime ( '<STR_LIT>' ) <EOL> def _update_marker ( self , batch ) : <EOL> timestamps = [ isotime . parse ( item . end_timestamp ) for item in batch ] <EOL> new_marker = max ( timestamps ) <EOL> if self . _persisted_marker and self . _persisted_marker > new_marker : <EOL> LOG . warn ( '<STR_LIT>' ) <EOL> try : <EOL> self . _write_marker_to_db ( new_marker ) <EOL> except : <EOL> LOG . exception ( '<STR_LIT>' ) <EOL> else : <EOL> self . _persisted_marker = new_marker <EOL> return self . _persisted_marker <EOL> def _write_marker_to_db ( self , new_marker ) : <EOL> LOG . info ( '<STR_LIT>' , new_marker ) <EOL> markers = DumperMarker . get_all ( ) <EOL> if len ( markers ) > <NUM_LIT:1> : <EOL> LOG . exception ( '<STR_LIT>' ) <EOL> marker = isotime . format ( new_marker , offset = False ) <EOL> updated_at = date_utils . get_datetime_utc_now ( ) <EOL> if markers : <EOL> marker_id = markers [ <NUM_LIT:0> ] [ '<STR_LIT:id>' ] <EOL> else : <EOL> marker_id = None <EOL> marker_db = DumperMarkerDB ( id = marker_id , marker = marker , updated_at = updated_at ) <EOL> return DumperMarker . add_or_update ( marker_db ) </s>
