From 76fa6c814344e19358f2b5cdd9fd104082bae293 Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Tue, 18 Aug 2026 23:13:18 -0400 Subject: [PATCH] ide: goto def for prepared transactions --- crates/squawk_ide/src/binder.rs | 59 +++++++++- crates/squawk_ide/src/classify.rs | 37 +++++- crates/squawk_ide/src/goto_definition.rs | 138 +++++++++++++++++++++++ crates/squawk_ide/src/hover.rs | 2 + crates/squawk_ide/src/location.rs | 1 + crates/squawk_ide/src/resolve.rs | 12 +- crates/squawk_ide/src/semantic_tokens.rs | 1 + 7 files changed, 247 insertions(+), 3 deletions(-) diff --git a/crates/squawk_ide/src/binder.rs b/crates/squawk_ide/src/binder.rs index 09317d84..80b6abb9 100644 --- a/crates/squawk_ide/src/binder.rs +++ b/crates/squawk_ide/src/binder.rs @@ -54,6 +54,8 @@ pub(crate) struct Binder { schema_regions: Vec<(TextRange, Schema)>, savepoint_stack: Vec<(Name, SyntaxNodePtr)>, savepoint_refs: FxHashMap, + prepared_transactions: FxHashMap, + prepared_transaction_refs: FxHashMap, } impl Binder { @@ -73,6 +75,8 @@ impl Binder { schema_regions: vec![], savepoint_stack: vec![], savepoint_refs: FxHashMap::default(), + prepared_transactions: FxHashMap::default(), + prepared_transaction_refs: FxHashMap::default(), } } @@ -98,6 +102,15 @@ impl Binder { .copied() } + pub(crate) fn lookup_prepared_transaction( + &self, + literal: &ast::Literal, + ) -> Option { + self.prepared_transaction_refs + .get(&SyntaxNodePtr::new(literal.syntax())) + .copied() + } + pub(crate) fn resolved_schemas( &self, position: TextSize, @@ -374,7 +387,9 @@ fn bind_stmt(b: &mut Binder, stmt: ast::Stmt) { ast::Stmt::SavepointCreate(savepoint) => bind_savepoint(b, savepoint), ast::Stmt::ReleaseSavepoint(release) => bind_release_savepoint(b, release), ast::Stmt::Rollback(rollback) => bind_rollback(b, rollback), - ast::Stmt::Begin(_) | ast::Stmt::Commit(_) => b.savepoint_stack.clear(), + ast::Stmt::PrepareTransaction(prepare) => bind_prepare_transaction(b, prepare), + ast::Stmt::Commit(commit) => bind_commit(b, commit), + ast::Stmt::Begin(_) => b.savepoint_stack.clear(), ast::Stmt::Select(select) => bind_select(b, select), ast::Stmt::Set(set) => bind_set(b, set), ast::Stmt::CreatePolicy(create_policy) => bind_create_policy(b, create_policy), @@ -1724,7 +1739,49 @@ fn bind_release_savepoint(b: &mut Binder, release: ast::ReleaseSavepoint) { } } +fn bind_commit(b: &mut Binder, commit: ast::Commit) { + if commit.prepared_token().is_some() { + bind_prepared_transaction_ref(b, commit.literal()); + } + + b.savepoint_stack.clear(); +} + +fn bind_prepare_transaction(b: &mut Binder, prepare: ast::PrepareTransaction) { + b.savepoint_stack.clear(); + + let Some(literal) = prepare.literal() else { + return; + }; + let Some(transaction_id) = literal_string_value(&literal) else { + return; + }; + + b.prepared_transactions + .insert(transaction_id, SyntaxNodePtr::new(literal.syntax())); +} + +fn bind_prepared_transaction_ref(b: &mut Binder, literal: Option) { + let Some(literal) = literal else { + return; + }; + let Some(transaction_id) = literal_string_value(&literal) else { + return; + }; + + let Some(ptr) = b.prepared_transactions.remove(&transaction_id) else { + return; + }; + + b.prepared_transaction_refs + .insert(SyntaxNodePtr::new(literal.syntax()), ptr); +} + fn bind_rollback(b: &mut Binder, rollback: ast::Rollback) { + if rollback.prepared_token().is_some() { + bind_prepared_transaction_ref(b, rollback.literal()); + } + let Some(savepoint_ref) = rollback.savepoint_ref() else { b.savepoint_stack.clear(); return; diff --git a/crates/squawk_ide/src/classify.rs b/crates/squawk_ide/src/classify.rs index 1ea57f10..a7058db7 100644 --- a/crates/squawk_ide/src/classify.rs +++ b/crates/squawk_ide/src/classify.rs @@ -1,7 +1,7 @@ use crate::{location::LocationKind, name, symbols::Name}; use squawk_syntax::{ SyntaxKind, SyntaxNode, - ast::{self, AstNode}, + ast::{self, AstNode, LitKind}, }; #[derive(Debug, Clone, Copy)] @@ -45,6 +45,7 @@ pub(crate) enum NameRefClass { PolicyColumn, PolicyQualifiedColumnTable, PreparedStatement, + PreparedTransaction, PrivilegeColumn, PrivilegeObjectTable, Procedure, @@ -392,12 +393,46 @@ pub(crate) fn classify_literal(node: &SyntaxNode) -> Option { if ast::SetSchemaValue::can_cast(parent.kind()) { return Some(NameRefClass::Schema); } + if is_prepared_transaction_id(node) { + return Some(NameRefClass::PreparedTransaction); + } if is_search_path_config_value(node) { return Some(NameRefClass::Schema); } None } +// commit prepared 'foo' | rollback prepared 'foo' +fn is_prepared_transaction_id(node: &SyntaxNode) -> bool { + let Some(literal) = ast::Literal::cast(node.clone()) else { + return false; + }; + if !matches!( + literal.kind(), + Some( + LitKind::String(_) + | LitKind::BitString(_) + | LitKind::ByteString(_) + | LitKind::EscString(_) + | LitKind::NationalString(_) + | LitKind::UnicodeEscString(_) + | LitKind::DollarQuotedString(_) + ) + ) { + return false; + } + let Some(parent) = node.parent() else { + return false; + }; + if let Some(commit) = ast::Commit::cast(parent.clone()) { + commit.prepared_token().is_some() + } else if let Some(rollback) = ast::Rollback::cast(parent) { + rollback.prepared_token().is_some() + } else { + false + } +} + // set search_path to ... fn is_search_path_config_value(node: &SyntaxNode) -> bool { let Some(to_config_value) = node.parent().and_then(ast::ToConfigValue::cast) else { diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 625d4007..294424f0 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -827,6 +827,18 @@ release savepoint sp$0; ); } + #[test] + fn goto_prepare_transaction_discards_savepoints() { + goto_not_found( + " +begin; +savepoint sp; +prepare transaction 'foo'; +release savepoint sp$0; +", + ); + } + #[test] fn goto_bare_rollback_discards_savepoints() { goto_not_found( @@ -6071,6 +6083,132 @@ commit; "); } + #[test] + fn commit_prepared_to_prepare_transaction() { + assert_snapshot!(goto( + " +prepare transaction 'foo'; +select 1; +commit prepared 'foo'$0; +", + ), @" + ╭▸ + 2 │ prepare transaction 'foo'; + │ ───── 2. destination + 3 │ select 1; + 4 │ commit prepared 'foo'; + ╰╴ ─ 1. source + "); + } + + #[test] + fn rollback_prepared_to_prepare_transaction() { + assert_snapshot!(goto( + " +prepare transaction 'foo'; +rollback prepared 'foo'$0; +", + ), @" + ╭▸ + 2 │ prepare transaction 'foo'; + │ ───── 2. destination + 3 │ rollback prepared 'foo'; + ╰╴ ─ 1. source + "); + } + + #[test] + fn commit_prepared_to_most_recent_prepare_transaction() { + assert_snapshot!(goto( + " +prepare transaction 'foo'; +commit prepared 'foo'; +prepare transaction 'foo'; +commit prepared 'foo'$0; +", + ), @" + ╭▸ + 4 │ prepare transaction 'foo'; + │ ───── 2. destination + 5 │ commit prepared 'foo'; + ╰╴ ─ 1. source + "); + } + + #[test] + fn commit_prepared_before_prepare_transaction() { + goto_not_found( + " +commit prepared 'foo'$0; +prepare transaction 'foo'; +", + ); + } + + #[test] + fn commit_prepared_frees_transaction_id() { + goto_not_found( + " +prepare transaction 'foo'; +commit prepared 'foo'; +commit prepared 'foo'$0; +", + ); + } + + #[test] + fn commit_prepared_matches_escaped_transaction_id() { + assert_snapshot!(goto( + r#" +prepare transaction e'fo\u006f'; +commit prepared 'foo'$0; +"#, + ), @r" + ╭▸ + 2 │ prepare transaction e'fo\u006f'; + │ ─────────── 2. destination + 3 │ commit prepared 'foo'; + ╰╴ ─ 1. source + "); + } + + #[test] + fn commit_prepared_matches_dollar_quoted_transaction_id() { + assert_snapshot!(goto( + " +prepare transaction $$foo$$; +commit prepared 'foo'$0; +", + ), @" + ╭▸ + 2 │ prepare transaction $$foo$$; + │ ─────── 2. destination + 3 │ commit prepared 'foo'; + ╰╴ ─ 1. source + "); + } + + #[test] + fn commit_prepared_with_unknown_transaction_id() { + goto_not_found( + " +begin; +prepare transaction 'foo'; +commit prepared 'bar'$0; +", + ); + } + + #[test] + fn commit_prepared_ignores_enclosing_begin() { + goto_not_found( + " +begin; +commit prepared 'foo'$0; +", + ); + } + #[test] fn goto_with_search_path() { assert_snapshot!(goto(r#" diff --git a/crates/squawk_ide/src/hover.rs b/crates/squawk_ide/src/hover.rs index 99994be0..572d5d3a 100644 --- a/crates/squawk_ide/src/hover.rs +++ b/crates/squawk_ide/src/hover.rs @@ -357,6 +357,7 @@ fn hover_name(db: &dyn Db, def: Location) -> Option { | LocationKind::CommitEnd | LocationKind::ElementTable | LocationKind::Label + | LocationKind::PreparedTransaction | LocationKind::Property => None, LocationKind::Channel => hover_channel(db, def), LocationKind::Column => hover_name_column(db, def), @@ -445,6 +446,7 @@ fn hover_position(db: &dyn Db, position: InFile) -> Option { | LocationKind::CommitEnd | LocationKind::ElementTable | LocationKind::Label + | LocationKind::PreparedTransaction | LocationKind::Property => None, LocationKind::Channel => hover_channel(db, def), LocationKind::Column => { diff --git a/crates/squawk_ide/src/location.rs b/crates/squawk_ide/src/location.rs index 43f4f240..7f8d95fb 100644 --- a/crates/squawk_ide/src/location.rs +++ b/crates/squawk_ide/src/location.rs @@ -37,6 +37,7 @@ pub enum LocationKind { OperatorFamily, Policy, PreparedStatement, + PreparedTransaction, Procedure, Property, PropertyGraph, diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 8bb4ddae..a998a74f 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -922,6 +922,7 @@ pub(crate) fn resolve_name_ref( let table_path = resolve_alter_column_relation_path(name_ref.syntax())?; resolve_column_for_path(db, InFile::new(file, &table_path), column_name) } + NameRefClass::PreparedTransaction => None, } .or_else(|| resolve_special_keyword_as_function(db, InFile::new(file, name_ref))) } @@ -948,7 +949,7 @@ pub(crate) fn resolve_config_value_name( } /// Resolves a string literal to its definition(s), e.g. the schema name in -/// `set schema 'app'` or `set search_path to 'app'`. +/// `set schema 'app'` or the transaction id in `commit prepared 'foo'`. pub(crate) fn resolve_literal( db: &dyn Db, literal: InFile<&ast::Literal>, @@ -977,6 +978,15 @@ pub(crate) fn resolve_literal( LocationKind::Schema )]) } + NameRefClass::PreparedTransaction => { + let binder = bind(db, file); + let ptr = binder.lookup_prepared_transaction(literal)?; + Some(smallvec![Location::new( + file, + ptr.text_range(), + LocationKind::PreparedTransaction + )]) + } _ => None, } } diff --git a/crates/squawk_ide/src/semantic_tokens.rs b/crates/squawk_ide/src/semantic_tokens.rs index 55be0209..e7a4ed99 100644 --- a/crates/squawk_ide/src/semantic_tokens.rs +++ b/crates/squawk_ide/src/semantic_tokens.rs @@ -220,6 +220,7 @@ impl TryFrom for SemanticTokenType { | LocationKind::OperatorFamily | LocationKind::Policy | LocationKind::PreparedStatement + | LocationKind::PreparedTransaction | LocationKind::Publication | LocationKind::Role | LocationKind::Rule