From 0837eacc4e1fa7916e48135e8ba43f54a7a64997 Mon Sep 17 00:00:00 2001 From: Michael Manfre Date: Wed, 8 Jan 2014 23:31:34 -0500 Subject: [PATCH] Made SQLCompiler.execute_sql(result_type) more explicit. Updated SQLUpdateCompiler.execute_sql to match the behavior described in the docstring; the 'first non-empty query' will now include all queries, not just the main and first related update. Added CURSOR and NO_RESULTS result_type constants to make the usages more self documenting and allow execute_sql to explicitly close the cursor when it is no longer needed. --- django/db/models/query.py | 5 +- django/db/models/sql/compiler.py | 78 ++++++++++++++++++++++-------- django/db/models/sql/constants.py | 2 + django/db/models/sql/subqueries.py | 8 +-- tests/backends/tests.py | 3 +- 5 files changed, 69 insertions(+), 27 deletions(-) diff --git a/django/db/models/query.py b/django/db/models/query.py index 48d295ccca..353dd95794 100644 --- a/django/db/models/query.py +++ b/django/db/models/query.py @@ -14,6 +14,7 @@ from django.db.models.fields import AutoField, Empty from django.db.models.query_utils import (Q, select_related_descend, deferred_class_factory, InvalidQuery) from django.db.models.deletion import Collector +from django.db.models.sql.constants import CURSOR from django.db.models import sql from django.utils.functional import partition from django.utils import six @@ -574,7 +575,7 @@ class QuerySet(object): query = self.query.clone(sql.UpdateQuery) query.add_update_values(kwargs) with transaction.commit_on_success_unless_managed(using=self.db): - rows = query.get_compiler(self.db).execute_sql(None) + rows = query.get_compiler(self.db).execute_sql(CURSOR) self._result_cache = None return rows update.alters_data = True @@ -591,7 +592,7 @@ class QuerySet(object): query = self.query.clone(sql.UpdateQuery) query.add_update_fields(values) self._result_cache = None - return query.get_compiler(self.db).execute_sql(None) + return query.get_compiler(self.db).execute_sql(CURSOR) _update.alters_data = True _update.queryset_only = False diff --git a/django/db/models/sql/compiler.py b/django/db/models/sql/compiler.py index 123427cf8b..536a66d139 100644 --- a/django/db/models/sql/compiler.py +++ b/django/db/models/sql/compiler.py @@ -5,8 +5,8 @@ from django.core.exceptions import FieldError from django.db.backends.utils import truncate_name from django.db.models.constants import LOOKUP_SEP from django.db.models.query_utils import select_related_descend, QueryWrapper -from django.db.models.sql.constants import (SINGLE, MULTI, ORDER_DIR, - GET_ITERATOR_CHUNK_SIZE, SelectInfo) +from django.db.models.sql.constants import (CURSOR, SINGLE, MULTI, NO_RESULTS, + ORDER_DIR, GET_ITERATOR_CHUNK_SIZE, SelectInfo) from django.db.models.sql.datastructures import EmptyResultSet from django.db.models.sql.expressions import SQLEvaluator from django.db.models.sql.query import get_order_dir, Query @@ -762,6 +762,8 @@ class SQLCompiler(object): is needed, as the filters describe an empty set. In that case, None is returned, to avoid any unnecessary database interaction. """ + if not result_type: + result_type = NO_RESULTS try: sql, params = self.as_sql() if not sql: @@ -773,27 +775,44 @@ class SQLCompiler(object): return cursor = self.connection.cursor() - cursor.execute(sql, params) + try: + cursor.execute(sql, params) + except: + cursor.close() + raise - if not result_type: + if result_type == CURSOR: + # Caller didn't specify a result_type, so just give them back the + # cursor to process (and close). return cursor if result_type == SINGLE: - if self.ordering_aliases: - return cursor.fetchone()[:-len(self.ordering_aliases)] - return cursor.fetchone() + try: + if self.ordering_aliases: + return cursor.fetchone()[:-len(self.ordering_aliases)] + return cursor.fetchone() + finally: + # done with the cursor + cursor.close() + if result_type == NO_RESULTS: + cursor.close() + return # The MULTI case. if self.ordering_aliases: result = order_modified_iter(cursor, len(self.ordering_aliases), self.connection.features.empty_fetchmany_value) else: - result = iter((lambda: cursor.fetchmany(GET_ITERATOR_CHUNK_SIZE)), - self.connection.features.empty_fetchmany_value) + result = cursor_iter(cursor, + self.connection.features.empty_fetchmany_value) if not self.connection.features.can_use_chunked_reads: - # If we are using non-chunked reads, we return the same data - # structure as normally, but ensure it is all read into memory - # before going any further. - return list(result) + try: + # If we are using non-chunked reads, we return the same data + # structure as normally, but ensure it is all read into memory + # before going any further. + return list(result) + finally: + # done with the cursor + cursor.close() return result def as_subquery_condition(self, alias, columns, qn): @@ -970,12 +989,15 @@ class SQLUpdateCompiler(SQLCompiler): related queries are not available. """ cursor = super(SQLUpdateCompiler, self).execute_sql(result_type) - rows = cursor.rowcount if cursor else 0 - is_empty = cursor is None - del cursor + try: + rows = cursor.rowcount if cursor else 0 + is_empty = cursor is None + finally: + if cursor: + cursor.close() for query in self.query.get_related_updates(): aux_rows = query.get_compiler(self.using).execute_sql(result_type) - if is_empty: + if is_empty and aux_rows: rows = aux_rows is_empty = False return rows @@ -1111,6 +1133,19 @@ class SQLDateTimeCompiler(SQLCompiler): yield datetime +def cursor_iter(cursor, sentinel): + """ + Yields blocks of rows from a cursor and ensures the cursor is closed when + done. + """ + try: + for rows in iter((lambda: cursor.fetchmany(GET_ITERATOR_CHUNK_SIZE)), + sentinel): + yield rows + finally: + cursor.close() + + def order_modified_iter(cursor, trim, sentinel): """ Yields blocks of rows from a cursor. We use this iterator in the special @@ -1118,6 +1153,9 @@ def order_modified_iter(cursor, trim, sentinel): requirements. We must trim those extra columns before anything else can use the results, since they're only needed to make the SQL valid. """ - for rows in iter((lambda: cursor.fetchmany(GET_ITERATOR_CHUNK_SIZE)), - sentinel): - yield [r[:-trim] for r in rows] + try: + for rows in iter((lambda: cursor.fetchmany(GET_ITERATOR_CHUNK_SIZE)), + sentinel): + yield [r[:-trim] for r in rows] + finally: + cursor.close() diff --git a/django/db/models/sql/constants.py b/django/db/models/sql/constants.py index 904f7b2c8b..36aab23bae 100644 --- a/django/db/models/sql/constants.py +++ b/django/db/models/sql/constants.py @@ -33,6 +33,8 @@ SelectInfo = namedtuple('SelectInfo', 'col field') # How many results to expect from a cursor.execute call MULTI = 'multi' SINGLE = 'single' +CURSOR = 'cursor' +NO_RESULTS = 'no results' ORDER_PATTERN = re.compile(r'\?|[-+]?[.\w]+$') ORDER_DIR = { diff --git a/django/db/models/sql/subqueries.py b/django/db/models/sql/subqueries.py index 86b1efd3f8..cfda1f552c 100644 --- a/django/db/models/sql/subqueries.py +++ b/django/db/models/sql/subqueries.py @@ -8,7 +8,7 @@ from django.db import connections from django.db.models.query_utils import Q from django.db.models.constants import LOOKUP_SEP from django.db.models.fields import DateField, DateTimeField, FieldDoesNotExist -from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, SelectInfo +from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, NO_RESULTS, SelectInfo from django.db.models.sql.datastructures import Date, DateTime from django.db.models.sql.query import Query from django.utils import six @@ -30,7 +30,7 @@ class DeleteQuery(Query): def do_query(self, table, where, using): self.tables = [table] self.where = where - self.get_compiler(using).execute_sql(None) + self.get_compiler(using).execute_sql(NO_RESULTS) def delete_batch(self, pk_list, using, field=None): """ @@ -82,7 +82,7 @@ class DeleteQuery(Query): values = innerq self.where = self.where_class() self.add_q(Q(pk__in=values)) - self.get_compiler(using).execute_sql(None) + self.get_compiler(using).execute_sql(NO_RESULTS) class UpdateQuery(Query): @@ -116,7 +116,7 @@ class UpdateQuery(Query): for offset in range(0, len(pk_list), GET_ITERATOR_CHUNK_SIZE): self.where = self.where_class() self.add_q(Q(pk__in=pk_list[offset: offset + GET_ITERATOR_CHUNK_SIZE])) - self.get_compiler(using).execute_sql(None) + self.get_compiler(using).execute_sql(NO_RESULTS) def add_update_values(self, values): """ diff --git a/tests/backends/tests.py b/tests/backends/tests.py index 0ff3ad0bba..4a3fc31b7a 100644 --- a/tests/backends/tests.py +++ b/tests/backends/tests.py @@ -20,6 +20,7 @@ from django.db.backends.utils import format_number, CursorWrapper from django.db.models import Sum, Avg, Variance, StdDev from django.db.models.fields import (AutoField, DateField, DateTimeField, DecimalField, IntegerField, TimeField) +from django.db.models.sql.constants import CURSOR from django.db.utils import ConnectionHandler from django.test import (TestCase, TransactionTestCase, override_settings, skipUnlessDBFeature, skipIfDBFeature) @@ -209,7 +210,7 @@ class LastExecutedQueryTest(TestCase): """ persons = models.Reporter.objects.filter(raw_data=b'\x00\x46 \xFE').extra(select={'föö': 1}) sql, params = persons.query.sql_with_params() - cursor = persons.query.get_compiler('default').execute_sql(None) + cursor = persons.query.get_compiler('default').execute_sql(CURSOR) last_sql = cursor.db.ops.last_executed_query(cursor, sql, params) self.assertIsInstance(last_sql, six.text_type)