ide: goto def for prepared transactions

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<SyntaxNodePtr, SyntaxNodePtr>,
+    prepared_transactions: FxHashMap<String, SyntaxNodePtr>,
+    prepared_transaction_refs: FxHashMap<SyntaxNodePtr, SyntaxNodePtr>,
 }
 
 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<SyntaxNodePtr> {
+        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<ast::Literal>) {
+    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;

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<NameRefClass> {
     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 {

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#"

crates/squawk_ide/src/hover.rs

@@ -357,6 +357,7 @@ fn hover_name(db: &dyn Db, def: Location) -> Option<Hover> {
         | 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<TextSize>) -> Option<Hover> {
         | LocationKind::CommitEnd
         | LocationKind::ElementTable
         | LocationKind::Label
+        | LocationKind::PreparedTransaction
         | LocationKind::Property => None,
         LocationKind::Channel => hover_channel(db, def),
         LocationKind::Column => {

crates/squawk_ide/src/location.rs

@@ -37,6 +37,7 @@ pub enum LocationKind {
     OperatorFamily,
     Policy,
     PreparedStatement,
+    PreparedTransaction,
     Procedure,
     Property,
     PropertyGraph,

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,
     }
 }

crates/squawk_ide/src/semantic_tokens.rs

@@ -220,6 +220,7 @@ impl TryFrom<LocationKind> for SemanticTokenType {
             | LocationKind::OperatorFamily
             | LocationKind::Policy
             | LocationKind::PreparedStatement
+            | LocationKind::PreparedTransaction
             | LocationKind::Publication
             | LocationKind::Role
             | LocationKind::Rule