<s> from collections import defaultdict <EOL> from alembic . autogenerate . api import compare_metadata <EOL> from datetime import datetime , timedelta <EOL> from corehq . apps . change_feed import topics <EOL> from corehq . apps . change_feed . consumer . feed import KafkaChangeFeed , MultiTopicCheckpointEventHandler <EOL> from corehq . apps . userreports . data_source_providers import DynamicDataSourceProvider , StaticDataSourceProvider <EOL> from corehq . apps . userreports . exceptions import TableRebuildError , StaleRebuildError <EOL> from corehq . apps . userreports . sql import IndicatorSqlAdapter , metadata <EOL> from corehq . apps . userreports . tasks import is_static , rebuild_indicators <EOL> from corehq . sql_db . connections import connection_manager <EOL> from corehq . util . soft_assert import soft_assert <EOL> from fluff . signals import get_migration_context , get_tables_to_rebuild <EOL> from pillowtop . checkpoints . manager import PillowCheckpoint <EOL> from pillowtop . pillow . interface import ConstructedPillow <EOL> from pillowtop . processors import PillowProcessor <EOL> REBUILD_CHECK_INTERVAL = <NUM_LIT:10> * <NUM_LIT> <EOL> UCR_CHECKPOINT_ID = '<STR_LIT>' <EOL> UCR_STATIC_CHECKPOINT_ID = '<STR_LIT>' <EOL> class ConfigurableReportTableManagerMixin ( object ) : <EOL> def __init__ ( self , data_source_provider , auto_repopulate_tables = False , * args , ** kwargs ) : <EOL> self . bootstrapped = False <EOL> self . last_bootstrapped = datetime . utcnow ( ) <EOL> self . data_source_provider = data_source_provider <EOL> self . auto_repopulate_tables = auto_repopulate_tables <EOL> super ( ConfigurableReportTableManagerMixin , self ) . __init__ ( * args , ** kwargs ) <EOL> def get_all_configs ( self ) : <EOL> return self . data_source_provider . get_data_sources ( ) <EOL> def needs_bootstrap ( self ) : <EOL> return ( <EOL> not self . bootstrapped <EOL> or datetime . utcnow ( ) - self . last_bootstrapped > timedelta ( seconds = REBUILD_CHECK_INTERVAL ) <EOL> ) <EOL> def bootstrap_if_needed ( self ) : <EOL> if self . needs_bootstrap ( ) : <EOL> self . bootstrap ( ) <EOL> def bootstrap ( self , configs = None ) : <EOL> if configs is None : <EOL> configs = self . get_all_configs ( ) <EOL> self . table_adapters = [ IndicatorSqlAdapter ( config ) for config in configs ] <EOL> self . rebuild_tables_if_necessary ( ) <EOL> self . bootstrapped = True <EOL> self . last_bootstrapped = datetime . utcnow ( ) <EOL> def rebuild_tables_if_necessary ( self ) : <EOL> tables_by_engine = defaultdict ( dict ) <EOL> for adapter in self . table_adapters : <EOL> tables_by_engine [ adapter . engine_id ] [ adapter . get_table ( ) . name ] = adapter <EOL> _assert = soft_assert ( to = '<STR_LIT:@>' . join ( [ '<STR_LIT>' , '<STR_LIT>' ] ) ) <EOL> _notify_cory = lambda msg , obj : _assert ( False , msg , obj ) <EOL> for engine_id , table_map in tables_by_engine . items ( ) : <EOL> engine = connection_manager . get_engine ( engine_id ) <EOL> with engine . begin ( ) as connection : <EOL> migration_context = get_migration_context ( connection , table_map . keys ( ) ) <EOL> diffs = compare_metadata ( migration_context , metadata ) <EOL> tables_to_rebuild = get_tables_to_rebuild ( diffs , table_map . keys ( ) ) <EOL> for table_name in tables_to_rebuild : <EOL> sql_adapter = table_map [ table_name ] <EOL> if not is_static ( sql_adapter . config . _id ) : <EOL> try : <EOL> rev_before_rebuild = sql_adapter . config . get_db ( ) . get_rev ( sql_adapter . config . _id ) <EOL> self . rebuild_table ( sql_adapter ) <EOL> except TableRebuildError , e : <EOL> _notify_cory ( unicode ( e ) , sql_adapter . config . to_json ( ) ) <EOL> else : <EOL> rev_after_rebuild = sql_adapter . config . get_db ( ) . get_rev ( sql_adapter . config . _id ) <EOL> _notify_cory ( <EOL> u'<STR_LIT>' . format ( <EOL> table_name , <EOL> u'<STR_LIT>' . format ( sql_adapter . config . display_name , sql_adapter . config . _id ) , <EOL> diffs , <EOL> rev_before_rebuild , <EOL> rev_after_rebuild , <EOL> ) , <EOL> sql_adapter . config . to_json ( ) , <EOL> ) <EOL> else : <EOL> self . rebuild_table ( sql_adapter ) <EOL> def rebuild_table ( self , sql_adapter ) : <EOL> config = sql_adapter . config <EOL> if not is_static ( config . _id ) : <EOL> latest_rev = config . get_db ( ) . get_rev ( config . _id ) <EOL> if config . _rev != latest_rev : <EOL> raise StaleRebuildError ( '<STR_LIT>' . format ( config ) ) <EOL> sql_adapter . rebuild_table ( ) <EOL> if self . auto_repopulate_tables : <EOL> rebuild_indicators . delay ( sql_adapter . config . get_id ) <EOL> class ConfigurableReportPillowProcessor ( ConfigurableReportTableManagerMixin , PillowProcessor ) : <EOL> def process_change ( self , pillow_instance , change , do_set_checkpoint ) : <EOL> self . bootstrap_if_needed ( ) <EOL> if change . deleted : <EOL> return <EOL> domain = change . metadata . domain <EOL> if not domain : <EOL> return <EOL> for table in self . table_adapters : <EOL> if table . config . domain == domain : <EOL> doc = change . get_document ( ) <EOL> if table . config . filter ( doc ) : <EOL> table . best_effort_save ( doc ) <EOL> elif table . config . deleted_filter ( doc ) : <EOL> table . delete ( doc ) <EOL> class ConfigurableReportKafkaPillow ( ConstructedPillow ) : <EOL> def __init__ ( self , processor , pillow_name ) : <EOL> change_feed = KafkaChangeFeed ( topics . ALL , group_id = pillow_name ) <EOL> checkpoint = PillowCheckpoint ( pillow_name ) <EOL> event_handler = MultiTopicCheckpointEventHandler ( <EOL> checkpoint = checkpoint , checkpoint_frequency = <NUM_LIT:1000> , change_feed = change_feed <EOL> ) <EOL> super ( ConfigurableReportKafkaPillow , self ) . __init__ ( <EOL> name = pillow_name , <EOL> change_feed = change_feed , <EOL> processor = processor , <EOL> checkpoint = checkpoint , <EOL> change_processed_event_handler = event_handler <EOL> ) <EOL> assert self . _processor is not None <EOL> assert self . _processor . bootstrapped is not None <EOL> def bootstrap ( self , configs = None ) : <EOL> self . _processor . bootstrap ( configs ) <EOL> def rebuild_table ( self , sql_adapter ) : <EOL> self . _processor . rebuild_table ( sql_adapter ) <EOL> def get_kafka_ucr_pillow ( pillow_id = '<STR_LIT>' ) : <EOL> return ConfigurableReportKafkaPillow ( <EOL> processor = ConfigurableReportPillowProcessor ( <EOL> data_source_provider = DynamicDataSourceProvider ( ) , <EOL> auto_repopulate_tables = False , <EOL> ) , <EOL> pillow_name = pillow_id , <EOL> ) <EOL> def get_kafka_ucr_static_pillow ( pillow_id = '<STR_LIT>' ) : <EOL> return ConfigurableReportKafkaPillow ( <EOL> processor = ConfigurableReportPillowProcessor ( <EOL> data_source_provider = StaticDataSourceProvider ( ) , <EOL> auto_repopulate_tables = True , <EOL> ) , <EOL> pillow_name = pillow_id , <EOL> ) </s>
