From 4531703f440e166b462fa0106f5429f7b904ce6a Mon Sep 17 00:00:00 2001 From: shaohuzhang1 Date: Thu, 27 Jun 2024 12:14:56 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=8D=87=E7=BA=A7django?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/common/db/compiler.py | 114 +++++++++++++++++++++++++++++++++---- pyproject.toml | 4 +- 2 files changed, 106 insertions(+), 12 deletions(-) diff --git a/apps/common/db/compiler.py b/apps/common/db/compiler.py index 9a65f93e1..69640c8a0 100644 --- a/apps/common/db/compiler.py +++ b/apps/common/db/compiler.py @@ -7,9 +7,10 @@ @desc: """ -from django.core.exceptions import EmptyResultSet +from django.core.exceptions import EmptyResultSet, FullResultSet from django.db import NotSupportedError from django.db.models.sql.compiler import SQLCompiler +from django.db.transaction import TransactionManagementError class AppSQLCompiler(SQLCompiler): @@ -19,15 +20,16 @@ class AppSQLCompiler(SQLCompiler): field_replace_dict = {} self.field_replace_dict = field_replace_dict - def get_query_str(self, with_limits=True, with_table_name=False): + def get_query_str(self, with_limits=True, with_table_name=False, with_col_aliases=False): refcounts_before = self.query.alias_refcount.copy() try: - extra_select, order_by, group_by = self.pre_sql_setup() + combinator = self.query.combinator + extra_select, order_by, group_by = self.pre_sql_setup( + with_col_aliases=with_col_aliases or bool(combinator), + ) for_update_part = None # Is a LIMIT/OFFSET clause needed? - with_limit_offset = with_limits and ( - self.query.high_mark is not None or self.query.low_mark - ) + with_limit_offset = with_limits and self.query.is_sliced combinator = self.query.combinator features = self.connection.features if combinator: @@ -40,8 +42,14 @@ class AppSQLCompiler(SQLCompiler): result, params = self.get_combinator_sql( combinator, self.query.combinator_all ) + elif self.qualify: + result, params = self.get_qualify_sql() + order_by = None else: distinct_fields, distinct_params = self.get_distinct() + # This must come after 'select', 'ordering', and 'distinct' + # (see docstring of get_from_clause() for details). + from_, f_params = self.get_from_clause() try: where, w_params = ( self.compile(self.where) if self.where is not None else ("", []) @@ -51,11 +59,92 @@ class AppSQLCompiler(SQLCompiler): raise # Use a predicate that's always False. where, w_params = "0 = 1", [] - having, h_params = ( - self.compile(self.having) if self.having is not None else ("", []) - ) + except FullResultSet: + where, w_params = "", [] + try: + having, h_params = ( + self.compile(self.having) + if self.having is not None + else ("", []) + ) + except FullResultSet: + having, h_params = "", [] result = [] params = [] + + if self.query.distinct: + distinct_result, distinct_params = self.connection.ops.distinct_sql( + distinct_fields, + distinct_params, + ) + result += distinct_result + params += distinct_params + + out_cols = [] + for _, (s_sql, s_params), alias in self.select + extra_select: + if alias: + s_sql = "%s AS %s" % ( + s_sql, + self.connection.ops.quote_name(alias), + ) + params.extend(s_params) + out_cols.append(s_sql) + + params.extend(f_params) + + if self.query.select_for_update and features.has_select_for_update: + if ( + self.connection.get_autocommit() + # Don't raise an exception when database doesn't + # support transactions, as it's a noop. + and features.supports_transactions + ): + raise TransactionManagementError( + "select_for_update cannot be used outside of a transaction." + ) + + if ( + with_limit_offset + and not features.supports_select_for_update_with_limit + ): + raise NotSupportedError( + "LIMIT/OFFSET is not supported with " + "select_for_update on this database backend." + ) + nowait = self.query.select_for_update_nowait + skip_locked = self.query.select_for_update_skip_locked + of = self.query.select_for_update_of + no_key = self.query.select_for_no_key_update + # If it's a NOWAIT/SKIP LOCKED/OF/NO KEY query but the + # backend doesn't support it, raise NotSupportedError to + # prevent a possible deadlock. + if nowait and not features.has_select_for_update_nowait: + raise NotSupportedError( + "NOWAIT is not supported on this database backend." + ) + elif skip_locked and not features.has_select_for_update_skip_locked: + raise NotSupportedError( + "SKIP LOCKED is not supported on this database backend." + ) + elif of and not features.has_select_for_update_of: + raise NotSupportedError( + "FOR UPDATE OF is not supported on this database backend." + ) + elif no_key and not features.has_select_for_no_key_update: + raise NotSupportedError( + "FOR NO KEY UPDATE is not supported on this " + "database backend." + ) + for_update_part = self.connection.ops.for_update_sql( + nowait=nowait, + skip_locked=skip_locked, + of=self.get_select_for_update_of_arguments(), + no_key=no_key, + ) + + if for_update_part and features.for_update_after_from: + result.append(for_update_part) + if where: result.append("WHERE %s" % where) params.extend(w_params) @@ -91,7 +180,11 @@ class AppSQLCompiler(SQLCompiler): for _, (o_sql, o_params, _) in order_by: ordering.append(o_sql) params.extend(o_params) - result.append("ORDER BY %s" % ", ".join(ordering)) + order_by_sql = "ORDER BY %s" % ", ".join(ordering) + if combinator and features.requires_compound_order_by_subquery: + result = ["SELECT * FROM (", *result, ")", order_by_sql] + else: + result.append(order_by_sql) if with_limit_offset: result.append( @@ -102,6 +195,7 @@ class AppSQLCompiler(SQLCompiler): if for_update_part and not features.for_update_after_from: result.append(for_update_part) + from_, f_params = self.get_from_clause() sql = " ".join(result) if not with_table_name: diff --git a/pyproject.toml b/pyproject.toml index fffd51cff..826ec71ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,8 +7,8 @@ readme = "README.md" [tool.poetry.dependencies] python = "^3.11" -django = "4.1.13" -djangorestframework = "3.14.0" +django = "4.2.13" +djangorestframework = "^3.15.2" drf-yasg = "1.21.7" django-filter = "23.2" langchain = "^0.2.3"