1use ahash::AHashMap;
2use itertools::{Itertools, enumerate};
3use sqruff_lib_core::dialects::syntax::{SyntaxKind, SyntaxSet};
4use sqruff_lib_core::lint_fix::LintFix;
5use sqruff_lib_core::parser::segments::{ErasedSegment, SegmentBuilder, Tables};
6use sqruff_lib_core::utils::functional::segments::Segments;
7
8use crate::core::config::Value;
9use crate::core::rules::context::RuleContext;
10use crate::core::rules::crawlers::{Crawler, SegmentSeekerCrawler};
11use crate::core::rules::{Erased, ErasedRule, LintResult, Rule, RuleGroups};
12use crate::utils::functional::context::FunctionalContext;
13
14struct SelectTargetsInfo {
15 select_idx: Option<usize>,
16 first_new_line_idx: Option<usize>,
17 first_select_target_idx: Option<usize>,
18
19 #[allow(dead_code)]
20 first_whitespace_idx: Option<usize>,
21 comment_after_select_idx: Option<usize>,
22 select_targets: Segments,
23 from_segment: Option<ErasedSegment>,
24 pre_from_whitespace: Segments,
25}
26
27#[derive(Debug, Clone)]
28pub struct RuleLT09 {
29 wildcard_policy: String,
30}
31
32impl Rule for RuleLT09 {
33 fn load_from_config(&self, _config: &AHashMap<String, Value>) -> Result<ErasedRule, String> {
34 Ok(RuleLT09 {
35 wildcard_policy: _config["wildcard_policy"].as_string().unwrap().to_owned(),
36 }
37 .erased())
38 }
39 fn name(&self) -> &'static str {
40 "layout.select_targets"
41 }
42
43 fn description(&self) -> &'static str {
44 "Select targets should be on a new line unless there is only one select target."
45 }
46
47 fn long_description(&self) -> &'static str {
48 r#"
49**Anti-pattern**
50
51Multiple select targets on the same line.
52
53```sql
54select a, b
55from foo;
56
57-- Single select target on its own line.
58
59SELECT
60 a
61FROM foo;
62```
63
64**Best practice**
65
66Multiple select targets each on their own line.
67
68```sql
69select
70 a,
71 b
72from foo;
73
74-- Single select target on the same line as the ``SELECT``
75-- keyword.
76
77SELECT a
78FROM foo;
79
80-- When select targets span multiple lines, however they
81-- can still be on a new line.
82
83SELECT
84 SUM(
85 1 + SUM(
86 2 + 3
87 )
88 ) AS col
89FROM test_table;
90```
91"#
92 }
93
94 fn groups(&self) -> &'static [RuleGroups] {
95 &[RuleGroups::All, RuleGroups::Layout]
96 }
97
98 fn eval(&self, context: &RuleContext) -> Vec<LintResult> {
99 let select_targets_info = Self::get_indexes(context);
100 let select_clause = FunctionalContext::new(context).segment();
101
102 let wildcards = select_clause
103 .children_where(|sp| sp.is_type(SyntaxKind::SelectClauseElement))
104 .children_where(|sp| sp.is_type(SyntaxKind::WildcardExpression));
105
106 let has_wildcard = !wildcards.is_empty();
107
108 if select_targets_info.select_targets.len() == 1
109 && (!has_wildcard || self.wildcard_policy == "single")
110 {
111 return self.eval_single_select_target_element(select_targets_info, context);
112 } else if !select_targets_info.select_targets.is_empty() {
113 return self.eval_multiple_select_target_elements(
114 context.tables,
115 select_targets_info,
116 context.segment.clone(),
117 );
118 }
119
120 Vec::new()
121 }
122
123 fn is_fix_compatible(&self) -> bool {
124 true
125 }
126
127 fn crawl_behaviour(&self) -> Crawler {
128 SegmentSeekerCrawler::new(const { SyntaxSet::new(&[SyntaxKind::SelectClause]) }).into()
129 }
130}
131
132impl RuleLT09 {
133 fn get_indexes(context: &RuleContext) -> SelectTargetsInfo {
134 let children = FunctionalContext::new(context).segment().children_all();
135
136 let select_targets = children
137 .filter(|segment: &ErasedSegment| segment.is_type(SyntaxKind::SelectClauseElement));
138
139 let first_select_target_idx = select_targets.first().and_then(|it| children.find(it));
140
141 let selects = children.filter(|segment: &ErasedSegment| segment.is_keyword("select"));
142
143 let select_idx =
144 (!selects.is_empty()).then(|| children.find(selects.first().unwrap()).unwrap());
145
146 let newlines = children.filter(|it: &ErasedSegment| it.is_type(SyntaxKind::Newline));
147
148 let first_new_line_idx =
149 (!newlines.is_empty()).then(|| children.find(newlines.first().unwrap()).unwrap());
150 let mut comment_after_select_idx = None;
151
152 if !newlines.is_empty() {
153 let select_head = selects.first().unwrap();
154 if let Some(first_comment) = children
155 .iter_after_while(select_head, |seg| {
156 seg.is_type(SyntaxKind::Comment)
157 | seg.is_type(SyntaxKind::Whitespace)
158 | seg.is_meta()
159 })
160 .find(|seg| seg.is_type(SyntaxKind::Comment))
161 {
162 comment_after_select_idx = children.find(first_comment);
163 }
164 }
165
166 let mut first_whitespace_idx = None;
167 if let Some(first_new_line_idx) = first_new_line_idx {
168 let segments_after_first_line = children
169 .after(&children[first_new_line_idx])
170 .filter(|seg: &ErasedSegment| seg.is_type(SyntaxKind::Whitespace));
171
172 if !segments_after_first_line.is_empty() {
173 first_whitespace_idx =
174 children.find(&segments_after_first_line.get(0, None).unwrap());
175 }
176 }
177
178 let siblings_post = FunctionalContext::new(context).siblings_post();
179 let from_segment = siblings_post
180 .find_first_where(|seg: &ErasedSegment| seg.is_type(SyntaxKind::FromClause))
181 .head()
182 .get(0, None);
183 let pre_from_whitespace = {
184 let range = if let Some(ref stop) = from_segment {
185 siblings_post.before(stop)
186 } else {
187 siblings_post.clone()
188 };
189 range.filter(|seg: &ErasedSegment| seg.is_type(SyntaxKind::Whitespace))
190 };
191
192 SelectTargetsInfo {
193 select_idx,
194 first_new_line_idx,
195 first_select_target_idx,
196 first_whitespace_idx,
197 comment_after_select_idx,
198 select_targets,
199 from_segment,
200 pre_from_whitespace,
201 }
202 }
203
204 fn eval_multiple_select_target_elements(
205 &self,
206 tables: &Tables,
207 select_targets_info: SelectTargetsInfo,
208 segment: ErasedSegment,
209 ) -> Vec<LintResult> {
210 let mut fixes = Vec::new();
211
212 for (i, select_target) in enumerate(select_targets_info.select_targets.iter()) {
213 let base_segment = if i == 0 {
214 segment.clone()
215 } else {
216 select_targets_info.select_targets[i - 1].clone()
217 };
218
219 if let Some((_, _)) = base_segment
220 .get_position_marker()
221 .zip(select_target.get_position_marker())
222 .filter(|(a, b)| a.working_line_no == b.working_line_no)
223 {
224 let mut start_seg = select_targets_info.select_idx.unwrap();
225 let modifier =
226 segment.child(const { &SyntaxSet::new(&[SyntaxKind::SelectClauseModifier]) });
227
228 if let Some(modifier) = modifier {
229 start_seg = segment
230 .segments()
231 .iter()
232 .position(|it| it == &modifier)
233 .unwrap();
234 }
235
236 let segments = segment.segments();
237
238 let start = if i == 0 {
239 &segments[start_seg]
240 } else {
241 &select_targets_info.select_targets[i - 1]
242 };
243
244 let start_position = segments.iter().position(|it| it == start).unwrap();
245 let ws_to_delete = segments[start_position + 1..]
246 .iter()
247 .take_while(|it| {
248 it.is_type(SyntaxKind::Whitespace)
249 | it.is_type(SyntaxKind::Comma)
250 | it.is_meta()
251 })
252 .filter(|it| it.is_type(SyntaxKind::Whitespace));
253
254 fixes.extend(ws_to_delete.cloned().map(LintFix::delete));
255 fixes.push(LintFix::create_before(
256 select_target.clone(),
257 vec![SegmentBuilder::newline(tables.next_id(), "\n")],
258 ));
259 }
260
261 if let Some(from_segment) = &select_targets_info.from_segment
262 && i + 1 == select_targets_info.select_targets.len()
263 && select_target.get_position_marker().unwrap().working_line_no
264 == from_segment.get_position_marker().unwrap().working_line_no
265 {
266 fixes.extend(
267 select_targets_info
268 .pre_from_whitespace
269 .clone()
270 .into_iter()
271 .map(LintFix::delete),
272 );
273
274 fixes.push(LintFix::create_before(
275 from_segment.clone(),
276 vec![SegmentBuilder::newline(tables.next_id(), "\n")],
277 ));
278 }
279 }
280
281 if !fixes.is_empty() {
282 return vec![LintResult::new(segment.into(), fixes, None, None)];
283 }
284
285 Vec::new()
286 }
287
288 fn eval_single_select_target_element(
289 &self,
290 select_targets_info: SelectTargetsInfo,
291 context: &RuleContext,
292 ) -> Vec<LintResult> {
293 let select_clause = FunctionalContext::new(context).segment();
294 let parent_stack = &context.parent_stack;
295
296 if !(select_targets_info.select_idx < select_targets_info.first_new_line_idx
297 && select_targets_info.first_new_line_idx < select_targets_info.first_select_target_idx)
298 {
299 return Vec::new();
300 }
301
302 let select_children = select_clause.children_all();
303 let mut modifier = select_children
304 .find_first_where(|seg: &ErasedSegment| seg.is_type(SyntaxKind::SelectClauseModifier));
305
306 if select_children[select_targets_info.first_select_target_idx.unwrap()]
307 .descendant_type_set()
308 .contains(SyntaxKind::Newline)
309 {
310 return Vec::new();
311 }
312
313 let mut insert_buff = vec![
314 SegmentBuilder::whitespace(context.tables.next_id(), " "),
315 select_children[select_targets_info.first_select_target_idx.unwrap()].clone(),
316 ];
317
318 if !modifier.is_empty()
319 && select_children.index(&modifier.get(0, None).unwrap())
320 < select_targets_info.first_new_line_idx
321 {
322 modifier = Segments::from_vec(Vec::new(), None);
323 }
324
325 let mut fixes = vec![LintFix::delete(
326 select_children[select_targets_info.first_select_target_idx.unwrap()].clone(),
327 )];
328 let target_idx = select_targets_info.first_select_target_idx.unwrap();
329 if target_idx > 0 && select_children[target_idx - 1].is_type(SyntaxKind::Whitespace) {
330 fixes.push(LintFix::delete(select_children[target_idx - 1].clone()));
331 }
332
333 let start_idx = if !modifier.is_empty() {
334 let buff = std::mem::take(&mut insert_buff);
335
336 insert_buff = vec![
337 SegmentBuilder::whitespace(context.tables.next_id(), " "),
338 modifier[0].clone(),
339 ];
340
341 insert_buff.extend(buff);
342
343 let modifier_idx = select_children
344 .index(&modifier.get(0, None).unwrap())
345 .unwrap();
346
347 if select_children.len() > modifier_idx + 1
348 && select_children[modifier_idx + 2].is_whitespace()
349 {
350 fixes.push(LintFix::delete(select_children[modifier_idx + 2].clone()));
351 }
352
353 fixes.push(LintFix::delete(modifier[0].clone()));
354
355 modifier_idx
356 } else {
357 select_targets_info.first_select_target_idx.unwrap()
358 };
359
360 if !parent_stack.is_empty()
361 && parent_stack
362 .last()
363 .unwrap()
364 .is_type(SyntaxKind::SelectStatement)
365 {
366 let select_stmt = parent_stack.last().unwrap();
367 let select_clause_idx = select_stmt
368 .segments()
369 .iter()
370 .position(|it| it.clone() == select_clause.get(0, None).unwrap())
371 .unwrap();
372 let after_select_clause_idx = select_clause_idx + 1;
373
374 let fixes_for_move_after_select_clause =
375 |fixes: &mut Vec<LintFix>,
376 stop_seg: ErasedSegment,
377 delete_segments: Option<Segments>,
378 add_newline: bool| {
379 let start_seg = if !modifier.is_empty() {
380 modifier[0].clone()
381 } else {
382 select_children[select_targets_info.first_new_line_idx.unwrap()].clone()
383 };
384
385 let move_after_select_clause =
386 select_children.between_exclusive(&start_seg, &stop_seg);
387 let mut local_fixes = Vec::new();
388 let mut all_deletes = fixes
389 .iter()
390 .filter(|fix| matches!(fix, LintFix::Delete { .. }))
391 .map(|fix| fix.anchor().clone())
392 .collect_vec();
393 for seg in delete_segments.unwrap_or_default() {
394 fixes.push(LintFix::delete(seg.clone()));
395 all_deletes.push(seg);
396 }
397
398 let new_fixes = move_after_select_clause
399 .iter()
400 .filter(|it| !all_deletes.contains(it))
401 .cloned()
402 .map(LintFix::delete);
403 local_fixes.extend(new_fixes);
404
405 if !move_after_select_clause.is_empty() || add_newline {
406 local_fixes.push(LintFix::create_after(
407 select_clause[0].clone(),
408 if add_newline {
409 vec![SegmentBuilder::newline(context.tables.next_id(), "\n")]
410 } else {
411 vec![]
412 }
413 .into_iter()
414 .chain(move_after_select_clause)
415 .collect_vec(),
416 None,
417 ));
418 }
419
420 local_fixes
421 };
422
423 if select_stmt.segments().len() > after_select_clause_idx {
424 if select_stmt.segments()[after_select_clause_idx].is_type(SyntaxKind::Newline) {
425 let to_delete = select_children
426 .reversed()
427 .after(&select_children[start_idx])
428 .take_while(|seg| seg.is_type(SyntaxKind::Whitespace));
429
430 if !to_delete.is_empty() {
431 let delete_last_newline = select_children[start_idx - to_delete.len() - 1]
432 .is_type(SyntaxKind::Newline);
433
434 if delete_last_newline {
435 fixes.push(LintFix::delete(
436 select_stmt.segments()[after_select_clause_idx].clone(),
437 ));
438 }
439
440 let new_fixes = fixes_for_move_after_select_clause(
441 &mut fixes,
442 to_delete.last().unwrap().clone(),
443 to_delete.into(),
444 true,
445 );
446 fixes.extend(new_fixes);
447 }
448 } else if select_stmt.segments()[after_select_clause_idx]
449 .is_type(SyntaxKind::Whitespace)
450 {
451 fixes.push(LintFix::delete(
452 select_stmt.segments()[after_select_clause_idx].clone(),
453 ));
454
455 let new_fixes = fixes_for_move_after_select_clause(
456 &mut fixes,
457 select_children[select_targets_info.first_select_target_idx.unwrap()]
458 .clone(),
459 None,
460 true,
461 );
462
463 fixes.extend(new_fixes);
464 } else if select_stmt.segments()[after_select_clause_idx]
465 .is_type(SyntaxKind::Dedent)
466 {
467 let start_seg = if select_clause_idx == 0 {
468 select_children.last().unwrap()
469 } else {
470 &select_children[select_clause_idx - 1]
471 };
472
473 let to_delete = select_children
474 .reversed()
475 .after(start_seg)
476 .take_while(|it| it.is_type(SyntaxKind::Whitespace));
477
478 if !to_delete.is_empty() {
479 let add_newline =
480 to_delete.iter().any(|it| it.is_type(SyntaxKind::Newline));
481 let local_fixes = fixes_for_move_after_select_clause(
482 &mut fixes,
483 to_delete.last().unwrap().clone(),
484 to_delete.into(),
485 add_newline,
486 );
487 fixes.extend(local_fixes);
488 }
489 } else {
490 let local_fixes = fixes_for_move_after_select_clause(
491 &mut fixes,
492 select_children[select_targets_info.first_select_target_idx.unwrap()]
493 .clone(),
494 None,
495 true,
496 );
497 fixes.extend(local_fixes);
498 }
499 }
500 }
501
502 if select_targets_info.comment_after_select_idx.is_none() {
503 fixes.push(LintFix::replace(
504 select_children[select_targets_info.first_new_line_idx.unwrap()].clone(),
505 insert_buff,
506 None,
507 ));
508 }
509
510 vec![LintResult::new(
511 select_clause.get(0, None).unwrap().clone().into(),
512 fixes,
513 None,
514 None,
515 )]
516 }
517}
518
519impl Default for RuleLT09 {
520 fn default() -> Self {
521 Self {
522 wildcard_policy: "single".into(),
523 }
524 }
525}