<s> from collections import namedtuple <EOL> from corehq . apps . change_feed . consumer . feed import KafkaChangeFeed <EOL> from corehq . apps . change_feed . document_types import GROUP <EOL> from corehq . apps . groups . models import Group <EOL> from corehq . elastic import stream_es_query , get_es_new , ES_META <EOL> from corehq . pillows . mappings . user_mapping import USER_INDEX <EOL> from pillowtop . checkpoints . manager import PillowCheckpoint , PillowCheckpointEventHandler <EOL> from pillowtop . pillow . interface import ConstructedPillow <EOL> from pillowtop . processors import PillowProcessor <EOL> from pillowtop . reindexer . change_providers . couch import CouchViewChangeProvider <EOL> from pillowtop . reindexer . reindexer import PillowReindexer <EOL> class GroupsToUsersProcessor ( PillowProcessor ) : <EOL> def __init__ ( self ) : <EOL> self . _es = get_es_new ( ) <EOL> def process_change ( self , pillow_instance , change , do_set_checkpoint ) : <EOL> if change . deleted : <EOL> remove_group_from_users ( change . get_document ( ) , self . _es ) <EOL> else : <EOL> update_es_user_with_groups ( change . get_document ( ) , self . _es ) <EOL> def get_group_to_user_pillow ( pillow_id = '<STR_LIT>' ) : <EOL> checkpoint = PillowCheckpoint ( <EOL> pillow_id , <EOL> ) <EOL> processor = GroupsToUsersProcessor ( ) <EOL> return ConstructedPillow ( <EOL> name = pillow_id , <EOL> checkpoint = checkpoint , <EOL> change_feed = KafkaChangeFeed ( topics = [ GROUP ] , group_id = '<STR_LIT>' ) , <EOL> processor = processor , <EOL> change_processed_event_handler = PillowCheckpointEventHandler ( <EOL> checkpoint = checkpoint , checkpoint_frequency = <NUM_LIT:100> , <EOL> ) , <EOL> ) <EOL> def remove_group_from_users ( group_doc , es_client ) : <EOL> for user_source in stream_user_sources ( group_doc . get ( "<STR_LIT>" , [ ] ) ) : <EOL> made_changes = False <EOL> if group_doc [ "<STR_LIT:name>" ] in user_source . group_names : <EOL> user_source . group_names . remove ( group_doc [ "<STR_LIT:name>" ] ) <EOL> made_changes = True <EOL> if group_doc [ "<STR_LIT>" ] in user_source . group_ids : <EOL> user_source . group_ids . remove ( group_doc [ "<STR_LIT>" ] ) <EOL> made_changes = True <EOL> if made_changes : <EOL> doc = { "<STR_LIT>" : list ( user_source . group_ids ) , "<STR_LIT>" : list ( user_source . group_names ) } <EOL> es_client . update ( USER_INDEX , ES_META [ '<STR_LIT>' ] . type , user_source . user_id , body = { "<STR_LIT>" : doc } ) <EOL> def update_es_user_with_groups ( group_doc , es_client = None ) : <EOL> if not es_client : <EOL> es_client = get_es_new ( ) <EOL> for user_source in stream_user_sources ( group_doc . get ( "<STR_LIT>" , [ ] ) ) : <EOL> if group_doc [ "<STR_LIT:name>" ] not in user_source . group_names or group_doc [ "<STR_LIT>" ] not in user_source . group_ids : <EOL> user_source . group_ids . add ( group_doc [ "<STR_LIT>" ] ) <EOL> user_source . group_names . add ( group_doc [ "<STR_LIT:name>" ] ) <EOL> doc = { "<STR_LIT>" : list ( user_source . group_ids ) , "<STR_LIT>" : list ( user_source . group_names ) } <EOL> es_client . update ( USER_INDEX , ES_META [ '<STR_LIT>' ] . type , user_source . user_id , body = { "<STR_LIT>" : doc } ) <EOL> UserSource = namedtuple ( '<STR_LIT>' , [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> def stream_user_sources ( user_ids ) : <EOL> q = { "<STR_LIT>" : { "<STR_LIT>" : [ { "<STR_LIT>" : { "<STR_LIT>" : user_ids } } ] } } <EOL> for result in stream_es_query ( es_index = '<STR_LIT>' , q = q , fields = [ "<STR_LIT>" , "<STR_LIT>" ] ) : <EOL> group_ids = result . get ( '<STR_LIT>' , { } ) . get ( "<STR_LIT>" , [ ] ) <EOL> group_ids = set ( group_ids ) if isinstance ( group_ids , list ) else { group_ids } <EOL> group_names = result . get ( '<STR_LIT>' , { } ) . get ( "<STR_LIT>" , [ ] ) <EOL> group_names = set ( group_names ) if isinstance ( group_names , list ) else { group_names } <EOL> yield UserSource ( result [ '<STR_LIT>' ] , group_ids , group_names ) <EOL> def get_groups_to_user_reindexer ( ) : <EOL> return PillowReindexer ( <EOL> pillow = get_group_to_user_pillow ( ) , <EOL> change_provider = CouchViewChangeProvider ( <EOL> couch_db = Group . get_db ( ) , <EOL> view_name = '<STR_LIT>' , <EOL> view_kwargs = { <EOL> '<STR_LIT>' : [ '<STR_LIT>' ] , <EOL> '<STR_LIT>' : [ '<STR_LIT>' , { } ] , <EOL> '<STR_LIT>' : True , <EOL> } <EOL> ) , <EOL> ) </s>
