add procedure to create tables for range_id

Author: kobayurii

database/src/base/state_indexer.rs

@@ -85,4 +85,6 @@ pub trait StateIndexerDbManager {
         block_height: u64,
         block_hash: near_primitives::hash::CryptoHash,
     ) -> anyhow::Result<()>;
+
+    async fn create_new_range_tables(&self, range_id: u64) -> anyhow::Result<()>;
 }

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

@@ -0,0 +1 @@
+-- Add down migration script here

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

@@ -0,0 +1,78 @@
+CREATE OR REPLACE PROCEDURE create_new_range_tables(range_id INT)
+    LANGUAGE plpgsql
+AS $$
+DECLARE
+    i INT;
+BEGIN
+    EXECUTE format('
+        CREATE TABLE IF NOT EXISTS state_changes_data_%s (
+            account_id text NOT NULL,
+            block_height numeric(20,0) NOT NULL,
+            block_hash text NOT NULL,
+            data_key text NOT NULL,
+            data_value bytea NULL,
+            PRIMARY KEY (account_id, data_key, block_height)
+        ) 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_%s PARTITION OF state_changes_data_%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_access_key_%s (
+            account_id text NOT NULL,
+            block_height numeric(20,0) NOT NULL,
+            block_hash text NOT NULL,
+            data_key text NOT NULL,
+            data_value bytea NULL,
+            PRIMARY KEY (account_id, data_key, block_height)
+        ) 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_%s PARTITION OF state_changes_access_key_%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_contract_%s (
+            account_id text NOT NULL,
+            block_height numeric(20,0) NOT NULL,
+            block_hash text NOT NULL,
+            data_value bytea NULL,
+            PRIMARY KEY (account_id, block_height)
+        ) 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_%s PARTITION OF state_changes_contract_%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_account_%s (
+            account_id text NOT NULL,
+            block_height numeric(20,0) NOT NULL,
+            block_hash text NOT NULL,
+            data_value bytea NULL,
+            PRIMARY KEY (account_id, block_height)
+        ) 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_%s PARTITION OF state_changes_account_%s FOR VALUES WITH (MODULUS 100, REMAINDER %s);
+            ', range_id, i, range_id, i);
+    END LOOP;
+END;
+$$;

database/src/postgres/state_indexer.rs

@@ -438,4 +438,14 @@ impl crate::StateIndexerDbManager for crate::PostgresDBManager {
             .await?;
         Ok(())
     }
+
+    async fn create_new_range_tables(&self, range_id: u64) -> anyhow::Result<()> {
+        for (_, pool) in self.shards_pool.iter() {
+            sqlx::query("CALL create_new_range_tables($1);")
+                .bind(range_id as i64)
+                .execute(pool)
+                .await?;
+        }
+        Ok(())
+    }
 }