<s> """<STR_LIT>""" <EOL> import os <EOL> import urllib2 <EOL> import csv <EOL> import re <EOL> import datetime <EOL> from pyesgf . exceptions import Error , DuplicateHashError <EOL> import logging <EOL> log = logging . getLogger ( __name__ ) <EOL> class Manifest ( object ) : <EOL> VERSION = '<STR_LIT>' <EOL> def __init__ ( self , drs_id ) : <EOL> self . drs_id = drs_id <EOL> self . _contents = { } <EOL> @ classmethod <EOL> def from_result ( cls , dataset_result ) : <EOL> """<STR_LIT>""" <EOL> ctx = dataset_result . file_context ( ) <EOL> obj = cls ( ) <EOL> for file_result in ctx . search ( ) : <EOL> filename = file_result . filename <EOL> filehash = '<STR_LIT>' % ( file_result . checksum_type . lower ( ) , <EOL> file_result . checksum ) <EOL> tracking_id = file_result . tracking_id <EOL> if tracking_id is None : <EOL> tracking_id = '<STR_LIT>' <EOL> size = file_result . size <EOL> obj . add ( filename , filehash , tracking_id , size ) <EOL> return obj <EOL> @ classmethod <EOL> def from_mapfile ( cls , mapfile ) : <EOL> raise NotImplementedError <EOL> def add ( self , filename , filehash , tracking_id , size ) : <EOL> if filename in self . _contents : <EOL> old_filehash , old_tracking_id , old_size = self . _contents [ filename ] <EOL> if old_filehash == filehash : <EOL> log . warn ( '<STR_LIT>' <EOL> '<STR_LIT>' . format ( filename , self . drs_id ) ) <EOL> elif old_tracking_id == tracking_id : <EOL> raise DuplicateHashError ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' . format ( filename , <EOL> self . drs_id ) ) <EOL> else : <EOL> raise DuplicateHashError ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' . format ( filename , self . drs_id ) ) <EOL> self . _contents [ filename ] = ( filehash , tracking_id , size ) <EOL> def write ( self , fh ) : <EOL> fh . write ( '<STR_LIT>' . format ( self . VERSION ) ) <EOL> for filename in sorted ( self . _contents ) : <EOL> filehash , tracking_id , size = self . _contents [ filename ] <EOL> fh . write ( '<STR_LIT>' . format ( filename , filehash , tracking_id , size ) + '<STR_LIT:\n>' ) <EOL> class ManifestExtractor ( object ) : <EOL> pass <EOL> class SolrManifestExtractor ( ManifestExtractor ) : <EOL> SOLR_BATCH_SIZE = <NUM_LIT:1000> <EOL> SOLR_FIELDS = [ '<STR_LIT>' , '<STR_LIT:title>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT:size>' ] <EOL> def __init__ ( self , endpoint , project , from_date = None ) : <EOL> self . endpoint = endpoint <EOL> self . project = project <EOL> self . from_date = from_date <EOL> def _query ( self , offset ) : <EOL> fq = '<STR_LIT>' . format ( self . project ) <EOL> if self . from_date : <EOL> fq += '<STR_LIT>' . format ( self . from_date ) <EOL> params = ( <EOL> ( '<STR_LIT:q>' , '<STR_LIT:*>' ) , <EOL> ( '<STR_LIT>' , fq ) , <EOL> ( '<STR_LIT>' , '<STR_LIT:U+002C>' . join ( self . SOLR_FIELDS ) ) , <EOL> ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> ( '<STR_LIT:start>' , offset ) , <EOL> ( '<STR_LIT>' , self . SOLR_BATCH_SIZE ) , <EOL> ) <EOL> param_str = '<STR_LIT:&>' . join ( ( '<STR_LIT>' . format ( k , v ) ) for ( k , v ) in params ) <EOL> url = '<STR_LIT>' . format ( self . endpoint , <EOL> param_str ) <EOL> log . info ( '<STR_LIT>' . format ( url ) ) <EOL> response = urllib2 . urlopen ( url ) <EOL> return response <EOL> def _check_header ( self , header ) : <EOL> if header != self . SOLR_FIELDS : <EOL> raise Error ( '<STR_LIT>' ) <EOL> def _init_manifest ( self , dataset_id ) : <EOL> drs_id , datanode = dataset_id . split ( '<STR_LIT:|>' ) <EOL> log . info ( '<STR_LIT>' . format ( drs_id ) ) <EOL> return Manifest ( drs_id ) <EOL> def __iter__ ( self ) : <EOL> """<STR_LIT>""" <EOL> offset = <NUM_LIT:0> <EOL> current_dataset_id = None <EOL> current_manifest = None <EOL> while <NUM_LIT:1> : <EOL> response = self . _query ( offset ) <EOL> reader = csv . reader ( response ) <EOL> header = reader . next ( ) <EOL> self . _check_header ( header ) <EOL> empty_batch = True <EOL> for row in reader : <EOL> empty_batch = False <EOL> log . debug ( '<STR_LIT>' . format ( row ) ) <EOL> dataset_id , filename , checksum_type , checksum , tracking_id , size = row <EOL> if dataset_id != current_dataset_id : <EOL> if current_manifest : <EOL> yield current_manifest <EOL> current_manifest = self . _init_manifest ( dataset_id ) <EOL> current_dataset_id = dataset_id <EOL> filehash = '<STR_LIT>' . format ( checksum_type . lower ( ) , checksum ) <EOL> size = int ( size ) <EOL> try : <EOL> current_manifest . add ( filename , filehash , tracking_id , size ) <EOL> except DuplicateHashError , e : <EOL> log . error ( e ) <EOL> if empty_batch : <EOL> break <EOL> else : <EOL> offset += self . SOLR_BATCH_SIZE <EOL> if current_manifest : <EOL> yield current_manifest <EOL> def cmip5_manifest_partitioner ( drs_id ) : <EOL> """<STR_LIT>""" <EOL> ( activity , product , institute , model , experiment , <EOL> frequency , realm , table , ensemble , version ) = drs_id . split ( '<STR_LIT:.>' ) <EOL> parts = ( '<STR_LIT:.>' . join ( ( activity , product , institute ) ) , <EOL> '<STR_LIT:.>' . join ( ( model , experiment , frequency ) ) ) <EOL> return parts <EOL> def extract_from_solr ( endpoint , project , target_dir , from_date = None ) : <EOL> solr_extractor = SolrManifestExtractor ( endpoint , project , from_date ) <EOL> for manifest in solr_extractor : <EOL> manifest_path = os . path . join ( target_dir , <EOL> * cmip5_manifest_partitioner ( manifest . drs_id ) ) <EOL> if not os . path . exists ( manifest_path ) : <EOL> log . info ( '<STR_LIT>' . format ( manifest_path ) ) <EOL> os . makedirs ( manifest_path ) <EOL> manifest_file = os . path . join ( manifest_path , manifest . drs_id ) <EOL> if os . path . exists ( manifest_file ) : <EOL> log . warn ( '<STR_LIT>' . format ( manifest_file ) ) <EOL> with open ( manifest_file , '<STR_LIT:w>' ) as fh : <EOL> log . info ( '<STR_LIT>' . format ( manifest_file ) ) <EOL> manifest . write ( fh ) <EOL> def parse_timestamp ( timestamp ) : <EOL> """<STR_LIT>""" <EOL> if re . match ( r'<STR_LIT>' , timestamp ) : <EOL> timestamp += '<STR_LIT>' <EOL> elif timestamp [ - <NUM_LIT:1> ] != '<STR_LIT>' : <EOL> timestamp += '<STR_LIT>' <EOL> dt = datetime . datetime . strptime ( timestamp , '<STR_LIT>' ) <EOL> return dt <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> import argparse <EOL> logging . basicConfig ( level = logging . INFO ) <EOL> parser = argparse . ArgumentParser ( description = '<STR_LIT>' ) <EOL> parser . add_argument ( '<STR_LIT>' , help = '<STR_LIT>' ) <EOL> parser . add_argument ( '<STR_LIT>' , help = '<STR_LIT>' ) <EOL> parser . add_argument ( '<STR_LIT>' , help = '<STR_LIT>' ) <EOL> parser . add_argument ( '<STR_LIT>' , dest = '<STR_LIT>' , action = '<STR_LIT:store>' , <EOL> help = '<STR_LIT>' ) <EOL> args = parser . parse_args ( ) <EOL> if args . from_date : <EOL> from_date = parse_timestamp ( args . from_date ) <EOL> else : <EOL> from_date is None <EOL> extract_from_solr ( args . endpoint , args . project , args . repository , from_date = from_date ) </s>
