Skip to content
This repository has been archived by the owner on Dec 26, 2024. It is now read-only.

feat(storage): add dump storage table to file utility #1239

Merged
merged 1 commit into from
Oct 3, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions crates/papyrus_storage/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@
pub mod base_layer;
pub mod body;
pub mod compiled_class;
pub mod utils;
// TODO(yair): Make the compression_utils module pub(crate) or extract it from the crate.
#[doc(hidden)]
pub mod compression_utils;
Expand Down Expand Up @@ -415,6 +416,10 @@ pub enum StorageError {
CompiledClassReWrite { class_hash: ClassHash },
#[error("The table {table_name} is unused under the {storage_scope:?} storage scope.")]
ScopeError { table_name: String, storage_scope: StorageScope },
#[error(transparent)]
IOError(#[from] std::io::Error),
#[error(transparent)]
SerdeError(#[from] serde_json::Error),
}

/// A type alias that maps to std::result::Result<T, StorageError>.
Expand Down
48 changes: 48 additions & 0 deletions crates/papyrus_storage/src/utils.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
//! module for external utils, such as dumping a storage table to a file
#[cfg(test)]
#[path = "utils_test.rs"]
mod utils_test;

use std::fs::File;
use std::io::{BufWriter, Write};

use crate::db::serialization::StorageSerde;
use crate::db::{DbIter, TableIdentifier, RO};
use crate::{open_storage, StorageConfig, StorageResult, StorageTxn};

/// Dumps a table from the storage to a file in JSON format.
fn dump_table_to_file<K, V>(
txn: &StorageTxn<'_, RO>,
table_id: &TableIdentifier<K, V>,
file_path: &str,
) -> StorageResult<()>
where
K: StorageSerde + serde::Serialize,
V: StorageSerde + serde::Serialize,
{
let table_handle = txn.txn.open_table(table_id)?;
let mut cursor = table_handle.cursor(&txn.txn)?;
let iter = DbIter::new(&mut cursor);
let file = File::create(file_path)?;
let mut writer = BufWriter::new(file);
writer.write_all(b"[")?;
let mut first = true;
for data in iter {
if !first {
writer.write_all(b",")?;
}
serde_json::to_writer(&mut writer, &data?)?;
first = false;
}
writer.write_all(b"]")?;
Ok(())
}

/// Dumps the declared_classes table from the storage to a file.
pub fn dump_declared_classes_table_to_file(file_path: &str) -> StorageResult<()> {
let storage_config = StorageConfig::default();
let (storage_reader, _) = open_storage(storage_config.clone())?;
let txn = storage_reader.begin_ro_txn()?;
dump_table_to_file(&txn, &txn.tables.declared_classes, file_path)?;
Ok(())
}
65 changes: 65 additions & 0 deletions crates/papyrus_storage/src/utils_test.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
use std::collections::HashMap;
use std::fs;

use indexmap::indexmap;
use starknet_api::block::BlockNumber;
use starknet_api::core::{ClassHash, CompiledClassHash};
use starknet_api::hash::{StarkFelt, StarkHash};
use starknet_api::state::{ContractClass, StateDiff};

use super::dump_table_to_file;
use crate::state::StateStorageWriter;
use crate::test_utils::get_test_storage;

#[test]
fn test_dump_table_to_file() {
let file_path = "tmp_test_dump_declared_classes_table.json";
let declared_class1 = (
ClassHash(1u128.into()),
ContractClass {
sierra_program: vec![StarkFelt::ONE, StarkFelt::TWO],
entry_point_by_type: HashMap::new(),
abi: "".to_string(),
},
);
let declared_class2 = (
ClassHash(2u128.into()),
ContractClass {
sierra_program: vec![StarkFelt::THREE, StarkFelt::ZERO],
entry_point_by_type: HashMap::new(),
abi: "".to_string(),
},
);
let compiled_class_hash = CompiledClassHash(StarkHash::default());
let declared_classes = vec![declared_class1.clone(), declared_class2.clone()];
let declared_classes_for_append_state = indexmap!(
declared_class1.0 =>
(compiled_class_hash, declared_class1.1.clone()),
declared_class2.0 =>
(compiled_class_hash, declared_class2.1.clone()),
);

let ((reader, mut writer), _temp_dir) = get_test_storage();
let txn = writer.begin_rw_txn().unwrap();
txn.append_state_diff(
BlockNumber(0),
StateDiff {
deployed_contracts: indexmap!(),
storage_diffs: indexmap!(),
declared_classes: declared_classes_for_append_state,
deprecated_declared_classes: indexmap!(),
nonces: indexmap!(),
replaced_classes: indexmap!(),
},
indexmap!(),
)
.unwrap()
.commit()
.unwrap();

let txn = reader.begin_ro_txn().unwrap();
dump_table_to_file(&txn, &txn.tables.declared_classes, file_path).unwrap();
let file_content = fs::read_to_string(file_path).unwrap();
let _ = fs::remove_file(file_path);
assert_eq!(file_content, serde_json::to_string(&declared_classes).unwrap());
}