<s> """<STR_LIT>""" <EOL> import logging <EOL> import datetime <EOL> import os <EOL> import re <EOL> import collections <EOL> import io <EOL> import bz2 <EOL> try : <EOL> import boto <EOL> except ImportError : <EOL> boto = None <EOL> else : <EOL> from boto . s3 . connection import OrdinaryCallingFormat <EOL> from . import errors <EOL> log = logging . getLogger ( __name__ ) <EOL> DATE_FORMATS = [ "<STR_LIT>" , "<STR_LIT>" ] <EOL> class InvalidDateError ( errors . Error ) : <EOL> pass <EOL> def parse_date_range_argument ( value ) : <EOL> for fmt in DATE_FORMATS : <EOL> try : <EOL> return datetime . datetime . strptime ( value , fmt ) <EOL> except ValueError : <EOL> pass <EOL> raise InvalidDateError ( ) <EOL> class LogFile ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , type_name , host = None , dt = None , date = None , bzip = False ) : <EOL> self . type_name = type_name <EOL> self . host = host <EOL> self . bzip = bzip <EOL> if ( dt , date ) == ( None , None ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . dt = dt <EOL> self . date = date or dt . date ( ) <EOL> @ property <EOL> def sort_dt ( self ) : <EOL> if self . dt : <EOL> return self . dt <EOL> else : <EOL> return datetime . datetime ( <EOL> self . date . year , self . date . month , self . date . day ) <EOL> @ property <EOL> def file_name ( self ) : <EOL> if self . dt : <EOL> date_name_str = self . dt . strftime ( '<STR_LIT>' ) <EOL> else : <EOL> date_name_str = self . date . strftime ( '<STR_LIT>' ) <EOL> bzip_str = "<STR_LIT>" <EOL> if self . bzip : <EOL> bzip_str = "<STR_LIT>" <EOL> host_str = "<STR_LIT>" <EOL> if self . host : <EOL> host_str = "<STR_LIT>" . format ( self . host ) <EOL> return "<STR_LIT>" . format ( <EOL> type = self . type_name , <EOL> host_str = host_str , <EOL> date_name = date_name_str , <EOL> bzip = bzip_str ) <EOL> @ property <EOL> def file_path ( self ) : <EOL> date_str = self . date . strftime ( '<STR_LIT>' ) <EOL> return "<STR_LIT>" . format ( <EOL> date = date_str , <EOL> filename = self . file_name ) <EOL> @ classmethod <EOL> def from_filename ( cls , filename ) : <EOL> basename = os . path . basename ( filename ) <EOL> match = re . match ( <EOL> r"<STR_LIT>" <EOL> r"<STR_LIT>" <EOL> r"<STR_LIT>" <EOL> r"<STR_LIT>" <EOL> r"<STR_LIT>" , basename ) <EOL> if match is None : <EOL> raise ValueError ( basename ) <EOL> match_info = match . groupdict ( ) <EOL> log_dt = log_date = None <EOL> if len ( match_info [ '<STR_LIT:date>' ] ) == len ( '<STR_LIT>' ) : <EOL> log_dt = datetime . datetime . strptime ( match_info [ "<STR_LIT:date>" ] , '<STR_LIT>' ) <EOL> else : <EOL> log_date = datetime . datetime . strptime ( match_info [ "<STR_LIT:date>" ] , <EOL> '<STR_LIT>' ) . date ( ) <EOL> return cls ( <EOL> match_info [ "<STR_LIT>" ] , <EOL> host = match_info . get ( '<STR_LIT:host>' ) , <EOL> dt = log_dt , <EOL> date = log_date , <EOL> bzip = bool ( match_info [ '<STR_LIT>' ] ) ) <EOL> class S3LogFile ( LogFile ) : <EOL> def s3_key ( self , bucket ) : <EOL> return boto . s3 . key . Key ( bucket , name = self . file_path ) <EOL> def open ( self , bucket ) : <EOL> """<STR_LIT>""" <EOL> def stream ( ) : <EOL> if self . bzip : <EOL> decompressor = bz2 . BZ2Decompressor ( ) <EOL> else : <EOL> decompressor = None <EOL> for data in self . s3_key ( bucket ) : <EOL> if decompressor : <EOL> r = decompressor . decompress ( data ) <EOL> else : <EOL> r = data <EOL> if r is not None : <EOL> yield r <EOL> return stream ( ) <EOL> @ classmethod <EOL> def from_s3_key ( cls , key ) : <EOL> return cls . from_filename ( key . name ) <EOL> class LocalLogFile ( LogFile ) : <EOL> def get_local_file_path ( self , log_path ) : <EOL> return os . path . join ( log_path , self . file_path ) <EOL> def open ( self , log_path ) : <EOL> """<STR_LIT>""" <EOL> def stream ( ) : <EOL> if self . bzip : <EOL> decompressor = bz2 . BZ2Decompressor ( ) <EOL> else : <EOL> decompressor = None <EOL> with io . open ( self . get_local_file_path ( log_path ) , "<STR_LIT:rb>" ) as f : <EOL> for data in f : <EOL> if decompressor : <EOL> r = decompressor . decompress ( data ) <EOL> else : <EOL> r = data <EOL> if r is not None : <EOL> yield r <EOL> return stream ( ) <EOL> def build_remote ( self , host ) : <EOL> return S3LogFile ( <EOL> self . type_name , <EOL> host = host , <EOL> dt = self . dt , <EOL> date = self . date , <EOL> bzip = self . bzip ) <EOL> def list_log_files ( log_path ) : <EOL> """<STR_LIT>""" <EOL> log_files = [ ] <EOL> for dirpath , dirnames , filenames in os . walk ( log_path ) : <EOL> for filename in filenames : <EOL> full_path = os . path . join ( dirpath , filename ) <EOL> try : <EOL> log_file = LocalLogFile . from_filename ( filename ) <EOL> except ValueError : <EOL> log . warning ( "<STR_LIT>" , full_path ) <EOL> continue <EOL> log_files . append ( log_file ) <EOL> return log_files <EOL> def filter_log_files_for_active ( log_files ) : <EOL> """<STR_LIT>""" <EOL> out_log_files = [ ] <EOL> files_by_type = collections . defaultdict ( list ) <EOL> for lf in log_files : <EOL> files_by_type [ lf . type_name ] . append ( lf ) <EOL> for type_files in files_by_type . values ( ) : <EOL> type_files . sort ( key = lambda f : f . sort_dt ) <EOL> last_lf = type_files . pop ( ) <EOL> out_log_files += type_files <EOL> cutoff_date = ( datetime . datetime . utcnow ( ) - datetime . timedelta ( hours = <NUM_LIT:1> ) <EOL> ) . date ( ) <EOL> if last_lf . date < cutoff_date : <EOL> out_log_files . append ( last_lf ) <EOL> return out_log_files <EOL> def filter_log_files_for_zipping ( log_files ) : <EOL> """<STR_LIT>""" <EOL> out_files = [ ] <EOL> for lf in filter_log_files_for_active ( log_files ) : <EOL> if lf . bzip : <EOL> continue <EOL> out_files . append ( lf ) <EOL> return out_files <EOL> def filter_log_files_for_uploading ( log_files , zipped_only ) : <EOL> """<STR_LIT>""" <EOL> out_files = [ ] <EOL> for lf in filter_log_files_for_active ( log_files ) : <EOL> if zipped_only and not lf . bzip : <EOL> continue <EOL> out_files . append ( lf ) <EOL> return out_files <EOL> def zip_log_file ( log_file , log_path ) : <EOL> orig_path = log_file . get_local_file_path ( log_path ) <EOL> log_file . bzip = True <EOL> zip_path = log_file . get_local_file_path ( log_path ) <EOL> zip_file = bz2 . BZ2File ( zip_path , '<STR_LIT:w>' , io . DEFAULT_BUFFER_SIZE ) <EOL> with io . open ( orig_path , "<STR_LIT:rb>" ) as fp : <EOL> for data in fp : <EOL> zip_file . write ( data ) <EOL> zip_file . close ( ) <EOL> os . unlink ( orig_path ) <EOL> def s3_prefix_for_date_and_type ( date , type_name ) : <EOL> date_str = date . strftime ( '<STR_LIT>' ) <EOL> return "<STR_LIT>" . format ( date_str , type_name ) <EOL> def inclusive_date_range ( start_dt , end_dt ) : <EOL> start_date = start_dt . date ( ) <EOL> end_date = end_dt . date ( ) <EOL> current_date = start_date <EOL> while current_date <= end_date : <EOL> yield current_date <EOL> current_date += datetime . timedelta ( days = <NUM_LIT:1> ) <EOL> def find_log_files_in_s3 ( bucket , type_name , start_dt , end_dt ) : <EOL> prefixes = [ ] <EOL> for dt in inclusive_date_range ( start_dt , end_dt ) : <EOL> prefixes . append ( s3_prefix_for_date_and_type ( dt , type_name ) ) <EOL> log_files = [ ] <EOL> for pf in prefixes : <EOL> for key in bucket . list ( pf ) : <EOL> try : <EOL> log_files . append ( S3LogFile . from_s3_key ( key ) ) <EOL> except ValueError as e : <EOL> log . warning ( "<STR_LIT>" , key , e ) <EOL> continue <EOL> out_log_files = [ ] <EOL> for lf in log_files : <EOL> assert lf . type_name == type_name <EOL> if lf . dt is None : <EOL> out_log_files . append ( lf ) <EOL> elif lf . dt >= start_dt and lf . dt <= end_dt : <EOL> out_log_files . append ( lf ) <EOL> return out_log_files <EOL> def find_log_files_in_path ( log_path , type_name , start_dt , end_dt ) : <EOL> log_files = list_log_files ( log_path ) <EOL> out_log_files = [ ] <EOL> for lf in log_files : <EOL> if lf . dt is None : <EOL> if lf . date < start_dt . date ( ) or lf . date > end_dt . date ( ) : <EOL> continue <EOL> else : <EOL> if lf . dt < start_dt or lf . dt > end_dt : <EOL> continue <EOL> out_log_files . append ( lf ) <EOL> return out_log_files <EOL> def open_bucket ( bucket_name ) : <EOL> region_name = os . environ . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> conn = boto . s3 . connect_to_region ( region_name , <EOL> calling_format = OrdinaryCallingFormat ( ) ) <EOL> return conn . get_bucket ( bucket_name ) </s>
