From 1469952b442324d0f4cbe5d9c925815878bb2b27 Mon Sep 17 00:00:00 2001 From: Raphael Gaschignard Date: Tue, 19 Nov 2024 14:50:24 +1000 Subject: [PATCH] Fixed #35918 -- Refactored execute_sql to reduce cursor management. This in particular adds support for execute_sql to return row counts directly --- django/db/models/query.py | 12 ++-- django/db/models/sql/compiler.py | 88 ++++++++++++++++++------------ django/db/models/sql/constants.py | 6 ++ django/db/models/sql/subqueries.py | 12 ++-- docs/releases/5.2.txt | 11 ++++ docs/spelling_wordlist | 2 + 6 files changed, 82 insertions(+), 49 deletions(-) diff --git a/django/db/models/query.py b/django/db/models/query.py index 21d5534cc9..62ea6d9288 100644 --- a/django/db/models/query.py +++ b/django/db/models/query.py @@ -26,7 +26,7 @@ from django.db.models.deletion import Collector from django.db.models.expressions import Case, F, Value, When from django.db.models.functions import Cast, Trunc from django.db.models.query_utils import FilteredRelation, Q -from django.db.models.sql.constants import CURSOR, GET_ITERATOR_CHUNK_SIZE +from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, ROW_COUNT from django.db.models.utils import ( AltersData, create_namedtuple_class, @@ -1206,11 +1206,7 @@ class QuerySet(AltersData): """ query = self.query.clone() query.__class__ = sql.DeleteQuery - cursor = query.get_compiler(using).execute_sql(CURSOR) - if cursor: - with cursor: - return cursor.rowcount - return 0 + return query.get_compiler(using).execute_sql(ROW_COUNT) _raw_delete.alters_data = True @@ -1249,7 +1245,7 @@ class QuerySet(AltersData): # Clear any annotations so that they won't be present in subqueries. query.annotations = {} with transaction.mark_for_rollback_on_error(using=self.db): - rows = query.get_compiler(self.db).execute_sql(CURSOR) + rows = query.get_compiler(self.db).execute_sql(ROW_COUNT) self._result_cache = None return rows @@ -1274,7 +1270,7 @@ class QuerySet(AltersData): # Clear any annotations so that they won't be present in subqueries. query.annotations = {} self._result_cache = None - return query.get_compiler(self.db).execute_sql(CURSOR) + return query.get_compiler(self.db).execute_sql(ROW_COUNT) _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 49263d5944..027f3946a8 100644 --- a/django/db/models/sql/compiler.py +++ b/django/db/models/sql/compiler.py @@ -17,6 +17,7 @@ from django.db.models.sql.constants import ( MULTI, NO_RESULTS, ORDER_DIR, + ROW_COUNT, SINGLE, ) from django.db.models.sql.query import Query, get_order_dir @@ -1560,12 +1561,19 @@ class SQLCompiler: return value is a single data item if result_type is SINGLE, or an iterator over the results if the result_type is MULTI. - result_type is either MULTI (use fetchmany() to retrieve all rows), - SINGLE (only retrieve a single row), or None. In this last case, the - cursor is returned if any query is executed, since it's used by - subclasses such as InsertQuery). It's possible, however, that no query - is needed, as the filters describe an empty set. In that case, None is - returned, to avoid any unnecessary database interaction. + result_type can be one of the following: + - MULTI + uses fetchmany() to retrieve all rows, potentially wrapping + it in an iterator for chunked reads for connections that + support it + - SINGLE + retrieves a single row using fetchone() + - ROW_COUNT + retrieve the number of rows in the result + - CURSOR + run the query, and then return the cursor object. In this case + it is the caller's responsibility to close the cursor when they + are done with it """ result_type = result_type or NO_RESULTS try: @@ -1588,10 +1596,15 @@ class SQLCompiler: cursor.close() raise - if result_type == CURSOR: - # Give the caller the cursor to process and close. + if result_type == ROW_COUNT: + try: + return cursor.rowcount + finally: + cursor.close() + elif result_type == CURSOR: + # here we are returning the cursor without closing it return cursor - if result_type == SINGLE: + elif result_type == SINGLE: try: val = cursor.fetchone() if val: @@ -1600,23 +1613,26 @@ class SQLCompiler: finally: # done with the cursor cursor.close() - if result_type == NO_RESULTS: + elif result_type == NO_RESULTS: cursor.close() return - - result = cursor_iter( - cursor, - self.connection.features.empty_fetchmany_value, - self.col_count if self.has_extra_select else None, - chunk_size, - ) - if not chunked_fetch or 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. Use chunked_fetch if requested, - # unless the database doesn't support it. - return list(result) - return result + else: + assert result_type == MULTI + # NB: cursor is now managed by cursor_iter, which + # will close the cursor if/when everything is consumed + result = cursor_iter( + cursor, + self.connection.features.empty_fetchmany_value, + self.col_count if self.has_extra_select else None, + chunk_size, + ) + if not chunked_fetch or 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. Use chunked_fetch if requested, + # unless the database doesn't support it. + return list(result) + return result def as_subquery_condition(self, alias, columns, compiler): qn = compiler.quote_name_unless_alias @@ -2012,19 +2028,21 @@ class SQLUpdateCompiler(SQLCompiler): non-empty query that is executed. Row counts for any subsequent, related queries are not available. """ - cursor = super().execute_sql(result_type) - try: - rows = cursor.rowcount if cursor else 0 - is_empty = cursor is None - finally: - if cursor: - cursor.close() + if result_type not in {ROW_COUNT, NO_RESULTS}: + raise ValueError(f"Unknown cursor type for update {repr(result_type)}") + row_count = super().execute_sql(result_type) + is_empty = row_count is None + row_count = row_count or 0 + for query in self.query.get_related_updates(): - aux_rows = query.get_compiler(self.using).execute_sql(result_type) - if is_empty and aux_rows: - rows = aux_rows + # NB: if result_type == NO_RESULTS then aux_row_count is None + aux_row_count = query.get_compiler(self.using).execute_sql(result_type) + if is_empty and aux_row_count: + # this will return the row count for any related updates as + # the number of rows updated + row_count = aux_row_count is_empty = False - return rows + return row_count def pre_sql_setup(self): """ diff --git a/django/db/models/sql/constants.py b/django/db/models/sql/constants.py index fdfb2ea891..12c3dc3915 100644 --- a/django/db/models/sql/constants.py +++ b/django/db/models/sql/constants.py @@ -9,9 +9,15 @@ GET_ITERATOR_CHUNK_SIZE = 100 # Namedtuples for sql.* internal use. # How many results to expect from a cursor.execute call +# multiple rows are expected MULTI = "multi" +# a single row is expected SINGLE = "single" +# do not return the rows, instead return the cursor +# used for the query CURSOR = "cursor" +# instead of returning the rows, return the row count +ROW_COUNT = "row count" NO_RESULTS = "no results" ORDER_DIR = { diff --git a/django/db/models/sql/subqueries.py b/django/db/models/sql/subqueries.py index f639eb8b82..b2810c8413 100644 --- a/django/db/models/sql/subqueries.py +++ b/django/db/models/sql/subqueries.py @@ -3,7 +3,11 @@ Query subclasses which provide extra functionality beyond simple data retrieval. """ from django.core.exceptions import FieldError -from django.db.models.sql.constants import CURSOR, GET_ITERATOR_CHUNK_SIZE, NO_RESULTS +from django.db.models.sql.constants import ( + GET_ITERATOR_CHUNK_SIZE, + NO_RESULTS, + ROW_COUNT, +) from django.db.models.sql.query import Query __all__ = ["DeleteQuery", "UpdateQuery", "InsertQuery", "AggregateQuery"] @@ -17,11 +21,7 @@ class DeleteQuery(Query): def do_query(self, table, where, using): self.alias_map = {table: self.alias_map[table]} self.where = where - cursor = self.get_compiler(using).execute_sql(CURSOR) - if cursor: - with cursor: - return cursor.rowcount - return 0 + return self.get_compiler(using).execute_sql(ROW_COUNT) def delete_batch(self, pk_list, using): """ diff --git a/docs/releases/5.2.txt b/docs/releases/5.2.txt index f1ffe07569..0d171bafe3 100644 --- a/docs/releases/5.2.txt +++ b/docs/releases/5.2.txt @@ -389,6 +389,17 @@ Dropped support for PostgreSQL 13 Upstream support for PostgreSQL 13 ends in November 2025. Django 5.2 supports PostgreSQL 14 and higher. +Models +------ + +* Multiple changes have been made to the undocumented ``django.db.models.sql.compiler.SQLCompiler.execute_sql`` + method. + + * ``ROW_COUNT`` has been added as a result type, which returns the number of rows + returned by the query directly, closing the cursor in the process. + * ``SQLUpdateCompiler.execute_sql`` now only accepts ``NO_RESULT`` and ``ROW_COUNT`` + as result types. + Changed MySQL connection character set default ---------------------------------------------- diff --git a/docs/spelling_wordlist b/docs/spelling_wordlist index 747a712a62..b2777e06b1 100644 --- a/docs/spelling_wordlist +++ b/docs/spelling_wordlist @@ -472,6 +472,8 @@ Sorani sortable Spectre Springmeyer +sql +SQLCompiler SSL stacktrace stateful