Ticket #14298: close-cursor-v2.diff

File close-cursor-v2.diff, 10.4 KB (added by xoror, 5 years ago)
  • django/db/models/sql/compiler.py

     
    1313        self.connection = connection
    1414        self.using = using
    1515        self.quote_cache = {}
     16        self.iter_cursor_map = {}
    1617
    1718    def pre_sql_setup(self):
    1819        """
     
    677676        resolve_columns = hasattr(self, 'resolve_columns')
    678677        fields = None
    679678        has_aggregate_select = bool(self.query.aggregate_select)
    680         for rows in self.execute_sql(MULTI):
    681             for row in rows:
    682                 if resolve_columns:
    683                     if fields is None:
    684                         # We only set this up here because
    685                         # related_select_fields isn't populated until
    686                         # execute_sql() has been called.
    687                         if self.query.select_fields:
    688                             fields = self.query.select_fields + self.query.related_select_fields
    689                         else:
    690                             fields = self.query.model._meta.fields
    691                         # If the field was deferred, exclude it from being passed
    692                         # into `resolve_columns` because it wasn't selected.
    693                         only_load = self.deferred_to_columns()
    694                         if only_load:
    695                             db_table = self.query.model._meta.db_table
    696                             fields = [f for f in fields if db_table in only_load and
    697                                       f.column in only_load[db_table]]
    698                     row = self.resolve_columns(row, fields)
     679        cursor_iter = self.execute_sql(MULTI)
     680        try:
     681            for rows in cursor_iter:
     682                for row in rows:
     683                    if resolve_columns:
     684                        if fields is None:
     685                            # We only set this up here because
     686                            # related_select_fields isn't populated until
     687                            # execute_sql() has been called.
     688                            if self.query.select_fields:
     689                                fields = self.query.select_fields + self.query.related_select_fields
     690                            else:
     691                                fields = self.query.model._meta.fields
     692                            # If the field was deferred, exclude it from being passed
     693                            # into `resolve_columns` because it wasn't selected.
     694                            only_load = self.deferred_to_columns()
     695                            if only_load:
     696                                db_table = self.query.model._meta.db_table
     697                                fields = [f for f in fields if db_table in only_load and
     698                                          f.column in only_load[db_table]]
     699                        row = self.resolve_columns(row, fields)
     700   
     701                    if has_aggregate_select:
     702                        aggregate_start = len(self.query.extra_select.keys()) + len(self.query.select)
     703                        aggregate_end = aggregate_start + len(self.query.aggregate_select)
     704                        row = tuple(row[:aggregate_start]) + tuple([
     705                            self.query.resolve_aggregate(value, aggregate, self.connection)
     706                            for (alias, aggregate), value
     707                            in zip(self.query.aggregate_select.items(), row[aggregate_start:aggregate_end])
     708                        ]) + tuple(row[aggregate_end:])
     709   
     710                    yield row
     711        finally:
     712            # iteration over cursor is completed.
     713            # it's now save to close the cursor and remove it from the lookup dict
     714            if cursor_iter in self.iter_cursor_map:
     715                self.close_cursor_by_iter(cursor_iter)
     716               
     717    def close_cursor_by_iter(self, iter_to_close):
     718        """
     719        Added method to close iterator from a subclass
     720        """
     721        if iter_to_close in self.iter_cursor_map:
     722            self.iter_cursor_map[iter_to_close].close()
     723            del self.iter_cursor_map[iter_to_close]
    699724
    700                 if has_aggregate_select:
    701                     aggregate_start = len(self.query.extra_select.keys()) + len(self.query.select)
    702                     aggregate_end = aggregate_start + len(self.query.aggregate_select)
    703                     row = tuple(row[:aggregate_start]) + tuple([
    704                         self.query.resolve_aggregate(value, aggregate, self.connection)
    705                         for (alias, aggregate), value
    706                         in zip(self.query.aggregate_select.items(), row[aggregate_start:aggregate_end])
    707                     ]) + tuple(row[aggregate_end:])
    708 
    709                 yield row
    710 
    711725    def execute_sql(self, result_type=MULTI):
    712726        """
    713727        Run the query against the database and returns the result(s). The
     
    734748        cursor = self.connection.cursor()
    735749        cursor.execute(sql, params)
    736750
     751        # close cursors for single results.
    737752        if not result_type:
    738753            return cursor
    739754        if result_type == SINGLE:
    740755            if self.query.ordering_aliases:
    741                 return cursor.fetchone()[:-len(self.query.ordering_aliases)]
    742             return cursor.fetchone()
     756                temp_result = cursor.fetchone()[:-len(self.query.ordering_aliases)]
     757                cursor.close()
     758                return temp_result
     759            temp_result = cursor.fetchone()
     760            cursor.close()
     761            return temp_result
    743762
    744763        # The MULTI case.
    745764        if self.query.ordering_aliases:
     
    748767        else:
    749768            result = iter((lambda: cursor.fetchmany(GET_ITERATOR_CHUNK_SIZE)),
    750769                    self.connection.features.empty_fetchmany_value)
     770        # store iterator only if we use chunks
     771        if self.connection.features.can_use_chunked_reads:
     772            self.iter_cursor_map[result] = cursor
    751773        if not self.connection.features.can_use_chunked_reads:
    752774            # If we are using non-chunked reads, we return the same data
    753775            # structure as normally, but ensure it is all read into memory
    754776            # before going any further.
    755             return list(result)
     777            temp_ressult = list(result)
     778            cursor.close()
     779            return temp_result
    756780        return result
    757781
    758 
    759782class SQLInsertCompiler(SQLCompiler):
    760783    def placeholder(self, field, val):
    761784        if field is None:
     
    790813        self.return_id = return_id
    791814        cursor = super(SQLInsertCompiler, self).execute_sql(None)
    792815        if not (return_id and cursor):
     816            if cursor:
     817                cursor.close()
    793818            return
    794819        if self.connection.features.can_return_id_from_insert:
    795             return self.connection.ops.fetch_returned_insert_id(cursor)
    796         return self.connection.ops.last_insert_id(cursor,
     820            lastid = self.connection.ops.fetch_returned_insert_id(cursor)
     821            cursor.close()
     822            return lastid
     823        result = self.connection.ops.last_insert_id(cursor,
    797824                self.query.model._meta.db_table, self.query.model._meta.pk.column)
     825        cursor.close()
     826        return result
    798827
    799828
    800829class SQLDeleteCompiler(SQLCompiler):
     
    869898        cursor = super(SQLUpdateCompiler, self).execute_sql(result_type)
    870899        rows = cursor and cursor.rowcount or 0
    871900        is_empty = cursor is None
     901        # explicit close, a Del prob won't trigger immediate GC in Jython
     902        cursor.close()
    872903        del cursor
    873904        for query in self.query.get_related_updates():
    874             aux_rows = query.get_compiler(self.using).execute_sql(result_type)
     905            qcompiler = query.get_compiler(self.using)
     906            aux_rows = qcompiler.execute_sql(result_type)
    875907            if is_empty:
    876908                rows = aux_rows
    877909                is_empty = False
     910            if (result_type is None and type(qcompiler).__name__ not in ['SQLInsertCompiler','SQLUpdateCompiler']):
     911                # Only insert and update compilers are cursor close-safe.
     912                # Other compilers using SQLCompiler.execute_sql(None) will return open cursor.
     913                aux_rows.close()
    878914        return rows
    879915
    880916    def pre_sql_setup(self):
     
    953989            needs_string_cast = self.connection.features.needs_datetime_string_cast
    954990
    955991        offset = len(self.query.extra_select)
    956         for rows in self.execute_sql(MULTI):
    957             for row in rows:
    958                 date = row[offset]
    959                 if resolve_columns:
    960                     date = self.resolve_columns(row, fields)[offset]
    961                 elif needs_string_cast:
    962                     date = typecast_timestamp(str(date))
    963                 yield date
     992        date_iter = self.execute_sql(MULTI)
     993        try:
     994            for rows in date_iter:
     995                for row in rows:
     996                    date = row[offset]
     997                    if resolve_columns:
     998                        date = self.resolve_columns(row, fields)[offset]
     999                    elif needs_string_cast:
     1000                        date = typecast_timestamp(str(date))
     1001                    yield date
     1002        finally:
     1003            # request parent to clean up cursor associated with date_iter
     1004            self.close_cursor_by_iter(date_iter)
    9641005
    965 
    9661006def empty_iter():
    9671007    """
    9681008    Returns an iterator containing no results.
  • django/db/models/sql/subqueries.py

     
    2626    def do_query(self, table, where, using):
    2727        self.tables = [table]
    2828        self.where = where
    29         self.get_compiler(using).execute_sql(None)
     29        cursor = self.get_compiler(using).execute_sql(None)
     30        #parent returns cursor for type None, close it
     31        cursor.close()
    3032
    3133    def delete_batch(self, pk_list, using, field=None):
    3234        """
     
    7880            self.where.add((Constraint(None, pk_field.column, pk_field), 'in',
    7981                    pk_list[offset : offset + GET_ITERATOR_CHUNK_SIZE]),
    8082                    AND)
    81             self.get_compiler(using).execute_sql(None)
     83            cursor = self.get_compiler(using).execute_sql(None)
     84            cursor.close()
    8285
    8386    def add_update_values(self, values):
    8487        """
Back to Top