Skip to content

Commit 4c38e8a

Browse files
liuruibinliuruibin
committed
feat: add dynamic SQL execution and pagination functionality with custom query compiler
1 parent fbb4e7d commit 4c38e8a

File tree

5 files changed

+481
-1
lines changed

5 files changed

+481
-1
lines changed

‎apps/common/db/__init__.py‎

Whitespace-only changes.

‎apps/common/db/compiler.py‎

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
# coding=utf-8
2+
"""
3+
@project: maxkb
4+
@Author：虎
5+
@file： compiler.py
6+
@date：2023/10/7 10:53
7+
@desc:
8+
"""
9+
10+
from django.core.exceptions import EmptyResultSet, FullResultSet
11+
from django.db import NotSupportedError
12+
from django.db.models.sql.compiler import SQLCompiler
13+
from django.db.transaction import TransactionManagementError
14+
15+
16+
class AppSQLCompiler(SQLCompiler):
17+
def __init__(self, query, connection, using, elide_empty=True, field_replace_dict=None):
18+
super().__init__(query, connection, using, elide_empty)
19+
if field_replace_dict is None:
20+
field_replace_dict = {}
21+
self.field_replace_dict = field_replace_dict
22+
23+
def get_query_str(self, with_limits=True, with_table_name=False, with_col_aliases=False):
24+
refcounts_before = self.query.alias_refcount.copy()
25+
try:
26+
combinator = self.query.combinator
27+
extra_select, order_by, group_by = self.pre_sql_setup(
28+
with_col_aliases=with_col_aliases or bool(combinator),
29+
)
30+
for_update_part = None
31+
# Is a LIMIT/OFFSET clause needed?
32+
with_limit_offset = with_limits and self.query.is_sliced
33+
combinator = self.query.combinator
34+
features = self.connection.features
35+
if combinator:
36+
if not getattr(features, "supports_select_{}".format(combinator)):
37+
raise NotSupportedError(
38+
"{} is not supported on this database backend.".format(
39+
combinator
40+
)
41+
)
42+
result, params = self.get_combinator_sql(
43+
combinator, self.query.combinator_all
44+
)
45+
elif self.qualify:
46+
result, params = self.get_qualify_sql()
47+
order_by = None
48+
else:
49+
distinct_fields, distinct_params = self.get_distinct()
50+
# This must come after 'select', 'ordering', and 'distinct'
51+
# (see docstring of get_from_clause() for details).
52+
from_, f_params = self.get_from_clause()
53+
try:
54+
where, w_params = (
55+
self.compile(self.where) if self.where is not None else ("", [])
56+
)
57+
except EmptyResultSet:
58+
if self.elide_empty:
59+
raise
60+
# Use a predicate that's always False.
61+
where, w_params = "0 = 1", []
62+
except FullResultSet:
63+
where, w_params = "", []
64+
try:
65+
having, h_params = (
66+
self.compile(self.having)
67+
if self.having is not None
68+
else ("", [])
69+
)
70+
except FullResultSet:
71+
having, h_params = "", []
72+
result = []
73+
params = []
74+
75+
if self.query.distinct:
76+
distinct_result, distinct_params = self.connection.ops.distinct_sql(
77+
distinct_fields,
78+
distinct_params,
79+
)
80+
result += distinct_result
81+
params += distinct_params
82+
83+
out_cols = []
84+
for _, (s_sql, s_params), alias in self.select + extra_select:
85+
if alias:
86+
s_sql = "%s AS %s" % (
87+
s_sql,
88+
self.connection.ops.quote_name(alias),
89+
)
90+
params.extend(s_params)
91+
out_cols.append(s_sql)
92+
93+
params.extend(f_params)
94+
95+
if self.query.select_for_update and features.has_select_for_update:
96+
if (
97+
self.connection.get_autocommit()
98+
# Don't raise an exception when database doesn't
99+
# support transactions, as it's a noop.
100+
and features.supports_transactions
101+
):
102+
raise TransactionManagementError(
103+
"select_for_update cannot be used outside of a transaction."
104+
)
105+
106+
if (
107+
with_limit_offset
108+
and not features.supports_select_for_update_with_limit
109+
):
110+
raise NotSupportedError(
111+
"LIMIT/OFFSET is not supported with "
112+
"select_for_update on this database backend."
113+
)
114+
nowait = self.query.select_for_update_nowait
115+
skip_locked = self.query.select_for_update_skip_locked
116+
of = self.query.select_for_update_of
117+
no_key = self.query.select_for_no_key_update
118+
# If it's a NOWAIT/SKIP LOCKED/OF/NO KEY query but the
119+
# backend doesn't support it, raise NotSupportedError to
120+
# prevent a possible deadlock.
121+
if nowait and not features.has_select_for_update_nowait:
122+
raise NotSupportedError(
123+
"NOWAIT is not supported on this database backend."
124+
)
125+
elif skip_locked and not features.has_select_for_update_skip_locked:
126+
raise NotSupportedError(
127+
"SKIP LOCKED is not supported on this database backend."
128+
)
129+
elif of and not features.has_select_for_update_of:
130+
raise NotSupportedError(
131+
"FOR UPDATE OF is not supported on this database backend."
132+
)
133+
elif no_key and not features.has_select_for_no_key_update:
134+
raise NotSupportedError(
135+
"FOR NO KEY UPDATE is not supported on this "
136+
"database backend."
137+
)
138+
for_update_part = self.connection.ops.for_update_sql(
139+
nowait=nowait,
140+
skip_locked=skip_locked,
141+
of=self.get_select_for_update_of_arguments(),
142+
no_key=no_key,
143+
)
144+
145+
if for_update_part and features.for_update_after_from:
146+
result.append(for_update_part)
147+
148+
if where:
149+
result.append("WHERE %s" % where)
150+
params.extend(w_params)
151+
152+
grouping = []
153+
for g_sql, g_params in group_by:
154+
grouping.append(g_sql)
155+
params.extend(g_params)
156+
if grouping:
157+
if distinct_fields:
158+
raise NotImplementedError(
159+
"annotate() + distinct(fields) is not implemented."
160+
)
161+
order_by = order_by or self.connection.ops.force_no_ordering()
162+
result.append("GROUP BY %s" % ", ".join(grouping))
163+
if self._meta_ordering:
164+
order_by = None
165+
if having:
166+
result.append("HAVING %s" % having)
167+
params.extend(h_params)
168+
169+
if self.query.explain_info:
170+
result.insert(
171+
0,
172+
self.connection.ops.explain_query_prefix(
173+
self.query.explain_info.format,
174+
**self.query.explain_info.options,
175+
),
176+
)
177+
178+
if order_by:
179+
ordering = []
180+
for _, (o_sql, o_params, _) in order_by:
181+
ordering.append(o_sql)
182+
params.extend(o_params)
183+
order_by_sql = "ORDER BY %s" % ", ".join(ordering)
184+
if combinator and features.requires_compound_order_by_subquery:
185+
result = ["SELECT * FROM (", *result, ")", order_by_sql]
186+
else:
187+
result.append(order_by_sql)
188+
189+
if with_limit_offset:
190+
result.append(
191+
self.connection.ops.limit_offset_sql(
192+
self.query.low_mark, self.query.high_mark
193+
)
194+
)
195+
196+
if for_update_part and not features.for_update_after_from:
197+
result.append(for_update_part)
198+
199+
from_, f_params = self.get_from_clause()
200+
sql = " ".join(result)
201+
if not with_table_name:
202+
for table_name in from_:
203+
sql = sql.replace(table_name + ".", "")
204+
for key in self.field_replace_dict.keys():
205+
value = self.field_replace_dict.get(key)
206+
sql = sql.replace(key, value)
207+
return sql, tuple(params)
208+
finally:
209+
# Finally do cleanup - get rid of the joins we created above.
210+
self.query.reset_refcounts(refcounts_before)
211+
212+
def as_sql(self, with_limits=True, with_col_aliases=False, select_string=None):
213+
if select_string is None:
214+
return super().as_sql(with_limits, with_col_aliases)
215+
else:
216+
sql, params = self.get_query_str(with_table_name=False)
217+
return (select_string + " " + sql), params

0 commit comments

Comments
 (0)