<s> from functools import wraps <EOL> import logging <EOL> from couchdbkit . exceptions import ResourceNotFound <EOL> from elasticsearch . exceptions import RequestError , ConnectionError , NotFoundError , ConflictError <EOL> from psycopg2 . _psycopg import InterfaceError as Psycopg2InterfaceError <EOL> from django . db . utils import InterfaceError as DjangoInterfaceError <EOL> from datetime import datetime , timedelta <EOL> import hashlib <EOL> import traceback <EOL> import math <EOL> import time <EOL> import simplejson <EOL> import sys <EOL> from dimagi . utils . decorators . memoized import memoized <EOL> from dimagi . utils . couch import LockManager <EOL> from pillow_retry . models import PillowError <EOL> from pillowtop . checkpoints . manager import PillowCheckpoint , get_default_django_checkpoint_for_legacy_pillow_class <EOL> from pillowtop . checkpoints . util import get_machine_id , construct_checkpoint_doc_id_from_name <EOL> from pillowtop . const import CHECKPOINT_FREQUENCY <EOL> from pillowtop . couchdb import CachedCouchDB <EOL> from django import db <EOL> from pillowtop . dao . couch import CouchDocumentStore <EOL> from pillowtop . es_utils import completely_initialize_pillow_index , doc_exists <EOL> from pillowtop . feed . couch import CouchChangeFeed <EOL> from pillowtop . logger import pillow_logging <EOL> from pillowtop . pillow . interface import PillowBase <EOL> from pillowtop . utils import prepare_bulk_payloads <EOL> try : <EOL> from corehq . util . soft_assert import soft_assert <EOL> _assert = soft_assert ( to = '<STR_LIT:@>' . join ( [ '<STR_LIT>' , '<STR_LIT>' ] ) , fail_if_debug = True ) <EOL> except ImportError : <EOL> _assert = lambda assertion , message : None <EOL> WAIT_HEARTBEAT = <NUM_LIT> <EOL> CHANGES_TIMEOUT = <NUM_LIT> <EOL> RETRY_INTERVAL = <NUM_LIT:2> <EOL> MAX_RETRIES = <NUM_LIT:4> <EOL> class PillowtopIndexingError ( Exception ) : <EOL> pass <EOL> class PillowtopNetworkError ( Exception ) : <EOL> pass <EOL> def ms_from_timedelta ( td ) : <EOL> """<STR_LIT>""" <EOL> return ( td . seconds * <NUM_LIT:1000> ) + ( td . microseconds / <NUM_LIT> ) <EOL> def lock_manager ( obj ) : <EOL> if isinstance ( obj , LockManager ) : <EOL> return obj <EOL> else : <EOL> return LockManager ( obj , None ) <EOL> class BasicPillow ( PillowBase ) : <EOL> """<STR_LIT>""" <EOL> checkpoint_frequency = CHECKPOINT_FREQUENCY <EOL> couch_filter = None <EOL> extra_args = { } <EOL> document_class = None <EOL> _couch_db = None <EOL> include_docs = True <EOL> use_locking = False <EOL> def __init__ ( self , couch_db = None , document_class = None , checkpoint = None , change_feed = None ) : <EOL> if document_class : <EOL> self . document_class = document_class <EOL> self . _couch_db = couch_db <EOL> self . _checkpoint = checkpoint <EOL> self . _change_feed = change_feed <EOL> if self . use_locking : <EOL> assert hasattr ( self . document_class , '<STR_LIT>' ) <EOL> @ property <EOL> def pillow_id ( self ) : <EOL> return self . __class__ . __name__ <EOL> def get_couch_db ( self ) : <EOL> if self . _couch_db is None : <EOL> self . _couch_db = self . get_default_couch_db ( ) <EOL> return self . _couch_db <EOL> def set_couch_db ( self , couch_db ) : <EOL> self . _couch_db = couch_db <EOL> def get_default_couch_db ( self ) : <EOL> return self . document_class . get_db ( ) if self . document_class else None <EOL> @ property <EOL> def couch_db ( self ) : <EOL> _assert ( False , '<STR_LIT>' ) <EOL> return self . get_couch_db ( ) <EOL> @ couch_db . setter <EOL> def couch_db ( self , value ) : <EOL> _assert ( False , '<STR_LIT>' ) <EOL> self . _couch_db = value <EOL> @ property <EOL> def document_store ( self ) : <EOL> return CouchDocumentStore ( self . get_couch_db ( ) ) <EOL> @ property <EOL> def checkpoint ( self ) : <EOL> if self . _checkpoint is None : <EOL> self . _checkpoint = self . _get_default_checkpoint ( ) <EOL> return self . _checkpoint <EOL> def _get_default_checkpoint ( self ) : <EOL> return PillowCheckpoint ( <EOL> construct_checkpoint_doc_id_from_name ( self . get_name ( ) ) , <EOL> ) <EOL> def get_change_feed ( self ) : <EOL> if self . _change_feed is None : <EOL> self . _change_feed = self . _get_default_change_feed ( ) <EOL> return self . _change_feed <EOL> def _get_default_change_feed ( self ) : <EOL> return CouchChangeFeed ( <EOL> couch_db = self . get_couch_db ( ) , <EOL> include_docs = self . include_docs , <EOL> couch_filter = self . couch_filter , <EOL> extra_couch_view_params = self . extra_args <EOL> ) <EOL> @ memoized <EOL> def get_name ( self ) : <EOL> return self . get_legacy_name ( ) <EOL> @ classmethod <EOL> def get_legacy_name ( cls ) : <EOL> return "<STR_LIT>" % ( cls . _get_base_name ( ) , cls . __name__ , get_machine_id ( ) ) <EOL> @ classmethod <EOL> def _get_base_name ( cls ) : <EOL> return cls . __module__ <EOL> def processor ( self , change , context ) : <EOL> """<STR_LIT>""" <EOL> self . process_change ( change ) <EOL> def fire_change_processed_event ( self , change , context ) : <EOL> if context . changes_seen % self . checkpoint_frequency == <NUM_LIT:0> and context . do_set_checkpoint : <EOL> self . set_checkpoint ( change ) <EOL> def process_change ( self , change , is_retry_attempt = False ) : <EOL> try : <EOL> with lock_manager ( self . change_trigger ( change ) ) as t : <EOL> if t is not None : <EOL> tr = self . change_transform ( t ) <EOL> if tr is not None : <EOL> self . change_transport ( tr ) <EOL> except Exception , ex : <EOL> if not is_retry_attempt : <EOL> self . _handle_pillow_error ( change , ex ) <EOL> else : <EOL> raise <EOL> def _handle_pillow_error ( self , change , exception ) : <EOL> try : <EOL> meta = self . get_couch_db ( ) . show ( '<STR_LIT>' , change [ '<STR_LIT:id>' ] ) <EOL> except ResourceNotFound : <EOL> meta = None <EOL> error = PillowError . get_or_create ( change , self , change_meta = meta ) <EOL> error . add_attempt ( exception , sys . exc_info ( ) [ <NUM_LIT:2> ] ) <EOL> error . save ( ) <EOL> pillow_logging . exception ( <EOL> "<STR_LIT>" % ( <EOL> self . get_name ( ) , <EOL> change [ '<STR_LIT:id>' ] , <EOL> exception , <EOL> error . id <EOL> ) <EOL> ) <EOL> def change_trigger ( self , changes_dict ) : <EOL> """<STR_LIT>""" <EOL> if changes_dict . get ( '<STR_LIT>' , False ) : <EOL> return None <EOL> id = changes_dict [ '<STR_LIT:id>' ] <EOL> if self . use_locking : <EOL> lock = self . document_class . get_obj_lock_by_id ( id ) <EOL> lock . acquire ( ) <EOL> return LockManager ( self . get_couch_db ( ) . open_doc ( id ) , lock ) <EOL> elif changes_dict . get ( '<STR_LIT>' , None ) is not None : <EOL> return changes_dict [ '<STR_LIT>' ] <EOL> elif hasattr ( changes_dict , '<STR_LIT>' ) and changes_dict . get_document ( ) : <EOL> return changes_dict . get_document ( ) <EOL> else : <EOL> return self . get_couch_db ( ) . open_doc ( id ) <EOL> def change_transform ( self , doc_dict ) : <EOL> """<STR_LIT>""" <EOL> return doc_dict <EOL> def change_transport ( self , doc_dict ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( <EOL> "<STR_LIT>" ) <EOL> PYTHONPILLOW_CHUNK_SIZE = <NUM_LIT> <EOL> PYTHONPILLOW_CHECKPOINT_FREQUENCY = CHECKPOINT_FREQUENCY * <NUM_LIT:10> <EOL> PYTHONPILLOW_MAX_WAIT_TIME = <NUM_LIT> <EOL> class PythonPillow ( BasicPillow ) : <EOL> """<STR_LIT>""" <EOL> process_deletions = False <EOL> def __init__ ( self , document_class = None , chunk_size = PYTHONPILLOW_CHUNK_SIZE , <EOL> checkpoint_frequency = PYTHONPILLOW_CHECKPOINT_FREQUENCY , <EOL> couch_db = None , checkpoint = None , change_feed = None , preload_docs = True ) : <EOL> """<STR_LIT>""" <EOL> super ( PythonPillow , self ) . __init__ ( <EOL> document_class = document_class , <EOL> checkpoint = checkpoint , <EOL> couch_db = couch_db , <EOL> change_feed = change_feed , <EOL> ) <EOL> self . change_queue = [ ] <EOL> self . chunk_size = chunk_size <EOL> self . use_chunking = chunk_size > <NUM_LIT:0> <EOL> self . checkpoint_frequency = checkpoint_frequency <EOL> self . include_docs = not self . use_chunking <EOL> self . last_processed_time = None <EOL> self . preload_docs = preload_docs <EOL> def get_default_couch_db ( self ) : <EOL> if self . document_class and self . use_chunking : <EOL> return CachedCouchDB ( self . document_class . get_db ( ) . uri , readonly = False ) <EOL> else : <EOL> return super ( PythonPillow , self ) . get_default_couch_db ( ) <EOL> def python_filter ( self , change ) : <EOL> """<STR_LIT>""" <EOL> return True <EOL> def process_chunk ( self ) : <EOL> def _assert_change_has_id ( change ) : <EOL> if '<STR_LIT:id>' not in change : <EOL> _assert ( False , "<STR_LIT>" . format ( <EOL> simplejson . dumps ( change ) <EOL> ) ) <EOL> return False <EOL> return True <EOL> changes_to_process = filter ( _assert_change_has_id , self . change_queue ) <EOL> if self . preload_docs : <EOL> self . get_couch_db ( ) . bulk_load ( [ change [ '<STR_LIT:id>' ] for change in changes_to_process ] , <EOL> purge_existing = True ) <EOL> for change in changes_to_process : <EOL> if self . preload_docs : <EOL> doc = self . get_couch_db ( ) . open_doc ( change [ '<STR_LIT:id>' ] , check_main = False ) <EOL> change . set_document ( doc ) <EOL> valid_change = ( not self . preload_docs or change . document ) and self . python_filter ( change ) <EOL> valid_deletion = self . process_deletions and change . get ( '<STR_LIT>' , None ) <EOL> if valid_change or valid_deletion : <EOL> try : <EOL> self . process_change ( change ) <EOL> except Exception : <EOL> logging . exception ( '<STR_LIT>' % <EOL> ( change . get ( '<STR_LIT>' , None ) , change [ '<STR_LIT:id>' ] ) ) <EOL> self . change_queue = [ ] <EOL> self . last_processed_time = datetime . utcnow ( ) <EOL> @ property <EOL> def queue_full ( self ) : <EOL> return len ( self . change_queue ) > self . chunk_size <EOL> @ property <EOL> def wait_expired ( self ) : <EOL> if not self . last_processed_time : <EOL> return False <EOL> wait_time = datetime . utcnow ( ) - self . last_processed_time <EOL> return wait_time > timedelta ( seconds = PYTHONPILLOW_MAX_WAIT_TIME ) <EOL> def processor ( self , change , context ) : <EOL> if self . use_chunking : <EOL> self . change_queue . append ( change ) <EOL> if self . queue_full or self . wait_expired : <EOL> self . process_chunk ( ) <EOL> elif self . python_filter ( change ) or ( change . get ( '<STR_LIT>' , None ) and self . process_deletions ) : <EOL> self . process_change ( change ) <EOL> def fire_change_processed_event ( self , change , context ) : <EOL> if context . changes_seen % self . checkpoint_frequency == <NUM_LIT:0> and context . do_set_checkpoint : <EOL> if self . use_chunking : <EOL> self . process_chunk ( ) <EOL> self . set_checkpoint ( change ) <EOL> def run ( self ) : <EOL> self . change_queue = [ ] <EOL> self . last_processed_time = datetime . utcnow ( ) <EOL> super ( PythonPillow , self ) . run ( ) <EOL> def send_to_elasticsearch ( index , doc_type , doc_id , es_getter , name , data = None , retries = MAX_RETRIES , <EOL> except_on_failure = False , update = False , delete = False ) : <EOL> """<STR_LIT>""" <EOL> data = data if data is not None else { } <EOL> current_tries = <NUM_LIT:0> <EOL> while current_tries < retries : <EOL> try : <EOL> if delete : <EOL> es_getter ( ) . delete ( index , doc_type , doc_id ) <EOL> elif update : <EOL> params = { '<STR_LIT>' : <NUM_LIT:2> } <EOL> es_getter ( ) . update ( index , doc_type , doc_id , body = { "<STR_LIT>" : data } , params = params ) <EOL> else : <EOL> es_getter ( ) . create ( index , doc_type , body = data , id = doc_id ) <EOL> break <EOL> except ConnectionError , ex : <EOL> current_tries += <NUM_LIT:1> <EOL> pillow_logging . error ( "<STR_LIT>" % ( <EOL> name , ex , current_tries , retries ) ) <EOL> if current_tries == retries : <EOL> message = "<STR_LIT>" % ( name , index , doc_type , doc_id ) <EOL> if except_on_failure : <EOL> raise PillowtopIndexingError ( message ) <EOL> else : <EOL> pillow_logging . error ( message ) <EOL> time . sleep ( math . pow ( RETRY_INTERVAL , current_tries ) ) <EOL> except RequestError as ex : <EOL> error_message = "<STR_LIT>" % ( <EOL> name , <EOL> ex . error or "<STR_LIT>" , <EOL> index , doc_type , doc_id , <EOL> data . keys ( ) ) <EOL> if except_on_failure : <EOL> raise PillowtopIndexingError ( error_message ) <EOL> else : <EOL> pillow_logging . error ( error_message ) <EOL> break <EOL> except ConflictError : <EOL> break <EOL> except NotFoundError : <EOL> break <EOL> class AliasedElasticPillow ( BasicPillow ) : <EOL> """<STR_LIT>""" <EOL> es_index = "<STR_LIT>" <EOL> es_type = "<STR_LIT>" <EOL> es_alias = '<STR_LIT>' <EOL> es_meta = { } <EOL> es_timeout = <NUM_LIT:3> <EOL> default_mapping = None <EOL> bulk = False <EOL> allow_updates = True <EOL> def __init__ ( self , online = True , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' not in kwargs : <EOL> kwargs [ '<STR_LIT>' ] = get_default_django_checkpoint_for_legacy_pillow_class ( self . __class__ ) <EOL> super ( AliasedElasticPillow , self ) . __init__ ( ** kwargs ) <EOL> if online : <EOL> completely_initialize_pillow_index ( self ) <EOL> def get_doc_path ( self , doc_id ) : <EOL> return "<STR_LIT>" % ( self . es_index , self . es_type , doc_id ) <EOL> @ memoized <EOL> def get_es_new ( self ) : <EOL> from corehq . elastic import get_es_new <EOL> return get_es_new ( timeout = self . es_timeout ) <EOL> def change_trigger ( self , changes_dict ) : <EOL> id = changes_dict [ '<STR_LIT:id>' ] <EOL> if changes_dict . get ( '<STR_LIT>' , False ) : <EOL> try : <EOL> if doc_exists ( self , id ) : <EOL> self . get_es_new ( ) . delete ( self . es_index , self . es_type , id ) <EOL> except Exception , ex : <EOL> pillow_logging . error ( <EOL> "<STR_LIT>" % ( <EOL> self . get_doc_path ( changes_dict [ '<STR_LIT:id>' ] ) , <EOL> ex , <EOL> ) <EOL> ) <EOL> return None <EOL> return super ( AliasedElasticPillow , self ) . change_trigger ( changes_dict ) <EOL> def send_robust ( self , doc_dict , retries = MAX_RETRIES , except_on_failure = False , update = False ) : <EOL> send_to_elasticsearch ( <EOL> index = self . es_index , <EOL> doc_type = self . es_type , <EOL> doc_id = doc_dict [ '<STR_LIT>' ] , <EOL> es_getter = self . get_es_new , <EOL> name = self . get_name ( ) , <EOL> data = doc_dict , <EOL> retries = retries , <EOL> except_on_failure = except_on_failure , <EOL> update = update , <EOL> ) <EOL> def change_transport ( self , doc_dict ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> if not self . bulk : <EOL> doc_exists_val = doc_exists ( self , doc_dict ) <EOL> if self . allow_updates : <EOL> can_put = True <EOL> else : <EOL> can_put = not doc_exists_val <EOL> if can_put and not self . bulk : <EOL> self . send_robust ( doc_dict , update = doc_exists_val ) <EOL> except Exception , ex : <EOL> tb = traceback . format_exc ( ) <EOL> pillow_logging . error ( <EOL> "<STR_LIT>" % <EOL> { <EOL> "<STR_LIT>" : self . get_name ( ) , <EOL> "<STR_LIT>" : doc_dict [ '<STR_LIT>' ] , <EOL> "<STR_LIT:error>" : ex , <EOL> "<STR_LIT>" : tb <EOL> } <EOL> ) <EOL> raise <EOL> def process_bulk ( self , changes ) : <EOL> if not changes : <EOL> return <EOL> self . allow_updates = False <EOL> self . bulk = True <EOL> bstart = datetime . utcnow ( ) <EOL> bulk_changes = self . bulk_builder ( changes ) <EOL> max_payload_size = pow ( <NUM_LIT:10> , <NUM_LIT:8> ) <EOL> payloads = prepare_bulk_payloads ( bulk_changes , max_payload_size ) <EOL> if len ( payloads ) > <NUM_LIT:1> : <EOL> pillow_logging . info ( "<STR_LIT>" % ( self . get_name ( ) , len ( payloads ) ) ) <EOL> pillow_logging . info ( <EOL> "<STR_LIT>" % ( self . get_name ( ) , str ( ms_from_timedelta ( datetime . utcnow ( ) - bstart ) / <NUM_LIT> ) ) ) <EOL> send_start = datetime . utcnow ( ) <EOL> for payload in payloads : <EOL> self . send_bulk ( payload ) <EOL> pillow_logging . info ( <EOL> "<STR_LIT>" % ( self . get_name ( ) , str ( ms_from_timedelta ( datetime . utcnow ( ) - send_start ) / <NUM_LIT> ) ) ) <EOL> def send_bulk ( self , payload ) : <EOL> self . get_es_new ( ) . bulk ( payload ) <EOL> @ staticmethod <EOL> def calc_mapping_hash ( mapping ) : <EOL> return hashlib . md5 ( simplejson . dumps ( mapping , sort_keys = True ) ) . hexdigest ( ) <EOL> @ classmethod <EOL> def get_unique_id ( cls ) : <EOL> """<STR_LIT>""" <EOL> return cls . calc_meta ( ) <EOL> @ classmethod <EOL> def calc_meta ( cls ) : <EOL> raise NotImplementedError ( "<STR_LIT>" ) <EOL> def bulk_builder ( self , changes ) : <EOL> """<STR_LIT>""" <EOL> for change in changes : <EOL> try : <EOL> with lock_manager ( self . change_trigger ( change ) ) as t : <EOL> if t is not None : <EOL> tr = self . change_transform ( t ) <EOL> if tr is not None : <EOL> self . change_transport ( tr ) <EOL> yield { <EOL> "<STR_LIT:index>" : { <EOL> "<STR_LIT>" : self . es_index , <EOL> "<STR_LIT>" : self . es_type , <EOL> "<STR_LIT>" : tr [ '<STR_LIT>' ] <EOL> } <EOL> } <EOL> yield tr <EOL> except Exception , ex : <EOL> pillow_logging . error ( <EOL> "<STR_LIT>" % ( change [ '<STR_LIT:id>' ] , ex ) <EOL> ) <EOL> @ memoized <EOL> def get_name ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . get_legacy_name ( ) <EOL> @ classmethod <EOL> def get_legacy_name ( cls ) : <EOL> return "<STR_LIT>" % ( <EOL> cls . __module__ , cls . __name__ , cls . get_unique_id ( ) , get_machine_id ( ) ) <EOL> def retry_on_connection_failure ( fn ) : <EOL> @ wraps ( fn ) <EOL> def _inner ( * args , ** kwargs ) : <EOL> retry = kwargs . pop ( '<STR_LIT>' , True ) <EOL> try : <EOL> return fn ( * args , ** kwargs ) <EOL> except db . utils . DatabaseError : <EOL> db . transaction . rollback ( ) <EOL> raise <EOL> except ( Psycopg2InterfaceError , DjangoInterfaceError ) : <EOL> db . connection . close ( ) <EOL> if retry : <EOL> _inner ( retry = False , * args , ** kwargs ) <EOL> else : <EOL> raise <EOL> return _inner <EOL> class SQLPillowMixIn ( object ) : <EOL> def change_trigger ( self , changes_dict ) : <EOL> if changes_dict . get ( '<STR_LIT>' , False ) : <EOL> self . change_transport ( { '<STR_LIT>' : changes_dict [ '<STR_LIT:id>' ] } , delete = True ) <EOL> return None <EOL> return super ( SQLPillowMixIn , self ) . change_trigger ( changes_dict ) <EOL> @ retry_on_connection_failure <EOL> @ db . transaction . atomic <EOL> def change_transport ( self , doc_dict , delete = False ) : <EOL> self . process_sql ( doc_dict , delete ) <EOL> def process_sql ( self , doc_dict , delete = False ) : <EOL> pass <EOL> class SQLPillow ( SQLPillowMixIn , BasicPillow ) : <EOL> def __init__ ( self ) : <EOL> checkpoint = get_default_django_checkpoint_for_legacy_pillow_class ( self . __class__ ) <EOL> super ( SQLPillow , self ) . __init__ ( checkpoint = checkpoint ) </s>
