#![cfg(feature = "state_machine")]
use arrow::record_batch::RecordBatch;
use bytes::Bytes;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::error::Error as StdError;
use std::fs;
use std::path::PathBuf;
use std::process::{Command, Output};
use tempfile::TempDir;
mod common;
const CLI_BINARY: &str = "cqlite";
fn run_cli_command(args: &[&str]) -> Output {
Command::new("cargo")
.args(["run", "--quiet", "--bin", CLI_BINARY, "--"])
.args(args)
.output()
.expect("Failed to execute CLI command")
}
fn get_test_data_root() -> PathBuf {
std::env::var("CQLITE_DATASETS_ROOT")
.map(PathBuf::from)
.unwrap_or_else(|_| {
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.parent()
.unwrap()
.join("test-data/datasets")
})
}
fn get_schemas_dir() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.parent()
.unwrap()
.join("test-data/schemas")
}
fn read_parquet_back(bytes: &[u8]) -> Result<RecordBatch, Box<dyn StdError>> {
let bytes = Bytes::copy_from_slice(bytes);
let builder = ParquetRecordBatchReaderBuilder::try_new(bytes)?;
let mut reader = builder.build()?;
match reader.next() {
Some(result) => result.map_err(|e| Box::new(e) as Box<dyn StdError>),
None => Err(Box::new(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"No batches in Parquet file",
)) as Box<dyn StdError>),
}
}
fn verify_parquet_magic(bytes: &[u8]) {
assert!(bytes.len() >= 8, "Parquet file too small");
assert_eq!(&bytes[0..4], b"PAR1", "Should start with PAR1 magic bytes");
assert_eq!(
&bytes[bytes.len() - 4..],
b"PAR1",
"Should end with PAR1 magic bytes"
);
}
fn assert_test_data_available() -> (PathBuf, PathBuf) {
let data_dir = get_test_data_root().join("sstables");
let schema_file = get_schemas_dir().join("basic-types.cql");
assert!(
data_dir.exists() && schema_file.exists(),
"Test requires full SSTable dataset. \n Set CQLITE_DATASETS_ROOT or run: bash test-data/scripts/fetch-datasets.sh\n data_dir={data_dir:?}, schema_file={schema_file:?}"
);
(data_dir, schema_file)
}
#[test]
fn test_export_csv_basic_types() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("export.csv");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"csv",
"--table",
"test_basic.simple_table",
]);
eprintln!("Exit status: {}", output.status);
eprintln!("STDOUT:\n{}", String::from_utf8_lossy(&output.stdout));
eprintln!("STDERR:\n{}", String::from_utf8_lossy(&output.stderr));
assert!(
output.status.success(),
"CSV export should succeed. Exit code: {:?}",
output.status.code()
);
assert!(output_file.exists(), "Output CSV file should exist");
let csv_content = fs::read_to_string(&output_file).expect("Failed to read CSV");
assert!(!csv_content.is_empty(), "CSV should not be empty");
let lines: Vec<&str> = csv_content.lines().collect();
assert!(lines.len() >= 2, "CSV should have header + data rows");
let header = lines[0];
assert!(
header.contains("id") || header.contains("name"),
"CSV header should contain column names: {header}"
);
}
#[test]
fn test_export_json_basic_types() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("export.json");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"json",
"--table",
"test_basic.simple_table",
]);
eprintln!("Exit status: {}", output.status);
eprintln!("STDERR:\n{}", String::from_utf8_lossy(&output.stderr));
assert!(
output.status.success(),
"JSON export should succeed. Exit code: {:?}",
output.status.code()
);
assert!(output_file.exists(), "Output JSON file should exist");
let json_content = fs::read_to_string(&output_file).expect("Failed to read JSON");
assert!(!json_content.is_empty(), "JSON should not be empty");
let parsed: serde_json::Value =
serde_json::from_str(&json_content).expect("Should be valid JSON");
assert!(parsed.is_array(), "JSON output should be an array");
let array = parsed.as_array().unwrap();
assert!(!array.is_empty(), "JSON array should have rows");
}
#[test]
fn test_export_csv_collections() {
let (data_dir, _) = assert_test_data_available();
let schema_file = get_schemas_dir().join("collections.cql");
if !schema_file.exists() {
eprintln!("Skipping test: collections schema not found");
return;
}
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("export_collections.csv");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"csv",
"--table",
"test_collections.collection_table",
]);
eprintln!("Exit status: {}", output.status);
eprintln!("STDERR:\n{}", String::from_utf8_lossy(&output.stderr));
assert!(
output.status.success(),
"CSV export of collections should succeed. Exit code: {:?}",
output.status.code()
);
assert!(output_file.exists(), "Output CSV file should exist");
let csv_content = fs::read_to_string(&output_file).expect("Failed to read CSV");
assert!(!csv_content.is_empty(), "CSV should not be empty");
}
#[test]
fn test_export_json_collections() {
let (data_dir, _) = assert_test_data_available();
let schema_file = get_schemas_dir().join("collections.cql");
if !schema_file.exists() {
eprintln!("Skipping test: collections schema not found");
return;
}
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("export_collections.json");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"json",
"--table",
"test_collections.collection_table",
]);
eprintln!("Exit status: {}", output.status);
eprintln!("STDERR:\n{}", String::from_utf8_lossy(&output.stderr));
assert!(
output.status.success(),
"JSON export of collections should succeed. Exit code: {:?}",
output.status.code()
);
assert!(output_file.exists(), "Output JSON file should exist");
let json_content = fs::read_to_string(&output_file).expect("Failed to read JSON");
let parsed: serde_json::Value =
serde_json::from_str(&json_content).expect("Should be valid JSON");
assert!(parsed.is_array(), "JSON output should be an array");
}
#[test]
fn test_export_parquet_basic() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("export.parquet");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"parquet",
"--table",
"test_basic.simple_table",
]);
eprintln!("Exit status: {}", output.status);
eprintln!("STDERR:\n{}", String::from_utf8_lossy(&output.stderr));
assert!(
output.status.success(),
"Parquet export should succeed. Exit code: {:?}",
output.status.code()
);
assert!(output_file.exists(), "Output Parquet file should exist");
let parquet_bytes = fs::read(&output_file).expect("Failed to read Parquet file");
verify_parquet_magic(&parquet_bytes);
}
#[test]
fn test_export_parquet_roundtrip() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("export_roundtrip.parquet");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"parquet",
"--table",
"test_basic.simple_table",
]);
assert!(
output.status.success(),
"Parquet export should succeed for roundtrip test"
);
let parquet_bytes = fs::read(&output_file).expect("Failed to read Parquet file");
let batch = read_parquet_back(&parquet_bytes).expect("Failed to read Parquet back");
assert!(batch.num_rows() > 0, "Should have rows in Parquet file");
assert!(
batch.num_columns() > 0,
"Should have columns in Parquet file"
);
eprintln!(
"Parquet roundtrip: {} rows, {} columns",
batch.num_rows(),
batch.num_columns()
);
}
#[test]
fn test_export_parquet_schema_matches() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("export_schema.parquet");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"parquet",
"--table",
"test_basic.simple_table",
]);
assert!(
output.status.success(),
"Parquet export should succeed for schema test"
);
let parquet_bytes = fs::read(&output_file).expect("Failed to read Parquet file");
let batch = read_parquet_back(&parquet_bytes).expect("Failed to read Parquet back");
let schema = batch.schema();
let column_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
assert!(
column_names.contains(&"id"),
"Parquet schema should contain 'id' column. Found: {column_names:?}"
);
eprintln!("Parquet columns: {column_names:?}");
}
#[test]
fn test_export_with_query_filter() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("export_filtered.csv");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"csv",
"--table",
"test_basic.simple_table",
"--query",
"active = true",
]);
eprintln!("Exit status: {}", output.status);
eprintln!("STDERR:\n{}", String::from_utf8_lossy(&output.stderr));
assert!(
output.status.success(),
"Export with --query filter should succeed. Exit code: {:?}\nSTDERR: {}",
output.status.code(),
String::from_utf8_lossy(&output.stderr)
);
assert!(output_file.exists(), "Filtered output should exist");
let csv_content = fs::read_to_string(&output_file).expect("Failed to read CSV");
let line_count = csv_content.lines().count();
assert!(line_count >= 1, "CSV should have at least a header row");
eprintln!("Filtered CSV rows (including header): {line_count}");
}
#[test]
fn test_export_row_count_matches_query() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let csv_file = temp_dir.path().join("export_count.csv");
let export_output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
csv_file.to_str().unwrap(),
"--format",
"csv",
"--table",
"test_basic.simple_table",
]);
assert!(
export_output.status.success(),
"Export should succeed for count test"
);
let query_output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"-e",
"SELECT * FROM test_basic.simple_table",
"--format",
"json",
]);
assert!(
query_output.status.success(),
"Direct query should succeed. Exit code: {:?}\nSTDERR: {}",
query_output.status.code(),
String::from_utf8_lossy(&query_output.stderr)
);
let query_stdout = String::from_utf8_lossy(&query_output.stdout);
let query_json: serde_json::Value =
serde_json::from_str(&query_stdout).expect("Query output should be valid JSON");
let query_row_count = query_json
.as_array()
.expect("Query output should be a JSON array")
.len();
let csv_content = fs::read_to_string(&csv_file).expect("Failed to read CSV");
let csv_row_count = csv_content.lines().count().saturating_sub(1);
eprintln!("Row counts - Query: {query_row_count}, CSV: {csv_row_count}");
assert_eq!(
csv_row_count, query_row_count,
"CSV export row count should match query result count"
);
}
#[test]
fn test_export_with_limit() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let csv_file = temp_dir.path().join("export_limit.csv");
const LIMIT: usize = 3;
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
csv_file.to_str().unwrap(),
"--format",
"csv",
"--table",
"test_basic.simple_table",
"--limit",
&LIMIT.to_string(),
]);
eprintln!("Exit status: {}", output.status);
eprintln!("STDERR:\n{}", String::from_utf8_lossy(&output.stderr));
assert!(
output.status.success(),
"Export with --limit should succeed. Exit code: {:?}\nSTDERR: {}",
output.status.code(),
String::from_utf8_lossy(&output.stderr)
);
assert!(csv_file.exists(), "Output CSV file should exist");
let csv_content = fs::read_to_string(&csv_file).expect("Failed to read CSV");
let lines: Vec<&str> = csv_content.lines().collect();
let data_row_count = lines.len().saturating_sub(1); assert_eq!(
data_row_count, LIMIT,
"CSV should have exactly {LIMIT} data rows (got {data_row_count}). Full content:\n{csv_content}"
);
eprintln!("Limit test passed: {LIMIT} rows exported as expected");
}
#[test]
fn test_export_csv_deterministic() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("deterministic.csv");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"csv",
"--table",
"test_basic.simple_table",
]);
assert!(output.status.success(), "CSV export should succeed");
let csv_content = fs::read_to_string(&output_file).expect("Failed to read CSV");
let lines: Vec<&str> = csv_content.lines().collect();
assert!(lines.len() > 1, "Should have header + data rows");
let header = lines[0];
assert!(header.contains("id"), "Header should contain 'id'");
assert!(header.contains("name"), "Header should contain 'name'");
assert!(header.contains("age"), "Header should contain 'age'");
let data_row_count = lines.len() - 1;
assert!(
data_row_count > 0,
"Should have at least one data row, got {data_row_count}"
);
eprintln!(
"CSV structure verified: {} columns, {} data rows",
header.split(',').count(),
data_row_count
);
}
#[test]
fn test_export_json_deterministic() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("deterministic.json");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"json",
"--table",
"test_basic.simple_table",
]);
assert!(output.status.success(), "JSON export should succeed");
let json_content = fs::read_to_string(&output_file).expect("Failed to read JSON");
let parsed: serde_json::Value =
serde_json::from_str(&json_content).expect("Should be valid JSON");
assert!(parsed.is_array(), "JSON output should be an array");
let array = parsed.as_array().unwrap();
assert!(!array.is_empty(), "JSON array should have rows");
let first_row = &array[0];
assert!(first_row.is_object(), "Each row should be an object");
let obj = first_row.as_object().unwrap();
assert!(obj.contains_key("id"), "Row should contain 'id'");
assert!(obj.contains_key("name"), "Row should contain 'name'");
assert!(obj.contains_key("age"), "Row should contain 'age'");
eprintln!(
"JSON structure verified: {} keys per row, {} rows",
obj.len(),
array.len()
);
}
#[test]
fn test_export_nonexistent_table_behavior() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("empty_output.csv");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"csv",
"--table",
"nonexistent_keyspace.nonexistent_table",
]);
eprintln!("Exit status: {}", output.status);
eprintln!("STDERR:\n{}", String::from_utf8_lossy(&output.stderr));
assert!(
!output.status.success(),
"Export command should fail for non-existent table (strict validation)"
);
let stderr = String::from_utf8_lossy(&output.stderr);
assert!(
stderr.contains("column") || stderr.contains("Could not determine"),
"Should indicate column metadata issue: {stderr}"
);
assert!(
!output_file.exists(),
"Output file should not be created when export fails"
);
}
#[test]
fn test_export_invalid_format_error() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("error_output.xyz");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"invalid_format",
"--table",
"test_basic.simple_table",
]);
eprintln!("Exit status: {}", output.status);
eprintln!("STDERR:\n{}", String::from_utf8_lossy(&output.stderr));
assert!(
!output.status.success(),
"Export should fail for invalid format"
);
let stderr = String::from_utf8_lossy(&output.stderr);
assert!(
stderr.to_lowercase().contains("format")
|| stderr.contains("invalid")
|| stderr.contains("possible values"),
"Error message should indicate format issue: {stderr}"
);
}
#[test]
fn test_export_missing_table_arg_error() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("error_output.csv");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"csv",
]);
eprintln!("Exit status: {}", output.status);
eprintln!("STDERR:\n{}", String::from_utf8_lossy(&output.stderr));
assert!(
!output.status.success(),
"Export should fail when --table is missing"
);
let stderr = String::from_utf8_lossy(&output.stderr);
assert!(
stderr.contains("--table") || stderr.contains("required") || stderr.contains("table"),
"Error message should indicate missing table argument: {stderr}"
);
}
#[test]
fn test_export_nonexistent_output_dir_error() {
let (data_dir, schema_file) = assert_test_data_available();
let output_file = PathBuf::from("/nonexistent_directory_12345/output.csv");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"csv",
"--table",
"test_basic.simple_table",
]);
eprintln!("Exit status: {}", output.status);
eprintln!("STDERR:\n{}", String::from_utf8_lossy(&output.stderr));
assert!(
!output.status.success(),
"Export should fail for nonexistent output directory"
);
let stderr = String::from_utf8_lossy(&output.stderr);
assert!(
stderr.to_lowercase().contains("directory")
|| stderr.to_lowercase().contains("path")
|| stderr.to_lowercase().contains("permission")
|| stderr.to_lowercase().contains("no such file"),
"Error message should indicate path/directory issue: {stderr}"
);
}
#[test]
fn test_export_csv_json_row_count_matches() {
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let csv_file = temp_dir.path().join("consistency.csv");
let json_file = temp_dir.path().join("consistency.json");
let csv_output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
csv_file.to_str().unwrap(),
"--format",
"csv",
"--table",
"test_basic.simple_table",
]);
let json_output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
json_file.to_str().unwrap(),
"--format",
"json",
"--table",
"test_basic.simple_table",
]);
assert!(csv_output.status.success(), "CSV export should succeed");
assert!(json_output.status.success(), "JSON export should succeed");
let csv_content = fs::read_to_string(&csv_file).expect("Failed to read CSV");
let csv_row_count = csv_content.lines().count().saturating_sub(1);
let json_content = fs::read_to_string(&json_file).expect("Failed to read JSON");
let json_parsed: serde_json::Value =
serde_json::from_str(&json_content).expect("Should be valid JSON");
let json_row_count = json_parsed.as_array().map(|a| a.len()).unwrap_or(0);
eprintln!("Cross-format row counts - CSV: {csv_row_count}, JSON: {json_row_count}");
assert_eq!(
csv_row_count, json_row_count,
"CSV and JSON exports should have same row count"
);
}
#[tokio::test]
async fn test_export_sstable_to_parquet() {
use cqlite_cli::cli::ExportFormat;
use cqlite_cli::commands::export_sstable;
use std::io::Write;
let (data_dir, _cql_schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let output_file = temp_dir.path().join("sstable_export.parquet");
let schema_content = r#"{
"keyspace": "test_basic",
"table": "simple_table",
"columns": {
"id": { "type": "uuid", "kind": "PartitionKey" },
"name": { "type": "text", "kind": "Regular" },
"age": { "type": "int", "kind": "Regular" },
"active": { "type": "boolean", "kind": "Regular" }
}
}"#;
let schema_file = temp_dir.path().join("test_schema.json");
{
let mut f = fs::File::create(&schema_file).expect("Failed to create schema file");
f.write_all(schema_content.as_bytes())
.expect("Failed to write schema");
}
let test_basic_dir = data_dir.join("test_basic");
let simple_table_dir = fs::read_dir(&test_basic_dir)
.expect("Failed to read test_basic directory")
.filter_map(Result::ok)
.find(|entry| {
entry
.file_name()
.to_string_lossy()
.starts_with("simple_table-")
})
.expect("No simple_table directory found");
let sstable_file = simple_table_dir.path().join("nb-1-big-Data.db");
assert!(
sstable_file.exists(),
"Test requires SSTable file: {sstable_file:?}"
);
eprintln!("Using SSTable: {sstable_file:?}");
eprintln!("Using schema: {schema_file:?}");
eprintln!("Output file: {output_file:?}");
let result = export_sstable(
&sstable_file,
&schema_file,
&output_file,
ExportFormat::Parquet,
)
.await;
assert!(
result.is_ok(),
"export_sstable to Parquet should succeed: {:?}",
result.err()
);
assert!(output_file.exists(), "Output Parquet file should exist");
let parquet_bytes = fs::read(&output_file).expect("Failed to read Parquet file");
verify_parquet_magic(&parquet_bytes);
let batch = read_parquet_back(&parquet_bytes).expect("Failed to read Parquet back");
eprintln!(
"SSTable to Parquet export verified: {} rows, {} columns",
batch.num_rows(),
batch.num_columns()
);
assert!(
!parquet_bytes.is_empty(),
"Parquet file should have content"
);
}
#[test]
#[ignore]
fn test_export_memory_efficiency() {
use sysinfo::{ProcessRefreshKind, System};
let (data_dir, schema_file) = assert_test_data_available();
let temp_dir = TempDir::new().expect("Failed to create temp dir");
let mut system = System::new();
let pid = sysinfo::get_current_pid().expect("Failed to get current PID");
system.refresh_process_specifics(pid, ProcessRefreshKind::new().with_memory());
let baseline_memory = system.process(pid).map(|p| p.memory()).unwrap_or(0);
eprintln!(
"Baseline memory: {} bytes ({:.1} MB)",
baseline_memory,
baseline_memory as f64 / (1024.0 * 1024.0)
);
let output_file = temp_dir.path().join("memory_test.parquet");
let output = run_cli_command(&[
"--schema",
schema_file.to_str().unwrap(),
"--data-dir",
data_dir.to_str().unwrap(),
"export",
output_file.to_str().unwrap(),
"--format",
"parquet",
"--table",
"test_basic.simple_table",
]);
assert!(
output.status.success(),
"Export should succeed for memory test. Exit code: {:?}\nSTDERR: {}",
output.status.code(),
String::from_utf8_lossy(&output.stderr)
);
system.refresh_process_specifics(pid, ProcessRefreshKind::new().with_memory());
let peak_memory = system.process(pid).map(|p| p.memory()).unwrap_or(0);
let memory_delta = peak_memory.saturating_sub(baseline_memory);
let memory_delta_mb = memory_delta as f64 / (1024.0 * 1024.0);
eprintln!(
"Peak memory: {} bytes ({:.1} MB)",
peak_memory,
peak_memory as f64 / (1024.0 * 1024.0)
);
eprintln!("Memory delta: {memory_delta} bytes ({memory_delta_mb:.1} MB)");
const MEMORY_LIMIT_MB: f64 = 128.0;
assert!(
memory_delta_mb < MEMORY_LIMIT_MB,
"Export memory usage ({memory_delta_mb:.1} MB) should stay under {MEMORY_LIMIT_MB} MB limit"
);
eprintln!(
"Memory efficiency test passed: {memory_delta_mb:.1} MB < {MEMORY_LIMIT_MB} MB limit"
);
}