Code

Ticket #14298: close-cursor-v2.diff

File close-cursor-v2.diff, 10.4 KB (added by xoror, 3 years ago)
Line 
1Index: django/db/models/sql/compiler.py
2===================================================================
3--- django/db/models/sql/compiler.py    (revision 15147)
4+++ django/db/models/sql/compiler.py    (working copy)
5@@ -13,6 +13,7 @@
6         self.connection = connection
7         self.using = using
8         self.quote_cache = {}
9+        self.iter_cursor_map = {}
10 
11     def pre_sql_setup(self):
12         """
13@@ -677,37 +676,52 @@
14         resolve_columns = hasattr(self, 'resolve_columns')
15         fields = None
16         has_aggregate_select = bool(self.query.aggregate_select)
17-        for rows in self.execute_sql(MULTI):
18-            for row in rows:
19-                if resolve_columns:
20-                    if fields is None:
21-                        # We only set this up here because
22-                        # related_select_fields isn't populated until
23-                        # execute_sql() has been called.
24-                        if self.query.select_fields:
25-                            fields = self.query.select_fields + self.query.related_select_fields
26-                        else:
27-                            fields = self.query.model._meta.fields
28-                        # If the field was deferred, exclude it from being passed
29-                        # into `resolve_columns` because it wasn't selected.
30-                        only_load = self.deferred_to_columns()
31-                        if only_load:
32-                            db_table = self.query.model._meta.db_table
33-                            fields = [f for f in fields if db_table in only_load and
34-                                      f.column in only_load[db_table]]
35-                    row = self.resolve_columns(row, fields)
36+        cursor_iter = self.execute_sql(MULTI)
37+        try:
38+            for rows in cursor_iter:
39+                for row in rows:
40+                    if resolve_columns:
41+                        if fields is None:
42+                            # We only set this up here because
43+                            # related_select_fields isn't populated until
44+                            # execute_sql() has been called.
45+                            if self.query.select_fields:
46+                                fields = self.query.select_fields + self.query.related_select_fields
47+                            else:
48+                                fields = self.query.model._meta.fields
49+                            # If the field was deferred, exclude it from being passed
50+                            # into `resolve_columns` because it wasn't selected.
51+                            only_load = self.deferred_to_columns()
52+                            if only_load:
53+                                db_table = self.query.model._meta.db_table
54+                                fields = [f for f in fields if db_table in only_load and
55+                                          f.column in only_load[db_table]]
56+                        row = self.resolve_columns(row, fields)
57+   
58+                    if has_aggregate_select:
59+                        aggregate_start = len(self.query.extra_select.keys()) + len(self.query.select)
60+                        aggregate_end = aggregate_start + len(self.query.aggregate_select)
61+                        row = tuple(row[:aggregate_start]) + tuple([
62+                            self.query.resolve_aggregate(value, aggregate, self.connection)
63+                            for (alias, aggregate), value
64+                            in zip(self.query.aggregate_select.items(), row[aggregate_start:aggregate_end])
65+                        ]) + tuple(row[aggregate_end:])
66+   
67+                    yield row
68+        finally:
69+            # iteration over cursor is completed.
70+            # it's now save to close the cursor and remove it from the lookup dict
71+            if cursor_iter in self.iter_cursor_map:
72+                self.close_cursor_by_iter(cursor_iter)
73+               
74+    def close_cursor_by_iter(self, iter_to_close):
75+        """
76+        Added method to close iterator from a subclass
77+        """
78+        if iter_to_close in self.iter_cursor_map:
79+            self.iter_cursor_map[iter_to_close].close()
80+            del self.iter_cursor_map[iter_to_close]
81 
82-                if has_aggregate_select:
83-                    aggregate_start = len(self.query.extra_select.keys()) + len(self.query.select)
84-                    aggregate_end = aggregate_start + len(self.query.aggregate_select)
85-                    row = tuple(row[:aggregate_start]) + tuple([
86-                        self.query.resolve_aggregate(value, aggregate, self.connection)
87-                        for (alias, aggregate), value
88-                        in zip(self.query.aggregate_select.items(), row[aggregate_start:aggregate_end])
89-                    ]) + tuple(row[aggregate_end:])
90-
91-                yield row
92-
93     def execute_sql(self, result_type=MULTI):
94         """
95         Run the query against the database and returns the result(s). The
96@@ -734,12 +748,17 @@
97         cursor = self.connection.cursor()
98         cursor.execute(sql, params)
99 
100+        # close cursors for single results.
101         if not result_type:
102             return cursor
103         if result_type == SINGLE:
104             if self.query.ordering_aliases:
105-                return cursor.fetchone()[:-len(self.query.ordering_aliases)]
106-            return cursor.fetchone()
107+                temp_result = cursor.fetchone()[:-len(self.query.ordering_aliases)]
108+                cursor.close()
109+                return temp_result
110+            temp_result = cursor.fetchone()
111+            cursor.close()
112+            return temp_result
113 
114         # The MULTI case.
115         if self.query.ordering_aliases:
116@@ -748,14 +767,18 @@
117         else:
118             result = iter((lambda: cursor.fetchmany(GET_ITERATOR_CHUNK_SIZE)),
119                     self.connection.features.empty_fetchmany_value)
120+        # store iterator only if we use chunks
121+        if self.connection.features.can_use_chunked_reads:
122+            self.iter_cursor_map[result] = cursor
123         if not self.connection.features.can_use_chunked_reads:
124             # If we are using non-chunked reads, we return the same data
125             # structure as normally, but ensure it is all read into memory
126             # before going any further.
127-            return list(result)
128+            temp_ressult = list(result)
129+            cursor.close()
130+            return temp_result
131         return result
132 
133-
134 class SQLInsertCompiler(SQLCompiler):
135     def placeholder(self, field, val):
136         if field is None:
137@@ -790,11 +813,17 @@
138         self.return_id = return_id
139         cursor = super(SQLInsertCompiler, self).execute_sql(None)
140         if not (return_id and cursor):
141+            if cursor:
142+                cursor.close()
143             return
144         if self.connection.features.can_return_id_from_insert:
145-            return self.connection.ops.fetch_returned_insert_id(cursor)
146-        return self.connection.ops.last_insert_id(cursor,
147+            lastid = self.connection.ops.fetch_returned_insert_id(cursor)
148+            cursor.close()
149+            return lastid
150+        result = self.connection.ops.last_insert_id(cursor,
151                 self.query.model._meta.db_table, self.query.model._meta.pk.column)
152+        cursor.close()
153+        return result
154 
155 
156 class SQLDeleteCompiler(SQLCompiler):
157@@ -869,12 +898,19 @@
158         cursor = super(SQLUpdateCompiler, self).execute_sql(result_type)
159         rows = cursor and cursor.rowcount or 0
160         is_empty = cursor is None
161+        # explicit close, a Del prob won't trigger immediate GC in Jython
162+        cursor.close()
163         del cursor
164         for query in self.query.get_related_updates():
165-            aux_rows = query.get_compiler(self.using).execute_sql(result_type)
166+            qcompiler = query.get_compiler(self.using)
167+            aux_rows = qcompiler.execute_sql(result_type)
168             if is_empty:
169                 rows = aux_rows
170                 is_empty = False
171+            if (result_type is None and type(qcompiler).__name__ not in ['SQLInsertCompiler','SQLUpdateCompiler']):
172+                # Only insert and update compilers are cursor close-safe.
173+                # Other compilers using SQLCompiler.execute_sql(None) will return open cursor.
174+                aux_rows.close()
175         return rows
176 
177     def pre_sql_setup(self):
178@@ -953,16 +989,20 @@
179             needs_string_cast = self.connection.features.needs_datetime_string_cast
180 
181         offset = len(self.query.extra_select)
182-        for rows in self.execute_sql(MULTI):
183-            for row in rows:
184-                date = row[offset]
185-                if resolve_columns:
186-                    date = self.resolve_columns(row, fields)[offset]
187-                elif needs_string_cast:
188-                    date = typecast_timestamp(str(date))
189-                yield date
190+        date_iter = self.execute_sql(MULTI)
191+        try:
192+            for rows in date_iter:
193+                for row in rows:
194+                    date = row[offset]
195+                    if resolve_columns:
196+                        date = self.resolve_columns(row, fields)[offset]
197+                    elif needs_string_cast:
198+                        date = typecast_timestamp(str(date))
199+                    yield date
200+        finally:
201+            # request parent to clean up cursor associated with date_iter
202+            self.close_cursor_by_iter(date_iter)
203 
204-
205 def empty_iter():
206     """
207     Returns an iterator containing no results.
208Index: django/db/models/sql/subqueries.py
209===================================================================
210--- django/db/models/sql/subqueries.py  (revision 15147)
211+++ django/db/models/sql/subqueries.py  (working copy)
212@@ -26,7 +26,9 @@
213     def do_query(self, table, where, using):
214         self.tables = [table]
215         self.where = where
216-        self.get_compiler(using).execute_sql(None)
217+        cursor = self.get_compiler(using).execute_sql(None)
218+        #parent returns cursor for type None, close it
219+        cursor.close()
220 
221     def delete_batch(self, pk_list, using, field=None):
222         """
223@@ -78,7 +80,8 @@
224             self.where.add((Constraint(None, pk_field.column, pk_field), 'in',
225                     pk_list[offset : offset + GET_ITERATOR_CHUNK_SIZE]),
226                     AND)
227-            self.get_compiler(using).execute_sql(None)
228+            cursor = self.get_compiler(using).execute_sql(None)
229+            cursor.close()
230 
231     def add_update_values(self, values):
232         """