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