spg_engine/

maintenance.rs

//! Table-maintenance executors: `ANALYZE` (re-stat) and
//! `COMPACT COLD SEGMENTS` (cold-segment merge). Split out of
//! `lib.rs` (cut 19); verbatim move, the only edit beyond
//! visibility is reuniting the `exec_compact_cold_segments` doc,
//! whose first half had drifted above `set_session_param`.

use alloc::string::{String, ToString};
use alloc::vec::Vec;

use spg_storage::{ColumnSchema, CompactReport, DataType, IndexKind, Row, StorageError, Value};

use crate::{
    COMPACTION_TARGET_DEFAULT_BYTES, Engine, EngineError, QueryResult, canonical_value_repr,
    is_internal_table_name, sort_values_for_histogram, statistics,
};

impl Engine {
    /// v6.2.0 — `ANALYZE [<table>]` runtime. Bare `ANALYZE` walks
    /// every user table; `ANALYZE <name>` re-stats one. For each
    /// target table, single-pass scan + per-column histogram +
    /// `null_frac` + `n_distinct`. Replaces the table's prior
    /// stats; resets the modified-row counter.
    ///
    /// v6.2.0 doesn't sample — it scans the full table. v6.2.x
    /// can add reservoir sampling at the > 100 K-row mark; not a
    /// scope blocker for the current commit since rows ≤ 100 K
    /// analyse in milliseconds.
    /// v7.39 (round 320, V53) — `DISCARD { ALL | PLANS | SEQUENCES | TEMP }`.
    /// pgbouncer sends `DISCARD ALL` between pooled client sessions so the
    /// next client sees a clean connection. It used to be swallowed as
    /// "dump noise" on the theory that SPG had no per-connection state
    /// worth discarding — untrue since round 279 gave every connection its
    /// own session bag (GUC overrides, prepared statements, large-object
    /// descriptors). A no-op meant one pooled client's state leaked to the
    /// next one.
    ///
    /// PG 18.4 measured: `SET application_name='x'; PREPARE p …;
    /// DISCARD ALL` leaves application_name back at its startup value and
    /// `EXECUTE p` failing with "prepared statement \"p\" does not
    /// exist"; inside a transaction block it is
    /// `ERROR: DISCARD ALL cannot run inside a transaction block`.
    ///
    /// v7.39 (round 321, V54) — cursors are closed here too, as PG's
    /// DISCARD ALL does. Round 320 had to leave them alone because the
    /// cursor table was process-wide: closing "all" from one connection
    /// would have closed another's. They live in the session bag now.
    pub(crate) fn exec_discard(
        &mut self,
        target: spg_sql::ast::DiscardTarget,
    ) -> Result<QueryResult, EngineError> {
        use spg_sql::ast::DiscardTarget;
        if matches!(target, DiscardTarget::All | DiscardTarget::Temp)
            && self.current_tx.is_some_and(|tx| self.is_tx_open(tx))
        {
            return Err(EngineError::Unsupported(alloc::format!(
                "DISCARD {target} cannot run inside a transaction block"
            )));
        }
        match target {
            DiscardTarget::All => {
                self.reset_all_gucs();
                self.prepared_statements.clear();
                self.lo_descriptors.clear();
                self.lo_next_fd = 0;
                self.cursors.clear();
                self.listen_channels.clear();
                self.plan_cache.clear();
                self.refresh_render_style();
            }
            DiscardTarget::Plans => self.plan_cache.clear(),
            // SPG has neither temp tables nor per-session sequence state
            // (`currval` reads the catalog), so there is nothing of either
            // kind to throw away. Accepted so a client's sequence of
            // DISCARDs runs, and tagged with the target it named.
            DiscardTarget::Sequences | DiscardTarget::Temp => {}
        }
        Ok(QueryResult::CommandOk {
            affected: 0,
            modified_catalog: false,
        })
    }

    pub(crate) fn exec_analyze(
        &mut self,
        target: Option<&str>,
    ) -> Result<QueryResult, EngineError> {
        // v7.38 元机制 D acceptor — `SPG_TEST_STATS_FROZEN=1` turns
        // ANALYZE into a no-op so a test's statistic-version snapshot
        // is stable across sessions. Pairs with the plan-cache
        // version-aware invalidation: freezing stats also freezes
        // plan reuse, which is what regression tests want when
        // proving "same SQL, same plan".
        if self.env_cfg().stats_frozen {
            return Ok(QueryResult::CommandOk {
                affected: 0,
                modified_catalog: false,
            });
        }
        let names: Vec<String> = if let Some(name) = target {
            // Verify the table exists; surface a clear error if not.
            // v7.38.18 (S10) — the ACTIVE catalog, which includes the
            // open transaction's shadow. A multi-statement simple query
            // is an implicit transaction, so
            //
            //   CREATE TABLE t (k INT); INSERT INTO t VALUES (1); ANALYZE t;
            //
            // sent as ONE string answered `relation "t" does not exist`
            // — while the INSERT in the same string had just succeeded.
            // Seven other statement kinds in that position were already
            // right; `ANALYZE` read the committed catalog alone.
            if self.active_catalog().get(name).is_none() {
                return Err(EngineError::Storage(StorageError::TableNotFound {
                    name: name.to_string(),
                }));
            }
            alloc::vec![name.to_string()]
        } else {
            // Same reason: a bare ANALYZE must cover the tables this
            // transaction created, not only the committed ones.
            self.active_catalog()
                .table_names()
                .into_iter()
                .filter(|n| !is_internal_table_name(n))
                .collect()
        };
        let mut analysed = 0usize;
        let now_us = self.clock.map(|f| f());
        for table_name in &names {
            self.analyze_one_table(table_name)?;
            // v7.39 (pg_stat knife C) — stamp last_analyze.
            if let Some(us) = now_us
                && let Some(t) = self.active_catalog_mut().get_mut(table_name)
            {
                t.stamp_analyze(us);
            }
            analysed += 1;
        }
        // v6.3.1 — plan cache invalidation. Bump stats version so
        // future lookups see the new generation, and selectively
        // evict every plan whose `source_tables` overlap with the
        // ANALYZE target set. Bare ANALYZE (all tables) clears the
        // whole cache.
        if analysed > 0 {
            self.statistics.bump_version();
            if target.is_some() {
                for t in &names {
                    self.plan_cache.evict_referencing(t);
                }
            } else {
                self.plan_cache.clear();
            }
        }
        Ok(QueryResult::CommandOk {
            affected: analysed,
            modified_catalog: true,
        })
    }

    /// Walk a single table's rows once and (re-)populate per-column
    /// stats. Drops the existing stats for `table` first so columns
    /// that have been DROP-ed between ANALYZEs don't leave stale
    /// rows.
    fn analyze_one_table(&mut self, table_name: &str) -> Result<(), EngineError> {
        // v7.38.18 (S10) — take what this pass needs and release the
        // catalog borrow. `active_catalog()` borrows the whole engine
        // (it may hand back the transaction's shadow), where the old
        // `self.catalog` borrowed one field, and the statistics store
        // below is mutated through `self`.
        let (schema, rows) = {
            let table = self.active_catalog().get(table_name).ok_or_else(|| {
                EngineError::Storage(StorageError::TableNotFound {
                    name: table_name.to_string(),
                })
            })?;
            (table.schema().clone(), table.rows().clone())
        };
        let row_count = rows.len();
        // For each column, collect (sorted) non-NULL textual values
        // + count NULLs; then ask `statistics::build_histogram` to
        // produce the 101 bounds and `estimate_n_distinct` the
        // distinct count.
        self.statistics.clear_table(table_name);
        for (col_pos, col_schema) in schema.columns.iter().enumerate() {
            // v6.2.0 skip: vector columns have their own stats
            // shape (HNSW graph topology). v6.2 deliberation #1.
            if matches!(col_schema.ty, DataType::Vector { .. }) {
                continue;
            }
            let mut non_null_values: Vec<Value<'static>> = Vec::with_capacity(row_count);
            let mut nulls: u64 = 0;
            for row in rows.iter() {
                match row.values.get(col_pos) {
                    Some(Value::Null) | None => nulls += 1,
                    Some(v) => non_null_values.push(v.clone()),
                }
            }
            // Sort by type-aware ordering (Int as int, Text as
            // lex, etc.) so histogram bounds reflect the column's
            // natural order — not lexicographic on the string
            // representation, which would put "9" after "49".
            non_null_values.sort_by(|a, b| sort_values_for_histogram(a, b));
            let non_null: Vec<String> = non_null_values.iter().map(canonical_value_repr).collect();
            let null_frac = if row_count == 0 {
                0.0
            } else {
                #[allow(clippy::cast_precision_loss)]
                let f = nulls as f32 / row_count as f32;
                f
            };
            let n_distinct = statistics::estimate_n_distinct(&non_null);
            let histogram_bounds = statistics::build_histogram(&non_null);
            self.statistics.set(
                table_name.to_string(),
                col_schema.name.clone(),
                statistics::ColumnStats {
                    null_frac,
                    n_distinct,
                    histogram_bounds,
                },
            );
        }
        self.statistics.reset_modified(table_name);
        // v6.7.0 — refresh the per-table cold_rows cache. Walk the
        // BTree indices and count Cold locators (MAX across
        // indices); store the result on the table. Surfaced via
        // `spg_statistic.cold_row_count` (new column) and
        // `spg_stat_segment.table_name` (new column).
        let cold_count = {
            let table = self
                .active_catalog()
                .get(table_name)
                .expect("table still present");
            table.count_cold_locators()
        };
        let table_mut = self
            .active_catalog_mut()
            .get_mut(table_name)
            .expect("table still present");
        table_mut.set_cold_row_count(cold_count);
        Ok(())
    }

    /// v6.7.3 — `COMPACT COLD SEGMENTS` runtime path. Drives the
    /// engine-layer compaction shim with the default
    /// 4 MiB segment-size threshold. spg-server intercepts the
    /// SQL before it reaches the engine on a server build —
    /// it reads `SPG_COMPACTION_TARGET_SEGMENT_BYTES`, calls
    /// `Engine::compact_cold_segments_with_target` directly with
    /// the env value, and persists every merged segment to
    /// `<db>.spg/segments/`. This arm only fires for engine-only
    /// callers (spg-embedded, lib tests); in that mode merged
    /// segments live in memory and are dropped at process exit.
    /// v7.37.17 (17.6 sibling) — `TRUNCATE [TABLE] <t>[, ...]
    /// [RESTART IDENTITY]`. Clears every row from each named
    /// table by dispatching to `Table::truncate()` (already exists
    /// for internal callers). RESTART IDENTITY additionally resets
    /// the table's associated sequence back to its start value.
    pub(crate) fn exec_truncate(
        &mut self,
        tables: &[String],
        _restart_identity: bool,
        only: bool,
    ) -> Result<QueryResult, EngineError> {
        // RESTART IDENTITY is parsed but not honored yet — the
        // SequenceDef doesn't expose a restart primitive on the
        // storage side today. Accepted-and-no-op for pg_dump compat;
        // real sequence reset lands with the sequence-lifecycle epic.
        //
        // v7.39 (round 647) — `ONLY` and the descent it turns off.
        // Measured on PG18: a plain TRUNCATE of a parent empties its
        // children too; `TRUNCATE ONLY <inheritance parent>` empties the
        // parent alone; and `TRUNCATE ONLY <partitioned parent>` is not
        // a no-op but an error, because a partitioned parent holds
        // nothing and the spelling can only mean a mistake.
        //
        // The whole expansion sits inside this branch rather than beside
        // it — rounds 641/643/644/646 each measured what an extra test
        // in a function's body costs when the row loop is downstream.
        let mut targets: Vec<String> = tables.to_vec();
        if !tables.is_empty() {
            for name in tables {
                if only {
                    if crate::partition::is_partition_parent(self.active_catalog(), name) {
                        return Err(EngineError::Unsupported(alloc::format!(
                            "cannot truncate only a partitioned table"
                        )));
                    }
                    continue;
                }
                let mut frontier = alloc::vec![name.clone()];
                while let Some(cur) = frontier.pop() {
                    for kid in crate::partition::children_of_parent(self.active_catalog(), &cur) {
                        frontier.push(kid.clone());
                        targets.push(kid);
                    }
                }
            }
        }
        let mut affected: usize = 0;
        for name in &targets {
            let cat = self.active_catalog_mut();
            let Some(t) = cat.get_mut(name) else {
                // v7.39 (read01 round 50) — PG says "relation" for TRUNCATE
                // (only DROP TABLE says "table"). 42P01 at the wire.
                return Err(EngineError::Storage(StorageError::Corrupt(alloc::format!(
                    "relation {name:?} does not exist"
                ))));
            };
            affected = affected.saturating_add(t.row_count());
            t.truncate();
        }
        Ok(QueryResult::CommandOk {
            affected,
            modified_catalog: false,
        })
    }

    pub(crate) fn exec_compact_cold_segments(&mut self) -> Result<QueryResult, EngineError> {
        let target = COMPACTION_TARGET_DEFAULT_BYTES;
        let reports = self.compact_cold_segments_with_target(target)?;
        let columns = alloc::vec![
            ColumnSchema::new("table_name", DataType::Text, false),
            ColumnSchema::new("index_name", DataType::Text, false),
            ColumnSchema::new("sources_merged", DataType::BigInt, false),
            ColumnSchema::new("merged_segment_id", DataType::BigInt, false),
            ColumnSchema::new("merged_rows", DataType::BigInt, false),
            ColumnSchema::new("deleted_rows_pruned", DataType::BigInt, false),
            ColumnSchema::new("bytes_reclaimed_estimate", DataType::BigInt, false),
        ];
        let rows: Vec<Row<'static>> = reports
            .into_iter()
            .map(|(tname, iname, report)| {
                Row::new(alloc::vec![
                    Value::text(tname),
                    Value::text(iname),
                    Value::BigInt(i64::try_from(report.sources.len()).unwrap_or(i64::MAX)),
                    Value::BigInt(i64::from(report.merged_segment_id.unwrap_or(0))),
                    Value::BigInt(i64::try_from(report.merged_rows).unwrap_or(i64::MAX)),
                    Value::BigInt(i64::try_from(report.deleted_rows_pruned).unwrap_or(i64::MAX),),
                    Value::BigInt(
                        i64::try_from(report.bytes_reclaimed_estimate).unwrap_or(i64::MAX),
                    ),
                ])
            })
            .collect();
        Ok(QueryResult::Rows { columns, rows })
    }

    /// v6.7.3 — public shim around `Catalog::compact_cold_segments`
    /// driving every BTree index on every user table. Returns one
    /// `(table, index, report)` triple for each merge that
    /// actually happened (no-op (table, index) pairs are filtered
    /// out so callers can size persist-side work to the live
    /// merges). Caller is responsible for persisting each
    /// `report.merged_segment_bytes` and updating the on-disk
    /// segment registry; engine layer is no_std and never
    /// touches disk.
    ///
    /// Marks every touched table's cached `cold_row_count` stale
    /// — compaction GC'd some shadowed rows, so the count must be
    /// re-derived on the next ANALYZE.
    pub fn compact_cold_segments_with_target(
        &mut self,
        target_segment_bytes: u64,
    ) -> Result<Vec<(String, String, CompactReport)>, EngineError> {
        let table_names = self.active_catalog().table_names();
        let mut reports: Vec<(String, String, CompactReport)> = Vec::new();
        for tname in table_names {
            if is_internal_table_name(&tname) {
                continue;
            }
            let idx_names: Vec<String> = {
                let Some(t) = self.active_catalog().get(&tname) else {
                    continue;
                };
                t.indices()
                    .iter()
                    .filter(|i| matches!(i.kind, IndexKind::BTree(_)))
                    .map(|i| i.name.clone())
                    .collect()
            };
            for iname in idx_names {
                let report = self
                    .active_catalog_mut()
                    .compact_cold_segments(&tname, &iname, target_segment_bytes)
                    .map_err(EngineError::Storage)?;
                if report.merged_segment_id.is_some() {
                    if let Some(t) = self.active_catalog_mut().get_mut(&tname) {
                        t.mark_cold_row_count_stale();
                    }
                    reports.push((tname.clone(), iname, report));
                }
            }
        }
        Ok(reports)
    }
}