<s> import collections <EOL> import datetime <EOL> import gc <EOL> import sys <EOL> from django . conf import settings <EOL> from django . http import Http404 <EOL> from elasticsearch import Elasticsearch , ElasticsearchException <EOL> from . import settings as es_settings <EOL> from . signals import post_indices_create , post_indices_rebuild <EOL> try : <EOL> from importlib import import_module <EOL> except ImportError : <EOL> from django . utils . importlib import import_module <EOL> _elasticsearch_indices = collections . defaultdict ( lambda : [ ] ) <EOL> def get_indices ( indices = [ ] ) : <EOL> if not _elasticsearch_indices : <EOL> type_classes = getattr ( settings , '<STR_LIT>' , ( ) ) <EOL> if not type_classes : <EOL> raise Exception ( '<STR_LIT>' ) <EOL> for type_class in type_classes : <EOL> package_name , klass_name = type_class . rsplit ( '<STR_LIT:.>' , <NUM_LIT:1> ) <EOL> try : <EOL> package = import_module ( package_name ) <EOL> klass = getattr ( package , klass_name ) <EOL> except ImportError : <EOL> sys . stderr . write ( '<STR_LIT>' . format ( type_class ) ) <EOL> continue <EOL> _elasticsearch_indices [ klass . get_index_name ( ) ] . append ( klass ) <EOL> if not indices : <EOL> return _elasticsearch_indices <EOL> else : <EOL> result = { } <EOL> for k , v in _elasticsearch_indices . items ( ) : <EOL> if k in indices : <EOL> result [ k ] = v <EOL> return result <EOL> def create_aliases ( es = None , indices = [ ] ) : <EOL> es = es or Elasticsearch ( ** es_settings . ELASTICSEARCH_CONNECTION_PARAMS ) <EOL> current_aliases = es . indices . get_aliases ( ) <EOL> aliases_for_removal = collections . defaultdict ( lambda : [ ] ) <EOL> for item , tmp in current_aliases . items ( ) : <EOL> for iname in list ( tmp . get ( '<STR_LIT>' , { } ) . keys ( ) ) : <EOL> aliases_for_removal [ iname ] . append ( item ) <EOL> actions = [ ] <EOL> for index_alias , index_name in indices : <EOL> for item in aliases_for_removal [ index_alias ] : <EOL> actions . append ( { <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT:index>' : item , <EOL> '<STR_LIT>' : index_alias <EOL> } <EOL> } ) <EOL> actions . append ( { <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT:index>' : index_name , <EOL> '<STR_LIT>' : index_alias <EOL> } <EOL> } ) <EOL> es . indices . update_aliases ( { '<STR_LIT>' : actions } ) <EOL> def create_indices ( es = None , indices = [ ] , set_aliases = True ) : <EOL> es = es or Elasticsearch ( ** es_settings . ELASTICSEARCH_CONNECTION_PARAMS ) <EOL> result = [ ] <EOL> aliases = [ ] <EOL> now = datetime . datetime . now ( ) . strftime ( "<STR_LIT>" ) <EOL> for index_alias , type_classes in get_indices ( indices ) . items ( ) : <EOL> index_settings = es_settings . ELASTICSEARCH_DEFAULT_INDEX_SETTINGS <EOL> index_settings = recursive_dict_update ( <EOL> index_settings , <EOL> es_settings . ELASTICSEARCH_CUSTOM_INDEX_SETTINGS . get ( index_alias , { } ) <EOL> ) <EOL> index_name = "<STR_LIT>" . format ( index_alias , now ) <EOL> aliases . append ( ( index_alias , index_name ) ) <EOL> type_mappings = { } <EOL> for type_class in type_classes : <EOL> tmp = type_class . get_type_mapping ( ) <EOL> if tmp : <EOL> type_mappings [ type_class . get_type_name ( ) ] = tmp <EOL> result . append ( ( <EOL> type_class , <EOL> index_alias , <EOL> index_name <EOL> ) ) <EOL> if type_mappings : <EOL> index_settings [ '<STR_LIT>' ] = type_mappings <EOL> es . indices . create ( index_name , index_settings ) <EOL> if set_aliases : <EOL> create_aliases ( es , aliases ) <EOL> post_indices_create . send ( None , indices = aliases , aliases_set = set_aliases ) <EOL> return result , aliases <EOL> def rebuild_indices ( es = None , indices = [ ] , set_aliases = True ) : <EOL> es = es or Elasticsearch ( ** es_settings . ELASTICSEARCH_CONNECTION_PARAMS ) <EOL> created_indices , aliases = create_indices ( es , indices , False ) <EOL> current_index_name = None <EOL> current_index_settings = { } <EOL> def change_index ( ) : <EOL> if current_index_name : <EOL> settings = { <EOL> '<STR_LIT>' : current_index_settings . get ( '<STR_LIT:index>' , { } ) . get ( '<STR_LIT>' , <NUM_LIT:1> ) , <EOL> '<STR_LIT>' : current_index_settings . get ( '<STR_LIT:index>' , { } ) . get ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> '<STR_LIT>' : current_index_settings . get ( '<STR_LIT:index>' , { } ) . get ( '<STR_LIT>' , <NUM_LIT:10> ) <EOL> } <EOL> es . indices . put_settings ( { '<STR_LIT:index>' : settings } , current_index_name ) <EOL> es . indices . refresh ( current_index_name ) <EOL> for type_class , index_alias , index_name in created_indices : <EOL> if index_name != current_index_name : <EOL> change_index ( ) <EOL> current_index_settings = es . indices . get_settings ( index_name ) . get ( index_name , { } ) . get ( '<STR_LIT>' , { } ) <EOL> current_index_name = index_name <EOL> es . indices . put_settings ( { '<STR_LIT:index>' : { <EOL> '<STR_LIT>' : <NUM_LIT:0> , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : <NUM_LIT:30> <EOL> } } , index = index_name ) <EOL> try : <EOL> type_class . bulk_index ( es , index_name ) <EOL> except NotImplementedError : <EOL> sys . stderr . write ( '<STR_LIT>' . format ( type_class . get_index_name ( ) ) ) <EOL> continue <EOL> else : <EOL> change_index ( ) <EOL> if set_aliases : <EOL> create_aliases ( es , aliases ) <EOL> post_indices_rebuild . send ( None , indices = aliases , aliases_set = set_aliases ) <EOL> return created_indices , aliases <EOL> def recursive_dict_update ( d , u ) : <EOL> for k , v in u . items ( ) : <EOL> if isinstance ( v , collections . Mapping ) : <EOL> r = recursive_dict_update ( d . get ( k , { } ) , v ) <EOL> d [ k ] = r <EOL> else : <EOL> d [ k ] = u [ k ] <EOL> return d <EOL> def queryset_iterator ( queryset , chunksize = <NUM_LIT:1000> ) : <EOL> try : <EOL> last_pk = queryset . order_by ( '<STR_LIT>' ) [ <NUM_LIT:0> ] . pk <EOL> except IndexError : <EOL> return <EOL> queryset = queryset . order_by ( '<STR_LIT>' ) <EOL> pk = queryset [ <NUM_LIT:0> ] . pk - <NUM_LIT:1> <EOL> while pk < last_pk : <EOL> for row in queryset . filter ( pk__gt = pk ) [ : chunksize ] : <EOL> pk = row . pk <EOL> yield row <EOL> gc . collect ( ) <EOL> def get_from_es_or_None ( index , type , id , ** kwargs ) : <EOL> es = kwargs . pop ( '<STR_LIT>' , Elasticsearch ( es_settings . ELASTICSEARCH_SERVER ) ) <EOL> try : <EOL> return es . get ( index , id , type , ** kwargs ) <EOL> except ElasticsearchException : <EOL> return None <EOL> def get_from_es_or_404 ( index , type , id , ** kwargs ) : <EOL> item = get_from_es_or_None ( index , type , id , ** kwargs ) <EOL> if not item : <EOL> raise Http404 ( '<STR_LIT>' . format ( type ) ) <EOL> return item </s>
