<s> """<STR_LIT>""" <EOL> from __future__ import with_statement <EOL> __all__ = [ <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" <EOL> ] <EOL> import cStringIO <EOL> import gc <EOL> import logging <EOL> import pickle <EOL> import random <EOL> import string <EOL> import time <EOL> from mapreduce import context <EOL> from mapreduce import errors <EOL> from mapreduce import json_util <EOL> from mapreduce import kv_pb <EOL> from mapreduce import model <EOL> from mapreduce import operation <EOL> from mapreduce import records <EOL> from mapreduce import shard_life_cycle <EOL> try : <EOL> cloudstorage = None <EOL> import cloudstorage <EOL> if hasattr ( cloudstorage , "<STR_LIT>" ) : <EOL> cloudstorage = None <EOL> if cloudstorage : <EOL> from cloudstorage import cloudstorage_api <EOL> from cloudstorage import errors as cloud_errors <EOL> except ImportError : <EOL> pass <EOL> if cloudstorage is None : <EOL> try : <EOL> import cloudstorage <EOL> from cloudstorage import cloudstorage_api <EOL> except ImportError : <EOL> pass <EOL> COUNTER_IO_WRITE_BYTES = "<STR_LIT>" <EOL> COUNTER_IO_WRITE_MSEC = "<STR_LIT>" <EOL> class OutputWriter ( json_util . JsonMixin ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def validate ( cls , mapper_spec ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( "<STR_LIT>" % cls ) <EOL> @ classmethod <EOL> def init_job ( cls , mapreduce_state ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> @ classmethod <EOL> def finalize_job ( cls , mapreduce_state ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> @ classmethod <EOL> def from_json ( cls , state ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( "<STR_LIT>" % cls ) <EOL> def to_json ( self ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( "<STR_LIT>" % <EOL> self . __class__ ) <EOL> @ classmethod <EOL> def create ( cls , mr_spec , shard_number , shard_attempt , _writer_state = None ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( "<STR_LIT>" % cls ) <EOL> def write ( self , data ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( "<STR_LIT>" % <EOL> self . __class__ ) <EOL> def finalize ( self , ctx , shard_state ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( "<STR_LIT>" % <EOL> self . __class__ ) <EOL> @ classmethod <EOL> def get_filenames ( cls , mapreduce_state ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( "<STR_LIT>" % cls ) <EOL> def _supports_shard_retry ( self , tstate ) : <EOL> """<STR_LIT>""" <EOL> return False <EOL> def _supports_slice_recovery ( self , mapper_spec ) : <EOL> """<STR_LIT>""" <EOL> return False <EOL> def _recover ( self , mr_spec , shard_number , shard_attempt ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> _FILE_POOL_FLUSH_SIZE = <NUM_LIT> * <NUM_LIT> <EOL> _FILE_POOL_MAX_SIZE = <NUM_LIT:1000> * <NUM_LIT> <EOL> def _get_params ( mapper_spec , allowed_keys = None , allow_old = True ) : <EOL> """<STR_LIT>""" <EOL> if "<STR_LIT>" not in mapper_spec . params : <EOL> message = ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> if not allow_old or allowed_keys : <EOL> raise errors . BadWriterParamsError ( message ) <EOL> params = mapper_spec . params <EOL> params = dict ( ( str ( n ) , v ) for n , v in params . iteritems ( ) ) <EOL> else : <EOL> if not isinstance ( mapper_spec . params . get ( "<STR_LIT>" ) , dict ) : <EOL> raise errors . BadWriterParamsError ( <EOL> "<STR_LIT>" ) <EOL> params = mapper_spec . params . get ( "<STR_LIT>" ) <EOL> params = dict ( ( str ( n ) , v ) for n , v in params . iteritems ( ) ) <EOL> if allowed_keys : <EOL> params_diff = set ( params . keys ( ) ) - allowed_keys <EOL> if params_diff : <EOL> raise errors . BadWriterParamsError ( <EOL> "<STR_LIT>" % "<STR_LIT:U+002C>" . join ( params_diff ) ) <EOL> return params <EOL> class _RecordsPoolBase ( context . Pool ) : <EOL> """<STR_LIT>""" <EOL> _RECORD_OVERHEAD_BYTES = <NUM_LIT:10> <EOL> def __init__ ( self , <EOL> flush_size_chars = _FILE_POOL_FLUSH_SIZE , <EOL> ctx = None , <EOL> exclusive = False ) : <EOL> """<STR_LIT>""" <EOL> self . _flush_size = flush_size_chars <EOL> self . _buffer = [ ] <EOL> self . _size = <NUM_LIT:0> <EOL> self . _ctx = ctx <EOL> self . _exclusive = exclusive <EOL> def append ( self , data ) : <EOL> """<STR_LIT>""" <EOL> data_length = len ( data ) <EOL> if self . _size + data_length > self . _flush_size : <EOL> self . flush ( ) <EOL> if not self . _exclusive and data_length > _FILE_POOL_MAX_SIZE : <EOL> raise errors . Error ( <EOL> "<STR_LIT>" % ( data_length , _FILE_POOL_MAX_SIZE ) ) <EOL> else : <EOL> self . _buffer . append ( data ) <EOL> self . _size += data_length <EOL> if self . _size > self . _flush_size : <EOL> self . flush ( ) <EOL> def flush ( self ) : <EOL> """<STR_LIT>""" <EOL> buf = cStringIO . StringIO ( ) <EOL> with records . RecordsWriter ( buf ) as w : <EOL> for record in self . _buffer : <EOL> w . write ( record ) <EOL> w . _pad_block ( ) <EOL> str_buf = buf . getvalue ( ) <EOL> buf . close ( ) <EOL> if not self . _exclusive and len ( str_buf ) > _FILE_POOL_MAX_SIZE : <EOL> raise errors . Error ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % <EOL> ( _FILE_POOL_MAX_SIZE , len ( str_buf ) ) ) <EOL> start_time = time . time ( ) <EOL> self . _write ( str_buf ) <EOL> if self . _ctx : <EOL> operation . counters . Increment ( <EOL> COUNTER_IO_WRITE_BYTES , len ( str_buf ) ) ( self . _ctx ) <EOL> operation . counters . Increment ( <EOL> COUNTER_IO_WRITE_MSEC , <EOL> int ( ( time . time ( ) - start_time ) * <NUM_LIT:1000> ) ) ( self . _ctx ) <EOL> self . _buffer = [ ] <EOL> self . _size = <NUM_LIT:0> <EOL> gc . collect ( ) <EOL> def _write ( self , str_buf ) : <EOL> raise NotImplementedError ( "<STR_LIT>" % type ( self ) ) <EOL> def __enter__ ( self ) : <EOL> return self <EOL> def __exit__ ( self , atype , value , traceback ) : <EOL> self . flush ( ) <EOL> class GCSRecordsPool ( _RecordsPoolBase ) : <EOL> """<STR_LIT>""" <EOL> _GCS_BLOCK_SIZE = <NUM_LIT> * <NUM_LIT> <EOL> def __init__ ( self , <EOL> filehandle , <EOL> flush_size_chars = _FILE_POOL_FLUSH_SIZE , <EOL> ctx = None , <EOL> exclusive = False ) : <EOL> """<STR_LIT>""" <EOL> super ( GCSRecordsPool , self ) . __init__ ( flush_size_chars , ctx , exclusive ) <EOL> self . _filehandle = filehandle <EOL> self . _buf_size = <NUM_LIT:0> <EOL> def _write ( self , str_buf ) : <EOL> """<STR_LIT>""" <EOL> self . _filehandle . write ( str_buf ) <EOL> self . _buf_size += len ( str_buf ) <EOL> def flush ( self , force = False ) : <EOL> """<STR_LIT>""" <EOL> super ( GCSRecordsPool , self ) . flush ( ) <EOL> if force : <EOL> extra_padding = self . _buf_size % self . _GCS_BLOCK_SIZE <EOL> if extra_padding > <NUM_LIT:0> : <EOL> self . _write ( "<STR_LIT:\x00>" * ( self . _GCS_BLOCK_SIZE - extra_padding ) ) <EOL> self . _filehandle . flush ( ) <EOL> class _GoogleCloudStorageBase ( shard_life_cycle . _ShardLifeCycle , <EOL> OutputWriter ) : <EOL> """<STR_LIT>""" <EOL> BUCKET_NAME_PARAM = "<STR_LIT>" <EOL> TMP_BUCKET_NAME_PARAM = "<STR_LIT>" <EOL> ACL_PARAM = "<STR_LIT>" <EOL> NAMING_FORMAT_PARAM = "<STR_LIT>" <EOL> CONTENT_TYPE_PARAM = "<STR_LIT>" <EOL> _ACCOUNT_ID_PARAM = "<STR_LIT>" <EOL> _TMP_ACCOUNT_ID_PARAM = "<STR_LIT>" <EOL> @ classmethod <EOL> def _get_gcs_bucket ( cls , writer_spec ) : <EOL> return writer_spec [ cls . BUCKET_NAME_PARAM ] <EOL> @ classmethod <EOL> def _get_account_id ( cls , writer_spec ) : <EOL> return writer_spec . get ( cls . _ACCOUNT_ID_PARAM , None ) <EOL> @ classmethod <EOL> def _get_tmp_gcs_bucket ( cls , writer_spec ) : <EOL> """<STR_LIT>""" <EOL> if cls . TMP_BUCKET_NAME_PARAM in writer_spec : <EOL> return writer_spec [ cls . TMP_BUCKET_NAME_PARAM ] <EOL> return cls . _get_gcs_bucket ( writer_spec ) <EOL> @ classmethod <EOL> def _get_tmp_account_id ( cls , writer_spec ) : <EOL> """<STR_LIT>""" <EOL> if cls . TMP_BUCKET_NAME_PARAM in writer_spec : <EOL> return writer_spec . get ( cls . _TMP_ACCOUNT_ID_PARAM , None ) <EOL> return cls . _get_account_id ( writer_spec ) <EOL> class _GoogleCloudStorageOutputWriterBase ( _GoogleCloudStorageBase ) : <EOL> """<STR_LIT>""" <EOL> _DEFAULT_NAMING_FORMAT = "<STR_LIT>" <EOL> _MR_TMP = "<STR_LIT>" <EOL> _TMP_FILE_NAMING_FORMAT = ( <EOL> _MR_TMP + "<STR_LIT>" ) <EOL> @ classmethod <EOL> def _generate_filename ( cls , writer_spec , name , job_id , num , <EOL> attempt = None , seg_index = None ) : <EOL> """<STR_LIT>""" <EOL> naming_format = cls . _TMP_FILE_NAMING_FORMAT <EOL> if seg_index is None : <EOL> naming_format = writer_spec . get ( cls . NAMING_FORMAT_PARAM , <EOL> cls . _DEFAULT_NAMING_FORMAT ) <EOL> template = string . Template ( naming_format ) <EOL> try : <EOL> if seg_index is None : <EOL> return template . substitute ( name = name , id = job_id , num = num ) <EOL> else : <EOL> return template . substitute ( name = name , id = job_id , num = num , <EOL> attempt = attempt , <EOL> seg = seg_index ) <EOL> except ValueError , error : <EOL> raise errors . BadWriterParamsError ( "<STR_LIT>" % ( error ) ) <EOL> except KeyError , error : <EOL> raise errors . BadWriterParamsError ( "<STR_LIT>" <EOL> "<STR_LIT>" % ( naming_format , error ) ) <EOL> @ classmethod <EOL> def get_params ( cls , mapper_spec , allowed_keys = None , allow_old = True ) : <EOL> params = _get_params ( mapper_spec , allowed_keys , allow_old ) <EOL> if ( mapper_spec . params . get ( cls . BUCKET_NAME_PARAM ) is not None and <EOL> params . get ( cls . BUCKET_NAME_PARAM ) is None ) : <EOL> params [ cls . BUCKET_NAME_PARAM ] = mapper_spec . params [ cls . BUCKET_NAME_PARAM ] <EOL> return params <EOL> @ classmethod <EOL> def validate ( cls , mapper_spec ) : <EOL> """<STR_LIT>""" <EOL> writer_spec = cls . get_params ( mapper_spec , allow_old = False ) <EOL> if cls . BUCKET_NAME_PARAM not in writer_spec : <EOL> raise errors . BadWriterParamsError ( <EOL> "<STR_LIT>" % <EOL> cls . BUCKET_NAME_PARAM ) <EOL> try : <EOL> cloudstorage . validate_bucket_name ( <EOL> writer_spec [ cls . BUCKET_NAME_PARAM ] ) <EOL> except ValueError , error : <EOL> raise errors . BadWriterParamsError ( "<STR_LIT>" % ( error ) ) <EOL> cls . _generate_filename ( writer_spec , "<STR_LIT:name>" , "<STR_LIT:id>" , <NUM_LIT:0> ) <EOL> cls . _generate_filename ( writer_spec , "<STR_LIT:name>" , "<STR_LIT:id>" , <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:0> ) <EOL> @ classmethod <EOL> def _open_file ( cls , writer_spec , filename_suffix , use_tmp_bucket = False ) : <EOL> """<STR_LIT>""" <EOL> if use_tmp_bucket : <EOL> bucket = cls . _get_tmp_gcs_bucket ( writer_spec ) <EOL> account_id = cls . _get_tmp_account_id ( writer_spec ) <EOL> else : <EOL> bucket = cls . _get_gcs_bucket ( writer_spec ) <EOL> account_id = cls . _get_account_id ( writer_spec ) <EOL> filename = "<STR_LIT>" % ( bucket , filename_suffix ) <EOL> content_type = writer_spec . get ( cls . CONTENT_TYPE_PARAM , None ) <EOL> options = { } <EOL> if cls . ACL_PARAM in writer_spec : <EOL> options [ "<STR_LIT>" ] = writer_spec . get ( cls . ACL_PARAM ) <EOL> return cloudstorage . open ( filename , mode = "<STR_LIT:w>" , content_type = content_type , <EOL> options = options , _account_id = account_id ) <EOL> @ classmethod <EOL> def _get_filename ( cls , shard_state ) : <EOL> return shard_state . writer_state [ "<STR_LIT:filename>" ] <EOL> @ classmethod <EOL> def get_filenames ( cls , mapreduce_state ) : <EOL> filenames = [ ] <EOL> for shard in model . ShardState . find_all_by_mapreduce_state ( mapreduce_state ) : <EOL> if shard . result_status == model . ShardState . RESULT_SUCCESS : <EOL> filenames . append ( cls . _get_filename ( shard ) ) <EOL> return filenames <EOL> def _get_write_buffer ( self ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> def write ( self , data ) : <EOL> """<STR_LIT>""" <EOL> start_time = time . time ( ) <EOL> self . _get_write_buffer ( ) . write ( data ) <EOL> ctx = context . get ( ) <EOL> operation . counters . Increment ( COUNTER_IO_WRITE_BYTES , len ( data ) ) ( ctx ) <EOL> operation . counters . Increment ( <EOL> COUNTER_IO_WRITE_MSEC , int ( ( time . time ( ) - start_time ) * <NUM_LIT:1000> ) ) ( ctx ) <EOL> def _supports_shard_retry ( self , tstate ) : <EOL> return True <EOL> class _GoogleCloudStorageOutputWriter ( _GoogleCloudStorageOutputWriterBase ) : <EOL> """<STR_LIT>""" <EOL> _SEG_PREFIX = "<STR_LIT>" <EOL> _LAST_SEG_INDEX = "<STR_LIT>" <EOL> _JSON_GCS_BUFFER = "<STR_LIT>" <EOL> _JSON_SEG_INDEX = "<STR_LIT>" <EOL> _JSON_NO_DUP = "<STR_LIT>" <EOL> _VALID_LENGTH = "<STR_LIT>" <EOL> _NO_DUPLICATE = "<STR_LIT>" <EOL> def __init__ ( self , streaming_buffer , writer_spec = None ) : <EOL> """<STR_LIT>""" <EOL> self . _streaming_buffer = streaming_buffer <EOL> self . _no_dup = False <EOL> if writer_spec : <EOL> self . _no_dup = writer_spec . get ( self . _NO_DUPLICATE , False ) <EOL> if self . _no_dup : <EOL> self . _seg_index = int ( streaming_buffer . name . rsplit ( "<STR_LIT:->" , <NUM_LIT:1> ) [ <NUM_LIT:1> ] ) <EOL> self . _seg_valid_length = <NUM_LIT:0> <EOL> @ classmethod <EOL> def validate ( cls , mapper_spec ) : <EOL> """<STR_LIT>""" <EOL> writer_spec = cls . get_params ( mapper_spec , allow_old = False ) <EOL> if writer_spec . get ( cls . _NO_DUPLICATE , False ) not in ( True , False ) : <EOL> raise errors . BadWriterParamsError ( "<STR_LIT>" ) <EOL> super ( _GoogleCloudStorageOutputWriter , cls ) . validate ( mapper_spec ) <EOL> def _get_write_buffer ( self ) : <EOL> return self . _streaming_buffer <EOL> @ classmethod <EOL> def create ( cls , mr_spec , shard_number , shard_attempt , _writer_state = None ) : <EOL> """<STR_LIT>""" <EOL> writer_spec = cls . get_params ( mr_spec . mapper , allow_old = False ) <EOL> seg_index = None <EOL> if writer_spec . get ( cls . _NO_DUPLICATE , False ) : <EOL> seg_index = <NUM_LIT:0> <EOL> key = cls . _generate_filename ( writer_spec , mr_spec . name , <EOL> mr_spec . mapreduce_id , <EOL> shard_number , shard_attempt , <EOL> seg_index ) <EOL> return cls . _create ( writer_spec , key ) <EOL> @ classmethod <EOL> def _create ( cls , writer_spec , filename_suffix ) : <EOL> """<STR_LIT>""" <EOL> writer = cls . _open_file ( writer_spec , filename_suffix ) <EOL> return cls ( writer , writer_spec = writer_spec ) <EOL> @ classmethod <EOL> def from_json ( cls , state ) : <EOL> writer = cls ( pickle . loads ( state [ cls . _JSON_GCS_BUFFER ] ) ) <EOL> no_dup = state . get ( cls . _JSON_NO_DUP , False ) <EOL> writer . _no_dup = no_dup <EOL> if no_dup : <EOL> writer . _seg_valid_length = state [ cls . _VALID_LENGTH ] <EOL> writer . _seg_index = state [ cls . _JSON_SEG_INDEX ] <EOL> return writer <EOL> def end_slice ( self , slice_ctx ) : <EOL> if not self . _streaming_buffer . closed : <EOL> self . _streaming_buffer . flush ( ) <EOL> def to_json ( self ) : <EOL> result = { self . _JSON_GCS_BUFFER : pickle . dumps ( self . _streaming_buffer ) , <EOL> self . _JSON_NO_DUP : self . _no_dup } <EOL> if self . _no_dup : <EOL> result . update ( { <EOL> self . _VALID_LENGTH : self . _streaming_buffer . tell ( ) , <EOL> self . _JSON_SEG_INDEX : self . _seg_index } ) <EOL> return result <EOL> def finalize ( self , ctx , shard_state ) : <EOL> self . _streaming_buffer . close ( ) <EOL> if self . _no_dup : <EOL> cloudstorage_api . copy2 ( <EOL> self . _streaming_buffer . name , <EOL> self . _streaming_buffer . name , <EOL> metadata = { self . _VALID_LENGTH : self . _streaming_buffer . tell ( ) } ) <EOL> mr_spec = ctx . mapreduce_spec <EOL> writer_spec = self . get_params ( mr_spec . mapper , allow_old = False ) <EOL> filename = self . _generate_filename ( writer_spec , <EOL> mr_spec . name , <EOL> mr_spec . mapreduce_id , <EOL> shard_state . shard_number ) <EOL> seg_filename = self . _streaming_buffer . name <EOL> prefix , last_index = seg_filename . rsplit ( "<STR_LIT:->" , <NUM_LIT:1> ) <EOL> shard_state . writer_state = { self . _SEG_PREFIX : prefix + "<STR_LIT:->" , <EOL> self . _LAST_SEG_INDEX : int ( last_index ) , <EOL> "<STR_LIT:filename>" : filename } <EOL> else : <EOL> shard_state . writer_state = { "<STR_LIT:filename>" : self . _streaming_buffer . name } <EOL> def _supports_slice_recovery ( self , mapper_spec ) : <EOL> writer_spec = self . get_params ( mapper_spec , allow_old = False ) <EOL> return writer_spec . get ( self . _NO_DUPLICATE , False ) <EOL> def _recover ( self , mr_spec , shard_number , shard_attempt ) : <EOL> next_seg_index = self . _seg_index <EOL> if self . _seg_valid_length != <NUM_LIT:0> : <EOL> try : <EOL> gcs_next_offset = self . _streaming_buffer . _get_offset_from_gcs ( ) + <NUM_LIT:1> <EOL> if gcs_next_offset > self . _streaming_buffer . tell ( ) : <EOL> self . _streaming_buffer . _force_close ( gcs_next_offset ) <EOL> else : <EOL> self . _streaming_buffer . close ( ) <EOL> except cloudstorage . FileClosedError : <EOL> pass <EOL> cloudstorage_api . copy2 ( <EOL> self . _streaming_buffer . name , <EOL> self . _streaming_buffer . name , <EOL> metadata = { self . _VALID_LENGTH : <EOL> self . _seg_valid_length } ) <EOL> next_seg_index = self . _seg_index + <NUM_LIT:1> <EOL> writer_spec = self . get_params ( mr_spec . mapper , allow_old = False ) <EOL> key = self . _generate_filename ( <EOL> writer_spec , mr_spec . name , <EOL> mr_spec . mapreduce_id , <EOL> shard_number , <EOL> shard_attempt , <EOL> next_seg_index ) <EOL> new_writer = self . _create ( writer_spec , key ) <EOL> new_writer . _seg_index = next_seg_index <EOL> return new_writer <EOL> def _get_filename_for_test ( self ) : <EOL> return self . _streaming_buffer . name <EOL> GoogleCloudStorageOutputWriter = _GoogleCloudStorageOutputWriter <EOL> class _ConsistentStatus ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self ) : <EOL> self . writer_spec = None <EOL> self . mapreduce_id = None <EOL> self . shard = None <EOL> self . mainfile = None <EOL> self . tmpfile = None <EOL> self . tmpfile_1ago = None <EOL> class GoogleCloudStorageConsistentOutputWriter ( <EOL> _GoogleCloudStorageOutputWriterBase ) : <EOL> """<STR_LIT>""" <EOL> _JSON_STATUS = "<STR_LIT:status>" <EOL> _RAND_BITS = <NUM_LIT> <EOL> _REWRITE_BLOCK_SIZE = <NUM_LIT> * <NUM_LIT> <EOL> _REWRITE_MR_TMP = "<STR_LIT>" <EOL> _TMPFILE_PATTERN = _REWRITE_MR_TMP + "<STR_LIT>" <EOL> _TMPFILE_PREFIX = _REWRITE_MR_TMP + "<STR_LIT>" <EOL> def __init__ ( self , status ) : <EOL> """<STR_LIT>""" <EOL> self . status = status <EOL> self . _data_written_to_slice = False <EOL> def _get_write_buffer ( self ) : <EOL> if not self . status . tmpfile : <EOL> raise errors . FailJobError ( <EOL> "<STR_LIT>" ) <EOL> return self . status . tmpfile <EOL> def _get_filename_for_test ( self ) : <EOL> return self . status . mainfile . name <EOL> @ classmethod <EOL> def create ( cls , mr_spec , shard_number , shard_attempt , _writer_state = None ) : <EOL> """<STR_LIT>""" <EOL> writer_spec = cls . get_params ( mr_spec . mapper , allow_old = False ) <EOL> key = cls . _generate_filename ( writer_spec , mr_spec . name , <EOL> mr_spec . mapreduce_id , <EOL> shard_number , shard_attempt ) <EOL> status = _ConsistentStatus ( ) <EOL> status . writer_spec = writer_spec <EOL> status . mainfile = cls . _open_file ( writer_spec , key ) <EOL> status . mapreduce_id = mr_spec . mapreduce_id <EOL> status . shard = shard_number <EOL> return cls ( status ) <EOL> def _remove_tmpfile ( self , filename , writer_spec ) : <EOL> if not filename : <EOL> return <EOL> account_id = self . _get_tmp_account_id ( writer_spec ) <EOL> try : <EOL> cloudstorage_api . delete ( filename , _account_id = account_id ) <EOL> except cloud_errors . NotFoundError : <EOL> pass <EOL> def _rewrite_tmpfile ( self , mainfile , tmpfile , writer_spec ) : <EOL> """<STR_LIT>""" <EOL> if mainfile . closed : <EOL> return <EOL> account_id = self . _get_tmp_account_id ( writer_spec ) <EOL> f = cloudstorage_api . open ( tmpfile , _account_id = account_id ) <EOL> data = f . read ( self . _REWRITE_BLOCK_SIZE ) <EOL> while data : <EOL> mainfile . write ( data ) <EOL> data = f . read ( self . _REWRITE_BLOCK_SIZE ) <EOL> f . close ( ) <EOL> mainfile . flush ( ) <EOL> @ classmethod <EOL> def _create_tmpfile ( cls , status ) : <EOL> """<STR_LIT>""" <EOL> tmpl = string . Template ( cls . _TMPFILE_PATTERN ) <EOL> filename = tmpl . substitute ( <EOL> id = status . mapreduce_id , shard = status . shard , <EOL> random = random . getrandbits ( cls . _RAND_BITS ) ) <EOL> return cls . _open_file ( status . writer_spec , filename , use_tmp_bucket = True ) <EOL> def begin_slice ( self , slice_ctx ) : <EOL> status = self . status <EOL> writer_spec = status . writer_spec <EOL> if status . tmpfile_1ago : <EOL> self . _remove_tmpfile ( status . tmpfile_1ago . name , writer_spec ) <EOL> files_to_keep = [ ] <EOL> if status . tmpfile : <EOL> self . _rewrite_tmpfile ( status . mainfile , status . tmpfile . name , writer_spec ) <EOL> files_to_keep . append ( status . tmpfile . name ) <EOL> self . _try_to_clean_garbage ( <EOL> writer_spec , exclude_list = files_to_keep ) <EOL> status . tmpfile_1ago = status . tmpfile <EOL> status . tmpfile = self . _create_tmpfile ( status ) <EOL> if status . mainfile . closed : <EOL> status . tmpfile . close ( ) <EOL> self . _remove_tmpfile ( status . tmpfile . name , writer_spec ) <EOL> @ classmethod <EOL> def from_json ( cls , state ) : <EOL> return cls ( pickle . loads ( state [ cls . _JSON_STATUS ] ) ) <EOL> def end_slice ( self , slice_ctx ) : <EOL> self . status . tmpfile . close ( ) <EOL> def to_json ( self ) : <EOL> return { self . _JSON_STATUS : pickle . dumps ( self . status ) } <EOL> def write ( self , data ) : <EOL> super ( GoogleCloudStorageConsistentOutputWriter , self ) . write ( data ) <EOL> self . _data_written_to_slice = True <EOL> def _try_to_clean_garbage ( self , writer_spec , exclude_list = ( ) ) : <EOL> """<STR_LIT>""" <EOL> tmpl = string . Template ( self . _TMPFILE_PREFIX ) <EOL> prefix = tmpl . substitute ( <EOL> id = self . status . mapreduce_id , shard = self . status . shard ) <EOL> bucket = self . _get_tmp_gcs_bucket ( writer_spec ) <EOL> account_id = self . _get_tmp_account_id ( writer_spec ) <EOL> for f in cloudstorage . listbucket ( "<STR_LIT>" % ( bucket , prefix ) , <EOL> _account_id = account_id ) : <EOL> if f . filename not in exclude_list : <EOL> self . _remove_tmpfile ( f . filename , self . status . writer_spec ) <EOL> def finalize ( self , ctx , shard_state ) : <EOL> if self . _data_written_to_slice : <EOL> raise errors . FailJobError ( <EOL> "<STR_LIT>" ) <EOL> if self . status . tmpfile : <EOL> self . status . tmpfile . close ( ) <EOL> self . status . mainfile . close ( ) <EOL> if self . status . tmpfile_1ago : <EOL> self . _remove_tmpfile ( self . status . tmpfile_1ago . name , <EOL> self . status . writer_spec ) <EOL> if self . status . tmpfile : <EOL> self . _remove_tmpfile ( self . status . tmpfile . name , <EOL> self . status . writer_spec ) <EOL> self . _try_to_clean_garbage ( self . status . writer_spec ) <EOL> shard_state . writer_state = { "<STR_LIT:filename>" : self . status . mainfile . name } <EOL> class _GoogleCloudStorageRecordOutputWriterBase ( _GoogleCloudStorageBase ) : <EOL> """<STR_LIT>""" <EOL> WRITER_CLS = None <EOL> def __init__ ( self , writer ) : <EOL> self . _writer = writer <EOL> self . _record_writer = records . RecordsWriter ( writer ) <EOL> @ classmethod <EOL> def validate ( cls , mapper_spec ) : <EOL> return cls . WRITER_CLS . validate ( mapper_spec ) <EOL> @ classmethod <EOL> def init_job ( cls , mapreduce_state ) : <EOL> return cls . WRITER_CLS . init_job ( mapreduce_state ) <EOL> @ classmethod <EOL> def finalize_job ( cls , mapreduce_state ) : <EOL> return cls . WRITER_CLS . finalize_job ( mapreduce_state ) <EOL> @ classmethod <EOL> def from_json ( cls , state ) : <EOL> return cls ( cls . WRITER_CLS . from_json ( state ) ) <EOL> def to_json ( self ) : <EOL> return self . _writer . to_json ( ) <EOL> @ classmethod <EOL> def create ( cls , mr_spec , shard_number , shard_attempt , _writer_state = None ) : <EOL> return cls ( cls . WRITER_CLS . create ( mr_spec , shard_number , shard_attempt , <EOL> _writer_state ) ) <EOL> def write ( self , data ) : <EOL> self . _record_writer . write ( data ) <EOL> def finalize ( self , ctx , shard_state ) : <EOL> return self . _writer . finalize ( ctx , shard_state ) <EOL> @ classmethod <EOL> def get_filenames ( cls , mapreduce_state ) : <EOL> return cls . WRITER_CLS . get_filenames ( mapreduce_state ) <EOL> def _supports_shard_retry ( self , tstate ) : <EOL> return self . _writer . _supports_shard_retry ( tstate ) <EOL> def _supports_slice_recovery ( self , mapper_spec ) : <EOL> return self . _writer . _supports_slice_recovery ( mapper_spec ) <EOL> def _recover ( self , mr_spec , shard_number , shard_attempt ) : <EOL> return self . _writer . _recover ( mr_spec , shard_number , shard_attempt ) <EOL> def begin_slice ( self , slice_ctx ) : <EOL> return self . _writer . begin_slice ( slice_ctx ) <EOL> def end_slice ( self , slice_ctx ) : <EOL> if not self . _writer . _get_write_buffer ( ) . closed : <EOL> self . _record_writer . _pad_block ( ) <EOL> return self . _writer . end_slice ( slice_ctx ) <EOL> class _GoogleCloudStorageRecordOutputWriter ( <EOL> _GoogleCloudStorageRecordOutputWriterBase ) : <EOL> WRITER_CLS = _GoogleCloudStorageOutputWriter <EOL> GoogleCloudStorageRecordOutputWriter = _GoogleCloudStorageRecordOutputWriter <EOL> class GoogleCloudStorageConsistentRecordOutputWriter ( <EOL> _GoogleCloudStorageRecordOutputWriterBase ) : <EOL> WRITER_CLS = GoogleCloudStorageConsistentOutputWriter <EOL> class _GoogleCloudStorageKeyValueOutputWriter ( <EOL> _GoogleCloudStorageRecordOutputWriter ) : <EOL> """<STR_LIT>""" <EOL> def write ( self , data ) : <EOL> if len ( data ) != <NUM_LIT:2> : <EOL> logging . error ( "<STR_LIT>" , <EOL> len ( data ) , data ) <EOL> try : <EOL> key = str ( data [ <NUM_LIT:0> ] ) <EOL> value = str ( data [ <NUM_LIT:1> ] ) <EOL> except TypeError : <EOL> logging . error ( "<STR_LIT>" , <EOL> data . __class__ . __name__ , data ) <EOL> proto = kv_pb . KeyValue ( ) <EOL> proto . set_key ( key ) <EOL> proto . set_value ( value ) <EOL> GoogleCloudStorageRecordOutputWriter . write ( self , proto . Encode ( ) ) <EOL> GoogleCloudStorageKeyValueOutputWriter = _GoogleCloudStorageKeyValueOutputWriter </s>
