Refs #28333 -- Moved SQLCompiler's forced column aliasing logic to get_select().

This extends query composability possibilities when dealing with
subqueries which is necessary to implement window function filtering.
This commit is contained in:
Simon Charette 2022-08-10 07:56:54 -04:00 committed by Mariusz Felisiak
parent e34dfad0a3
commit 8c3046daad
1 changed files with 14 additions and 15 deletions

View File

@ -57,19 +57,21 @@ class SQLCompiler:
f"connection={self.connection!r} using={self.using!r}>" f"connection={self.connection!r} using={self.using!r}>"
) )
def setup_query(self): def setup_query(self, with_col_aliases=False):
if all(self.query.alias_refcount[a] == 0 for a in self.query.alias_map): if all(self.query.alias_refcount[a] == 0 for a in self.query.alias_map):
self.query.get_initial_alias() self.query.get_initial_alias()
self.select, self.klass_info, self.annotation_col_map = self.get_select() self.select, self.klass_info, self.annotation_col_map = self.get_select(
with_col_aliases=with_col_aliases,
)
self.col_count = len(self.select) self.col_count = len(self.select)
def pre_sql_setup(self): def pre_sql_setup(self, with_col_aliases=False):
""" """
Do any necessary class setup immediately prior to producing SQL. This Do any necessary class setup immediately prior to producing SQL. This
is for things that can't necessarily be done in __init__ because we is for things that can't necessarily be done in __init__ because we
might not have all the pieces in place at that time. might not have all the pieces in place at that time.
""" """
self.setup_query() self.setup_query(with_col_aliases=with_col_aliases)
order_by = self.get_order_by() order_by = self.get_order_by()
self.where, self.having = self.query.where.split_having() self.where, self.having = self.query.where.split_having()
extra_select = self.get_extra_select(order_by, self.select) extra_select = self.get_extra_select(order_by, self.select)
@ -224,7 +226,7 @@ class SQLCompiler:
] ]
return expressions return expressions
def get_select(self): def get_select(self, with_col_aliases=False):
""" """
Return three values: Return three values:
- a list of 3-tuples of (expression, (sql, params), alias) - a list of 3-tuples of (expression, (sql, params), alias)
@ -287,6 +289,7 @@ class SQLCompiler:
get_select_from_parent(klass_info) get_select_from_parent(klass_info)
ret = [] ret = []
col_idx = 1
for col, alias in select: for col, alias in select:
try: try:
sql, params = self.compile(col) sql, params = self.compile(col)
@ -301,6 +304,9 @@ class SQLCompiler:
sql, params = self.compile(Value(empty_result_set_value)) sql, params = self.compile(Value(empty_result_set_value))
else: else:
sql, params = col.select_format(self, sql, params) sql, params = col.select_format(self, sql, params)
if alias is None and with_col_aliases:
alias = f"col{col_idx}"
col_idx += 1
ret.append((col, (sql, params), alias)) ret.append((col, (sql, params), alias))
return ret, klass_info, annotations return ret, klass_info, annotations
@ -588,7 +594,9 @@ class SQLCompiler:
""" """
refcounts_before = self.query.alias_refcount.copy() refcounts_before = self.query.alias_refcount.copy()
try: try:
extra_select, order_by, group_by = self.pre_sql_setup() extra_select, order_by, group_by = self.pre_sql_setup(
with_col_aliases=with_col_aliases,
)
for_update_part = None for_update_part = None
# Is a LIMIT/OFFSET clause needed? # Is a LIMIT/OFFSET clause needed?
with_limit_offset = with_limits and ( with_limit_offset = with_limits and (
@ -635,19 +643,12 @@ class SQLCompiler:
params += distinct_params params += distinct_params
out_cols = [] out_cols = []
col_idx = 1
for _, (s_sql, s_params), alias in self.select + extra_select: for _, (s_sql, s_params), alias in self.select + extra_select:
if alias: if alias:
s_sql = "%s AS %s" % ( s_sql = "%s AS %s" % (
s_sql, s_sql,
self.connection.ops.quote_name(alias), self.connection.ops.quote_name(alias),
) )
elif with_col_aliases:
s_sql = "%s AS %s" % (
s_sql,
self.connection.ops.quote_name("col%d" % col_idx),
)
col_idx += 1
params.extend(s_params) params.extend(s_params)
out_cols.append(s_sql) out_cols.append(s_sql)
@ -768,8 +769,6 @@ class SQLCompiler:
sub_selects = [] sub_selects = []
sub_params = [] sub_params = []
for index, (select, _, alias) in enumerate(self.select, start=1): for index, (select, _, alias) in enumerate(self.select, start=1):
if not alias and with_col_aliases:
alias = "col%d" % index
if alias: if alias:
sub_selects.append( sub_selects.append(
"%s.%s" "%s.%s"