<s> from __future__ import print_function , unicode_literals <EOL> import operator <EOL> import sys <EOL> import time <EOL> from copy import copy <EOL> from subprocess import PIPE , Popen <EOL> from django . db import connections , models <EOL> from django . db . models . sql . where import ExtraWhere <EOL> from django . db . transaction import atomic <EOL> from django . test . utils import CaptureQueriesContext <EOL> from django . utils import six <EOL> from django . utils . functional import cached_property <EOL> from django . utils . translation import ugettext as _ <EOL> from django_mysql . models . handler import Handler <EOL> from django_mysql . rewrite_query import REWRITE_MARKER <EOL> from django_mysql . status import GlobalStatus <EOL> from django_mysql . utils import ( <EOL> StopWatch , WeightedAverageRate , format_duration , have_program , <EOL> noop_context , settings_to_cmd_args <EOL> ) <EOL> class QuerySetMixin ( object ) : <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> super ( QuerySetMixin , self ) . __init__ ( * args , ** kwargs ) <EOL> self . _count_tries_approx = False <EOL> def _clone ( self , * args , ** kwargs ) : <EOL> clone = super ( QuerySetMixin , self ) . _clone ( * args , ** kwargs ) <EOL> clone . _count_tries_approx = copy ( <EOL> getattr ( self , '<STR_LIT>' , False ) <EOL> ) <EOL> if hasattr ( self , '<STR_LIT>' ) : <EOL> clone . _found_rows = None <EOL> return clone <EOL> def count ( self ) : <EOL> if self . _count_tries_approx : <EOL> return self . approx_count ( ** self . _count_tries_approx ) <EOL> return super ( QuerySetMixin , self ) . count ( ) <EOL> def count_tries_approx ( self , activate = True , fall_back = True , <EOL> return_approx_int = True , min_size = <NUM_LIT:1000> ) : <EOL> clone = self . _clone ( ) <EOL> if activate : <EOL> clone . _count_tries_approx = { <EOL> '<STR_LIT>' : fall_back , <EOL> '<STR_LIT>' : return_approx_int , <EOL> '<STR_LIT>' : min_size , <EOL> } <EOL> else : <EOL> clone . _count_tries_approx = False <EOL> return clone <EOL> def approx_count ( self , fall_back = True , return_approx_int = True , <EOL> min_size = <NUM_LIT:1000> ) : <EOL> try : <EOL> num = approx_count ( self ) <EOL> except ValueError : <EOL> if not fall_back : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> return super ( QuerySetMixin , self ) . count ( ) <EOL> if min_size and num < min_size : <EOL> return super ( QuerySetMixin , self ) . count ( ) <EOL> if return_approx_int : <EOL> return ApproximateInt ( num ) <EOL> else : <EOL> return num <EOL> def label ( self , string ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in string : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> return self . extra ( where = [ "<STR_LIT>" . format ( string ) ] ) <EOL> def straight_join ( self ) : <EOL> return self . extra ( where = [ "<STR_LIT>" ] ) <EOL> def sql_small_result ( self ) : <EOL> return self . extra ( where = [ "<STR_LIT>" ] ) <EOL> def sql_big_result ( self ) : <EOL> return self . extra ( where = [ "<STR_LIT>" ] ) <EOL> def sql_buffer_result ( self ) : <EOL> return self . extra ( where = [ "<STR_LIT>" ] ) <EOL> def sql_cache ( self ) : <EOL> return self . extra ( where = [ "<STR_LIT>" ] ) <EOL> def sql_no_cache ( self ) : <EOL> return self . extra ( where = [ "<STR_LIT>" ] ) <EOL> def sql_calc_found_rows ( self ) : <EOL> qs = self . extra ( where = [ "<STR_LIT>" ] ) <EOL> qs . _found_rows = None <EOL> return qs <EOL> @ property <EOL> def found_rows ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> raise ValueError ( <EOL> "<STR_LIT>" <EOL> ) <EOL> if self . _found_rows is None : <EOL> raise RuntimeError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> ) <EOL> return self . _found_rows <EOL> def iterator ( self ) : <EOL> for row in super ( QuerySetMixin , self ) . iterator ( ) : <EOL> yield row <EOL> if getattr ( self , '<STR_LIT>' , <NUM_LIT:0> ) is None : <EOL> with connections [ self . db ] . cursor ( ) as cursor : <EOL> cursor . execute ( "<STR_LIT>" ) <EOL> self . _found_rows = cursor . fetchone ( ) [ <NUM_LIT:0> ] <EOL> def use_index ( self , * index_names , ** kwargs ) : <EOL> kwargs [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> return self . _index_hint ( * index_names , ** kwargs ) <EOL> def force_index ( self , * index_names , ** kwargs ) : <EOL> kwargs [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> return self . _index_hint ( * index_names , ** kwargs ) <EOL> def ignore_index ( self , * index_names , ** kwargs ) : <EOL> kwargs [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> return self . _index_hint ( * index_names , ** kwargs ) <EOL> def _index_hint ( self , * index_names , ** kwargs ) : <EOL> hint = kwargs . pop ( '<STR_LIT>' ) <EOL> table_name = kwargs . pop ( '<STR_LIT>' , None ) <EOL> for_ = kwargs . pop ( '<STR_LIT>' , None ) <EOL> if kwargs : <EOL> raise ValueError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> . format ( hint . lower ( ) ) <EOL> ) <EOL> if hint != '<STR_LIT>' and not len ( index_names ) : <EOL> raise ValueError ( <EOL> "<STR_LIT>" <EOL> . format ( hint . lower ( ) ) <EOL> ) <EOL> if table_name is None : <EOL> table_name = self . model . _meta . db_table <EOL> if for_ in ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) : <EOL> for_bit = '<STR_LIT>' . format ( for_ ) <EOL> elif for_ is None : <EOL> for_bit = '<STR_LIT>' <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> if len ( index_names ) == <NUM_LIT:0> : <EOL> indexes = "<STR_LIT>" <EOL> else : <EOL> indexes = "<STR_LIT>" + "<STR_LIT>" . join ( index_names ) + "<STR_LIT>" <EOL> hint = ( <EOL> "<STR_LIT>" <EOL> . format ( table_name = table_name , hint = hint , for_bit = for_bit , <EOL> indexes = indexes ) <EOL> ) <EOL> return self . extra ( where = [ hint ] ) <EOL> def iter_smart ( self , ** kwargs ) : <EOL> assert '<STR_LIT>' not in kwargs , "<STR_LIT>" <EOL> return SmartIterator ( queryset = self , ** kwargs ) <EOL> def iter_smart_chunks ( self , ** kwargs ) : <EOL> assert '<STR_LIT>' not in kwargs , "<STR_LIT>" <EOL> return SmartChunkedIterator ( queryset = self , ** kwargs ) <EOL> def iter_smart_pk_ranges ( self , ** kwargs ) : <EOL> assert '<STR_LIT>' not in kwargs , "<STR_LIT>" <EOL> return SmartPKRangeIterator ( queryset = self , ** kwargs ) <EOL> def pt_visual_explain ( self , display = True ) : <EOL> return pt_visual_explain ( self , display ) <EOL> def handler ( self ) : <EOL> return Handler ( self ) <EOL> class QuerySet ( QuerySetMixin , models . QuerySet ) : <EOL> pass <EOL> def add_QuerySetMixin ( queryset ) : <EOL> queryset2 = queryset . _clone ( ) <EOL> queryset2 . __class__ = _make_mixin_class ( queryset . __class__ ) <EOL> return queryset2 <EOL> _mixin_classes = { } <EOL> def _make_mixin_class ( klass ) : <EOL> global _mixin_classes <EOL> if klass not in _mixin_classes : <EOL> class MixedInQuerySet ( QuerySetMixin , klass ) : <EOL> pass <EOL> if six . PY2 : <EOL> MixedInQuerySet . __name__ = b'<STR_LIT>' + klass . __name__ <EOL> else : <EOL> MixedInQuerySet . __name__ = '<STR_LIT>' + klass . __name__ <EOL> _mixin_classes [ klass ] = MixedInQuerySet <EOL> return _mixin_classes [ klass ] <EOL> @ six . python_2_unicode_compatible <EOL> class ApproximateInt ( int ) : <EOL> """<STR_LIT>""" <EOL> def __str__ ( self ) : <EOL> return _ ( "<STR_LIT>" ) % { <EOL> '<STR_LIT>' : super ( ApproximateInt , self ) . __str__ ( ) <EOL> } <EOL> class SmartChunkedIterator ( object ) : <EOL> def __init__ ( self , queryset , atomically = True , status_thresholds = None , <EOL> pk_range = None , chunk_time = <NUM_LIT:0.5> , chunk_size = <NUM_LIT:2> , chunk_min = <NUM_LIT:1> , <EOL> chunk_max = <NUM_LIT> , report_progress = False , total = None ) : <EOL> self . queryset = self . sanitize_queryset ( queryset ) <EOL> if atomically : <EOL> self . maybe_atomic = atomic <EOL> else : <EOL> self . maybe_atomic = noop_context <EOL> self . status_thresholds = status_thresholds <EOL> self . pk_range = pk_range <EOL> self . rate = WeightedAverageRate ( chunk_time ) <EOL> assert <NUM_LIT:0> < chunk_min <= chunk_max , "<STR_LIT>" <EOL> self . chunk_min = chunk_min <EOL> self . chunk_max = chunk_max <EOL> self . chunk_size = self . constrain_size ( chunk_size ) <EOL> self . report_progress = report_progress <EOL> self . total = total <EOL> def __iter__ ( self ) : <EOL> first_pk , last_pk = self . get_first_and_last ( ) <EOL> if first_pk <= last_pk : <EOL> comp = operator . le <EOL> direction = <NUM_LIT:1> <EOL> else : <EOL> comp = operator . ge <EOL> direction = - <NUM_LIT:1> <EOL> current_pk = first_pk <EOL> db_alias = self . queryset . db <EOL> status = GlobalStatus ( db_alias ) <EOL> self . init_progress ( direction ) <EOL> while comp ( current_pk , last_pk ) : <EOL> status . wait_until_load_low ( self . status_thresholds ) <EOL> start_pk = current_pk <EOL> current_pk = current_pk + self . chunk_size * direction <EOL> if direction == <NUM_LIT:1> : <EOL> end_pk = min ( current_pk , last_pk + <NUM_LIT:1> ) <EOL> else : <EOL> end_pk = max ( current_pk , last_pk - <NUM_LIT:1> ) <EOL> with StopWatch ( ) as timer , self . maybe_atomic ( using = db_alias ) : <EOL> if direction == <NUM_LIT:1> : <EOL> chunk = self . queryset . filter ( pk__gte = start_pk , <EOL> pk__lt = end_pk ) <EOL> else : <EOL> chunk = self . queryset . filter ( pk__lte = start_pk , <EOL> pk__gt = end_pk ) <EOL> chunk . _smart_iterator_pks = ( start_pk , end_pk ) <EOL> yield chunk <EOL> self . update_progress ( direction , chunk , end_pk ) <EOL> self . adjust_chunk_size ( chunk , timer . total_time ) <EOL> self . end_progress ( ) <EOL> def sanitize_queryset ( self , queryset ) : <EOL> if queryset . ordered : <EOL> raise ValueError ( <EOL> "<STR_LIT>" % <EOL> self . __class__ . __name__ <EOL> ) <EOL> if queryset . query . low_mark or queryset . query . high_mark : <EOL> raise ValueError ( <EOL> "<STR_LIT>" % <EOL> self . __class__ . __name__ <EOL> ) <EOL> pk = queryset . model . _meta . pk <EOL> allowed_field = ( <EOL> isinstance ( pk , self . ALLOWED_PK_FIELD_CLASSES ) or <EOL> ( isinstance ( pk , models . ForeignKey ) and <EOL> isinstance ( pk . foreign_related_fields [ <NUM_LIT:0> ] , <EOL> self . ALLOWED_PK_FIELD_CLASSES ) ) <EOL> ) <EOL> if not allowed_field : <EOL> raise ValueError ( <EOL> "<STR_LIT>" % <EOL> self . __class__ . __name__ <EOL> ) <EOL> return queryset . order_by ( '<STR_LIT>' ) <EOL> ALLOWED_PK_FIELD_CLASSES = ( <EOL> models . IntegerField , <EOL> models . AutoField , <EOL> ) <EOL> def get_first_and_last ( self ) : <EOL> if isinstance ( self . pk_range , tuple ) and len ( self . pk_range ) == <NUM_LIT:2> : <EOL> should_be_reversed = ( <EOL> self . pk_range [ <NUM_LIT:1> ] < self . pk_range [ <NUM_LIT:0> ] and <EOL> self . queryset . query . standard_ordering <EOL> ) <EOL> if should_be_reversed : <EOL> self . queryset = self . queryset . reverse ( ) <EOL> return self . pk_range <EOL> elif self . pk_range == '<STR_LIT:all>' : <EOL> base_qs = self . queryset . model . objects . using ( self . queryset . db ) . all ( ) <EOL> elif self . pk_range is None : <EOL> base_qs = self . queryset <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" <EOL> . format ( self . pk_range ) ) <EOL> if not base_qs . query . standard_ordering : <EOL> base_qs = base_qs . reverse ( ) <EOL> min_qs = base_qs . order_by ( '<STR_LIT>' ) . values_list ( '<STR_LIT>' , flat = True ) <EOL> max_qs = base_qs . order_by ( '<STR_LIT>' ) . values_list ( '<STR_LIT>' , flat = True ) <EOL> try : <EOL> min_pk = min_qs [ <NUM_LIT:0> ] <EOL> except IndexError : <EOL> max_pk = min_pk = <NUM_LIT:0> <EOL> else : <EOL> try : <EOL> max_pk = max_qs [ <NUM_LIT:0> ] <EOL> except IndexError : <EOL> max_pk = min_pk <EOL> if self . queryset . query . standard_ordering : <EOL> return ( min_pk , max_pk ) <EOL> else : <EOL> return ( max_pk , min_pk ) <EOL> def constrain_size ( self , chunk_size ) : <EOL> return max ( min ( chunk_size , self . chunk_max ) , self . chunk_min ) <EOL> def adjust_chunk_size ( self , chunk , chunk_time ) : <EOL> if chunk . _result_cache is None : <EOL> num_processed = self . chunk_size <EOL> else : <EOL> num_processed = len ( chunk ) <EOL> if num_processed > <NUM_LIT:0> : <EOL> new_chunk_size = self . rate . update ( num_processed , chunk_time ) <EOL> else : <EOL> new_chunk_size = self . chunk_size <EOL> if new_chunk_size < <NUM_LIT:1> : <EOL> new_chunk_size = <NUM_LIT:1> <EOL> self . chunk_size = self . constrain_size ( new_chunk_size ) <EOL> def init_progress ( self , direction ) : <EOL> if not self . report_progress : <EOL> return <EOL> self . start_time = time . time ( ) <EOL> self . old_report = "<STR_LIT>" <EOL> self . objects_done = <NUM_LIT:0> <EOL> self . chunks_done = <NUM_LIT:0> <EOL> if self . total is None : <EOL> try : <EOL> self . total = approx_count ( self . queryset ) <EOL> if self . total < <NUM_LIT:1000> : <EOL> self . total = self . queryset . count ( ) <EOL> except ValueError : <EOL> self . total = self . queryset . count ( ) <EOL> self . update_progress ( direction ) <EOL> def update_progress ( self , direction , chunk = None , end_pk = None ) : <EOL> if not self . report_progress : <EOL> return <EOL> if chunk is not None : <EOL> self . chunks_done += <NUM_LIT:1> <EOL> if self . objects_done != "<STR_LIT>" : <EOL> if chunk . _result_cache is None : <EOL> self . objects_done = "<STR_LIT>" <EOL> else : <EOL> self . objects_done += len ( chunk ) <EOL> try : <EOL> percent_complete = <NUM_LIT:100> * ( float ( self . objects_done ) / self . total ) <EOL> except ( ZeroDivisionError , ValueError ) : <EOL> percent_complete = <NUM_LIT:0.0> <EOL> report = "<STR_LIT>" . format ( <EOL> self . model_name , <EOL> self . __class__ . __name__ , <EOL> self . objects_done , <EOL> self . total , <EOL> percent_complete , <EOL> self . chunks_done , <EOL> ) <EOL> if end_pk is not None : <EOL> report += "<STR_LIT>" . format ( <EOL> dir = "<STR_LIT>" if direction == <NUM_LIT:1> else "<STR_LIT>" , <EOL> end_pk = end_pk , <EOL> ) <EOL> if self . objects_done != '<STR_LIT>' and self . rate . avg_rate : <EOL> n_remaining = self . total - self . objects_done <EOL> s_remaining = max ( <NUM_LIT:0> , int ( n_remaining // self . rate . avg_rate ) ) <EOL> report += '<STR_LIT>' . format ( <EOL> format_duration ( s_remaining ) <EOL> ) <EOL> spacing = "<STR_LIT:U+0020>" * max ( <NUM_LIT:0> , len ( self . old_report ) - len ( report ) ) <EOL> if self . old_report : <EOL> sys . stdout . write ( "<STR_LIT:\r>" ) <EOL> sys . stdout . write ( report ) <EOL> sys . stdout . write ( spacing ) <EOL> sys . stdout . flush ( ) <EOL> self . old_report = report <EOL> def end_progress ( self ) : <EOL> if not self . report_progress : <EOL> return <EOL> total_time = int ( time . time ( ) - self . start_time ) <EOL> sys . stdout . write ( <EOL> "<STR_LIT>" . format ( <EOL> n = self . objects_done , <EOL> s = '<STR_LIT:s>' if self . objects_done != <NUM_LIT:1> else '<STR_LIT>' , <EOL> duration = format_duration ( total_time ) <EOL> ) <EOL> ) <EOL> @ cached_property <EOL> def model_name ( self ) : <EOL> return self . queryset . model . __name__ <EOL> class SmartIterator ( SmartChunkedIterator ) : <EOL> """<STR_LIT>""" <EOL> def __iter__ ( self ) : <EOL> for chunk in super ( SmartIterator , self ) . __iter__ ( ) : <EOL> for obj in chunk : <EOL> yield obj <EOL> class SmartPKRangeIterator ( SmartChunkedIterator ) : <EOL> def __iter__ ( self ) : <EOL> for chunk in super ( SmartPKRangeIterator , self ) . __iter__ ( ) : <EOL> start_pk , end_pk = chunk . _smart_iterator_pks <EOL> yield start_pk , end_pk <EOL> def approx_count ( queryset ) : <EOL> if not can_approx_count ( queryset ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> connection = connections [ queryset . db ] <EOL> with connection . cursor ( ) as cursor : <EOL> table_name = queryset . model . _meta . db_table <EOL> cursor . execute ( <EOL> """<STR_LIT>""" , <EOL> ( table_name , ) <EOL> ) <EOL> approx_count = cursor . fetchone ( ) [ <NUM_LIT:0> ] <EOL> return approx_count <EOL> def can_approx_count ( queryset ) : <EOL> query = queryset . query <EOL> if query . select or query . group_by or query . distinct : <EOL> return False <EOL> elif query . high_mark is not None or query . low_mark != <NUM_LIT:0> : <EOL> return False <EOL> for child in query . where . children : <EOL> if not isinstance ( child , ExtraWhere ) : <EOL> return False <EOL> elif not all ( ( REWRITE_MARKER in sql ) for sql in child . sqls ) : <EOL> return False <EOL> if hasattr ( query , '<STR_LIT>' ) and query . having : <EOL> return False <EOL> return True <EOL> def pt_visual_explain ( queryset , display = True ) : <EOL> if not have_program ( '<STR_LIT>' ) : <EOL> raise OSError ( "<STR_LIT>" ) <EOL> connection = connections [ queryset . db ] <EOL> capturer = CaptureQueriesContext ( connection ) <EOL> with capturer , connection . cursor ( ) as cursor : <EOL> sql , params = queryset . query . sql_with_params ( ) <EOL> cursor . execute ( '<STR_LIT>' + sql , params ) <EOL> queries = [ q [ '<STR_LIT>' ] for q in capturer . captured_queries ] <EOL> explain_query = queries [ - <NUM_LIT:1> ] <EOL> mysql_command = ( <EOL> settings_to_cmd_args ( connection . settings_dict ) + <EOL> [ '<STR_LIT>' , explain_query ] <EOL> ) <EOL> mysql = Popen ( mysql_command , stdout = PIPE ) <EOL> visual_explain = Popen ( <EOL> [ '<STR_LIT>' , '<STR_LIT:->' ] , <EOL> stdin = mysql . stdout , <EOL> stdout = PIPE <EOL> ) <EOL> mysql . stdout . close ( ) <EOL> explanation = visual_explain . communicate ( ) [ <NUM_LIT:0> ] . decode ( encoding = "<STR_LIT:utf-8>" ) <EOL> if display : <EOL> print ( explanation ) <EOL> else : <EOL> return explanation </s>
