<s> """<STR_LIT>""" <EOL> from __future__ import division <EOL> import datetime <EOL> import itertools <EOL> import logging <EOL> import os <EOL> import os . path <EOL> import time <EOL> from multiprocessing . pool import ThreadPool <EOL> try : <EOL> from urlparse import urlsplit <EOL> except ImportError : <EOL> from urllib . parse import urlsplit <EOL> import warnings <EOL> try : <EOL> from ConfigParser import NoSectionError <EOL> except ImportError : <EOL> from configparser import NoSectionError <EOL> from luigi import six <EOL> from luigi . six . moves import range <EOL> from luigi import configuration <EOL> from luigi . format import get_default_format <EOL> from luigi . parameter import Parameter <EOL> from luigi . target import FileAlreadyExists , FileSystem , FileSystemException , FileSystemTarget , AtomicLocalFile , MissingParentDirectory <EOL> from luigi . task import ExternalTask <EOL> logger = logging . getLogger ( '<STR_LIT>' ) <EOL> S3_DIRECTORY_MARKER_SUFFIX_0 = '<STR_LIT>' <EOL> S3_DIRECTORY_MARKER_SUFFIX_1 = '<STR_LIT:/>' <EOL> class InvalidDeleteException ( FileSystemException ) : <EOL> pass <EOL> class FileNotFoundException ( FileSystemException ) : <EOL> pass <EOL> class S3Client ( FileSystem ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , aws_access_key_id = None , aws_secret_access_key = None , <EOL> ** kwargs ) : <EOL> import boto <EOL> import boto . s3 . connection <EOL> from boto . s3 . key import Key <EOL> options = self . _get_s3_config ( ) <EOL> options . update ( kwargs ) <EOL> role_arn = options . get ( '<STR_LIT>' ) <EOL> role_session_name = options . get ( '<STR_LIT>' ) <EOL> aws_session_token = None <EOL> if role_arn and role_session_name : <EOL> from boto import sts <EOL> sts_client = sts . STSConnection ( ) <EOL> assumed_role = sts_client . assume_role ( role_arn , role_session_name ) <EOL> aws_secret_access_key = assumed_role . credentials . secret_key <EOL> aws_access_key_id = assumed_role . credentials . access_key <EOL> aws_session_token = assumed_role . credentials . session_token <EOL> else : <EOL> if not aws_access_key_id : <EOL> aws_access_key_id = options . get ( '<STR_LIT>' ) <EOL> if not aws_secret_access_key : <EOL> aws_secret_access_key = options . get ( '<STR_LIT>' ) <EOL> for key in [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] : <EOL> if key in options : <EOL> options . pop ( key ) <EOL> self . s3 = boto . s3 . connection . S3Connection ( aws_access_key_id , <EOL> aws_secret_access_key , <EOL> security_token = aws_session_token , <EOL> ** options ) <EOL> self . Key = Key <EOL> def exists ( self , path ) : <EOL> """<STR_LIT>""" <EOL> ( bucket , key ) = self . _path_to_bucket_and_key ( path ) <EOL> s3_bucket = self . s3 . get_bucket ( bucket , validate = True ) <EOL> if self . _is_root ( key ) : <EOL> return True <EOL> s3_key = s3_bucket . get_key ( key ) <EOL> if s3_key : <EOL> return True <EOL> if self . isdir ( path ) : <EOL> return True <EOL> logger . debug ( '<STR_LIT>' , path ) <EOL> return False <EOL> def remove ( self , path , recursive = True ) : <EOL> """<STR_LIT>""" <EOL> if not self . exists ( path ) : <EOL> logger . debug ( '<STR_LIT>' , path ) <EOL> return False <EOL> ( bucket , key ) = self . _path_to_bucket_and_key ( path ) <EOL> if self . _is_root ( key ) : <EOL> raise InvalidDeleteException ( '<STR_LIT>' % path ) <EOL> s3_bucket = self . s3 . get_bucket ( bucket , validate = True ) <EOL> s3_key = s3_bucket . get_key ( key ) <EOL> if s3_key : <EOL> s3_bucket . delete_key ( s3_key ) <EOL> logger . debug ( '<STR_LIT>' , key , bucket ) <EOL> return True <EOL> if self . isdir ( path ) and not recursive : <EOL> raise InvalidDeleteException ( '<STR_LIT>' % path ) <EOL> delete_key_list = [ <EOL> k for k in s3_bucket . list ( self . _add_path_delimiter ( key ) ) ] <EOL> s3_dir_with_suffix_key = s3_bucket . get_key ( key + S3_DIRECTORY_MARKER_SUFFIX_0 ) <EOL> if s3_dir_with_suffix_key : <EOL> delete_key_list . append ( s3_dir_with_suffix_key ) <EOL> if len ( delete_key_list ) > <NUM_LIT:0> : <EOL> for k in delete_key_list : <EOL> logger . debug ( '<STR_LIT>' , k , bucket ) <EOL> s3_bucket . delete_keys ( delete_key_list ) <EOL> return True <EOL> return False <EOL> def get_key ( self , path ) : <EOL> ( bucket , key ) = self . _path_to_bucket_and_key ( path ) <EOL> s3_bucket = self . s3 . get_bucket ( bucket , validate = True ) <EOL> return s3_bucket . get_key ( key ) <EOL> def put ( self , local_path , destination_s3_path , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> ( bucket , key ) = self . _path_to_bucket_and_key ( destination_s3_path ) <EOL> s3_bucket = self . s3 . get_bucket ( bucket , validate = True ) <EOL> s3_key = self . Key ( s3_bucket ) <EOL> s3_key . key = key <EOL> s3_key . set_contents_from_filename ( local_path , ** kwargs ) <EOL> def put_string ( self , content , destination_s3_path , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> ( bucket , key ) = self . _path_to_bucket_and_key ( destination_s3_path ) <EOL> s3_bucket = self . s3 . get_bucket ( bucket , validate = True ) <EOL> s3_key = self . Key ( s3_bucket ) <EOL> s3_key . key = key <EOL> s3_key . set_contents_from_string ( content , ** kwargs ) <EOL> def put_multipart ( self , local_path , destination_s3_path , part_size = <NUM_LIT> , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> source_size = os . stat ( local_path ) . st_size <EOL> if source_size <= part_size : <EOL> return self . put ( local_path , destination_s3_path , ** kwargs ) <EOL> ( bucket , key ) = self . _path_to_bucket_and_key ( destination_s3_path ) <EOL> s3_bucket = self . s3 . get_bucket ( bucket , validate = True ) <EOL> num_parts = ( source_size + part_size - <NUM_LIT:1> ) // part_size <EOL> mp = None <EOL> try : <EOL> mp = s3_bucket . initiate_multipart_upload ( key , ** kwargs ) <EOL> for i in range ( num_parts ) : <EOL> offset = part_size * i <EOL> bytes = min ( part_size , source_size - offset ) <EOL> with open ( local_path , '<STR_LIT:rb>' ) as fp : <EOL> part_num = i + <NUM_LIT:1> <EOL> logger . info ( '<STR_LIT>' , part_num , num_parts , destination_s3_path ) <EOL> fp . seek ( offset ) <EOL> mp . upload_part_from_file ( fp , part_num = part_num , size = bytes ) <EOL> mp . complete_upload ( ) <EOL> except BaseException : <EOL> if mp : <EOL> logger . info ( '<STR_LIT>' , destination_s3_path ) <EOL> mp . cancel_upload ( ) <EOL> raise <EOL> def get ( self , s3_path , destination_local_path ) : <EOL> """<STR_LIT>""" <EOL> ( bucket , key ) = self . _path_to_bucket_and_key ( s3_path ) <EOL> s3_bucket = self . s3 . get_bucket ( bucket , validate = True ) <EOL> s3_key = self . Key ( s3_bucket ) <EOL> s3_key . key = key <EOL> s3_key . get_contents_to_filename ( destination_local_path ) <EOL> def get_as_string ( self , s3_path ) : <EOL> """<STR_LIT>""" <EOL> ( bucket , key ) = self . _path_to_bucket_and_key ( s3_path ) <EOL> s3_bucket = self . s3 . get_bucket ( bucket , validate = True ) <EOL> s3_key = self . Key ( s3_bucket ) <EOL> s3_key . key = key <EOL> contents = s3_key . get_contents_as_string ( ) <EOL> return contents <EOL> def copy ( self , source_path , destination_path , threads = <NUM_LIT:100> , start_time = None , end_time = None , part_size = <NUM_LIT> , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> start = datetime . datetime . now ( ) <EOL> ( src_bucket , src_key ) = self . _path_to_bucket_and_key ( source_path ) <EOL> ( dst_bucket , dst_key ) = self . _path_to_bucket_and_key ( destination_path ) <EOL> threads = <NUM_LIT:3> if threads < <NUM_LIT:3> else threads <EOL> total_keys = <NUM_LIT:0> <EOL> copy_pool = ThreadPool ( processes = threads ) <EOL> if self . isdir ( source_path ) : <EOL> copy_jobs = [ ] <EOL> management_pool = ThreadPool ( processes = threads ) <EOL> src_prefix = self . _add_path_delimiter ( src_key ) <EOL> dst_prefix = self . _add_path_delimiter ( dst_key ) <EOL> for key in self . list ( source_path , start_time = start_time , end_time = end_time ) : <EOL> if key != '<STR_LIT>' and key != '<STR_LIT:/>' : <EOL> total_keys += <NUM_LIT:1> <EOL> job = management_pool . apply_async ( self . __copy_multipart , <EOL> args = ( copy_pool , <EOL> src_bucket , src_prefix + key , <EOL> dst_bucket , dst_prefix + key , <EOL> part_size ) , <EOL> kwds = kwargs ) <EOL> copy_jobs . append ( job ) <EOL> management_pool . close ( ) <EOL> management_pool . join ( ) <EOL> copy_pool . close ( ) <EOL> copy_pool . join ( ) <EOL> for result in copy_jobs : <EOL> result . get ( ) <EOL> end = datetime . datetime . now ( ) <EOL> duration = end - start <EOL> logger . info ( '<STR_LIT>' % <EOL> ( datetime . datetime . now ( ) , total_keys , duration ) ) <EOL> else : <EOL> self . __copy_multipart ( copy_pool , src_bucket , src_key , dst_bucket , dst_key , part_size , ** kwargs ) <EOL> copy_pool . close ( ) <EOL> copy_pool . join ( ) <EOL> def __copy_multipart ( self , pool , src_bucket , src_key , dst_bucket , dst_key , part_size , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> source_bucket = self . s3 . get_bucket ( src_bucket , validate = True ) <EOL> dest_bucket = self . s3 . get_bucket ( dst_bucket , validate = True ) <EOL> key_size = source_bucket . lookup ( src_key ) . size <EOL> if key_size == <NUM_LIT:0> or key_size <= part_size : <EOL> result = pool . apply_async ( dest_bucket . copy_key , args = ( dst_key , src_bucket , src_key ) , kwds = kwargs ) <EOL> return result . get ( ) <EOL> mp = None <EOL> try : <EOL> mp = dest_bucket . initiate_multipart_upload ( dst_key , ** kwargs ) <EOL> cur_pos = <NUM_LIT:0> <EOL> results = [ ] <EOL> num_parts = ( key_size + part_size - <NUM_LIT:1> ) // part_size <EOL> for i in range ( num_parts ) : <EOL> part_start = cur_pos <EOL> cur_pos += part_size <EOL> part_end = min ( cur_pos - <NUM_LIT:1> , key_size - <NUM_LIT:1> ) <EOL> part_num = i + <NUM_LIT:1> <EOL> results . append ( pool . apply_async ( mp . copy_part_from_key , args = ( src_bucket , src_key , part_num , part_start , part_end ) ) ) <EOL> logger . info ( '<STR_LIT>' , part_num , num_parts , dst_bucket , dst_key ) <EOL> logger . info ( '<STR_LIT>' , dst_bucket , dst_key ) <EOL> for result in results : <EOL> result . get ( ) <EOL> mp . complete_upload ( ) <EOL> return mp . key_name <EOL> except : <EOL> logger . info ( '<STR_LIT>' , src_bucket , src_key , dst_bucket , dst_key ) <EOL> if mp : <EOL> mp . cancel_upload ( ) <EOL> raise <EOL> def rename ( self , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> self . move ( * args , ** kwargs ) <EOL> def move ( self , source_path , destination_path , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> self . copy ( source_path , destination_path , ** kwargs ) <EOL> self . remove ( source_path ) <EOL> def listdir ( self , path , start_time = None , end_time = None ) : <EOL> """<STR_LIT>""" <EOL> ( bucket , key ) = self . _path_to_bucket_and_key ( path ) <EOL> s3_bucket = self . s3 . get_bucket ( bucket , validate = True ) <EOL> key_path = self . _add_path_delimiter ( key ) <EOL> key_path_len = len ( key_path ) <EOL> for item in s3_bucket . list ( prefix = key_path ) : <EOL> last_modified_date = time . strptime ( item . last_modified , "<STR_LIT>" ) <EOL> if ( <EOL> ( not start_time and not end_time ) or <EOL> ( start_time and not end_time and start_time < last_modified_date ) or <EOL> ( not start_time and end_time and last_modified_date < end_time ) or <EOL> ( start_time and end_time and start_time < last_modified_date < end_time ) <EOL> ) : <EOL> yield self . _add_path_delimiter ( path ) + item . key [ key_path_len : ] <EOL> def list ( self , path , start_time = None , end_time = None ) : <EOL> key_path_len = len ( self . _add_path_delimiter ( path ) ) <EOL> for item in self . listdir ( path , start_time = start_time , end_time = end_time ) : <EOL> yield item [ key_path_len : ] <EOL> def isdir ( self , path ) : <EOL> """<STR_LIT>""" <EOL> ( bucket , key ) = self . _path_to_bucket_and_key ( path ) <EOL> s3_bucket = self . s3 . get_bucket ( bucket , validate = True ) <EOL> if self . _is_root ( key ) : <EOL> return True <EOL> for suffix in ( S3_DIRECTORY_MARKER_SUFFIX_0 , <EOL> S3_DIRECTORY_MARKER_SUFFIX_1 ) : <EOL> s3_dir_with_suffix_key = s3_bucket . get_key ( key + suffix ) <EOL> if s3_dir_with_suffix_key : <EOL> return True <EOL> key_path = self . _add_path_delimiter ( key ) <EOL> s3_bucket_list_result = list ( itertools . islice ( s3_bucket . list ( prefix = key_path ) , <NUM_LIT:1> ) ) <EOL> if s3_bucket_list_result : <EOL> return True <EOL> return False <EOL> is_dir = isdir <EOL> def mkdir ( self , path , parents = True , raise_if_exists = False ) : <EOL> if raise_if_exists and self . isdir ( path ) : <EOL> raise FileAlreadyExists ( ) <EOL> _ , key = self . _path_to_bucket_and_key ( path ) <EOL> if self . _is_root ( key ) : <EOL> return <EOL> key = self . _add_path_delimiter ( key ) <EOL> if not parents and not self . isdir ( os . path . dirname ( key ) ) : <EOL> raise MissingParentDirectory ( ) <EOL> return self . put_string ( "<STR_LIT>" , self . _add_path_delimiter ( path ) ) <EOL> def _get_s3_config ( self , key = None ) : <EOL> try : <EOL> config = dict ( configuration . get_config ( ) . items ( '<STR_LIT>' ) ) <EOL> except NoSectionError : <EOL> return { } <EOL> for k , v in six . iteritems ( config ) : <EOL> try : <EOL> config [ k ] = int ( v ) <EOL> except ValueError : <EOL> pass <EOL> if key : <EOL> return config . get ( key ) <EOL> return config <EOL> def _path_to_bucket_and_key ( self , path ) : <EOL> ( scheme , netloc , path , query , fragment ) = urlsplit ( path ) <EOL> path_without_initial_slash = path [ <NUM_LIT:1> : ] <EOL> return netloc , path_without_initial_slash <EOL> def _is_root ( self , key ) : <EOL> return ( len ( key ) == <NUM_LIT:0> ) or ( key == '<STR_LIT:/>' ) <EOL> def _add_path_delimiter ( self , key ) : <EOL> return key if key [ - <NUM_LIT:1> : ] == '<STR_LIT:/>' or key == '<STR_LIT>' else key + '<STR_LIT:/>' <EOL> class AtomicS3File ( AtomicLocalFile ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , path , s3_client , ** kwargs ) : <EOL> self . s3_client = s3_client <EOL> super ( AtomicS3File , self ) . __init__ ( path ) <EOL> self . s3_options = kwargs <EOL> def move_to_final_destination ( self ) : <EOL> self . s3_client . put_multipart ( self . tmp_path , self . path , ** self . s3_options ) <EOL> class ReadableS3File ( object ) : <EOL> def __init__ ( self , s3_key ) : <EOL> self . s3_key = s3_key <EOL> self . buffer = [ ] <EOL> self . closed = False <EOL> self . finished = False <EOL> def read ( self , size = <NUM_LIT:0> ) : <EOL> f = self . s3_key . read ( size = size ) <EOL> if f == b'<STR_LIT>' : <EOL> self . finished = True <EOL> if self . finished : <EOL> return b'<STR_LIT>' <EOL> return f <EOL> def close ( self ) : <EOL> self . s3_key . close ( ) <EOL> self . closed = True <EOL> def __del__ ( self ) : <EOL> self . close ( ) <EOL> def __exit__ ( self , exc_type , exc , traceback ) : <EOL> self . close ( ) <EOL> def __enter__ ( self ) : <EOL> return self <EOL> def _add_to_buffer ( self , line ) : <EOL> self . buffer . append ( line ) <EOL> def _flush_buffer ( self ) : <EOL> output = b'<STR_LIT>' . join ( self . buffer ) <EOL> self . buffer = [ ] <EOL> return output <EOL> def readable ( self ) : <EOL> return True <EOL> def writable ( self ) : <EOL> return False <EOL> def seekable ( self ) : <EOL> return False <EOL> def __iter__ ( self ) : <EOL> key_iter = self . s3_key . __iter__ ( ) <EOL> has_next = True <EOL> while has_next : <EOL> try : <EOL> chunk = next ( key_iter ) <EOL> for line in chunk . splitlines ( True ) : <EOL> if not line . endswith ( os . linesep ) : <EOL> self . _add_to_buffer ( line ) <EOL> else : <EOL> if self . buffer : <EOL> self . _add_to_buffer ( line ) <EOL> yield self . _flush_buffer ( ) <EOL> else : <EOL> yield line <EOL> except StopIteration : <EOL> output = self . _flush_buffer ( ) <EOL> if output : <EOL> yield output <EOL> has_next = False <EOL> self . close ( ) <EOL> class S3Target ( FileSystemTarget ) : <EOL> """<STR_LIT>""" <EOL> fs = None <EOL> def __init__ ( self , path , format = None , client = None , ** kwargs ) : <EOL> super ( S3Target , self ) . __init__ ( path ) <EOL> if format is None : <EOL> format = get_default_format ( ) <EOL> self . path = path <EOL> self . format = format <EOL> self . fs = client or S3Client ( ) <EOL> self . s3_options = kwargs <EOL> def open ( self , mode = '<STR_LIT:r>' ) : <EOL> """<STR_LIT:U+0020>""" <EOL> if mode not in ( '<STR_LIT:r>' , '<STR_LIT:w>' ) : <EOL> raise ValueError ( "<STR_LIT>" % mode ) <EOL> if mode == '<STR_LIT:r>' : <EOL> s3_key = self . fs . get_key ( self . path ) <EOL> if not s3_key : <EOL> raise FileNotFoundException ( "<STR_LIT>" % self . path ) <EOL> fileobj = ReadableS3File ( s3_key ) <EOL> return self . format . pipe_reader ( fileobj ) <EOL> else : <EOL> return self . format . pipe_writer ( AtomicS3File ( self . path , self . fs , ** self . s3_options ) ) <EOL> class S3FlagTarget ( S3Target ) : <EOL> """<STR_LIT>""" <EOL> fs = None <EOL> def __init__ ( self , path , format = None , client = None , flag = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> if format is None : <EOL> format = get_default_format ( ) <EOL> if path [ - <NUM_LIT:1> ] != "<STR_LIT:/>" : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> super ( S3FlagTarget , self ) . __init__ ( path , format , client ) <EOL> self . flag = flag <EOL> def exists ( self ) : <EOL> hadoopSemaphore = self . path + self . flag <EOL> return self . fs . exists ( hadoopSemaphore ) <EOL> class S3EmrTarget ( S3FlagTarget ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> warnings . warn ( "<STR_LIT>" ) <EOL> super ( S3EmrTarget , self ) . __init__ ( * args , ** kwargs ) <EOL> class S3PathTask ( ExternalTask ) : <EOL> """<STR_LIT>""" <EOL> path = Parameter ( ) <EOL> def output ( self ) : <EOL> return S3Target ( self . path ) <EOL> class S3EmrTask ( ExternalTask ) : <EOL> """<STR_LIT>""" <EOL> path = Parameter ( ) <EOL> def output ( self ) : <EOL> return S3EmrTarget ( self . path ) <EOL> class S3FlagTask ( ExternalTask ) : <EOL> """<STR_LIT>""" <EOL> path = Parameter ( ) <EOL> flag = Parameter ( default = None ) <EOL> def output ( self ) : <EOL> return S3FlagTarget ( self . path , flag = self . flag ) </s>
