From ab01ad2e0273a40834c0615e7d731327b38c991f Mon Sep 17 00:00:00 2001
From: psteinroe <philipp@steinroetter.com>
Date: Mon, 26 Oct 2026 15:18:34 +0000
Subject: [PATCH 1/3] fix(workspace): one walk for lint, hover, completions and
 inlay hints

pgls_analyser::walk_statements is the only loop that follows the catalog
and the session through a file: sections, unparsable statements, function
bodies and identifier parameters are handled once. The analyser runs its
rules in it, and the workspace resolves hover, completions, inlay hints,
semantic tokens and code actions in it. The catalog that a statement at
the end of the walk sees is the one of its section, so completions in the
first statement after a down section no longer see its changes.

Statements share their parse trees with the workspace's cache instead of
copying them.
---
 crates/pgls_analyser/src/lib.rs               | 358 ++++--------------
 .../src/lint/style/require_table_schema.rs    |   4 +-
 crates/pgls_analyser/src/linter_context.rs    |   8 +-
 crates/pgls_analyser/src/walk.rs              | 289 ++++++++++++++
 crates/pgls_workspace/src/workspace/server.rs | 100 ++---
 .../src/workspace/server.tests.rs             | 139 +++++++
 .../src/workspace/server/document.rs          |  15 +-
 .../src/workspace/server/formatting.rs        |   8 +-
 .../src/workspace/server/pg_query.rs          |  15 +-
 9 files changed, 578 insertions(+), 358 deletions(-)
 create mode 100644 crates/pgls_analyser/src/walk.rs

diff --git a/crates/pgls_analyser/src/lib.rs b/crates/pgls_analyser/src/lib.rs
index 80fca05fa..123635b1e 100644
--- a/crates/pgls_analyser/src/lib.rs
+++ b/crates/pgls_analyser/src/lib.rs
@@ -1,10 +1,7 @@
-use std::{
-    ops::Range,
-    sync::{Arc, LazyLock},
-};
+use std::sync::{Arc, LazyLock};
 
 use pgls_analyse::{AnalysisFilter, AppliesTo, MetadataRegistry};
-use pgls_catalog::{Catalog, CatalogBase, Session, TransactionMode, resolve::FunctionContext};
+use pgls_catalog::{CatalogBase, resolve::FunctionContext};
 use pgls_diagnostics::DiagnosticTags;
 use pgls_query_ext::function_body::{EmbeddedStatement, OffsetMap};
 use pgls_text_size::{TextRange, TextSize};
@@ -19,6 +16,7 @@ mod linter_rule;
 mod migration;
 pub mod options;
 mod registry;
+mod walk;
 
 // Re-export linter-specific types
 pub use linter_context::{AnalysedFileContext, FileKind, LinterRuleContext, StatementContext};
@@ -28,6 +26,8 @@ pub use linter_registry::{
 };
 pub use linter_rule::{FixEdit, LinterDiagnostic, LinterRule, RuleFix};
 pub use migration::{Direction, MigrationSection};
+use walk::is_plain_query;
+pub use walk::{WalkParams, walk_statements};
 
 // For convenience in macros and rule files - keep these shorter names
 pub use LinterDiagnostic as RuleDiagnostic;
@@ -79,7 +79,9 @@ pub fn replacement_of_removed_rule(name: &str) -> Option<&'static str> {
 
 #[derive(Debug)]
 pub struct AnalysableStatement {
-    pub root: pgls_query::NodeEnum,
+    /// Shared with the workspace's cache of parsed statements, so that analysing a file doesn't
+    /// copy every statement's tree.
+    pub root: Arc<pgls_query::NodeEnum>,
     pub range: TextRange,
     /// The statement text, used for precise spans of typecheck findings.
     pub sql: Option<String>,
@@ -95,9 +97,9 @@ pub struct AnalysableStatement {
 }
 
 impl AnalysableStatement {
-    pub fn new(root: pgls_query::NodeEnum, range: TextRange) -> Self {
+    pub fn new(root: impl Into<Arc<pgls_query::NodeEnum>>, range: TextRange) -> Self {
         Self {
-            root,
+            root: root.into(),
             range,
             sql: None,
             function: None,
@@ -109,7 +111,7 @@ impl AnalysableStatement {
     /// The statement `embedded`, like a statement of a SQL function body, in the statement that
     /// starts at `parent_start` in the file. `root` is the parsed `embedded.sql`.
     pub fn embedded(
-        root: pgls_query::NodeEnum,
+        root: impl Into<Arc<pgls_query::NodeEnum>>,
         parent_start: TextSize,
         embedded: &EmbeddedStatement,
     ) -> Self {
@@ -219,7 +221,7 @@ impl<'a> Analyser<'a> {
     }
 
     /// Runs the rules statement by statement, while the catalog and the session follow the
-    /// effects of the statements.
+    /// effects of the statements, see [`walk_statements`].
     pub fn analyse(&self, params: AnalyserParams) -> AnalysisResult {
         let mut result = AnalysisResult::default();
 
@@ -228,8 +230,6 @@ impl<'a> Analyser<'a> {
             .catalog_base
             .as_ref()
             .map(|base| base.snapshot().as_ref());
-        let roots: Vec<pgls_query::NodeEnum> =
-            params.stmts.iter().map(|s| s.root.clone()).collect();
 
         let rules: Vec<_> = self
             .registry
@@ -243,172 +243,82 @@ impl<'a> Analyser<'a> {
             })
             .collect();
 
-        // The catalog after the up sections so far. A down section runs after its up section,
-        // and its changes don't carry into the next section.
-        let mut catalog = Catalog::new(params.catalog_base.clone());
-        let mut unparsable = params.unparsable.iter().peekable();
-        for unit in migration_units(&params.stmts, &params.sections) {
-            let is_down = unit
-                .section
-                .is_some_and(|section| section.direction == Direction::Down);
-            let unit_catalog = if is_down {
-                catalog.clone()
-            } else {
-                std::mem::replace(&mut catalog, Catalog::new(None))
-            };
-            // Each section runs in a session and transaction of its own.
-            let session = Session::new(params.search_path.clone()).with_transaction_mode(
-                unit.section
-                    .map_or(TransactionMode::Unknown, MigrationSection::transaction_mode),
-            );
-            let stmts = &params.stmts[unit.stmts.clone()];
-            let end = params
-                .unparsable
-                .iter()
-                .filter(|skipped| skipped.range.start() < unit.range.end())
-                .map(|skipped| skipped.range.end())
-                .chain(stmts.last().map(|stmt| stmt.range.end()))
-                .max()
-                .unwrap_or_default();
-            // The statements the unit runs, including those that don't parse, but not those in
-            // function bodies.
-            let stmt_count = |range: TextRange| {
-                params.stmts[unit.stmts.clone()]
-                    .iter()
-                    .filter(|stmt| stmt.function.is_none() && range.contains(stmt.range.start()))
-                    .count()
-                    + params
-                        .unparsable
-                        .iter()
-                        .filter(|skipped| range.contains(skipped.range.start()))
-                        .count()
+        let walk = WalkParams {
+            stmts: &params.stmts,
+            unparsable: &params.unparsable,
+            catalog_base: params.catalog_base.clone(),
+            search_path: &params.search_path,
+            sections: &params.sections,
+            end: TextSize::from(u32::MAX),
+            best_effort: false,
+        };
+        walk_statements(walk, |stmt, file_context| {
+            let root = &*stmt.root;
+            // The catalog can't follow statements on placeholder names, like `drop table :name`,
+            // so the rest of the file isn't typechecked.
+            if stmt.has_identifier_parameters && !is_plain_query(root) {
+                typecheck = false;
+            }
+            let typecheck_statement = typecheck && !stmt.has_identifier_parameters;
+
+            let statement =
+                StatementContext::new(root, stmt.sql.as_deref(), stmt.function.as_ref());
+            let mut has_typecheck_findings = false;
+
+            let rule_params = LinterRegistryRuleParams {
+                root,
+                options: self.options,
+                analysed_file_context: file_context,
+                statement: &statement,
+                snapshot,
             };
-            // Outside a transaction, the queries with more than one statement run in an implicit
-            // transaction block.
-            let implicit_blocks: Vec<TextRange> = unit
-                .section
-                .filter(|section| !section.transaction)
-                .map(|section| {
-                    section
-                        .queries
-                        .iter()
-                        .copied()
-                        .filter(|query| stmt_count(*query) > 1)
-                        .collect()
-                })
-                .unwrap_or_default();
-            let mut file_context =
-                AnalysedFileContext::new(&roots[unit.stmts.clone()], unit_catalog, session)
-                    .with_positions(stmts.iter().map(|stmt| stmt.range.start()).collect(), end)
-                    .with_stmt_count(stmt_count(unit.range));
-            let mut implicit_block = None;
-
-            for (stmt, root) in params.stmts[unit.stmts.clone()]
-                .iter()
-                .zip(&roots[unit.stmts.clone()])
-            {
-                if stmt.function.is_none() {
-                    let block = implicit_blocks
-                        .iter()
-                        .position(|block| block.contains(stmt.range.start()));
-                    if block != implicit_block {
-                        if implicit_block.is_some() {
-                            file_context.end_implicit_block();
+            for rule in &rules {
+                if rule.group == TYPECHECK_GROUP && !typecheck_statement {
+                    continue;
+                }
+                let diagnostics = (rule.run)(&rule_params);
+                has_typecheck_findings |= rule.group == TYPECHECK_GROUP && !diagnostics.is_empty();
+
+                result
+                    .diagnostics
+                    .extend(diagnostics.into_iter().map(|mut diagnostic| {
+                        diagnostic.span = Some(match diagnostic.span {
+                            Some(span) => stmt.document_range(span),
+                            None => stmt.range,
+                        });
+                        // Fixes are computed on the statement text, which only matches the source
+                        // if it is the source as written. A fix must keep the statement valid.
+                        let fix_is_valid = stmt.sql.as_ref().is_some_and(|sql| {
+                            stmt.offsets.is_verbatim()
+                                && TextSize::of(sql.as_str()) == stmt.range.len()
+                                && diagnostic
+                                    .fix
+                                    .as_ref()
+                                    .is_none_or(|fix| keeps_statement_valid(fix, sql))
+                        });
+                        if !fix_is_valid {
+                            diagnostic.fix = None;
                         }
-                        if block.is_some() {
-                            file_context.begin_implicit_block();
+                        if let Some(fix) = &mut diagnostic.fix {
+                            fix.shift(stmt.range.start());
+                            diagnostic.tags |= DiagnosticTags::FIXABLE;
                         }
-                        implicit_block = block;
-                    }
-                }
-                while let Some(skipped) =
-                    unparsable.next_if(|skipped| skipped.range.start() < stmt.range.start())
-                {
-                    if may_change_catalog(&skipped.sql) {
-                        file_context.taint_catalog();
-                    }
-                }
-                // The catalog and session can't follow statements on placeholder names, like
-                // `drop table :name`, so the rest of the file isn't typechecked.
-                if stmt.has_identifier_parameters && !is_plain_query(root) {
-                    typecheck = false;
-                }
-                let typecheck_statement = typecheck && !stmt.has_identifier_parameters;
-
-                let statement =
-                    StatementContext::new(root, stmt.sql.as_deref(), stmt.function.as_ref());
-                let mut has_typecheck_findings = false;
-
-                let rule_params = LinterRegistryRuleParams {
-                    root,
-                    options: self.options,
-                    analysed_file_context: &file_context,
-                    statement: &statement,
-                    snapshot,
-                };
-                for rule in &rules {
-                    if rule.group == TYPECHECK_GROUP && !typecheck_statement {
-                        continue;
-                    }
-                    let diagnostics = (rule.run)(&rule_params);
-                    has_typecheck_findings |=
-                        rule.group == TYPECHECK_GROUP && !diagnostics.is_empty();
-
-                    result
-                        .diagnostics
-                        .extend(diagnostics.into_iter().map(|mut diagnostic| {
-                            diagnostic.span = Some(match diagnostic.span {
-                                Some(span) => stmt.document_range(span),
-                                None => stmt.range,
-                            });
-                            // Fixes are computed on the statement text, which only matches the
-                            // source if it is the source as written. A fix must keep the statement
-                            // valid.
-                            let fix_is_valid = stmt.sql.as_ref().is_some_and(|sql| {
-                                stmt.offsets.is_verbatim()
-                                    && TextSize::of(sql.as_str()) == stmt.range.len()
-                                    && diagnostic
-                                        .fix
-                                        .as_ref()
-                                        .is_none_or(|fix| keeps_statement_valid(fix, sql))
-                            });
-                            if !fix_is_valid {
-                                diagnostic.fix = None;
-                            }
-                            if let Some(fix) = &mut diagnostic.fix {
-                                fix.shift(stmt.range.start());
-                                diagnostic.tags |= DiagnosticTags::FIXABLE;
-                            }
-                            diagnostic
-                        }));
-                }
-
-                let database_only = typecheck_statement
-                    && statement
-                        .resolution(&file_context)
-                        .is_some_and(|resolution| resolution.database_only);
-
-                result.statements.push(StatementAnalysis {
-                    range: stmt.range,
-                    database_only,
-                    has_typecheck_findings,
-                    search_path: file_context.session().search_path().to_vec(),
-                });
-
-                file_context.next(stmt.function.is_none());
+                        diagnostic
+                    }));
             }
 
-            while let Some(skipped) =
-                unparsable.next_if(|skipped| skipped.range.start() < unit.range.end())
-            {
-                if may_change_catalog(&skipped.sql) {
-                    file_context.taint_catalog();
-                }
-            }
-            if !is_down {
-                catalog = file_context.into_catalog();
-            }
-        }
+            let database_only = typecheck_statement
+                && statement
+                    .resolution(file_context)
+                    .is_some_and(|resolution| resolution.database_only);
+
+            result.statements.push(StatementAnalysis {
+                range: stmt.range,
+                database_only,
+                has_typecheck_findings,
+                search_path: file_context.session().search_path().to_vec(),
+            });
+        });
 
         result
     }
@@ -425,108 +335,6 @@ fn keeps_statement_valid(fix: &RuleFix, sql: &str) -> bool {
     pgls_query::parse(&fix.apply(sql)).is_ok()
 }
 
-/// Statements that run together: those before the first section marker, then each section.
-struct MigrationUnit<'a> {
-    section: Option<&'a MigrationSection>,
-    /// The indices of the statements.
-    stmts: Range<usize>,
-    range: TextRange,
-}
-
-fn migration_units<'a>(
-    stmts: &[AnalysableStatement],
-    sections: &'a [MigrationSection],
-) -> Vec<MigrationUnit<'a>> {
-    let before_sections = sections
-        .first()
-        .map_or(TextSize::from(u32::MAX), |section| section.range.start());
-    let ranges = std::iter::once((None, TextRange::up_to(before_sections))).chain(
-        sections
-            .iter()
-            .map(|section| (Some(section), section.range)),
-    );
-
-    let mut first = 0;
-    ranges
-        .map(|(section, range)| {
-            let last =
-                first + stmts[first..].partition_point(|stmt| stmt.range.start() < range.end());
-            let unit = MigrationUnit {
-                section,
-                stmts: first..last,
-                range,
-            };
-            first = last;
-            unit
-        })
-        .collect()
-}
-
-/// Whether a statement that doesn't parse may have changed the catalog. Queries, `COPY`,
-/// maintenance, cursors, prepared statements, comments and privileges can't create, drop or
-/// rename objects, unless they are `SELECT ... INTO` or `EXPLAIN ANALYZE CREATE TABLE AS`. For
-/// anything else we can't tell.
-pub fn may_change_catalog(sql: &str) -> bool {
-    use pgls_query::protobuf::Token;
-    let Ok(scan) = pgls_query::scan(sql) else {
-        return true;
-    };
-    let mut tokens = scan
-        .tokens
-        .iter()
-        .map(|token| token.token())
-        .filter(|token| !matches!(token, Token::SqlComment | Token::CComment | Token::Ascii40));
-    match tokens.next() {
-        Some(
-            Token::Insert
-            | Token::Update
-            | Token::DeleteP
-            | Token::Merge
-            | Token::Copy
-            | Token::Analyse
-            | Token::Analyze
-            | Token::Vacuum
-            | Token::Reindex
-            | Token::Cluster
-            | Token::Checkpoint
-            | Token::LockP
-            | Token::Listen
-            | Token::Notify
-            | Token::Unlisten
-            | Token::Show
-            | Token::Declare
-            | Token::Fetch
-            | Token::Move
-            | Token::Close
-            | Token::Prepare
-            | Token::Execute
-            | Token::Deallocate
-            | Token::Comment
-            | Token::Grant
-            | Token::Revoke,
-        ) => false,
-        Some(Token::Select | Token::With | Token::Values | Token::Table | Token::Explain) => {
-            tokens.any(|token| matches!(token, Token::Into | Token::Create))
-        }
-        _ => true,
-    }
-}
-
-/// Whether the statement only reads or writes rows, without changing the catalog or the
-/// session.
-pub fn is_plain_query(root: &pgls_query::NodeEnum) -> bool {
-    use pgls_query::NodeEnum;
-    match root {
-        NodeEnum::SelectStmt(select) => select.into_clause.is_none(),
-        NodeEnum::InsertStmt(_)
-        | NodeEnum::UpdateStmt(_)
-        | NodeEnum::DeleteStmt(_)
-        | NodeEnum::MergeStmt(_) => true,
-        _ => false,
-    }
-}
-
-/// Runs a single rule on a statement and returns the text each diagnostic points at.
 #[cfg(test)]
 pub(crate) fn spanned_texts(group: &str, rule: &str, sql: &str) -> Vec<String> {
     let rule_filter = pgls_analyse::RuleFilter::Rule(group, rule);
diff --git a/crates/pgls_analyser/src/lint/style/require_table_schema.rs b/crates/pgls_analyser/src/lint/style/require_table_schema.rs
index 13ce4dddb..a83b0dab4 100644
--- a/crates/pgls_analyser/src/lint/style/require_table_schema.rs
+++ b/crates/pgls_analyser/src/lint/style/require_table_schema.rs
@@ -223,10 +223,10 @@ fn dropped_names(tokens: StatementTokens) -> Vec<TextRange> {
 
 /// The unqualified names of the temporary tables created by earlier statements of the file.
 /// An unqualified name refers to them, since `pg_temp` comes first in the search path.
-fn temporary_tables(previous: &[NodeEnum]) -> Vec<&str> {
+fn temporary_tables<'a>(previous: &[&'a NodeEnum]) -> Vec<&'a str> {
     let mut tables = Vec::new();
     for stmt in previous {
-        let created = match stmt {
+        let created = match *stmt {
             NodeEnum::CreateStmt(stmt) => stmt.relation.as_ref(),
             NodeEnum::CreateTableAsStmt(stmt) => {
                 stmt.into.as_ref().and_then(|into| into.rel.as_ref())
diff --git a/crates/pgls_analyser/src/linter_context.rs b/crates/pgls_analyser/src/linter_context.rs
index 8ce990126..52807100c 100644
--- a/crates/pgls_analyser/src/linter_context.rs
+++ b/crates/pgls_analyser/src/linter_context.rs
@@ -207,7 +207,7 @@ impl<'a> StatementContext<'a> {
 ///
 /// In a file with migration sections, the file context is that of the current section.
 pub struct AnalysedFileContext<'a> {
-    pub stmts: &'a [pgls_query::NodeEnum],
+    pub stmts: &'a [&'a pgls_query::NodeEnum],
     pos: usize,
     stmt_count: usize,
     catalog: Catalog,
@@ -219,7 +219,7 @@ pub struct AnalysedFileContext<'a> {
 }
 
 impl<'a> AnalysedFileContext<'a> {
-    pub fn new(stmts: &'a [pgls_query::NodeEnum], catalog: Catalog, session: Session) -> Self {
+    pub fn new(stmts: &'a [&'a pgls_query::NodeEnum], catalog: Catalog, session: Session) -> Self {
         Self {
             stmts,
             pos: 0,
@@ -251,7 +251,7 @@ impl<'a> AnalysedFileContext<'a> {
         self
     }
 
-    pub fn previous_stmts(&self) -> &[pgls_query::NodeEnum] {
+    pub fn previous_stmts(&self) -> &[&'a pgls_query::NodeEnum] {
         &self.stmts[0..self.pos]
     }
 
@@ -306,7 +306,7 @@ impl<'a> AnalysedFileContext<'a> {
     /// in the bodies of SQL functions are only analysed, not executed.
     pub fn next(&mut self, applies_effects: bool) {
         if applies_effects && self.pos < self.stmts.len() {
-            let stmt = &self.stmts[self.pos];
+            let stmt = self.stmts[self.pos];
             self.catalog.apply(stmt, self.session.search_path());
             self.session.apply(stmt);
         }
diff --git a/crates/pgls_analyser/src/walk.rs b/crates/pgls_analyser/src/walk.rs
new file mode 100644
index 000000000..4973b3277
--- /dev/null
+++ b/crates/pgls_analyser/src/walk.rs
@@ -0,0 +1,289 @@
+//! The walk over the statements of a file. Lint and the workspace's editor features (hover,
+//! completions, inlay hints, semantic tokens, code actions) share it, so that they see the same
+//! catalog and session at each statement.
+
+use std::{ops::Range, sync::Arc};
+
+use pgls_catalog::{Catalog, CatalogBase, Session, TransactionMode};
+use pgls_text_size::{TextRange, TextSize};
+
+use crate::{
+    AnalysableStatement, AnalysedFileContext, Direction, MigrationSection, UnparsableStatement,
+};
+
+pub struct WalkParams<'a> {
+    /// The statements that parse, in file order, including the statements of SQL function
+    /// bodies after their function.
+    pub stmts: &'a [AnalysableStatement],
+    /// The top-level statements that don't parse, in file order.
+    pub unparsable: &'a [UnparsableStatement],
+    /// The database snapshot the catalog starts from.
+    pub catalog_base: Option<Arc<CatalogBase>>,
+    /// Explicit search path at the start of the file.
+    pub search_path: &'a [String],
+    /// The sections of a migration file, like dbmate's up and down migrations. Without sections,
+    /// the whole file is one migration.
+    pub sections: &'a [MigrationSection],
+    /// The walk visits the statements that start before `end`.
+    pub end: TextSize,
+    /// Whether the catalog resolves names as well as it can after statements it can't follow,
+    /// like a `DO` block, see [`Catalog::best_effort`]. Editor features show what names most
+    /// likely refer to; lint only reports what is certain.
+    pub best_effort: bool,
+}
+
+/// Walks the statements like the database runs them: the catalog and the session follow the
+/// statements, so that each statement sees the objects that the statements before it create,
+/// change or drop. `visit` gets each statement with the file context as it is before it.
+/// Returns the catalog as a statement at [`WalkParams::end`] sees it.
+///
+/// - Each migration section runs in a session and transaction of its own. A down section runs
+///   after the up sections before it, and its changes don't carry into the next section.
+/// - A top-level statement that doesn't parse taints the catalog if it may have changed it.
+/// - The statements of SQL function bodies are visited, but they don't run.
+/// - The names of a statement with identifier parameters, like psql's `:name`, are placeholders.
+///   Such a statement is visited, but if it may change the catalog, it taints it instead of being
+///   applied.
+///
+/// The file context only knows the statements that the walk is given. When they stop before the
+/// end of the file, [`AnalysedFileContext::stmt_count`] and the implicit transaction blocks of
+/// the last query are those of the statements given, which only rules rely on.
+pub fn walk_statements(
+    params: WalkParams<'_>,
+    mut visit: impl FnMut(&AnalysableStatement, &AnalysedFileContext),
+) -> Catalog {
+    let roots: Vec<&pgls_query::NodeEnum> = params.stmts.iter().map(|s| &*s.root).collect();
+
+    // The catalog after the up sections so far.
+    let mut catalog = Catalog::new(params.catalog_base);
+    if params.best_effort {
+        catalog = catalog.best_effort();
+    }
+    let mut unparsable = params.unparsable.iter().peekable();
+    for unit in migration_units(params.stmts, params.sections) {
+        if unit.start >= params.end {
+            break;
+        }
+        let is_down = unit
+            .section
+            .is_some_and(|section| section.direction == Direction::Down);
+        let unit_catalog = if is_down {
+            catalog.clone()
+        } else {
+            std::mem::replace(&mut catalog, Catalog::new(None))
+        };
+        // Each section runs in a session and transaction of its own.
+        let session = Session::new(params.search_path.to_vec()).with_transaction_mode(
+            unit.section
+                .map_or(TransactionMode::Unknown, MigrationSection::transaction_mode),
+        );
+        let stmts = &params.stmts[unit.stmts.clone()];
+        let unit_range = TextRange::new(unit.start, unit.end);
+        let last_end = params
+            .unparsable
+            .iter()
+            .filter(|skipped| unit_range.contains(skipped.range.start()))
+            .map(|skipped| skipped.range.end())
+            .chain(stmts.last().map(|stmt| stmt.range.end()))
+            .max()
+            .unwrap_or_default();
+        // The statements the unit runs, including those that don't parse, but not those in
+        // function bodies.
+        let stmt_count = |range: TextRange| {
+            stmts
+                .iter()
+                .filter(|stmt| stmt.function.is_none() && range.contains(stmt.range.start()))
+                .count()
+                + params
+                    .unparsable
+                    .iter()
+                    .filter(|skipped| range.contains(skipped.range.start()))
+                    .count()
+        };
+        // Outside a transaction, the queries with more than one statement run in an implicit
+        // transaction block.
+        let implicit_blocks: Vec<TextRange> = unit
+            .section
+            .filter(|section| !section.transaction)
+            .map(|section| {
+                section
+                    .queries
+                    .iter()
+                    .copied()
+                    .filter(|query| stmt_count(*query) > 1)
+                    .collect()
+            })
+            .unwrap_or_default();
+        let mut file_context =
+            AnalysedFileContext::new(&roots[unit.stmts.clone()], unit_catalog, session)
+                .with_positions(
+                    stmts.iter().map(|stmt| stmt.range.start()).collect(),
+                    last_end,
+                )
+                .with_stmt_count(stmt_count(unit_range));
+        let mut implicit_block = None;
+
+        for stmt in stmts {
+            if stmt.range.start() >= params.end {
+                break;
+            }
+            let runs = stmt.function.is_none();
+            if runs {
+                let block = implicit_blocks
+                    .iter()
+                    .position(|block| block.contains(stmt.range.start()));
+                if block != implicit_block {
+                    if implicit_block.is_some() {
+                        file_context.end_implicit_block();
+                    }
+                    if block.is_some() {
+                        file_context.begin_implicit_block();
+                    }
+                    implicit_block = block;
+                }
+            }
+            skip_unparsable(&mut unparsable, stmt.range.start(), &mut file_context);
+
+            visit(stmt, &file_context);
+
+            let placeholders = stmt.has_identifier_parameters && !is_plain_query(&stmt.root);
+            if runs && placeholders {
+                file_context.taint_catalog();
+            }
+            file_context.next(runs && !placeholders);
+        }
+
+        skip_unparsable(&mut unparsable, unit.end.min(params.end), &mut file_context);
+        // The section that contains `end` decides what a statement there sees.
+        if params.end < unit.end {
+            return file_context.into_catalog();
+        }
+        if !is_down {
+            catalog = file_context.into_catalog();
+        }
+    }
+    catalog
+}
+
+/// Moves past the statements that don't parse and start before `before`. Those that may have
+/// changed the catalog taint it.
+fn skip_unparsable<'a>(
+    unparsable: &mut std::iter::Peekable<impl Iterator<Item = &'a UnparsableStatement>>,
+    before: TextSize,
+    file_context: &mut AnalysedFileContext,
+) {
+    while let Some(skipped) = unparsable.next_if(|skipped| skipped.range.start() < before) {
+        if may_change_catalog(&skipped.sql) {
+            file_context.taint_catalog();
+        }
+    }
+}
+
+/// Statements that run together: those before the first section marker, then each section.
+struct MigrationUnit<'a> {
+    section: Option<&'a MigrationSection>,
+    /// The indices of the statements.
+    stmts: Range<usize>,
+    start: TextSize,
+    /// The last unit ends at the end of all text, so that it contains the end of the file.
+    end: TextSize,
+}
+
+fn migration_units<'a>(
+    stmts: &[AnalysableStatement],
+    sections: &'a [MigrationSection],
+) -> Vec<MigrationUnit<'a>> {
+    let end_of_text = TextSize::from(u32::MAX);
+    let before_sections = sections
+        .first()
+        .map_or(end_of_text, |section| section.range.start());
+    let units = std::iter::once((None, TextSize::from(0), before_sections)).chain(
+        sections.iter().enumerate().map(|(i, section)| {
+            let end = if i + 1 == sections.len() {
+                end_of_text
+            } else {
+                section.range.end()
+            };
+            (Some(section), section.range.start(), end)
+        }),
+    );
+
+    let mut first = 0;
+    units
+        .map(|(section, start, end)| {
+            let last = first + stmts[first..].partition_point(|stmt| stmt.range.start() < end);
+            let unit = MigrationUnit {
+                section,
+                stmts: first..last,
+                start,
+                end,
+            };
+            first = last;
+            unit
+        })
+        .collect()
+}
+
+/// Whether a statement that doesn't parse may have changed the catalog. Queries, `COPY`,
+/// maintenance, cursors, prepared statements, comments and privileges can't create, drop or
+/// rename objects, unless they are `SELECT ... INTO` or `EXPLAIN ANALYZE CREATE TABLE AS`. For
+/// anything else we can't tell.
+fn may_change_catalog(sql: &str) -> bool {
+    use pgls_query::protobuf::Token;
+    let Ok(scan) = pgls_query::scan(sql) else {
+        return true;
+    };
+    let mut tokens = scan
+        .tokens
+        .iter()
+        .map(|token| token.token())
+        .filter(|token| !matches!(token, Token::SqlComment | Token::CComment | Token::Ascii40));
+    match tokens.next() {
+        Some(
+            Token::Insert
+            | Token::Update
+            | Token::DeleteP
+            | Token::Merge
+            | Token::Copy
+            | Token::Analyse
+            | Token::Analyze
+            | Token::Vacuum
+            | Token::Reindex
+            | Token::Cluster
+            | Token::Checkpoint
+            | Token::LockP
+            | Token::Listen
+            | Token::Notify
+            | Token::Unlisten
+            | Token::Show
+            | Token::Declare
+            | Token::Fetch
+            | Token::Move
+            | Token::Close
+            | Token::Prepare
+            | Token::Execute
+            | Token::Deallocate
+            | Token::Comment
+            | Token::Grant
+            | Token::Revoke,
+        ) => false,
+        Some(Token::Select | Token::With | Token::Values | Token::Table | Token::Explain) => {
+            tokens.any(|token| matches!(token, Token::Into | Token::Create))
+        }
+        _ => true,
+    }
+}
+
+/// Whether the statement only reads or writes rows, without changing the catalog or the
+/// session.
+pub(crate) fn is_plain_query(root: &pgls_query::NodeEnum) -> bool {
+    use pgls_query::NodeEnum;
+    match root {
+        NodeEnum::SelectStmt(select) => select.into_clause.is_none(),
+        NodeEnum::InsertStmt(_)
+        | NodeEnum::UpdateStmt(_)
+        | NodeEnum::DeleteStmt(_)
+        | NodeEnum::MergeStmt(_) => true,
+        _ => false,
+    }
+}
diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs
index 98e8ab48e..0f60170f6 100644
--- a/crates/pgls_workspace/src/workspace/server.rs
+++ b/crates/pgls_workspace/src/workspace/server.rs
@@ -27,10 +27,10 @@ use pgls_analyse::{
 use pgls_analyser::StatementAnalysis;
 use pgls_analyser::{
     Analyser, AnalyserConfig, AnalyserParams, AnalysisResult, Direction, FileKind,
-    LinterDiagnostic, LinterOptions, MigrationSection, RuleFix,
+    LinterDiagnostic, LinterOptions, MigrationSection, RuleFix, WalkParams,
 };
 use pgls_catalog::resolve::{ColumnSource, Reference, ResolveParams, resolve};
-use pgls_catalog::{Catalog, CatalogBase, Origin, Session, TransactionMode};
+use pgls_catalog::{Catalog, CatalogBase, Origin, Session};
 use std::borrow::Cow;
 
 use pgls_catalog::Snapshot;
@@ -458,12 +458,10 @@ impl WorkspaceServer {
             .collect()
     }
 
-    /// Walks the statements of `doc` that start before `end` like the analyser does: the catalog
-    /// and the session follow the statements, so that each statement sees the objects created
-    /// earlier in the file. Each migration section gets a session of its own, and the changes of
-    /// a down section don't carry into the next section. `visit` gets every parsed statement with
-    /// the catalog and the session as they are before it. `settings` are the ones of the file.
-    /// Returns the catalog after the walked statements.
+    /// Walks the statements of `doc` that start before `end` like the analyser does, see
+    /// [`pgls_analyser::walk_statements`]. `visit` gets every parsed statement, except those with
+    /// identifier parameters, with the catalog and the session as they are before it. `settings`
+    /// are the ones of the file. Returns the catalog as a statement at `end` sees it.
     fn walk_statements(
         &self,
         settings: Option<&Settings>,
@@ -480,69 +478,43 @@ impl WorkspaceServer {
                 )
             })
             .unwrap_or_default();
-        let search_path = snapshot.expand_search_path(&search_path);
-        let new_session = |section: Option<&MigrationSection>| {
-            Session::new(search_path.clone()).with_transaction_mode(
-                section.map_or(TransactionMode::Unknown, MigrationSection::transaction_mode),
-            )
-        };
-
-        // Editor features show what names most likely refer to, so a `DO` block or a statement
-        // that doesn't parse doesn't turn them off for the rest of the file.
-        let mut catalog = Catalog::new(Some(self.catalog_base(snapshot))).best_effort();
-        let mut session = new_session(None);
-        let mut section: Option<&MigrationSection> = None;
-        // The catalog after the up sections so far, while walking a down section.
-        let mut up_catalog: Option<Catalog> = None;
-        for (stmt, _, unparsable) in doc.iter(AnalyserDiagnosticsMapper) {
-            let Some(start) = stmt.as_ref().map(|stmt| stmt.range.start()).or_else(|| {
-                unparsable
-                    .as_ref()
-                    .map(|unparsable| unparsable.range.start())
-            }) else {
+        let mut stmts = Vec::new();
+        let mut unparsable = Vec::new();
+        for (stmt, _, skipped) in doc.iter(AnalyserDiagnosticsMapper) {
+            let Some(start) = stmt
+                .as_ref()
+                .map(|stmt| stmt.range.start())
+                .or_else(|| skipped.as_ref().map(|skipped| skipped.range.start()))
+            else {
                 continue;
             };
             if start >= end {
                 break;
             }
-            let current = sections
-                .iter()
-                .find(|section| section.range.contains(start));
-            if current != section {
-                if let Some(up_catalog) = up_catalog.take() {
-                    catalog = up_catalog;
-                }
-                if current.is_some_and(|section| section.direction == Direction::Down) {
-                    up_catalog = Some(catalog.clone());
-                }
-                session = new_session(current);
-                section = current;
-            }
+            stmts.extend(stmt);
+            unparsable.extend(skipped);
+        }
 
-            if let Some(unparsable) = unparsable {
-                if pgls_analyser::may_change_catalog(&unparsable.sql) {
-                    catalog.taint();
-                }
-                continue;
-            }
-            let Some(stmt) = stmt else {
-                continue;
-            };
-            // The names of statements with identifier parameters like `:name` are placeholders.
-            if stmt.has_identifier_parameters {
-                if !pgls_analyser::is_plain_query(&stmt.root) {
-                    catalog.taint();
+        pgls_analyser::walk_statements(
+            WalkParams {
+                stmts: &stmts,
+                unparsable: &unparsable,
+                catalog_base: Some(self.catalog_base(snapshot)),
+                search_path: &snapshot.expand_search_path(&search_path),
+                sections: &sections,
+                end,
+                // Editor features show what names most likely refer to, so a `DO` block or a
+                // statement that doesn't parse doesn't turn them off for the rest of the file.
+                best_effort: true,
+            },
+            |stmt, file_context| {
+                // The names of statements with identifier parameters like `:name` are
+                // placeholders.
+                if !stmt.has_identifier_parameters {
+                    visit(stmt, file_context.catalog(), file_context.session());
                 }
-                continue;
-            }
-            visit(&stmt, &catalog, &session);
-            // Statements in the bodies of SQL functions are not executed.
-            if stmt.function.is_none() {
-                catalog.apply(&stmt.root, session.search_path());
-                session.apply(&stmt.root);
-            }
-        }
-        catalog
+            },
+        )
     }
 
     /// For hover on the statement `id` at `range`: the references whose span contains
diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs
index 5680d18fa..f3fad69a4 100644
--- a/crates/pgls_workspace/src/workspace/server.tests.rs
+++ b/crates/pgls_workspace/src/workspace/server.tests.rs
@@ -4413,3 +4413,142 @@ async fn test_sql_function_bodies(test_db: PgPool) {
         .collect::<Vec<_>>();
     assert_eq!(unknown_columns, ["nope", "gone"], "{diagnostics:#?}");
 }
+
+/// Lint, hover and completions see the same catalog at each statement: after a down section,
+/// after a statement that doesn't parse, and after a statement on identifier parameters.
+#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")]
+async fn lint_hover_and_completions_see_the_same_catalog(test_db: PgPool) {
+    let mut conf = PartialConfiguration::init();
+    conf.merge_with(PartialConfiguration {
+        db: Some(PartialDatabaseConfiguration {
+            database: Some(
+                test_db
+                    .connect_options()
+                    .get_database()
+                    .unwrap()
+                    .to_string(),
+            ),
+            ..Default::default()
+        }),
+        migrations: Some(PartialMigrationsConfiguration {
+            section_markers: Some(SectionMarkers::Dbmate),
+            ..Default::default()
+        }),
+        ..Default::default()
+    });
+    let (workspace, project_key) =
+        get_test_workspace(Some(conf)).expect("Unable to create test workspace");
+    let path = PgLSPath::new("20250301120000_migration.sql");
+
+    // Opens `content` without the `|`, and returns the typecheck diagnostics of the file, the
+    // completions at the `|`, and the hover on each needle.
+    let check = |content: &str, needles: &[&str]| {
+        let cursor = content.find('|').expect("cursor in content");
+        let content = content.replace('|', "");
+        workspace
+            .open_file(OpenFileParams {
+                project_key,
+                path: path.clone(),
+                content: content.clone(),
+                version: 1,
+            })
+            .expect("Unable to open test file");
+        let typecheck: Vec<_> = workspace
+            .pull_file_diagnostics(crate::workspace::PullFileDiagnosticsParams {
+                project_key,
+                path: path.clone(),
+                categories: RuleCategories::all(),
+                max_diagnostics: 100,
+                only: vec![],
+                skip: vec![],
+            })
+            .expect("Unable to pull diagnostics")
+            .diagnostics
+            .into_iter()
+            .filter(is_typecheck)
+            .map(|diagnostic| {
+                let span = diagnostic.location().span.expect("span");
+                content[span].to_string()
+            })
+            .collect();
+        let completions: Vec<_> = workspace
+            .get_completions(crate::workspace::GetCompletionsParams {
+                project_key,
+                path: path.clone(),
+                position: u32::try_from(cursor).unwrap().into(),
+            })
+            .expect("Unable to request completions")
+            .items
+            .into_iter()
+            .map(|item| item.label)
+            .collect();
+        let hovers: Vec<_> = needles
+            .iter()
+            .map(|needle| {
+                let offset = content.rfind(needle).expect("needle in content") + 2;
+                workspace
+                    .on_hover(OnHoverParams {
+                        project_key,
+                        path: path.clone(),
+                        position: u32::try_from(offset).unwrap().into(),
+                    })
+                    .expect("Unable to hover")
+                    .markdown_blocks
+                    .join("\n")
+            })
+            .collect();
+        workspace
+            .close_file(crate::workspace::CloseFileParams {
+                project_key,
+                path: path.clone(),
+            })
+            .expect("Unable to close test file");
+        (typecheck, completions, hovers)
+    };
+
+    // The changes of a down section don't carry into the next up section, also for its first
+    // statement.
+    let (typecheck, completions, hovers) = check(
+        "-- migrate:up\n\
+         create table walk_items (id int);\n\
+         -- migrate:down\n\
+         alter table walk_items add column down_only int;\n\
+         -- migrate:up\n\
+         select | from walk_items;\n\
+         select down_only from walk_items;\n",
+        &["walk_items;"],
+    );
+    assert_eq!(typecheck, ["down_only"]);
+    assert!(completions.contains(&"id".to_string()), "{completions:?}");
+    assert!(
+        !completions.contains(&"down_only".to_string()),
+        "{completions:?}"
+    );
+    assert!(hovers[0].contains("- id: int4"), "{hovers:?}");
+    assert!(!hovers[0].contains("down_only"), "{hovers:?}");
+
+    // A statement that doesn't parse may have created anything, but the objects created before
+    // it are still known.
+    let (typecheck, completions, hovers) = check(
+        "create table walk_parsed (id int);\n\
+         create tabel walk_unparsed (id int);\n\
+         select id from walk_unparsed;\n\
+         select | from walk_parsed;\n",
+        &["walk_parsed;"],
+    );
+    assert!(typecheck.is_empty(), "{typecheck:?}");
+    assert!(completions.contains(&"id".to_string()), "{completions:?}");
+    assert!(hovers[0].contains("- id: int4"), "{hovers:?}");
+
+    // The names of a statement with identifier parameters are placeholders: `:tbl` doesn't
+    // create a table `tbl`.
+    let (typecheck, completions, hovers) = check(
+        "create table :tbl (id int);\n\
+         select id from tbl;\n\
+         select | from tbl;\n",
+        &["tbl;"],
+    );
+    assert!(typecheck.is_empty(), "{typecheck:?}");
+    assert!(!completions.contains(&"id".to_string()), "{completions:?}");
+    assert_eq!(hovers, [""]);
+}
diff --git a/crates/pgls_workspace/src/workspace/server/document.rs b/crates/pgls_workspace/src/workspace/server/document.rs
index b64297f31..f44b2c35f 100644
--- a/crates/pgls_workspace/src/workspace/server/document.rs
+++ b/crates/pgls_workspace/src/workspace/server/document.rs
@@ -247,7 +247,12 @@ impl<'a> StatementMapper<'a> for DefaultMapper {
 
 pub struct ExecuteStatementMapper;
 impl<'a> StatementMapper<'a> for ExecuteStatementMapper {
-    type Output = (StatementId, TextRange, String, Option<pgls_query::NodeEnum>);
+    type Output = (
+        StatementId,
+        TextRange,
+        String,
+        Option<Arc<pgls_query::NodeEnum>>,
+    );
 
     fn map(&self, doc: &'a Document, stmt: DocumentStatement) -> Self::Output {
         let ast_result = doc.ast_db.get_or_cache_ast(&stmt.id);
@@ -267,7 +272,7 @@ impl<'a> StatementMapper<'a> for FormatStatementMapper {
         StatementId,
         TextRange,
         String,
-        Result<pgls_query::NodeEnum, pgls_query_ext::diagnostics::SyntaxDiagnostic>,
+        Result<Arc<pgls_query::NodeEnum>, pgls_query_ext::diagnostics::SyntaxDiagnostic>,
     );
 
     fn map(&self, doc: &'a Document, stmt: DocumentStatement) -> Self::Output {
@@ -287,7 +292,7 @@ impl<'a> StatementMapper<'a> for TypecheckDiagnosticsMapper {
     type Output = (
         StatementId,
         TextRange,
-        Option<pgls_query::NodeEnum>,
+        Option<Arc<pgls_query::NodeEnum>>,
         Arc<tree_sitter::Tree>,
     );
 
@@ -341,7 +346,9 @@ impl<'a> StatementMapper<'a> for AnalyserDiagnosticsMapper {
         // Statements in the body of a SQL function can reference its parameters.
         let function = id.parent().and_then(|root| {
             let parent = doc.ast_db.get_or_cache_ast(&root);
-            let Ok(pgls_query::NodeEnum::CreateFunctionStmt(create_fn)) = parent.as_ref() else {
+            let Ok(pgls_query::NodeEnum::CreateFunctionStmt(create_fn)) =
+                parent.as_ref().as_deref()
+            else {
                 return None;
             };
             pgls_catalog::resolve::FunctionContext::from_create_function(create_fn)
diff --git a/crates/pgls_workspace/src/workspace/server/formatting.rs b/crates/pgls_workspace/src/workspace/server/formatting.rs
index 1e448c0e2..97b754603 100644
--- a/crates/pgls_workspace/src/workspace/server/formatting.rs
+++ b/crates/pgls_workspace/src/workspace/server/formatting.rs
@@ -1,7 +1,7 @@
 //! Formatting of a document. Each statement is printed on its own, and the text between the
 //! statements is kept, since the parser drops comments.
 
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 
 use pgls_diagnostics::{DiagnosticExt, serde::Diagnostic as SDiagnostic};
 use pgls_pretty_print::FormatConfig;
@@ -141,7 +141,11 @@ pub(super) fn format_document(
             Ok(ast) => {
                 let mut ast = ast;
                 if let Some(formatted_sql_fn_body) = formatted_sql_fn_bodies.get(&id) {
-                    sql_function::set_formatted_sql_fn_body(&mut ast, formatted_sql_fn_body);
+                    // Copies the tree only for the functions whose body is replaced.
+                    sql_function::set_formatted_sql_fn_body(
+                        Arc::make_mut(&mut ast),
+                        formatted_sql_fn_body,
+                    );
                 }
 
                 match pgls_pretty_print::format_statement(&ast, &text, &config) {
diff --git a/crates/pgls_workspace/src/workspace/server/pg_query.rs b/crates/pgls_workspace/src/workspace/server/pg_query.rs
index c8987b1e2..97e2d4eb3 100644
--- a/crates/pgls_workspace/src/workspace/server/pg_query.rs
+++ b/crates/pgls_workspace/src/workspace/server/pg_query.rs
@@ -10,10 +10,14 @@ use rustc_hash::FxHashMap;
 
 use super::statement_identifier::StatementId;
 
+/// The tree of a statement, or why it doesn't parse.
+pub type ParseResult = Result<Arc<pgls_query::NodeEnum>, SyntaxDiagnostic>;
+
 /// The parse results of the statements of a document. They are kept until the statement is
 /// no longer in the document, see [`PgQueryStore::retain`].
 pub struct PgQueryStore {
-    ast_db: Mutex<FxHashMap<StatementId, Arc<Result<pgls_query::NodeEnum, SyntaxDiagnostic>>>>,
+    /// The tree is in an `Arc` of its own, so that the analyser can share it.
+    ast_db: Mutex<FxHashMap<StatementId, Arc<ParseResult>>>,
     plpgsql_db: Mutex<FxHashMap<StatementId, Result<(), SyntaxDiagnostic>>>,
     /// The statements embedded in each top-level statement, like those of a SQL function body.
     embedded_db: Mutex<FxHashMap<StatementId, Arc<[EmbeddedStatement]>>>,
@@ -53,10 +57,7 @@ impl PgQueryStore {
         embedded
     }
 
-    pub fn get_or_cache_ast(
-        &self,
-        statement: &StatementId,
-    ) -> Arc<Result<pgls_query::NodeEnum, SyntaxDiagnostic>> {
+    pub fn get_or_cache_ast(&self, statement: &StatementId) -> Arc<ParseResult> {
         let mut cache = self.ast_db.lock().unwrap();
 
         if let Some(existing) = cache.get(statement) {
@@ -67,7 +68,7 @@ impl PgQueryStore {
             pgls_query::parse(&convert_to_positional_params(statement.content()))
                 .map_err(SyntaxDiagnostic::from)
                 .and_then(|ast| {
-                    ast.into_root().ok_or_else(|| {
+                    ast.into_root().map(Arc::new).ok_or_else(|| {
                         SyntaxDiagnostic::new("No root node found in parse result", None)
                     })
                 }),
@@ -82,7 +83,7 @@ impl PgQueryStore {
     ) -> Option<Result<(), SyntaxDiagnostic>> {
         let ast = self.get_or_cache_ast(statement);
 
-        let create_fn = match ast.as_ref() {
+        let create_fn = match ast.as_ref().as_deref() {
             Ok(pgls_query::NodeEnum::CreateFunctionStmt(node)) => node,
             _ => return None,
         };

From 3495f0d6ebb49212d08b6ab268c56cb09afbcd1b Mon Sep 17 00:00:00 2001
From: psteinroe <philipp@steinroetter.com>
Date: Mon, 26 Oct 2026 16:48:23 +0000
Subject: [PATCH 2/3] fix(analyser): keep the session effects of statements
 with identifier parameters

A statement whose names are psql variables still takes its locks and
changes the transaction, so the walk applies it to the session and only
taints the catalog. The rules spec harness converts psql variables like
the workspace does, and three specs check that the lock rules see such
statements.
---
 crates/pgls_analyser/src/linter_context.rs    | 15 +++++++------
 crates/pgls_analyser/src/walk.rs              |  9 +++++---
 crates/pgls_analyser/tests/rules_tests.rs     | 10 +++++++--
 .../identifier_parameter.sql                  |  4 ++++
 .../identifier_parameter.sql.snap             | 19 +++++++++++++++++
 .../valid_new_table_identifier_parameter.sql  |  9 ++++++++
 ...id_new_table_identifier_parameter.sql.snap | 17 +++++++++++++++
 .../identifier_parameter.sql                  |  4 ++++
 .../identifier_parameter.sql.snap             | 21 +++++++++++++++++++
 9 files changed, 97 insertions(+), 11 deletions(-)
 create mode 100644 crates/pgls_analyser/tests/specs/safety/avoidWideLockWindow/identifier_parameter.sql
 create mode 100644 crates/pgls_analyser/tests/specs/safety/avoidWideLockWindow/identifier_parameter.sql.snap
 create mode 100644 crates/pgls_analyser/tests/specs/safety/lockTimeoutWarning/valid_new_table_identifier_parameter.sql
 create mode 100644 crates/pgls_analyser/tests/specs/safety/lockTimeoutWarning/valid_new_table_identifier_parameter.sql.snap
 create mode 100644 crates/pgls_analyser/tests/specs/safety/runningStatementWhileHoldingAccessExclusive/identifier_parameter.sql
 create mode 100644 crates/pgls_analyser/tests/specs/safety/runningStatementWhileHoldingAccessExclusive/identifier_parameter.sql.snap

diff --git a/crates/pgls_analyser/src/linter_context.rs b/crates/pgls_analyser/src/linter_context.rs
index 52807100c..14f484d4c 100644
--- a/crates/pgls_analyser/src/linter_context.rs
+++ b/crates/pgls_analyser/src/linter_context.rs
@@ -302,12 +302,15 @@ impl<'a> AnalysedFileContext<'a> {
 
     /// Moves to the next statement.
     ///
-    /// Executed statements (`applies_effects`) update the catalog and the session. Statements
-    /// in the bodies of SQL functions are only analysed, not executed.
-    pub fn next(&mut self, applies_effects: bool) {
-        if applies_effects && self.pos < self.stmts.len() {
-            let stmt = self.stmts[self.pos];
-            self.catalog.apply(stmt, self.session.search_path());
+    /// A statement that `runs` updates the session, and the catalog too if `updates_catalog`.
+    /// Statements in the bodies of SQL functions are only analysed, not executed. A statement on
+    /// placeholder names, like psql's `:name`, still takes its locks and changes the transaction,
+    /// but the catalog can't follow it.
+    pub fn next(&mut self, runs: bool, updates_catalog: bool) {
+        if runs && let Some(stmt) = self.stmts.get(self.pos) {
+            if updates_catalog {
+                self.catalog.apply(stmt, self.session.search_path());
+            }
             self.session.apply(stmt);
         }
         self.pos += 1;
diff --git a/crates/pgls_analyser/src/walk.rs b/crates/pgls_analyser/src/walk.rs
index 4973b3277..2cb9c5ba2 100644
--- a/crates/pgls_analyser/src/walk.rs
+++ b/crates/pgls_analyser/src/walk.rs
@@ -42,8 +42,8 @@ pub struct WalkParams<'a> {
 /// - A top-level statement that doesn't parse taints the catalog if it may have changed it.
 /// - The statements of SQL function bodies are visited, but they don't run.
 /// - The names of a statement with identifier parameters, like psql's `:name`, are placeholders.
-///   Such a statement is visited, but if it may change the catalog, it taints it instead of being
-///   applied.
+///   Such a statement is visited and applied to the session (its locks and its effect on the
+///   transaction), but if it may change the catalog, it taints it instead of being applied to it.
 ///
 /// The file context only knows the statements that the walk is given. When they stop before the
 /// end of the file, [`AnalysedFileContext::stmt_count`] and the implicit transaction blocks of
@@ -146,11 +146,14 @@ pub fn walk_statements(
 
             visit(stmt, &file_context);
 
+            // The names of a statement with identifier parameters are placeholders, so the
+            // catalog can't follow its changes. Its locks and its effect on the transaction are
+            // still those of the statement.
             let placeholders = stmt.has_identifier_parameters && !is_plain_query(&stmt.root);
             if runs && placeholders {
                 file_context.taint_catalog();
             }
-            file_context.next(runs && !placeholders);
+            file_context.next(runs, !placeholders);
         }
 
         skip_unparsable(&mut unparsable, unit.end.min(params.end), &mut file_context);
diff --git a/crates/pgls_analyser/tests/rules_tests.rs b/crates/pgls_analyser/tests/rules_tests.rs
index 1a1dbf88d..4d8682d24 100644
--- a/crates/pgls_analyser/tests/rules_tests.rs
+++ b/crates/pgls_analyser/tests/rules_tests.rs
@@ -142,11 +142,17 @@ fn analyse(analyser: &Analyser, query: &str, is_typecheck: bool) -> Vec<LinterDi
     let mut unparsable = Vec::new();
     for r in &split.ranges {
         let text = &query[*r];
-        match pgls_query::parse(text) {
+        // psql variables (`:name`, `:'name'`) become positional parameters, as in the workspace.
+        let conversion = pgls_lexer::convert_to_positional_params_with_metadata(text);
+        match pgls_query::parse(&conversion.sql) {
             Ok(ast) => {
                 let root = ast.into_root().expect("Failed to convert AST to root node");
                 let body = sql_function_body(&root, text, *r);
-                stmts.push(AnalysableStatement::new(root, *r).with_sql(text));
+                stmts.push(
+                    AnalysableStatement::new(root, *r)
+                        .with_sql(text)
+                        .with_identifier_parameters(conversion.has_identifier_parameters),
+                );
                 stmts.extend(body);
             }
             Err(_) => unparsable.push(UnparsableStatement {
diff --git a/crates/pgls_analyser/tests/specs/safety/avoidWideLockWindow/identifier_parameter.sql b/crates/pgls_analyser/tests/specs/safety/avoidWideLockWindow/identifier_parameter.sql
new file mode 100644
index 000000000..e8e470c1f
--- /dev/null
+++ b/crates/pgls_analyser/tests/specs/safety/avoidWideLockWindow/identifier_parameter.sql
@@ -0,0 +1,4 @@
+-- expect_lint/avoidWideLockWindow
+-- The name of the column is a psql variable, but the table is known, so its lock is held.
+ALTER TABLE users ADD COLUMN :"column" TEXT;
+ALTER TABLE orders ADD COLUMN total NUMERIC;
diff --git a/crates/pgls_analyser/tests/specs/safety/avoidWideLockWindow/identifier_parameter.sql.snap b/crates/pgls_analyser/tests/specs/safety/avoidWideLockWindow/identifier_parameter.sql.snap
new file mode 100644
index 000000000..72f6dde52
--- /dev/null
+++ b/crates/pgls_analyser/tests/specs/safety/avoidWideLockWindow/identifier_parameter.sql.snap
@@ -0,0 +1,19 @@
+---
+source: crates/pgls_analyser/tests/rules_tests.rs
+expression: snapshot
+---
+# Input
+```
+-- expect_lint/avoidWideLockWindow
+-- The name of the column is a psql variable, but the table is known, so its lock is held.
+ALTER TABLE users ADD COLUMN :"column" TEXT;
+ALTER TABLE orders ADD COLUMN total NUMERIC;
+
+```
+
+# Diagnostics
+lint/avoidWideLockWindow ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+  × Acquiring a lock on public.orders while already holding ACCESS EXCLUSIVE locks on other tables.
+  
+  i This widens the lock window. Split into separate transactions to minimize lock duration.
diff --git a/crates/pgls_analyser/tests/specs/safety/lockTimeoutWarning/valid_new_table_identifier_parameter.sql b/crates/pgls_analyser/tests/specs/safety/lockTimeoutWarning/valid_new_table_identifier_parameter.sql
new file mode 100644
index 000000000..c8ef243ac
--- /dev/null
+++ b/crates/pgls_analyser/tests/specs/safety/lockTimeoutWarning/valid_new_table_identifier_parameter.sql
@@ -0,0 +1,9 @@
+-- expect_no_diagnostics
+-- The table is created in the same transaction, even though the name of a column is a psql
+-- variable.
+CREATE TABLE users (
+    id INT PRIMARY KEY,
+    :"column" TEXT
+);
+
+ALTER TABLE users ADD COLUMN email TEXT;
diff --git a/crates/pgls_analyser/tests/specs/safety/lockTimeoutWarning/valid_new_table_identifier_parameter.sql.snap b/crates/pgls_analyser/tests/specs/safety/lockTimeoutWarning/valid_new_table_identifier_parameter.sql.snap
new file mode 100644
index 000000000..66325544e
--- /dev/null
+++ b/crates/pgls_analyser/tests/specs/safety/lockTimeoutWarning/valid_new_table_identifier_parameter.sql.snap
@@ -0,0 +1,17 @@
+---
+source: crates/pgls_analyser/tests/rules_tests.rs
+expression: snapshot
+---
+# Input
+```
+-- expect_no_diagnostics
+-- The table is created in the same transaction, even though the name of a column is a psql
+-- variable.
+CREATE TABLE users (
+    id INT PRIMARY KEY,
+    :"column" TEXT
+);
+
+ALTER TABLE users ADD COLUMN email TEXT;
+
+```
diff --git a/crates/pgls_analyser/tests/specs/safety/runningStatementWhileHoldingAccessExclusive/identifier_parameter.sql b/crates/pgls_analyser/tests/specs/safety/runningStatementWhileHoldingAccessExclusive/identifier_parameter.sql
new file mode 100644
index 000000000..856a99ac6
--- /dev/null
+++ b/crates/pgls_analyser/tests/specs/safety/runningStatementWhileHoldingAccessExclusive/identifier_parameter.sql
@@ -0,0 +1,4 @@
+-- expect_only_lint/runningStatementWhileHoldingAccessExclusive
+-- The name of the column is a psql variable, but the table is known, so its lock is held.
+ALTER TABLE authors ADD COLUMN :"column" TEXT;
+SELECT pg_sleep(60);
diff --git a/crates/pgls_analyser/tests/specs/safety/runningStatementWhileHoldingAccessExclusive/identifier_parameter.sql.snap b/crates/pgls_analyser/tests/specs/safety/runningStatementWhileHoldingAccessExclusive/identifier_parameter.sql.snap
new file mode 100644
index 000000000..1c030bb75
--- /dev/null
+++ b/crates/pgls_analyser/tests/specs/safety/runningStatementWhileHoldingAccessExclusive/identifier_parameter.sql.snap
@@ -0,0 +1,21 @@
+---
+source: crates/pgls_analyser/tests/rules_tests.rs
+expression: snapshot
+---
+# Input
+```
+-- expect_only_lint/runningStatementWhileHoldingAccessExclusive
+-- The name of the column is a psql variable, but the table is known, so its lock is held.
+ALTER TABLE authors ADD COLUMN :"column" TEXT;
+SELECT pg_sleep(60);
+
+```
+
+# Diagnostics
+lint/runningStatementWhileHoldingAccessExclusive ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+  × Running statement while holding ACCESS EXCLUSIVE lock.
+  
+  i This blocks all access to the table for the duration of this statement.
+  
+  i Run this statement in a separate transaction to minimize lock duration.

From b220118b91cb3c41d9ffcfd72e6c27805d0b4bc3 Mon Sep 17 00:00:00 2001
From: psteinroe <philipp@steinroetter.com>
Date: Sat, 31 Oct 2026 09:17:30 +0000
Subject: [PATCH 3/3] perf(analyser): find the implicit blocks of a walk in one
 pass

Counting the statements of each query scanned all statements of the
section, and each statement scanned the blocks for its own, which made
the walk quadratic in a goose NO TRANSACTION migration with many
StatementBegin blocks. The statement starts are now sorted once and
counted with binary searches, and the statements move through the
blocks with a cursor.
---
 crates/pgls_analyser/src/migration.rs |   2 +-
 crates/pgls_analyser/src/walk.rs      | 148 ++++++++++++++++++++++----
 2 files changed, 131 insertions(+), 19 deletions(-)

diff --git a/crates/pgls_analyser/src/migration.rs b/crates/pgls_analyser/src/migration.rs
index 22a1ae543..6156e9447 100644
--- a/crates/pgls_analyser/src/migration.rs
+++ b/crates/pgls_analyser/src/migration.rs
@@ -21,7 +21,7 @@ pub struct MigrationSection {
     /// The parts of the section that the migration tool sends to Postgres as one query, like a
     /// whole dbmate section or a goose `StatementBegin` block. Postgres runs a query with more
     /// than one statement in a transaction block, also in a section without a transaction. The
-    /// other statements of the section are sent one by one.
+    /// other statements of the section are sent one by one. In order, and they don't overlap.
     pub queries: Vec<TextRange>,
 }
 
diff --git a/crates/pgls_analyser/src/walk.rs b/crates/pgls_analyser/src/walk.rs
index 2cb9c5ba2..cf2993ba6 100644
--- a/crates/pgls_analyser/src/walk.rs
+++ b/crates/pgls_analyser/src/walk.rs
@@ -78,27 +78,32 @@ pub fn walk_statements(
                 .map_or(TransactionMode::Unknown, MigrationSection::transaction_mode),
         );
         let stmts = &params.stmts[unit.stmts.clone()];
-        let unit_range = TextRange::new(unit.start, unit.end);
-        let last_end = params
-            .unparsable
+        let unit_unparsable = {
+            let starts = |end: TextSize| {
+                params
+                    .unparsable
+                    .partition_point(|skipped| skipped.range.start() < end)
+            };
+            &params.unparsable[starts(unit.start)..starts(unit.end)]
+        };
+        let last_end = unit_unparsable
             .iter()
-            .filter(|skipped| unit_range.contains(skipped.range.start()))
             .map(|skipped| skipped.range.end())
             .chain(stmts.last().map(|stmt| stmt.range.end()))
             .max()
             .unwrap_or_default();
-        // The statements the unit runs, including those that don't parse, but not those in
-        // function bodies.
+        // The starts of the statements the unit runs, in order, including those that don't
+        // parse, but not those in function bodies.
+        let mut starts: Vec<TextSize> = stmts
+            .iter()
+            .filter(|stmt| stmt.function.is_none())
+            .map(|stmt| stmt.range.start())
+            .chain(unit_unparsable.iter().map(|skipped| skipped.range.start()))
+            .collect();
+        starts.sort_unstable();
         let stmt_count = |range: TextRange| {
-            stmts
-                .iter()
-                .filter(|stmt| stmt.function.is_none() && range.contains(stmt.range.start()))
-                .count()
-                + params
-                    .unparsable
-                    .iter()
-                    .filter(|skipped| range.contains(skipped.range.start()))
-                    .count()
+            starts.partition_point(|start| *start < range.end())
+                - starts.partition_point(|start| *start < range.start())
         };
         // Outside a transaction, the queries with more than one statement run in an implicit
         // transaction block.
@@ -114,13 +119,15 @@ pub fn walk_statements(
                     .collect()
             })
             .unwrap_or_default();
+        // The blocks that don't end before the statement, which comes after those before it.
+        let mut implicit_blocks = implicit_blocks.into_iter().peekable();
         let mut file_context =
             AnalysedFileContext::new(&roots[unit.stmts.clone()], unit_catalog, session)
                 .with_positions(
                     stmts.iter().map(|stmt| stmt.range.start()).collect(),
                     last_end,
                 )
-                .with_stmt_count(stmt_count(unit_range));
+                .with_stmt_count(stmt_count(TextRange::new(unit.start, unit.end)));
         let mut implicit_block = None;
 
         for stmt in stmts {
@@ -129,9 +136,14 @@ pub fn walk_statements(
             }
             let runs = stmt.function.is_none();
             if runs {
+                while implicit_blocks
+                    .next_if(|block| block.end() <= stmt.range.start())
+                    .is_some()
+                {}
                 let block = implicit_blocks
-                    .iter()
-                    .position(|block| block.contains(stmt.range.start()));
+                    .peek()
+                    .copied()
+                    .filter(|block| block.contains(stmt.range.start()));
                 if block != implicit_block {
                     if implicit_block.is_some() {
                         file_context.end_implicit_block();
@@ -290,3 +302,103 @@ pub(crate) fn is_plain_query(root: &pgls_query::NodeEnum) -> bool {
         _ => false,
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::time::{Duration, Instant};
+
+    use pgls_text_size::{TextRange, TextSize};
+
+    use super::{WalkParams, walk_statements};
+    use crate::{AnalysableStatement, Direction, MigrationSection, UnparsableStatement};
+
+    /// Walks the statements of `sql` in a section without a transaction, with the given queries.
+    /// Returns each statement that parses with whether it runs in an implicit transaction block,
+    /// and how long the walk took.
+    fn walk_queries(sql: &str, queries: Vec<TextRange>) -> (Vec<(String, bool)>, Duration) {
+        let mut stmts = Vec::new();
+        let mut unparsable = Vec::new();
+        for range in pgls_statement_splitter::split(sql).ranges {
+            let text = &sql[range];
+            match pgls_query::parse(text) {
+                Ok(parsed) => stmts.push(
+                    AnalysableStatement::new(parsed.into_root().unwrap(), range).with_sql(text),
+                ),
+                Err(_) => unparsable.push(UnparsableStatement {
+                    range,
+                    sql: text.to_string(),
+                }),
+            }
+        }
+        let sections = [MigrationSection {
+            range: TextRange::up_to(TextSize::of(sql)),
+            direction: Direction::Up,
+            transaction: false,
+            queries,
+        }];
+        let mut visited = Vec::new();
+        let start = Instant::now();
+        walk_statements(
+            WalkParams {
+                stmts: &stmts,
+                unparsable: &unparsable,
+                catalog_base: None,
+                search_path: &[],
+                sections: &sections,
+                end: TextSize::of(sql),
+                best_effort: false,
+            },
+            |stmt, file_context| {
+                visited.push((
+                    stmt.sql.clone().unwrap_or_default(),
+                    file_context.session().in_implicit_block(),
+                ))
+            },
+        );
+        (visited, start.elapsed())
+    }
+
+    #[test]
+    fn queries_with_several_statements_run_in_implicit_blocks() {
+        let sql = "select 1; select 2; select 3; selec 4; select 5; select 6;";
+        let query = |from: &str, to: &str| {
+            TextRange::new(
+                TextSize::try_from(sql.find(from).unwrap()).unwrap(),
+                TextSize::try_from(sql.find(to).unwrap() + to.len()).unwrap(),
+            )
+        };
+        let (visited, _) = walk_queries(
+            sql,
+            vec![
+                query("select 1", "select 2;"),
+                // A single statement, and one with a statement that doesn't parse.
+                query("select 3", "select 3;"),
+                query("selec 4", "select 5;"),
+            ],
+        );
+        assert_eq!(
+            visited,
+            [
+                ("select 1;".to_string(), true),
+                ("select 2;".to_string(), true),
+                ("select 3;".to_string(), false),
+                ("select 5;".to_string(), true),
+                ("select 6;".to_string(), false),
+            ]
+        );
+    }
+
+    #[test]
+    fn many_queries_are_walked_in_linear_time() {
+        // Two statements per query, like goose `StatementBegin` blocks. Looking up the
+        // statements of each query among all statements took seconds.
+        let queries = 8_000;
+        let sql = "select 1; select 2;\n".repeat(queries);
+        let blocks = (0..queries)
+            .map(|i| TextRange::at(TextSize::try_from(i * 20).unwrap(), 19.into()))
+            .collect();
+        let (visited, elapsed) = walk_queries(&sql, blocks);
+        assert!(visited.iter().all(|(_, in_block)| *in_block));
+        assert!(elapsed < Duration::from_secs(1), "{elapsed:?}");
+    }
+}