1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
|
# frozen_string_literal: true
module Gitlab
module Pagination
module Keyset
module InOperatorOptimization
# rubocop: disable CodeReuse/ActiveRecord
class QueryBuilder
UnsupportedScopeOrder = Class.new(StandardError)
RECURSIVE_CTE_NAME = 'recursive_keyset_cte'
RECORDS_COLUMN = 'records'
# This class optimizes slow database queries (PostgreSQL specific) where the
# IN SQL operator is used with sorting.
#
# Arguments:
# scope - ActiveRecord::Relation supporting keyset pagination
# array_scope - ActiveRecord::Relation for the `IN` subselect
# array_mapping_scope - Lambda for connecting scope with array_scope
# finder_query - ActiveRecord::Relation for finding one row by the passed in cursor values
# values - keyset cursor values (optional)
#
# Example ActiveRecord query: Issues in the namespace hierarchy
# > scope = Issue
# > .where(project_id: Group.find(9970).all_projects.select(:id))
# > .order(:created_at, :id)
# > .limit(20);
#
# Optimized version:
#
# > scope = Issue.where({}).order(:created_at, :id) # base scope
# > array_scope = Group.find(9970).all_projects.select(:id)
# > array_mapping_scope = -> (id_expression) { Issue.where(Issue.arel_table[:project_id].eq(id_expression)) }
#
# # finding the record by id is good enough, we can ignore the created_at_expression
# > finder_query = -> (created_at_expression, id_expression) { Issue.where(Issue.arel_table[:id].eq(id_expression)) }
#
# > Gitlab::Pagination::Keyset::InOperatorOptimization::QueryBuilder.new(
# > scope: scope,
# > array_scope: array_scope,
# > array_mapping_scope: array_mapping_scope,
# > finder_query: finder_query
# > ).execute.limit(20)
def initialize(scope:, array_scope:, array_mapping_scope:, finder_query:, values: {})
@scope, success = Gitlab::Pagination::Keyset::SimpleOrderBuilder.build(scope)
unless success
error_message = <<~MSG
The order on the scope does not support keyset pagination. You might need to define a custom Order object.\n
See https://docs.gitlab.com/ee/development/database/keyset_pagination.html#complex-order-configuration\n
Or the Gitlab::Pagination::Keyset::Order class for examples
MSG
raise(UnsupportedScopeOrder, error_message)
end
@order = Gitlab::Pagination::Keyset::Order.extract_keyset_order_object(scope)
@array_scope = array_scope
@array_mapping_scope = array_mapping_scope
@finder_query = finder_query
@values = values
@model = @scope.model
@table_name = @model.table_name
@arel_table = @model.arel_table
end
def execute
selector_cte = Gitlab::SQL::CTE.new(:array_cte, array_scope)
cte = Gitlab::SQL::RecursiveCTE.new(RECURSIVE_CTE_NAME, union_args: { remove_duplicates: false, remove_order: false })
cte << initializer_query
cte << data_collector_query
q = cte
.apply_to(model.where({})
.with(selector_cte.to_arel))
.select(result_collector_final_projections)
.where("count <> 0") # filter out the initializer row
model.from(q.arel.as(table_name))
end
private
attr_reader :array_scope, :scope, :order, :array_mapping_scope, :finder_query, :values, :model, :table_name, :arel_table
def initializer_query
array_column_names = array_scope_columns.array_aggregated_column_names + order_by_columns.array_aggregated_column_names
projections = [
*result_collector_initializer_columns,
*array_column_names,
'0::bigint AS count'
]
model.select(projections).from(build_column_arrays_query).limit(1)
end
# This query finds the first cursor values for each item in the array CTE.
#
# array_cte:
#
# |project_id|
# |----------|
# | 1|
# | 2|
# | 3|
# | 4|
#
# For each project_id, find the first issues row by respecting the created_at, id order.
#
# The `array_mapping_scope` parameter defines how the `array_scope` and the `scope` can be combined.
#
# scope = Issue.where({}) # empty scope
# array_mapping_scope = Issue.where(project_id: X)
#
# scope.merge(array_mapping_scope) # Issue.where(project_id: X)
#
# X will be replaced with a value from the `array_cte` temporary table.
#
# |created_at|id|
# |----------|--|
# |2020-01-15| 2|
# |2020-01-07| 3|
# |2020-01-07| 4|
# |2020-01-10| 5|
def build_column_arrays_query
q = Arel::SelectManager.new
.project(array_scope_columns.array_aggregated_columns + order_by_columns.array_aggregated_columns)
.from(array_cte)
.join(Arel.sql("LEFT JOIN LATERAL (#{initial_keyset_query.to_sql}) #{table_name} ON TRUE"))
order_by_columns.each { |column| q.where(column.arel_column.not_eq(nil)) }
q.as('array_scope_lateral_query')
end
def array_cte
Arel::SelectManager.new
.project(array_scope_columns.arel_columns)
.from(Arel.sql(array_scope_columns.array_scope_cte_name))
.as(array_scope_columns.array_scope_cte_name)
end
def initial_keyset_query
keyset_scope = scope.merge(array_mapping_scope.call(*array_scope_columns.arel_columns))
order
.apply_cursor_conditions(keyset_scope, values, use_union_optimization: true)
.reselect(*order_by_columns.arel_columns)
.limit(1)
end
def data_collector_query
array_column_list = array_scope_columns.array_aggregated_column_names
order_column_value_arrays = order_by_columns.replace_value_in_array_by_position_expressions
select = [
*result_collector_columns,
*array_column_list,
*order_column_value_arrays,
"#{RECURSIVE_CTE_NAME}.count + 1"
]
from = <<~SQL
#{RECURSIVE_CTE_NAME},
#{array_order_query.lateral.as('position_query').to_sql},
#{ensure_one_row(next_cursor_values_query).lateral.as('next_cursor_values').to_sql}
SQL
model.select(select).from(from)
end
# NULL guard. This method ensures that NULL values are returned when the passed in scope returns 0 rows.
# Example query: returns issues.id or NULL
#
# SELECT issues.id FROM (VALUES (NULL)) nulls (id)
# LEFT JOIN (SELECT id FROM issues WHERE id = 1 LIMIT 1) issues ON TRUE
# LIMIT 1
def ensure_one_row(query)
q = Arel::SelectManager.new
q.projections = order_by_columns.original_column_names_as_tmp_tamble
null_values = [nil] * order_by_columns.count
from = Arel::Nodes::Grouping.new(Arel::Nodes::ValuesList.new([null_values])).as('nulls')
q.from(from)
q.join(Arel.sql("LEFT JOIN (#{query.to_sql}) record ON TRUE"))
q.limit = 1
q
end
# This subquery finds the cursor values for the next record by sorting the generated cursor arrays in memory and taking the first element.
# It combines the cursor arrays (UNNEST) together and sorts them according to the originally defined ORDER BY clause.
#
# Example: issues in the group hierarchy with ORDER BY created_at, id
#
# |project_id| |created_at|id| # 2 arrays combined: issues_created_at_array, issues_id_array
# |----------| |----------|--|
# | 1| |2020-01-15| 2|
# | 2| |2020-01-07| 3|
# | 3| |2020-01-07| 4|
# | 4| |2020-01-10| 5|
#
# The query will return the cursor values: (2020-01-07, 3) and the array position: 1
# From the position, we can tell that the record belongs to the project with id 2.
def array_order_query
q = Arel::SelectManager.new
.project([*order_by_columns.original_column_names_as_arel_string, Arel.sql('position')])
.from("UNNEST(#{list(order_by_columns.array_aggregated_column_names)}) WITH ORDINALITY AS u(#{list(order_by_columns.original_column_names)}, position)")
order_by_columns.each { |column| q.where(Arel.sql(column.original_column_name).not_eq(nil)) } # ignore rows where all columns are NULL
q.order(Arel.sql(order_by_without_table_references)).take(1)
end
# This subquery finds the next cursor values after the previously determined position (from array_order_query).
# The current cursor values are passed in as SQL literals since the actual values are encoded into SQL arrays.
#
# Example: issues in the group hierarchy with ORDER BY created_at, id
#
# |project_id| |created_at|id| # 2 arrays combined: issues_created_at_array, issues_id_array
# |----------| |----------|--|
# | 1| |2020-01-15| 2|
# | 2| |2020-01-07| 3|
# | 3| |2020-01-07| 4|
# | 4| |2020-01-10| 5|
#
# Assuming that the determined position is 1, the cursor values will be the following:
# - Filter: project_id = 2
# - created_at = 2020-01-07
# - id = 3
def next_cursor_values_query
cursor_values = order_by_columns.cursor_values(RECURSIVE_CTE_NAME)
array_mapping_scope_columns = array_scope_columns.array_lookup_expressions_by_position(RECURSIVE_CTE_NAME)
keyset_scope = scope
.reselect(*order_by_columns.arel_columns)
.merge(array_mapping_scope.call(*array_mapping_scope_columns))
order
.apply_cursor_conditions(keyset_scope, cursor_values, use_union_optimization: true)
.reselect(*order_by_columns.arel_columns)
.limit(1)
end
# Generates an ORDER BY clause by using the column position index and the original order clauses.
# This method is used to sort the collected arrays in SQL.
# Example: "issues".created_at DESC , "issues".id ASC => 1 DESC, 2 ASC
def order_by_without_table_references
order.column_definitions.each_with_index.map do |column_definition, i|
"#{i + 1} #{column_definition.order_direction_as_sql_string}"
end.join(", ")
end
def result_collector_initializer_columns
["NULL::#{table_name} AS #{RECORDS_COLUMN}"]
end
def result_collector_columns
query = finder_query
.call(*order_by_columns.array_lookup_expressions_by_position(RECURSIVE_CTE_NAME))
.select("#{table_name}")
.limit(1)
["(#{query.to_sql})"]
end
def result_collector_final_projections
["(#{RECORDS_COLUMN}).*"]
end
def array_scope_columns
@array_scope_columns ||= ArrayScopeColumns.new(array_scope.select_values)
end
def order_by_columns
@order_by_columns ||= OrderByColumns.new(order.column_definitions, arel_table)
end
def list(array)
array.join(', ')
end
end
# rubocop: enable CodeReuse/ActiveRecord
end
end
end
end
|