add script to snapshot state data

Author: kobayurii

configuration/src/configs/database.rs

@@ -5,7 +5,7 @@ use crate::configs::{deserialize_data_or_env, deserialize_optional_data_or_env};
 
 // Database connection URL
 // Example: "postgres://user:password@localhost:5432/dbname"
-type DatabaseConnectUrl = String;
+pub type DatabaseConnectUrl = String;
 
 #[derive(Validate, serde_derive::Deserialize, Debug, Clone, Default)]
 pub struct ShardDatabaseConfig {

configuration/src/lib.rs

@@ -9,7 +9,7 @@ mod configs;
 mod default_env_configs;
 pub mod utils;
 
-pub use crate::configs::database::DatabaseConfig;
+pub use crate::configs::database::{DatabaseConfig, DatabaseConnectUrl};
 pub use crate::configs::general::ChainId;
 pub use crate::configs::general::StorageProvider;
 pub use crate::configs::{

database/src/postgres/docker-compose.yml

@@ -1,5 +1,3 @@
-version: '3.7'
-
 services:
   metadata:
     image: postgres:15.5
@@ -10,6 +8,11 @@ services:
       POSTGRES_PASSWORD: password
     ports:
       - "5422:5432"
+    command: ["postgres", "-c", "max_connections=4000"]
+    ulimits:
+      nofile:
+        soft: 1048576
+        hard: 1048576
 
   shard_0:
     image: postgres:15.5
@@ -20,6 +23,11 @@ services:
       POSTGRES_PASSWORD: password
     ports:
       - "5430:5432"
+    command: ["postgres", "-c", "max_connections=4000"]
+    ulimits:
+      nofile:
+        soft: 1048576
+        hard: 1048576
 
   shard_1:
     image: postgres:15.5
@@ -30,6 +38,11 @@ services:
       POSTGRES_PASSWORD: password
     ports:
       - "5431:5432"
+    command: ["postgres", "-c", "max_connections=4000"]
+    ulimits:
+      nofile:
+        soft: 1048576
+        hard: 1048576
 
   shard_2:
     image: postgres:15.5
@@ -40,6 +53,7 @@ services:
       POSTGRES_PASSWORD: password
     ports:
       - "5432:5432"
+    command: ["postgres", "-c", "max_connections=4000"]
 
   shard_3:
     image: postgres:15.5
@@ -50,6 +64,11 @@ services:
       POSTGRES_PASSWORD: password
     ports:
       - "5433:5432"
+    command: ["postgres", "-c", "max_connections=4000"]
+    ulimits:
+      nofile:
+        soft: 1048576
+        hard: 1048576
 
   shard_4:
     image: postgres:15.5
@@ -60,6 +79,11 @@ services:
       POSTGRES_PASSWORD: password
     ports:
       - "5434:5432"
+    command: ["postgres", "-c", "max_connections=4000"]
+    ulimits:
+      nofile:
+        soft: 1048576
+        hard: 1048576
 
   shard_5:
     image: postgres:15.5
@@ -70,3 +94,8 @@ services:
       POSTGRES_PASSWORD: password
     ports:
       - "5435:5432"
+    command: ["postgres", "-c", "max_connections=4000"]
+    ulimits:
+      nofile:
+        soft: 1048576
+        hard: 1048576

database/src/postgres/migrations/shard_db/20250430152143_procedure_create_new_range_tables.up.sql

@@ -70,5 +70,77 @@ BEGIN
             CREATE TABLE IF NOT EXISTS state_changes_account_%s_%s PARTITION OF state_changes_account_%s FOR VALUES WITH (MODULUS 100, REMAINDER %s);
             ', range_id, i, range_id, i);
     END LOOP;
+
+
+    EXECUTE format('
+        CREATE TABLE IF NOT EXISTS state_changes_data_%s_compact (
+            account_id text NOT NULL,
+            block_height_from numeric(20,0) NOT NULL,
+            block_height_to numeric(20,0) NULL,
+            data_key text NOT NULL,
+            data_value bytea NOT NULL,
+            PRIMARY KEY (account_id, data_key, block_height_from)
+        ) PARTITION BY HASH (account_id);
+    ', range_id);
+
+    -- Create 100 partitions
+    FOR i IN 0..99 LOOP
+        EXECUTE format('
+            CREATE TABLE IF NOT EXISTS state_changes_data_%s_compact_%s PARTITION OF state_changes_data_%s_compact FOR VALUES WITH (MODULUS 100, REMAINDER %s);
+            ', range_id, i, range_id, i);
+    END LOOP;
+
+
+    EXECUTE format('
+        CREATE TABLE IF NOT EXISTS state_changes_access_key_%s_compact (
+            account_id text NOT NULL,
+            block_height_from numeric(20,0) NOT NULL,
+            block_height_to numeric(20,0) NULL,
+            data_key text NOT NULL,
+            data_value bytea NOT NULL,
+            PRIMARY KEY (account_id, data_key, block_height_from)
+        ) PARTITION BY HASH (account_id);
+    ', range_id);
+
+    -- Create 100 partitions
+    FOR i IN 0..99 LOOP
+        EXECUTE format('
+            CREATE TABLE IF NOT EXISTS state_changes_access_key_%s_compact_%s PARTITION OF state_changes_access_key_%s_compact FOR VALUES WITH (MODULUS 100, REMAINDER %s);
+            ', range_id, i, range_id, i);
+    END LOOP;
+
+    EXECUTE format('
+        CREATE TABLE IF NOT EXISTS state_changes_contract_%s_compact (
+            account_id text NOT NULL,
+            block_height_from numeric(20,0) NOT NULL,
+            block_height_to numeric(20,0) NULL,
+            data_value bytea,
+            PRIMARY KEY (account_id, block_height_from)
+        ) PARTITION BY HASH (account_id);
+    ', range_id);
+
+    -- Create 100 partitions
+    FOR i IN 0..99 LOOP
+        EXECUTE format('
+            CREATE TABLE IF NOT EXISTS state_changes_contract_%s_compact_%s PARTITION OF state_changes_contract_%s_compact FOR VALUES WITH (MODULUS 100, REMAINDER %s);
+            ', range_id, i, range_id, i);
+    END LOOP;
+
+    EXECUTE format('
+        CREATE TABLE IF NOT EXISTS state_changes_account_%s_compact (
+            account_id text NOT NULL,
+            block_height_from numeric(20,0) NOT NULL,
+            block_height_to numeric(20,0) NULL,
+            data_value bytea,
+            PRIMARY KEY (account_id, block_height_from)
+        ) PARTITION BY HASH (account_id);
+    ', range_id);
+
+    -- Create 100 partitions
+    FOR i IN 0..99 LOOP
+        EXECUTE format('
+            CREATE TABLE IF NOT EXISTS state_changes_account_%s_compact_%s PARTITION OF state_changes_account_%s_compact FOR VALUES WITH (MODULUS 100, REMAINDER %s);
+            ', range_id, i, range_id, i);
+    END LOOP;
 END;
 $$;

database/src/postgres/migrations/shard_db/20250521171609_drop_block_hash_from_state_changes.up.sql

@@ -3,3 +3,9 @@ ALTER TABLE state_changes_data DROP COLUMN IF EXISTS block_hash;
 ALTER TABLE state_changes_access_key DROP COLUMN IF EXISTS block_hash;
 ALTER TABLE state_changes_contract DROP COLUMN IF EXISTS block_hash;
 ALTER TABLE state_changes_account DROP COLUMN IF EXISTS block_hash;
+
+-- TODO: REMOVE ALL DEPRECATED TABLES.
+DROP TABLE IF EXISTS state_changes_data;
+DROP TABLE IF EXISTS state_changes_access_key;
+DROP TABLE IF EXISTS state_changes_contract;
+DROP TABLE IF EXISTS state_changes_account;

logic-state-indexer/src/configs.rs

@@ -20,6 +20,7 @@ pub enum StartOptions {
         height: Option<u64>,
     },
     FromLatest,
+    FromGenesis,
 }
 
 pub async fn get_start_block_height(
@@ -28,7 +29,9 @@ pub async fn get_start_block_height(
     start_options: &StartOptions,
     indexer_id: &str,
 ) -> anyhow::Result<u64> {
+    let genesis_config = configuration::read_genesis_config_from_root()?;
     let start_block_height = match start_options {
+        StartOptions::FromGenesis => return Ok(genesis_config.genesis_height),
         StartOptions::FromBlock { height } => *height,
         StartOptions::FromInterruption { height } => {
             if let Ok(block_height) = db_manager.get_last_processed_block_height(indexer_id).await {
@@ -41,7 +44,11 @@ pub async fn get_start_block_height(
         }
         StartOptions::FromLatest => final_block_height(near_client).await?,
     };
-    Ok(start_block_height - 100) // Start just a bit earlier to overlap indexed blocks to ensure we don't miss anything in-between
+    if start_block_height - 100 < genesis_config.genesis_height {
+        Ok(genesis_config.genesis_height)
+    } else {
+        Ok(start_block_height - 100) // Start just a bit earlier to overlap indexed blocks to ensure we don't miss anything in-between
+    }
 }
 
 pub(crate) async fn final_block_height(

logic-state-indexer/src/lib.rs

@@ -160,13 +160,32 @@ impl StateChangesToStore {
     }
 }
 
+async fn task_migrate_and_compact_db(
+    shard_db_config: HashMap<ShardId, configuration::DatabaseConnectUrl>,
+    previous_range_id: u64,
+    current_range_id: u64,
+    current_range_start_block_height: u64,
+) {
+    for (_, shard_database_url) in shard_db_config {
+        let output = std::process::Command::new("migrate_compact_with_state/shard_migration.sh")
+            .arg(shard_database_url)
+            .arg(previous_range_id.to_string())
+            .arg(current_range_id.to_string())
+            .arg(current_range_start_block_height.to_string())
+            .output();
+
+        println!("Task_migrate_and_compact_db exited with status: {:?}", output);
+    }
+}
+
 #[cfg_attr(
     feature = "tracing-instrumentation",
     tracing::instrument(skip(streamer_message, db_manager))
 )]
 pub async fn handle_streamer_message(
     streamer_message: near_indexer_primitives::StreamerMessage,
     db_manager: &(impl database::StateIndexerDbManager + Sync + Send + 'static),
+    shard_db_config: HashMap<ShardId, configuration::DatabaseConnectUrl>,
     near_client: &(impl NearClient + std::fmt::Debug + Sync),
     indexer_config: impl configuration::RightsizingConfig
         + configuration::IndexerConfig
@@ -179,14 +198,33 @@ pub async fn handle_streamer_message(
     let block_hash = streamer_message.block.header.hash;
 
     let range_id = configuration::utils::get_data_range_id(&block_height).await?;
-    if stats.read().await.current_range_id < range_id {
-        db_manager.create_new_range_tables(range_id).await?;
+    let current_range_id = stats.read().await.current_range_id;
+    if current_range_id < range_id {
+        match db_manager.create_new_range_tables(range_id).await {
+            Ok(_) => {
+                tracing::info!(target: INDEXER, "Created new range tables for range {}", range_id);
+            }
+            Err(e) => {
+                tracing::error!(target: INDEXER, "Failed to create new range tables for range {}: {}", range_id, e);
+                return Err(e);
+            }
+        };
+        if current_range_id > 0 {
+            // We spawn a task to do this in the background
+            tracing::info!(target: INDEXER, "Migrating and compacting DB for range {}", current_range_id);
+            tokio::spawn(async move { task_migrate_and_compact_db(
+                shard_db_config,
+                current_range_id,
+                range_id,
+                block_height,
+            ).await });
+        }
         stats.write().await.current_range_id = range_id;
     }
 
     let current_epoch_id = streamer_message.block.header.epoch_id;
     let next_epoch_id = streamer_message.block.header.next_epoch_id;
-
+    println!("current_epoch_id {}, next_epoch_id {}", current_epoch_id, next_epoch_id);
     tracing::debug!(target: INDEXER, "Block height {}", block_height,);
 
     stats

state-indexer/migrate_compact_with_state/migrate_access_keys.sh

@@ -0,0 +1,65 @@
+#!/bin/bash
+
+DATABASE_URL=$1
+PREVIOUS_RANGE_ID=$2
+CURRENT_RANGE_ID=$3
+CURRENT_RANGE_BLOCK_HEIGHT_START=$4
+MAX_PARALLEL=${MAX_PARALLEL_JOBS}
+
+# Function to migrate a single partition
+migrate_partition() {
+    local partition=$1
+
+    psql "$DATABASE_URL" -c "
+        WITH ordered_data AS (
+            SELECT
+                account_id,
+                data_key,
+                data_value,
+                block_height AS block_height_from,
+                LAG(block_height) OVER (
+                    PARTITION BY account_id, data_key
+                    ORDER BY block_height DESC
+                ) AS block_height_to
+            FROM state_changes_access_key_${PREVIOUS_RANGE_ID}_$partition
+        ),
+        insert_compact AS (
+            INSERT INTO state_changes_access_key_${PREVIOUS_RANGE_ID}_compact_$partition (
+                account_id, data_key, data_value, block_height_from, block_height_to
+            )
+            SELECT
+                account_id, data_key, data_value, block_height_from, block_height_to
+            FROM ordered_data
+            WHERE data_value IS NOT NULL
+            ON CONFLICT (account_id, data_key, block_height_from) DO NOTHING
+        )
+        INSERT INTO state_changes_access_key_${CURRENT_RANGE_ID}_$partition (
+            account_id, block_height, data_key, data_value
+        )
+        SELECT
+            account_id,
+            block_height_from AS block_height,
+            data_key,
+            data_value
+        FROM ordered_data
+        WHERE data_value IS NOT NULL
+          AND block_height_from <= ${CURRENT_RANGE_BLOCK_HEIGHT_START}
+          AND (block_height_to IS NULL OR block_height_to > ${CURRENT_RANGE_BLOCK_HEIGHT_START})
+        ON CONFLICT (account_id, data_key, block_height) DO NOTHING;
+    "
+}
+
+# Run migrations in parallel for partitions 0 to 99
+running=0
+for i in $(seq 0 99); do
+    migrate_partition "$i" &
+
+    ((running+=1))
+    if ((running >= MAX_PARALLEL)); then
+        wait -n  # Wait for at least one job to finish
+        ((running-=1))
+    fi
+done
+
+# Wait for all background jobs to finish
+wait