Skip to content

ide: add merge & returning clause support #825

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kodiakhq merged 2 commits into from
Jan 6, 2026
Merged

ide: add merge & returning clause support #825

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0fed51d
ide: basic returning support
sbdchd Jan 5, 2026
8ec07a8
ide: add merge & returning clause support
sbdchd Jan 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
156 changes: 156 additions & 0 deletions crates/squawk_ide/src/classify.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,12 +47,29 @@ pub(crate) enum NameRefClass {
CompositeTypeField,
InsertTable,
InsertColumn,
InsertQualifiedColumnTable,
DeleteTable,
DeleteWhereColumn,
DeleteQualifiedColumnTable,
DeleteUsingTable,
UpdateTable,
UpdateWhereColumn,
UpdateSetColumn,
UpdateFromTable,
UpdateSetQualifiedColumnTable,
UpdateReturningColumn,
InsertReturningColumn,
DeleteReturningColumn,
MergeReturningColumn,
MergeWhenColumn,
MergeOnColumn,
MergeQualifiedColumnTable,
MergeUsingTable,
MergeTable,
UpdateReturningQualifiedColumnTable,
InsertReturningQualifiedColumnTable,
DeleteReturningQualifiedColumnTable,
MergeReturningQualifiedColumnTable,
JoinUsingColumn,
SchemaQualifier,
TypeReference,
Expand All @@ -61,6 +78,7 @@ pub(crate) enum NameRefClass {
VacuumTable,
AlterTable,
AlterTableColumn,
AlterTableDropColumn,
RefreshMaterializedView,
ReindexTable,
ReindexIndex,
Expand All @@ -82,6 +100,9 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
let mut in_constraint_where_clause = false;
let mut in_partition_item = false;
let mut in_set_null_columns = false;
let mut in_using_clause = false;
let mut in_returning_clause = false;
let mut in_when_clause = false;

// TODO: can we combine this if and the one that follows?
if let Some(parent) = name_ref.syntax().parent()
Expand All @@ -105,13 +126,89 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>

let mut in_from_clause = false;
let mut in_on_clause = false;
let mut in_returning_clause = false;
let mut in_set_clause = false;
let mut in_where_clause = false;
let mut in_when_clause = false;
for ancestor in parent.ancestors() {
if ast::OnClause::can_cast(ancestor.kind()) {
in_on_clause = true;
}
if ast::FromClause::can_cast(ancestor.kind()) {
in_from_clause = true;
}
if ast::ReturningClause::can_cast(ancestor.kind()) {
in_returning_clause = true;
}
if ast::SetClause::can_cast(ancestor.kind()) {
in_set_clause = true;
}
if ast::WhereClause::can_cast(ancestor.kind()) {
in_where_clause = true;
}
if ast::MergeWhenMatched::can_cast(ancestor.kind()) {
in_when_clause = true;
}
if ast::Update::can_cast(ancestor.kind()) {
if in_returning_clause {
if is_function_call || is_schema_table_col {
return Some(NameRefClass::SchemaQualifier);
} else {
return Some(NameRefClass::UpdateReturningQualifiedColumnTable);
}
} else if in_set_clause || in_where_clause || in_from_clause {
if is_function_call || is_schema_table_col {
return Some(NameRefClass::SchemaQualifier);
} else {
return Some(NameRefClass::UpdateSetQualifiedColumnTable);
}
}
}
if ast::Insert::can_cast(ancestor.kind()) {
if in_returning_clause {
if is_function_call || is_schema_table_col {
return Some(NameRefClass::SchemaQualifier);
} else {
return Some(NameRefClass::InsertReturningQualifiedColumnTable);
}
} else if !in_from_clause && !in_on_clause {
if is_function_call || is_schema_table_col {
return Some(NameRefClass::SchemaQualifier);
} else {
return Some(NameRefClass::InsertQualifiedColumnTable);
}
}
}
if ast::Delete::can_cast(ancestor.kind()) {
if in_returning_clause {
if is_function_call || is_schema_table_col {
return Some(NameRefClass::SchemaQualifier);
} else {
return Some(NameRefClass::DeleteReturningQualifiedColumnTable);
}
} else if in_where_clause || in_using_clause {
if is_function_call || is_schema_table_col {
return Some(NameRefClass::SchemaQualifier);
} else {
return Some(NameRefClass::DeleteQualifiedColumnTable);
}
}
}
if ast::Merge::can_cast(ancestor.kind()) {
if in_returning_clause {
if is_function_call || is_schema_table_col {
return Some(NameRefClass::SchemaQualifier);
} else {
return Some(NameRefClass::MergeReturningQualifiedColumnTable);
}
} else if in_on_clause || in_when_clause {
if is_function_call || is_schema_table_col {
return Some(NameRefClass::SchemaQualifier);
} else {
return Some(NameRefClass::MergeQualifiedColumnTable);
}
}
}
if ast::Select::can_cast(ancestor.kind()) && (!in_from_clause || in_on_clause) {
if is_function_call || is_schema_table_col {
return Some(NameRefClass::SchemaQualifier);
Expand Down Expand Up @@ -147,6 +244,8 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
let mut in_from_clause = false;
let mut in_on_clause = false;
let mut in_cast_expr = false;
let mut in_when_clause = false;
let mut in_returning_clause = false;
for ancestor in parent.ancestors() {
if ast::OnClause::can_cast(ancestor.kind()) {
in_on_clause = true;
Expand All @@ -157,6 +256,21 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
if ast::FromClause::can_cast(ancestor.kind()) {
in_from_clause = true;
}
if ast::MergeWhenMatched::can_cast(ancestor.kind()) {
in_when_clause = true;
}
if ast::ReturningClause::can_cast(ancestor.kind()) {
in_returning_clause = true;
}
if ast::Merge::can_cast(ancestor.kind())
&& (in_on_clause || in_when_clause || in_returning_clause)
{
if let Some(base) = field_expr.base()
&& matches!(base, ast::Expr::NameRef(_) | ast::Expr::FieldExpr(_))
{
return Some(NameRefClass::SelectQualifiedColumn);
}
}
if ast::Select::can_cast(ancestor.kind()) && (!in_from_clause || in_on_clause) {
if in_cast_expr {
return Some(NameRefClass::TypeReference);
Expand Down Expand Up @@ -215,6 +329,9 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
if ast::AlterColumn::can_cast(ancestor.kind()) {
return Some(NameRefClass::AlterTableColumn);
}
if ast::DropColumn::can_cast(ancestor.kind()) {
return Some(NameRefClass::AlterTableDropColumn);
}
if ast::AlterTable::can_cast(ancestor.kind()) {
return Some(NameRefClass::AlterTable);
}
Expand Down Expand Up @@ -443,6 +560,9 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
in_partition_item = true;
}
if ast::Insert::can_cast(ancestor.kind()) {
if in_returning_clause {
return Some(NameRefClass::InsertReturningColumn);
}
if in_column_list {
return Some(NameRefClass::InsertColumn);
}
Expand All @@ -457,13 +577,31 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
if ast::SetClause::can_cast(ancestor.kind()) {
in_set_clause = true;
}
if ast::UsingClause::can_cast(ancestor.kind()) {
in_using_clause = true;
}
if ast::UsingOnClause::can_cast(ancestor.kind()) {
in_using_clause = true;
}
if ast::ReturningClause::can_cast(ancestor.kind()) {
in_returning_clause = true;
}
if ast::Delete::can_cast(ancestor.kind()) {
if in_returning_clause {
return Some(NameRefClass::DeleteReturningColumn);
}
if in_where_clause {
return Some(NameRefClass::DeleteWhereColumn);
}
if in_using_clause {
return Some(NameRefClass::DeleteUsingTable);
}
return Some(NameRefClass::DeleteTable);
}
if ast::Update::can_cast(ancestor.kind()) {
if in_returning_clause {
return Some(NameRefClass::UpdateReturningColumn);
}
if in_where_clause {
return Some(NameRefClass::UpdateWhereColumn);
}
Expand All @@ -475,6 +613,24 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
}
return Some(NameRefClass::UpdateTable);
}
if ast::MergeWhenMatched::can_cast(ancestor.kind()) {
in_when_clause = true;
}
if ast::Merge::can_cast(ancestor.kind()) {
if in_returning_clause {
return Some(NameRefClass::MergeReturningColumn);
}
if in_when_clause {
return Some(NameRefClass::MergeWhenColumn);
}
if in_on_clause {
return Some(NameRefClass::MergeOnColumn);
}
if in_using_clause {
return Some(NameRefClass::MergeUsingTable);
}
return Some(NameRefClass::MergeTable);
}
}

None
Expand Down
Loading
Loading