<s> from memsql_loader . loader_db . storage import LoaderStorage <EOL> from memsql_loader . util . attr_dict import AttrDict <EOL> from memsql_loader . util import super_json as json <EOL> from memsql_loader . util import apsw_sql_utility , apsw_helpers , log , schema <EOL> from memsql_loader . util . apsw_sql_step_queue . time_helpers import unix_timestamp <EOL> from memsql_loader . vendor import glob2 <EOL> from boto . exception import S3ResponseError <EOL> from pywebhdfs . webhdfs import PyWebHdfsClient <EOL> import uuid <EOL> import datetime <EOL> import os <EOL> import hashlib <EOL> PRIMARY_TABLE = apsw_sql_utility . TableDefinition ( '<STR_LIT>' , """<STR_LIT>""" , index_columns = ( '<STR_LIT>' , ) ) <EOL> def hash_64_bit ( value ) : <EOL> result = hashlib . sha256 ( value . encode ( '<STR_LIT:utf-8>' ) ) <EOL> return int ( result . hexdigest ( ) [ : <NUM_LIT:16> ] , <NUM_LIT:16> ) <EOL> class Jobs ( apsw_sql_utility . APSWSQLUtility ) : <EOL> def __init__ ( self ) : <EOL> super ( Jobs , self ) . __init__ ( LoaderStorage ( ) ) <EOL> self . _define_table ( PRIMARY_TABLE ) <EOL> def save ( self , job ) : <EOL> assert isinstance ( job , Job ) , '<STR_LIT>' <EOL> with self . storage . transaction ( ) as cursor : <EOL> cursor . execute ( '''<STR_LIT>''' , ( job . id , unix_timestamp ( datetime . datetime . utcnow ( ) ) , job . json_spec ( ) ) ) <EOL> def delete ( self , job ) : <EOL> assert isinstance ( job , Job ) , '<STR_LIT>' <EOL> with self . storage . transaction ( ) as cursor : <EOL> cursor . execute ( '<STR_LIT>' , ( job . id , ) ) <EOL> def get ( self , job_id ) : <EOL> with self . storage . transaction ( ) as cursor : <EOL> job = apsw_helpers . get ( <EOL> cursor , '<STR_LIT>' , job_id ) <EOL> if job is not None : <EOL> job [ '<STR_LIT>' ] = json . loads ( job . spec ) <EOL> return Job ( job . spec , job . id ) <EOL> def all ( self ) : <EOL> with self . storage . cursor ( ) as cursor : <EOL> result = apsw_helpers . query ( <EOL> cursor , '<STR_LIT>' ) <EOL> return [ Job ( json . loads ( job . spec ) , job . id ) for job in result ] <EOL> def query_target ( self , host , port , database , table ) : <EOL> with self . storage . cursor ( ) as cursor : <EOL> result = apsw_helpers . query ( cursor , '<STR_LIT>' ) <EOL> ret = [ ] <EOL> for job in result : <EOL> spec = json . loads ( job . spec ) <EOL> if spec [ '<STR_LIT>' ] [ '<STR_LIT:host>' ] != host : <EOL> continue <EOL> if spec [ '<STR_LIT>' ] [ '<STR_LIT:port>' ] != port : <EOL> continue <EOL> if spec [ '<STR_LIT:target>' ] [ '<STR_LIT>' ] != database : <EOL> continue <EOL> if spec [ '<STR_LIT:target>' ] [ '<STR_LIT>' ] != table : <EOL> continue <EOL> ret . append ( Job ( spec , job . id ) ) <EOL> return ret <EOL> class Job ( object ) : <EOL> def __init__ ( self , spec , job_id = None ) : <EOL> """<STR_LIT>""" <EOL> self . id = job_id if job_id is not None else uuid . uuid1 ( ) . hex <EOL> self . spec = schema . validate_spec ( spec ) <EOL> self . paths = [ schema . LoadPath ( path ) for path in self . spec . source . paths ] <EOL> def json_spec ( self ) : <EOL> return json . dumps ( self . spec ) <EOL> def get_file_id ( self , key ) : <EOL> """<STR_LIT>""" <EOL> bucket_name = '<STR_LIT>' <EOL> if key . bucket is not None : <EOL> bucket_name = key . bucket . name <EOL> return hash_64_bit ( bucket_name + key . name ) <EOL> def has_file_id ( self ) : <EOL> assert '<STR_LIT>' in self . spec . options <EOL> return self . spec . options . file_id_column is not None <EOL> def get_files ( self , s3_conn = None ) : <EOL> logger = log . get_logger ( '<STR_LIT>' ) <EOL> for load_path in self . paths : <EOL> if load_path . scheme == '<STR_LIT>' : <EOL> bucket = s3_conn . get_bucket ( load_path . bucket ) <EOL> s3_globber = glob2 . S3Globber ( bucket ) <EOL> for keyname in s3_globber . glob ( load_path . pattern ) : <EOL> if not s3_globber . isdir ( keyname ) : <EOL> try : <EOL> key = s3_globber . get_key ( keyname ) <EOL> if key is not None : <EOL> yield AttrDict ( { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT:name>' : key . name , <EOL> '<STR_LIT>' : key . etag , <EOL> '<STR_LIT:size>' : key . size , <EOL> '<STR_LIT>' : bucket <EOL> } ) <EOL> else : <EOL> logger . warning ( "<STR_LIT>" , keyname ) <EOL> except S3ResponseError as e : <EOL> logger . warning ( "<STR_LIT>" , e . status , e . reason , keyname ) <EOL> elif load_path . scheme == '<STR_LIT:file>' : <EOL> fs_globber = glob2 . Globber ( ) <EOL> for fname in fs_globber . glob ( load_path . pattern ) : <EOL> if not fs_globber . isdir ( fname ) : <EOL> yield AttrDict ( { <EOL> '<STR_LIT>' : '<STR_LIT:file>' , <EOL> '<STR_LIT:name>' : fname , <EOL> '<STR_LIT>' : None , <EOL> '<STR_LIT:size>' : os . path . getsize ( fs_globber . _normalize_string ( fname ) ) , <EOL> '<STR_LIT>' : None <EOL> } ) <EOL> elif load_path . scheme == '<STR_LIT>' : <EOL> hdfs_host = self . spec . source . hdfs_host <EOL> webhdfs_port = self . spec . source . webhdfs_port <EOL> hdfs_user = self . spec . source . hdfs_user <EOL> client = PyWebHdfsClient ( <EOL> hdfs_host , webhdfs_port , user_name = hdfs_user ) <EOL> hdfs_globber = glob2 . HDFSGlobber ( client ) <EOL> for fname in hdfs_globber . glob ( load_path . pattern ) : <EOL> if not hdfs_globber . isdir ( fname ) : <EOL> fileinfo = hdfs_globber . get_fileinfo ( fname ) <EOL> yield AttrDict ( { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT:name>' : fileinfo [ '<STR_LIT:path>' ] , <EOL> '<STR_LIT>' : fileinfo [ '<STR_LIT>' ] , <EOL> '<STR_LIT:size>' : fileinfo [ '<STR_LIT>' ] , <EOL> '<STR_LIT>' : None <EOL> } ) <EOL> else : <EOL> assert False , "<STR_LIT>" % load_path . scheme </s>
