<s> """<STR_LIT>""" <EOL> import copy <EOL> from mapreduce . lib . pipeline import common <EOL> from src . clients import gcs <EOL> from src . clients import s3 <EOL> from src . pipelines import pipeline <EOL> from src . pipelines import shardstage <EOL> from src . pipelines . stages import gcscompositor <EOL> class S3Input ( shardstage . ShardStage ) : <EOL> """<STR_LIT>""" <EOL> CHUNK_SIZE_8MB = <NUM_LIT:1> << <NUM_LIT> <EOL> CHUNK_SIZE_32MB = <NUM_LIT:1> << <NUM_LIT> <EOL> REQUEST_CHUNK_SIZE = CHUNK_SIZE_8MB <EOL> MAX_SHARD_SIZE = CHUNK_SIZE_32MB <EOL> DEFAULT_CONTENT_TYPE = '<STR_LIT>' <EOL> @ staticmethod <EOL> def GetHelp ( ) : <EOL> return """<STR_LIT>""" <EOL> def run ( self , config ) : <EOL> """<STR_LIT>""" <EOL> storage = s3 . S3 ( config = config . get ( '<STR_LIT>' ) ) <EOL> s3_objects = [ ] <EOL> if '<STR_LIT:object>' in config : <EOL> s3_objects . append ( config [ '<STR_LIT:object>' ] ) <EOL> if '<STR_LIT>' in config : <EOL> objects = config [ '<STR_LIT>' ] <EOL> for s3_obj in storage . ListBucket ( objects [ '<STR_LIT>' ] , <EOL> objects . get ( '<STR_LIT>' ) ) : <EOL> s3_objects . append ( s3 . S3 . MakeUrl ( objects [ '<STR_LIT>' ] , s3_obj ) ) <EOL> s3_objects = zip ( s3_objects , config [ '<STR_LIT>' ] ) <EOL> ( s3_obj , gcs_obj ) = s3_objects . pop ( ) <EOL> sub_stages = [ ] <EOL> if '<STR_LIT>' in config : <EOL> config . pop ( '<STR_LIT>' ) <EOL> for ( next_s3_obj , next_gcs_obj ) in s3_objects : <EOL> cfg = copy . deepcopy ( config ) <EOL> cfg [ '<STR_LIT:object>' ] = next_s3_obj <EOL> cfg [ '<STR_LIT>' ] = [ next_gcs_obj ] <EOL> s = yield S3Input ( cfg ) <EOL> sub_stages . append ( s ) <EOL> start = config . get ( '<STR_LIT:start>' ) <EOL> if not start : <EOL> start = <NUM_LIT:0> <EOL> config [ '<STR_LIT:start>' ] = <NUM_LIT:0> <EOL> length = config . get ( '<STR_LIT>' ) <EOL> if not length : <EOL> length = storage . StatObject ( s3_obj ) [ '<STR_LIT:size>' ] <EOL> config [ '<STR_LIT>' ] = length <EOL> if '<STR_LIT>' not in config : <EOL> config [ '<STR_LIT>' ] = self . REQUEST_CHUNK_SIZE <EOL> ( shards , compositors ) = self . ShardStage ( config ) <EOL> if shards and compositors : <EOL> with pipeline . After ( * [ ( yield shard ) for shard in shards ] ) : <EOL> _ = [ ( yield compositor ) for compositor in compositors ] <EOL> else : <EOL> handler = _S3ReadBufferHandler ( s3_obj , <EOL> gcs_obj , <EOL> config . get ( '<STR_LIT>' ) ) <EOL> storage . ReadObject ( url = s3_obj , <EOL> handler = handler . Handle , <EOL> start = start , <EOL> length = length ) <EOL> comp_stage = common . Ignore ( ) <EOL> if handler . chunk_urls : <EOL> comp_config = { '<STR_LIT>' : handler . content_type , <EOL> '<STR_LIT>' : handler . chunk_urls , <EOL> '<STR_LIT>' : [ gcs_obj ] } <EOL> comp_stage = yield gcscompositor . GcsCompositor ( comp_config ) <EOL> sub_stages . append ( comp_stage ) <EOL> yield common . Append ( * sub_stages ) <EOL> def Lint ( self , linter ) : <EOL> """<STR_LIT>""" <EOL> linter . AtLeastOneFieldRequiredCheck ( [ '<STR_LIT:object>' , '<STR_LIT>' ] ) <EOL> linter . FieldCheck ( '<STR_LIT:object>' , validator = s3 . S3 . UrlToBucketAndName ) <EOL> linter . FieldCheck ( '<STR_LIT>' , validator = s3 . S3 . UrlToBucketAndName ) <EOL> linter . FieldCheck ( '<STR_LIT>' , field_type = dict , required = True ) <EOL> linter . FieldCheck ( '<STR_LIT>' , required = True ) <EOL> linter . FieldCheck ( '<STR_LIT>' , required = True ) <EOL> linter . FieldCheck ( '<STR_LIT>' , validator = self . ValidateShardSize ) <EOL> def ValidateShardSize ( self , size ) : <EOL> if size > self . MAX_SHARD_SIZE : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> class _S3ReadBufferHandler ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , s3_obj , gcs_obj , shard_prefix ) : <EOL> self . s3_obj = s3_obj <EOL> self . gcs_obj = gcs_obj <EOL> self . gcs_storage = gcs . Gcs ( ) <EOL> ( gcs_bucket , _ ) = gcs . Gcs . UrlToBucketAndName ( gcs_obj ) <EOL> self . url_gen = gcs . Gcs . UrlCreator ( gcs_bucket , shard_prefix ) <EOL> self . chunk_urls = [ ] <EOL> self . content_type = None <EOL> self . once = True <EOL> def Handle ( self , buf , unused_bytes_read , bytes_remaining , ct ) : <EOL> """<STR_LIT>""" <EOL> self . content_type = ct <EOL> if bytes_remaining == <NUM_LIT:0> and self . once : <EOL> self . gcs_storage . InsertObject ( buf , self . gcs_obj ) <EOL> else : <EOL> tmp = self . url_gen ( ) <EOL> self . chunk_urls . append ( tmp ) <EOL> self . gcs_storage . InsertObject ( buf , tmp ) <EOL> self . once = False </s>
