<s> import base <EOL> from brewery import dq <EOL> import time <EOL> from brewery . metadata import expand_record <EOL> try : <EOL> from pyes . es import ES <EOL> except ImportError : <EOL> from brewery . utils import MissingPackage <EOL> pyes = MissingPackage ( "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ) <EOL> class ESDataSource ( base . DataSource ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , document_type , database = None , host = None , port = None , <EOL> expand = False , ** elasticsearch_args ) : <EOL> """<STR_LIT>""" <EOL> self . document_type = document_type <EOL> self . database_name = database <EOL> self . host = host <EOL> self . port = port <EOL> self . elasticsearch_args = elasticsearch_args <EOL> self . expand = expand <EOL> self . connection = None <EOL> self . _fields = None <EOL> def initialize ( self ) : <EOL> """<STR_LIT>""" <EOL> args = self . elasticsearch_args . copy ( ) <EOL> server = "<STR_LIT>" <EOL> if self . host : <EOL> server = self . host <EOL> if self . port : <EOL> server += "<STR_LIT::>" + self . port <EOL> self . connection = ES ( server , ** args ) <EOL> self . connection . default_indices = self . database_name <EOL> self . connection . default_types = self . document_type <EOL> def read_fields ( self , limit = <NUM_LIT:0> ) : <EOL> keys = [ ] <EOL> probes = { } <EOL> def probe_record ( record , parent = None ) : <EOL> for key , value in record . items ( ) : <EOL> if parent : <EOL> full_key = parent + "<STR_LIT:.>" + key <EOL> else : <EOL> full_key = key <EOL> if self . expand and type ( value ) == dict : <EOL> probe_record ( value , full_key ) <EOL> continue <EOL> if not full_key in probes : <EOL> probe = dq . FieldTypeProbe ( full_key ) <EOL> probes [ full_key ] = probe <EOL> keys . append ( full_key ) <EOL> else : <EOL> probe = probes [ full_key ] <EOL> probe . probe ( value ) <EOL> for record in self . document_type . find ( limit = limit ) : <EOL> probe_record ( record ) <EOL> fields = [ ] <EOL> for key in keys : <EOL> probe = probes [ key ] <EOL> field = base . Field ( probe . field ) <EOL> storage_type = probe . unique_storage_type <EOL> if not storage_type : <EOL> field . storage_type = "<STR_LIT>" <EOL> elif storage_type == "<STR_LIT>" : <EOL> field . storage_type = "<STR_LIT:string>" <EOL> else : <EOL> field . storage_type = "<STR_LIT>" <EOL> field . concrete_storage_type = storage_type <EOL> fields . append ( field ) <EOL> self . fields = list ( fields ) <EOL> return self . fields <EOL> def rows ( self ) : <EOL> if not self . connection : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> from pyes . query import MatchAllQuery <EOL> fields = self . fields . names ( ) <EOL> results = self . connection . search ( MatchAllQuery ( ) , search_type = "<STR_LIT>" , timeout = "<STR_LIT>" , size = "<STR_LIT>" ) <EOL> return ESRowIterator ( results , fields ) <EOL> def records ( self ) : <EOL> if not self . connection : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> from pyes . query import MatchAllQuery <EOL> results = self . connection . search ( MatchAllQuery ( ) , search_type = "<STR_LIT>" , timeout = "<STR_LIT>" , size = "<STR_LIT>" ) <EOL> return ESRecordIterator ( results , self . expand ) <EOL> class ESRowIterator ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , resultset , field_names ) : <EOL> self . resultset = resultset <EOL> self . field_names = field_names <EOL> def __getitem__ ( self , index ) : <EOL> record = self . resultset . __getitem__ ( index ) <EOL> array = [ ] <EOL> for field in self . field_names : <EOL> value = record <EOL> for key in field . split ( '<STR_LIT:.>' ) : <EOL> if key in value : <EOL> value = value [ key ] <EOL> else : <EOL> break <EOL> array . append ( value ) <EOL> return tuple ( array ) <EOL> class ESRecordIterator ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , resultset , expand = False ) : <EOL> self . resultset = resultset <EOL> self . expand = expand <EOL> def __getitem__ ( self , index ) : <EOL> def expand_record ( record , parent = None ) : <EOL> ret = { } <EOL> for key , value in record . items ( ) : <EOL> if parent : <EOL> full_key = parent + "<STR_LIT:.>" + key <EOL> else : <EOL> full_key = key <EOL> if type ( value ) == dict : <EOL> expanded = expand_record ( value , full_key ) <EOL> ret . update ( expanded ) <EOL> else : <EOL> ret [ full_key ] = value <EOL> return ret <EOL> record = self . resultset . __getitem__ ( index ) <EOL> if not self . expand : <EOL> return record <EOL> else : <EOL> return expand_record ( record ) <EOL> class ESDataTarget ( base . DataTarget ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , document_type , database = "<STR_LIT:test>" , host = "<STR_LIT:127.0.0.1>" , port = "<STR_LIT>" , <EOL> truncate = False , expand = False , ** elasticsearch_args ) : <EOL> """<STR_LIT>""" <EOL> self . document_type = document_type <EOL> self . database_name = database <EOL> self . host = host <EOL> self . port = port <EOL> self . elasticsearch_args = elasticsearch_args <EOL> self . expand = expand <EOL> self . truncate = truncate <EOL> self . _fields = None <EOL> def initialize ( self ) : <EOL> """<STR_LIT>""" <EOL> from pyes . es import ES <EOL> from pyes . exceptions import IndexAlreadyExistsException <EOL> args = self . elasticsearch_args . copy ( ) <EOL> server = "<STR_LIT>" <EOL> if self . host : <EOL> server = self . host <EOL> if self . port : <EOL> server += "<STR_LIT::>" + self . port <EOL> create = args . pop ( "<STR_LIT>" , False ) <EOL> replace = args . pop ( "<STR_LIT:replace>" , False ) <EOL> self . connection = ES ( server , ** args ) <EOL> self . connection . default_indices = self . database_name <EOL> self . connection . default_types = self . document_type <EOL> created = False <EOL> if create : <EOL> try : <EOL> self . connection . create_index ( self . database_name ) <EOL> self . connection . refresh ( self . database_name ) <EOL> created = True <EOL> except IndexAlreadyExistsException : <EOL> pass <EOL> if replace and not created : <EOL> self . connection . delete_index_if_exists ( self . database_name ) <EOL> time . sleep ( <NUM_LIT:2> ) <EOL> self . connection . create_index ( self . database_name ) <EOL> self . connection . refresh ( self . database_name ) <EOL> if self . truncate : <EOL> self . connection . delete_mapping ( self . database_name , self . document_type ) <EOL> self . connection . refresh ( self . database_name ) <EOL> def append ( self , obj ) : <EOL> record = obj <EOL> if not isinstance ( obj , dict ) : <EOL> record = dict ( zip ( self . fields . names ( ) , obj ) ) <EOL> if self . expand : <EOL> record = expand_record ( record ) <EOL> id = record . get ( '<STR_LIT:id>' ) or record . get ( '<STR_LIT>' ) <EOL> self . connection . index ( record , self . database_name , self . document_type , id , bulk = True ) <EOL> def finalize ( self ) : <EOL> self . connection . flush_bulk ( forced = True ) </s>
