PruningOptions -> RetentionPolicies

make life easier, requires epochs_to_keep with optional overrides. otherwise, no retention policy at all

strumming the pruning config. pruner.rs file will define prunable tables. the toml file when parsed into Rust config will warn if it tries to name any variants not supported by indexer code. additionally, there's also a check to make sure that prunable tables actually exist in the db

crates/sui-graphql-rpc/src/test_infra/cluster.rs

@@ -14,7 +14,7 @@ use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::Duration;
 use sui_graphql_rpc_client::simple_client::SimpleClient;
-use sui_indexer::config::PruningOptions;
+pub use sui_indexer::config::RetentionConfig;
 pub use sui_indexer::config::SnapshotLagConfig;
 use sui_indexer::errors::IndexerError;
 use sui_indexer::store::PgIndexerStore;
@@ -151,7 +151,7 @@ pub async fn start_network_cluster() -> NetworkCluster {
 pub async fn serve_executor(
     executor: Arc<dyn RestStateReader + Send + Sync>,
     snapshot_config: Option<SnapshotLagConfig>,
-    epochs_to_keep: Option<u64>,
+    retention_config: Option<RetentionConfig>,
     data_ingestion_path: PathBuf,
 ) -> ExecutorCluster {
     let database = TempDb::new().unwrap();
@@ -184,7 +184,7 @@ pub async fn serve_executor(
     let (pg_store, pg_handle, _) = start_indexer_writer_for_testing(
         db_url,
         Some(snapshot_config.clone()),
-        Some(PruningOptions { epochs_to_keep }),
+        retention_config,
         Some(data_ingestion_path),
         Some(cancellation_token.clone()),
     )

crates/sui-indexer/Cargo.toml

@@ -31,11 +31,14 @@ regex.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 serde_with.workspace = true
+strum.workspace = true
+strum_macros.workspace = true
 tap.workspace = true
 tempfile.workspace = true
 thiserror.workspace = true
 tokio = { workspace = true, features = ["full"] }
 tokio-util = { workspace = true, features = ["rt"] }
+toml.workspace = true
 tracing.workspace = true
 url.workspace = true
 

crates/sui-indexer/README.md

@@ -32,7 +32,7 @@ See the [docs](https://docs.sui.io/guides/developer/getting-started/local-networ
 
 Start a local network using the `sui` binary:
 ```sh
-cargo run --bin sui -- start --with-faucet --force-regenesis 
+cargo run --bin sui -- start --with-faucet --force-regenesis
 ```
 
 If you want to run a local network with the indexer enabled (note that `libpq` is required), you can run the following command after following the steps in the next section to set up an indexer DB:
@@ -65,11 +65,12 @@ cargo run --bin sui-indexer -- --db-url "<DATABASE_URL>" --rpc-client-url "https
 ```
 cargo run --bin sui-indexer -- --db-url "<DATABASE_URL>" --rpc-client-url "https://fullnode.devnet.sui.io:443" --rpc-server-worker
 ```
-More flags info can be found in this [file](https://github.com/MystenLabs/sui/blob/main/crates/sui-indexer/src/lib.rs#L83-L123).
+More flags info can be found in this [file](src/main.rs#L41).
+
 ### DB reset
-Run this command under `sui/crates/sui-indexer`, which will wipe DB; In case of schema changes in `.sql` files, this will also update corresponding `schema.rs` file.
+When making db-related changes, you may find yourself having to run migrations and reset dbs often. The commands below are how you can invoke these actions.
 ```sh
-diesel database reset --database-url="<DATABASE_URL>"
+cargo run --bin sui-indexer -- --database-url "<DATABASE_URL>" reset-database --force
 ```
 
 ## Steps to run locally (TiDB)

crates/sui-indexer/src/config.rs

@@ -1,14 +1,20 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::backfill::BackfillTaskKind;
 use crate::db::ConnectionPoolConfig;
+use crate::{backfill::BackfillTaskKind, handlers::pruner::PrunableTable};
 use clap::{Args, Parser, Subcommand};
-use std::{net::SocketAddr, path::PathBuf};
+use serde::{Deserialize, Serialize};
+use std::{collections::HashMap, net::SocketAddr, path::PathBuf};
+use strum::IntoEnumIterator;
 use sui_json_rpc::name_service::NameServiceConfig;
 use sui_types::base_types::{ObjectID, SuiAddress};
 use url::Url;
 
+/// The primary purpose of objects_history is to serve consistency query.
+/// A short retention is sufficient.
+const OBJECTS_HISTORY_EPOCHS_TO_KEEP: u64 = 2;
+
 #[derive(Parser, Clone, Debug)]
 #[clap(
     name = "Sui indexer",
@@ -208,8 +214,93 @@ pub enum Command {
 
 #[derive(Args, Default, Debug, Clone)]
 pub struct PruningOptions {
-    #[arg(long, env = "EPOCHS_TO_KEEP")]
-    pub epochs_to_keep: Option<u64>,
+    /// Path to TOML file containing configuration for retention policies.
+    #[arg(long)]
+    pub pruning_config_path: Option<PathBuf>,
+}
+
+/// Represents the default retention policy and overrides for prunable tables. Instantiated only if
+/// `PruningOptions` is provided on indexer start.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RetentionConfig {
+    /// Default retention policy for all tables.
+    pub epochs_to_keep: u64,
+    /// A map of tables to their respective retention policies that will override the default.
+    /// Prunable tables not named here will use the default retention policy.
+    #[serde(default)]
+    pub overrides: HashMap<PrunableTable, u64>,
+}
+
+impl PruningOptions {
+    /// Load default retention policy and overrides from file.
+    pub fn load_from_file(&self) -> Option<RetentionConfig> {
+        let config_path = self.pruning_config_path.as_ref()?;
+
+        let contents = std::fs::read_to_string(config_path)
+            .expect("Failed to read default retention policy and overrides from file");
+        let retention_with_overrides = toml::de::from_str::<RetentionConfig>(&contents)
+            .expect("Failed to parse into RetentionConfig struct");
+
+        let default_retention = retention_with_overrides.epochs_to_keep;
+
+        assert!(
+            default_retention > 0,
+            "Default retention must be greater than 0"
+        );
+        assert!(
+            retention_with_overrides
+                .overrides
+                .values()
+                .all(|&policy| policy > 0),
+            "All retention overrides must be greater than 0"
+        );
+
+        Some(retention_with_overrides)
+    }
+}
+
+impl RetentionConfig {
+    /// Create a new `RetentionConfig` with the specified default retention and overrides. Call
+    /// `finalize()` on the instance to update the `policies` field with the default retention
+    /// policy for all tables that do not have an override specified.
+    pub fn new(epochs_to_keep: u64, overrides: HashMap<PrunableTable, u64>) -> Self {
+        Self {
+            epochs_to_keep,
+            overrides,
+        }
+    }
+
+    pub fn new_with_default_retention_only_for_testing(epochs_to_keep: u64) -> Self {
+        let mut overrides = HashMap::new();
+        overrides.insert(
+            PrunableTable::ObjectsHistory,
+            OBJECTS_HISTORY_EPOCHS_TO_KEEP,
+        );
+
+        Self::new(epochs_to_keep, HashMap::new())
+    }
+
+    /// Consumes this struct to produce a full mapping of every prunable table and its retention
+    /// policy. By default, every prunable table will have the default retention policy from
+    /// `epochs_to_keep`. Some tables like `objects_history` will observe a different default
+    /// retention policy. These default values are overridden by any entries in `overrides`.
+    pub fn retention_policies(self) -> HashMap<PrunableTable, u64> {
+        let RetentionConfig {
+            epochs_to_keep,
+            mut overrides,
+        } = self;
+
+        for table in PrunableTable::iter() {
+            let default_retention = match table {
+                PrunableTable::ObjectsHistory => OBJECTS_HISTORY_EPOCHS_TO_KEEP,
+                _ => epochs_to_keep,
+            };
+
+            overrides.entry(table).or_insert(default_retention);
+        }
+
+        overrides
+    }
 }
 
 #[derive(Args, Debug, Clone)]
@@ -290,7 +381,9 @@ impl Default for RestoreConfig {
 #[cfg(test)]
 mod test {
     use super::*;
+    use std::io::Write;
     use tap::Pipe;
+    use tempfile::NamedTempFile;
 
     fn parse_args<'a, T>(args: impl IntoIterator<Item = &'a str>) -> Result<T, clap::error::Error>
     where
@@ -354,4 +447,129 @@ mod test {
         // fullnode rpc url must be present
         parse_args::<JsonRpcConfig>([]).unwrap_err();
     }
+
+    #[test]
+    fn pruning_options_with_objects_history_override() {
+        let mut temp_file = NamedTempFile::new().unwrap();
+        let toml_content = r#"
+        epochs_to_keep = 5
+
+        [overrides]
+        objects_history = 10
+        transactions = 20
+        "#;
+        temp_file.write_all(toml_content.as_bytes()).unwrap();
+        let temp_path: PathBuf = temp_file.path().to_path_buf();
+        let pruning_options = PruningOptions {
+            pruning_config_path: Some(temp_path.clone()),
+        };
+        let retention_config = pruning_options.load_from_file().unwrap();
+
+        // Assert the parsed values
+        assert_eq!(retention_config.epochs_to_keep, 5);
+        assert_eq!(
+            retention_config
+                .overrides
+                .get(&PrunableTable::ObjectsHistory)
+                .copied(),
+            Some(10)
+        );
+        assert_eq!(
+            retention_config
+                .overrides
+                .get(&PrunableTable::Transactions)
+                .copied(),
+            Some(20)
+        );
+        assert_eq!(retention_config.overrides.len(), 2);
+
+        let retention_policies = retention_config.retention_policies();
+
+        for table in PrunableTable::iter() {
+            let Some(retention) = retention_policies.get(&table).copied() else {
+                panic!("Expected a retention policy for table {:?}", table);
+            };
+
+            match table {
+                PrunableTable::ObjectsHistory => assert_eq!(retention, 10),
+                PrunableTable::Transactions => assert_eq!(retention, 20),
+                _ => assert_eq!(retention, 5),
+            };
+        }
+    }
+
+    #[test]
+    fn pruning_options_no_objects_history_override() {
+        let mut temp_file = NamedTempFile::new().unwrap();
+        let toml_content = r#"
+        epochs_to_keep = 5
+
+        [overrides]
+        tx_affected_addresses = 10
+        transactions = 20
+        "#;
+        temp_file.write_all(toml_content.as_bytes()).unwrap();
+        let temp_path: PathBuf = temp_file.path().to_path_buf();
+        let pruning_options = PruningOptions {
+            pruning_config_path: Some(temp_path.clone()),
+        };
+        let retention_config = pruning_options.load_from_file().unwrap();
+
+        // Assert the parsed values
+        assert_eq!(retention_config.epochs_to_keep, 5);
+        assert_eq!(
+            retention_config
+                .overrides
+                .get(&PrunableTable::TxAffectedAddresses)
+                .copied(),
+            Some(10)
+        );
+        assert_eq!(
+            retention_config
+                .overrides
+                .get(&PrunableTable::Transactions)
+                .copied(),
+            Some(20)
+        );
+        assert_eq!(retention_config.overrides.len(), 2);
+
+        let retention_policies = retention_config.retention_policies();
+
+        for table in PrunableTable::iter() {
+            let Some(retention) = retention_policies.get(&table).copied() else {
+                panic!("Expected a retention policy for table {:?}", table);
+            };
+
+            match table {
+                PrunableTable::ObjectsHistory => {
+                    assert_eq!(retention, OBJECTS_HISTORY_EPOCHS_TO_KEEP)
+                }
+                PrunableTable::TxAffectedAddresses => assert_eq!(retention, 10),
+                PrunableTable::Transactions => assert_eq!(retention, 20),
+                _ => assert_eq!(retention, 5),
+            };
+        }
+    }
+
+    #[test]
+    fn test_invalid_pruning_config_file() {
+        let toml_str = r#"
+        epochs_to_keep = 5
+
+        [overrides]
+        objects_history = 10
+        transactions = 20
+        invalid_table = 30
+        "#;
+
+        let result = toml::from_str::<RetentionConfig>(toml_str);
+        assert!(result.is_err(), "Expected an error, but parsing succeeded");
+
+        if let Err(e) = result {
+            assert!(
+                e.to_string().contains("unknown variant `invalid_table`"),
+                "Error message doesn't mention the invalid table"
+            );
+        }
+    }
 }