sql_insight/extractor/
helper.rs

1use crate::TableReference;
2use std::collections::HashMap;
3
4pub(crate) fn resolve_aliased_tables(
5    possibly_aliased_tables: Vec<TableReference>,
6    original_tables: Vec<TableReference>,
7) -> Vec<TableReference> {
8    possibly_aliased_tables
9        .iter()
10        .map(|possibly_aliased_table| {
11            if possibly_aliased_table.has_qualifiers() || possibly_aliased_table.has_alias() {
12                return possibly_aliased_table.clone();
13            }
14            if let Some(resolved_table) = original_tables.iter().find_map(|original_table| {
15                original_table.alias.as_ref().and_then(|alias| {
16                    if *alias == possibly_aliased_table.name {
17                        Some(original_table.clone())
18                    } else {
19                        None
20                    }
21                })
22            }) {
23                return resolved_table;
24            }
25            possibly_aliased_table.clone()
26        })
27        .collect()
28}
29
30pub(crate) fn calc_difference_of_tables(
31    base_tables: Vec<TableReference>,
32    exclude_tables: Vec<TableReference>,
33) -> Vec<TableReference> {
34    let mut exclude_tables_count = HashMap::new();
35    for exclude_table in exclude_tables.iter() {
36        *exclude_tables_count.entry(exclude_table).or_insert(0) += 1;
37    }
38    base_tables
39        .into_iter()
40        .filter(|base_table| {
41            if let Some(count) = exclude_tables_count.get_mut(base_table) {
42                if *count > 0 {
43                    *count -= 1;
44                    return false;
45                }
46            }
47            true
48        })
49        .collect()
50}
51
52#[cfg(test)]
53mod tests {
54    use super::*;
55    use sqlparser::ast::Ident;
56
57    mod resolve_aliased_tables {
58        use super::*;
59
60        #[test]
61        fn test_single_aliased_table() {
62            let possibly_aliased_tables = vec![TableReference {
63                catalog: None,
64                schema: None,
65                name: Ident::new("t1_alias"),
66                alias: None,
67            }];
68            let original_tables = vec![TableReference {
69                catalog: None,
70                schema: None,
71                name: Ident::new("t1"),
72                alias: Some(Ident::new("t1_alias")),
73            }];
74            let expected_resolved_tables = vec![TableReference {
75                catalog: None,
76                schema: None,
77                name: Ident::new("t1"),
78                alias: Some(Ident::new("t1_alias")),
79            }];
80            let result = resolve_aliased_tables(possibly_aliased_tables, original_tables);
81            assert_eq!(result, expected_resolved_tables);
82        }
83
84        #[test]
85        fn test_multiple_aliased_tables() {
86            let possibly_aliased_tables = vec![
87                TableReference {
88                    catalog: None,
89                    schema: None,
90                    name: Ident::new("t1_alias"),
91                    alias: None,
92                },
93                TableReference {
94                    catalog: None,
95                    schema: None,
96                    name: Ident::new("t2_alias"),
97                    alias: None,
98                },
99            ];
100            let original_tables = vec![
101                TableReference {
102                    catalog: None,
103                    schema: None,
104                    name: Ident::new("t1"),
105                    alias: Some(Ident::new("t1_alias")),
106                },
107                TableReference {
108                    catalog: None,
109                    schema: None,
110                    name: Ident::new("t2"),
111                    alias: Some(Ident::new("t2_alias")),
112                },
113            ];
114            let expected_resolved_tables = vec![
115                TableReference {
116                    catalog: None,
117                    schema: None,
118                    name: Ident::new("t1"),
119                    alias: Some(Ident::new("t1_alias")),
120                },
121                TableReference {
122                    catalog: None,
123                    schema: None,
124                    name: Ident::new("t2"),
125                    alias: Some(Ident::new("t2_alias")),
126                },
127            ];
128            let result = resolve_aliased_tables(possibly_aliased_tables, original_tables);
129            assert_eq!(result, expected_resolved_tables);
130        }
131
132        #[test]
133        fn test_catalog_and_schema_qualified_table_in_original_tables() {
134            let possibly_aliased_tables = vec![
135                TableReference {
136                    catalog: None,
137                    schema: None,
138                    name: Ident::new("t1_alias"),
139                    alias: None,
140                },
141                TableReference {
142                    catalog: None,
143                    schema: None,
144                    name: Ident::new("t2_alias"),
145                    alias: None,
146                },
147            ];
148            let original_tables = vec![
149                TableReference {
150                    catalog: Some(Ident::new("c1")),
151                    schema: Some(Ident::new("s1")),
152                    name: Ident::new("t1"),
153                    alias: Some(Ident::new("t1_alias")),
154                },
155                TableReference {
156                    catalog: None,
157                    schema: Some(Ident::new("s2")),
158                    name: Ident::new("t2"),
159                    alias: Some(Ident::new("t2_alias")),
160                },
161            ];
162            let expected_resolved_tables = vec![
163                TableReference {
164                    catalog: Some(Ident::new("c1")),
165                    schema: Some(Ident::new("s1")),
166                    name: Ident::new("t1"),
167                    alias: Some(Ident::new("t1_alias")),
168                },
169                TableReference {
170                    catalog: None,
171                    schema: Some(Ident::new("s2")),
172                    name: Ident::new("t2"),
173                    alias: Some(Ident::new("t2_alias")),
174                },
175            ];
176            let result = resolve_aliased_tables(possibly_aliased_tables, original_tables);
177            assert_eq!(result, expected_resolved_tables);
178        }
179
180        #[test]
181        fn test_catalog_and_schema_qualified_table_in_possible_aliased_tables() {
182            // qualified alias is not valid syntax in standard SQL,
183            // so qualified tables are not regarded as aliased tables, hence they are not resolved.
184            let possibly_aliased_tables = vec![
185                TableReference {
186                    catalog: Some(Ident::new("c1")),
187                    schema: Some(Ident::new("s1")),
188                    name: Ident::new("t1_alias"),
189                    alias: None,
190                },
191                TableReference {
192                    catalog: None,
193                    schema: Some(Ident::new("s2")),
194                    name: Ident::new("t2_alias"),
195                    alias: None,
196                },
197            ];
198            let original_tables = vec![
199                TableReference {
200                    catalog: Some(Ident::new("c1")),
201                    schema: Some(Ident::new("s1")),
202                    name: Ident::new("t1"),
203                    alias: Some(Ident::new("t1_alias")),
204                },
205                TableReference {
206                    catalog: None,
207                    schema: Some(Ident::new("s2")),
208                    name: Ident::new("t2"),
209                    alias: Some(Ident::new("t2_alias")),
210                },
211            ];
212            let expected_resolved_tables = vec![
213                TableReference {
214                    catalog: Some(Ident::new("c1")),
215                    schema: Some(Ident::new("s1")),
216                    name: Ident::new("t1_alias"),
217                    alias: None,
218                },
219                TableReference {
220                    catalog: None,
221                    schema: Some(Ident::new("s2")),
222                    name: Ident::new("t2_alias"),
223                    alias: None,
224                },
225            ];
226            let result = resolve_aliased_tables(possibly_aliased_tables, original_tables);
227            assert_eq!(result, expected_resolved_tables);
228        }
229    }
230
231    mod calc_difference_of_tables {
232        use super::*;
233
234        #[test]
235        fn test_single_table() {
236            let base_tables = vec![TableReference {
237                catalog: None,
238                schema: None,
239                name: Ident::new("t1"),
240                alias: None,
241            }];
242            let exclude_tables = vec![TableReference {
243                catalog: None,
244                schema: None,
245                name: Ident::new("t1"),
246                alias: None,
247            }];
248            let expected_result = vec![];
249            let result = calc_difference_of_tables(base_tables, exclude_tables);
250            assert_eq!(result, expected_result);
251        }
252
253        #[test]
254        fn test_multiple_unique_tables() {
255            let base_tables = vec![
256                TableReference {
257                    catalog: None,
258                    schema: None,
259                    name: Ident::new("t1"),
260                    alias: None,
261                },
262                TableReference {
263                    catalog: None,
264                    schema: None,
265                    name: Ident::new("t2"),
266                    alias: None,
267                },
268            ];
269            let exclude_tables = vec![
270                TableReference {
271                    catalog: None,
272                    schema: None,
273                    name: Ident::new("t1"),
274                    alias: None,
275                },
276                TableReference {
277                    catalog: None,
278                    schema: None,
279                    name: Ident::new("t2"),
280                    alias: None,
281                },
282            ];
283            let expected_result = vec![];
284            let result = calc_difference_of_tables(base_tables, exclude_tables);
285            assert_eq!(result, expected_result);
286        }
287
288        #[test]
289        fn test_multiple_tables_with_duplicates() {
290            let base_tables = vec![
291                TableReference {
292                    catalog: None,
293                    schema: None,
294                    name: Ident::new("t1"),
295                    alias: None,
296                },
297                TableReference {
298                    catalog: None,
299                    schema: None,
300                    name: Ident::new("t1"),
301                    alias: None,
302                },
303                TableReference {
304                    catalog: None,
305                    schema: None,
306                    name: Ident::new("t2"),
307                    alias: None,
308                },
309            ];
310            let exclude_tables = vec![
311                TableReference {
312                    catalog: None,
313                    schema: None,
314                    name: Ident::new("t1"),
315                    alias: None,
316                },
317                TableReference {
318                    catalog: None,
319                    schema: None,
320                    name: Ident::new("t2"),
321                    alias: None,
322                },
323            ];
324            let expected_result = vec![TableReference {
325                catalog: None,
326                schema: None,
327                name: Ident::new("t1"),
328                alias: None,
329            }];
330            let result = calc_difference_of_tables(base_tables, exclude_tables);
331            assert_eq!(result, expected_result);
332        }
333    }
334}