Skip to content

Commit 4260c95

Browse files
authored
ide: cte improvements, subqueries, values, union select (#787)
1 parent 2b64421 commit 4260c95

File tree

5 files changed

+443
-11
lines changed

5 files changed

+443
-11
lines changed

crates/squawk_ide/src/goto_definition.rs

Lines changed: 227 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1423,6 +1423,44 @@ select a$0 from t;
14231423
");
14241424
}
14251425

1426+
#[test]
1427+
fn goto_cte_with_partial_column_list() {
1428+
assert_snapshot!(goto("
1429+
with t(x) as (select 1 as a, 2 as b)
1430+
select b$0 from t;
1431+
"), @r"
1432+
╭▸
1433+
2 │ with t(x) as (select 1 as a, 2 as b)
1434+
│ ─ 2. destination
1435+
3 │ select b from t;
1436+
╰╴ ─ 1. source
1437+
");
1438+
}
1439+
1440+
#[test]
1441+
fn goto_cte_with_partial_column_list_renamed() {
1442+
assert_snapshot!(goto("
1443+
with t(x) as (select 1 as a, 2 as b)
1444+
select x$0 from t;
1445+
"), @r"
1446+
╭▸
1447+
2 │ with t(x) as (select 1 as a, 2 as b)
1448+
│ ─ 2. destination
1449+
3 │ select x from t;
1450+
╰╴ ─ 1. source
1451+
");
1452+
}
1453+
1454+
#[test]
1455+
fn goto_cte_column_list_overwrites_column() {
1456+
goto_not_found(
1457+
"
1458+
with t(x) as (select 1 as a)
1459+
select a$0 from t;
1460+
",
1461+
);
1462+
}
1463+
14261464
#[test]
14271465
fn goto_cte_shadows_table() {
14281466
assert_snapshot!(goto("
@@ -1438,6 +1476,50 @@ select a from t;
14381476
");
14391477
}
14401478

1479+
#[test]
1480+
fn goto_subquery_column() {
1481+
assert_snapshot!(goto("
1482+
select a$0 from (select 1 a);
1483+
"), @r"
1484+
╭▸
1485+
2 │ select a from (select 1 a);
1486+
╰╴ ─ 1. source ─ 2. destination
1487+
");
1488+
}
1489+
1490+
#[test]
1491+
fn goto_subquery_column_with_as() {
1492+
assert_snapshot!(goto("
1493+
select a$0 from (select 1 as a);
1494+
"), @r"
1495+
╭▸
1496+
2 │ select a from (select 1 as a);
1497+
╰╴ ─ 1. source ─ 2. destination
1498+
");
1499+
}
1500+
1501+
#[test]
1502+
fn goto_subquery_column_multiple_columns() {
1503+
assert_snapshot!(goto("
1504+
select b$0 from (select 1 a, 2 b);
1505+
"), @r"
1506+
╭▸
1507+
2 │ select b from (select 1 a, 2 b);
1508+
╰╴ ─ 1. source ─ 2. destination
1509+
");
1510+
}
1511+
1512+
#[test]
1513+
fn goto_subquery_column_nested_parens() {
1514+
assert_snapshot!(goto("
1515+
select a$0 from ((select 1 a));
1516+
"), @r"
1517+
╭▸
1518+
2 │ select a from ((select 1 a));
1519+
╰╴ ─ 1. source ─ 2. destination
1520+
");
1521+
}
1522+
14411523
#[test]
14421524
fn goto_insert_table() {
14431525
assert_snapshot!(goto("
@@ -2046,6 +2128,94 @@ select foo.t.a$0 from t;
20462128
");
20472129
}
20482130

2131+
#[test]
2132+
fn goto_cte_values_column1() {
2133+
assert_snapshot!(goto("
2134+
with t as (
2135+
values (1, 2), (3, 4)
2136+
)
2137+
select column1$0, column2 from t;
2138+
"), @r"
2139+
╭▸
2140+
3 │ values (1, 2), (3, 4)
2141+
│ ─ 2. destination
2142+
4 │ )
2143+
5 │ select column1, column2 from t;
2144+
╰╴ ─ 1. source
2145+
");
2146+
}
2147+
2148+
#[test]
2149+
fn goto_cte_values_column2() {
2150+
assert_snapshot!(goto("
2151+
with t as (
2152+
values (1, 2), (3, 4)
2153+
)
2154+
select column1, column2$0 from t;
2155+
"), @r"
2156+
╭▸
2157+
3 │ values (1, 2), (3, 4)
2158+
│ ─ 2. destination
2159+
4 │ )
2160+
5 │ select column1, column2 from t;
2161+
╰╴ ─ 1. source
2162+
");
2163+
}
2164+
2165+
#[test]
2166+
fn goto_cte_values_single_column() {
2167+
assert_snapshot!(goto("
2168+
with t as (
2169+
values (1), (2), (3)
2170+
)
2171+
select column1$0 from t;
2172+
"), @r"
2173+
╭▸
2174+
3 │ values (1), (2), (3)
2175+
│ ─ 2. destination
2176+
4 │ )
2177+
5 │ select column1 from t;
2178+
╰╴ ─ 1. source
2179+
");
2180+
}
2181+
2182+
#[test]
2183+
fn goto_cte_values_multiple_rows() {
2184+
assert_snapshot!(goto("
2185+
with t as (
2186+
values
2187+
(1, 2, 3),
2188+
(4, 5, 6),
2189+
(7, 8, 9)
2190+
)
2191+
select column3$0 from t;
2192+
"), @r"
2193+
╭▸
2194+
4 │ (1, 2, 3),
2195+
│ ─ 2. destination
2196+
2197+
8 │ select column3 from t;
2198+
╰╴ ─ 1. source
2199+
");
2200+
}
2201+
2202+
#[test]
2203+
fn goto_cte_values_uppercase_column_names() {
2204+
assert_snapshot!(goto("
2205+
with t as (
2206+
values (1, 2), (3, 4)
2207+
)
2208+
select COLUMN1$0, COLUMN2 from t;
2209+
"), @r"
2210+
╭▸
2211+
3 │ values (1, 2), (3, 4)
2212+
│ ─ 2. destination
2213+
4 │ )
2214+
5 │ select COLUMN1, COLUMN2 from t;
2215+
╰╴ ─ 1. source
2216+
");
2217+
}
2218+
20492219
#[test]
20502220
fn goto_qualified_column_with_schema_in_from_table() {
20512221
assert_snapshot!(goto("
@@ -2073,4 +2243,61 @@ select t.a$0 from foo.t;
20732243
╰╴ ─ 1. source
20742244
");
20752245
}
2246+
2247+
#[test]
2248+
fn goto_cte_union_all_column() {
2249+
assert_snapshot!(goto("
2250+
with t as (
2251+
select 1 as a, 2 as b
2252+
union all
2253+
select 3, 4
2254+
)
2255+
select a$0, b from t;
2256+
"), @r"
2257+
╭▸
2258+
3 │ select 1 as a, 2 as b
2259+
│ ─ 2. destination
2260+
2261+
7 │ select a, b from t;
2262+
╰╴ ─ 1. source
2263+
");
2264+
}
2265+
2266+
#[test]
2267+
fn goto_cte_union_all_column_second() {
2268+
assert_snapshot!(goto("
2269+
with t as (
2270+
select 1 as a, 2 as b
2271+
union all
2272+
select 3, 4
2273+
)
2274+
select a, b$0 from t;
2275+
"), @r"
2276+
╭▸
2277+
3 │ select 1 as a, 2 as b
2278+
│ ─ 2. destination
2279+
2280+
7 │ select a, b from t;
2281+
╰╴ ─ 1. source
2282+
");
2283+
}
2284+
2285+
#[test]
2286+
fn goto_cte_union_column() {
2287+
assert_snapshot!(goto("
2288+
with t as (
2289+
select 1 as a, 2 as b
2290+
union
2291+
select 3, 4
2292+
)
2293+
select a$0 from t;
2294+
"), @r"
2295+
╭▸
2296+
3 │ select 1 as a, 2 as b
2297+
│ ─ 2. destination
2298+
2299+
7 │ select a from t;
2300+
╰╴ ─ 1. source
2301+
");
2302+
}
20762303
}

crates/squawk_ide/src/hover.rs

Lines changed: 86 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1-
use crate::binder;
21
use crate::offsets::token_from_offset;
32
use crate::resolve;
3+
use crate::{binder, symbols::Name};
44
use rowan::TextSize;
55
use squawk_syntax::ast::{self, AstNode};
66

@@ -69,6 +69,10 @@ pub fn hover(file: &ast::SourceFile, offset: TextSize) -> Option<String> {
6969
return format_create_table(&create_table, &binder);
7070
}
7171

72+
if let Some(with_table) = name.syntax().parent().and_then(ast::WithTable::cast) {
73+
return format_with_table(&with_table);
74+
}
75+
7276
if let Some(create_index) = name.syntax().ancestors().find_map(ast::CreateIndex::cast) {
7377
return format_create_index(&create_index, &binder);
7478
}
@@ -110,7 +114,14 @@ fn hover_column(
110114

111115
if let Some(with_table) = column_name_node.ancestors().find_map(ast::WithTable::cast) {
112116
let cte_name = with_table.name()?.syntax().text().to_string();
113-
let column_name = column_name_node.text().to_string();
117+
let column_name = if column_name_node
118+
.ancestors()
119+
.any(|a| ast::Values::can_cast(a.kind()))
120+
{
121+
Name::new(name_ref.syntax().text().to_string())
122+
} else {
123+
Name::new(column_name_node.text().to_string())
124+
};
114125
return Some(format!("column {}.{}", cte_name, column_name));
115126
}
116127

@@ -1595,4 +1606,77 @@ select a$0 from t;
15951606
╰╴ ─ hover
15961607
");
15971608
}
1609+
1610+
#[test]
1611+
fn hover_on_cte_definition() {
1612+
assert_snapshot!(check_hover("
1613+
with t$0 as (select 1 a)
1614+
select a from t;
1615+
"), @r"
1616+
hover: with t as (select 1 a)
1617+
╭▸
1618+
2 │ with t as (select 1 a)
1619+
╰╴ ─ hover
1620+
");
1621+
}
1622+
1623+
#[test]
1624+
fn hover_on_cte_values_column1() {
1625+
assert_snapshot!(check_hover("
1626+
with t as (
1627+
values (1, 2), (3, 4)
1628+
)
1629+
select column1$0, column2 from t;
1630+
"), @r"
1631+
hover: column t.column1
1632+
╭▸
1633+
5 │ select column1, column2 from t;
1634+
╰╴ ─ hover
1635+
");
1636+
}
1637+
1638+
#[test]
1639+
fn hover_on_cte_values_column2() {
1640+
assert_snapshot!(check_hover("
1641+
with t as (
1642+
values (1, 2), (3, 4)
1643+
)
1644+
select column1, column2$0 from t;
1645+
"), @r"
1646+
hover: column t.column2
1647+
╭▸
1648+
5 │ select column1, column2 from t;
1649+
╰╴ ─ hover
1650+
");
1651+
}
1652+
1653+
#[test]
1654+
fn hover_on_cte_values_single_column() {
1655+
assert_snapshot!(check_hover("
1656+
with t as (
1657+
values (1), (2), (3)
1658+
)
1659+
select column1$0 from t;
1660+
"), @r"
1661+
hover: column t.column1
1662+
╭▸
1663+
5 │ select column1 from t;
1664+
╰╴ ─ hover
1665+
");
1666+
}
1667+
1668+
#[test]
1669+
fn hover_on_cte_values_uppercase_column_names() {
1670+
assert_snapshot!(check_hover("
1671+
with t as (
1672+
values (1, 2), (3, 4)
1673+
)
1674+
select COLUMN1$0, COLUMN2 from t;
1675+
"), @r"
1676+
hover: column t.column1
1677+
╭▸
1678+
5 │ select COLUMN1, COLUMN2 from t;
1679+
╰╴ ─ hover
1680+
");
1681+
}
15981682
}

0 commit comments

Comments
 (0)