use dbnexus::DbConfig;
use std::collections::HashMap;
use tempfile::TempDir;
#[allow(dead_code)]
pub async fn make_sqlite_memory_pool() -> dbnexus::DbPool {
dbnexus::DbPool::new("sqlite::memory:").await.expect("Failed to create test pool")
}
#[cfg(feature = "permission-engine")]
#[allow(dead_code)]
async fn join_all<F>(handles: Vec<tokio::task::JoinHandle<F>>) -> Vec<F>
where
F: Send,
{
let mut results = Vec::new();
for handle in handles {
results.push(handle.await.unwrap());
}
results
}
fn sanitize_url(url: &str) -> String {
if let Some(at_pos) = url.find('@') {
if let Some(proto_end) = url.find("://") {
let proto = &url[..proto_end + 3];
let rest = &url[at_pos..];
if let Some(colon_pos) = url[proto_end + 3..at_pos].find(':') {
let user = &url[proto_end + 3..proto_end + 3 + colon_pos];
return format!("{}{}:****{}", proto, user, rest);
}
}
}
url.to_string()
}
static TEST_PERMISSIONS_CONTENT: &str = r#"
roles:
admin:
tables:
- name: "*"
operations:
- select
- insert
- update
- delete
user:
tables:
- name: "users"
operations:
- select
- insert
test_role:
tables:
- name: "*"
operations:
- select
"#;
pub fn get_test_database_url() -> String {
let test_db_type = std::env::var("TEST_DB_TYPE").unwrap_or_else(|_| "sqlite".to_string());
let database_url = std::env::var("DATABASE_URL").ok();
match test_db_type.as_str() {
"postgres" => database_url.unwrap_or_else(|| {
let password = std::env::var("TEST_DB_PASSWORD").unwrap_or_else(|_| "dbnexus_password".to_string());
format!("postgres://dbnexus:{}@localhost:15433/dbnexus_test", password)
}),
"mysql" => database_url.unwrap_or_else(|| {
let password = std::env::var("TEST_DB_PASSWORD").unwrap_or_else(|_| "dbnexus_password".to_string());
format!("mysql://dbnexus:{}@localhost:13308/dbnexus_test", password)
}),
_ => database_url.unwrap_or_else(|| "sqlite::memory:".to_string()),
}
}
#[allow(dead_code)]
pub fn get_test_config() -> (DbConfig, Option<TempDir>) {
get_test_config_with_permissions(false)
}
#[allow(dead_code)]
pub fn get_test_config_with_permissions(with_permissions: bool) -> (DbConfig, Option<TempDir>) {
let url = get_test_database_url();
if with_permissions {
let temp_dir = TempDir::new().expect("Failed to create temp directory");
let perm_file = temp_dir.path().join("test_permissions.yaml");
std::fs::write(&perm_file, TEST_PERMISSIONS_CONTENT).expect("Failed to write test permissions file");
let perm_path = perm_file.to_string_lossy().to_string();
let config = dbnexus::DbConfig {
url,
max_connections: 5,
min_connections: 1,
idle_timeout: 300,
acquire_timeout: 5000,
admin_role: "admin".to_string(),
permissions_path: Some(perm_path),
..Default::default()
};
return (config, Some(temp_dir));
}
let config = dbnexus::DbConfig {
url,
max_connections: 5,
min_connections: 1,
idle_timeout: 300,
acquire_timeout: 5000,
admin_role: "admin".to_string(),
..Default::default()
};
(config, None)
}
#[allow(dead_code)]
pub fn generate_test_table_name(prefix: &str) -> String {
use std::sync::atomic::{AtomicU64, Ordering};
static COUNTER: AtomicU64 = AtomicU64::new(0);
let n = COUNTER.fetch_add(1, Ordering::SeqCst);
format!("{}_test_{}", prefix, n)
}
#[allow(dead_code)]
pub async fn cleanup_test_table(session: &mut dbnexus::Session, table_name: &str) {
let _ = session
.execute_raw_ddl(&format!("DROP TABLE IF EXISTS {}", table_name))
.await;
}
#[allow(dead_code)]
pub async fn create_test_table(session: &mut dbnexus::Session, table_name: &str) {
session
.execute_raw_ddl(&format!(
"CREATE TABLE IF NOT EXISTS {} (id INTEGER PRIMARY KEY, data TEXT)",
table_name
))
.await
.expect("Failed to create test table");
}
#[allow(dead_code)]
pub fn assert_pool_healthy(pool: &dbnexus::DbPool) {
let status = pool.status();
assert!(status.active <= status.total, "Active should not exceed total");
assert_eq!(
status.total,
status.active + status.idle,
"Total should equal active + idle"
);
}
#[allow(dead_code)]
pub fn assert_session_valid(session: &mut dbnexus::Session) {
assert!(!session.role().is_empty(), "Session should have a role");
}
#[cfg(feature = "permission-engine")]
#[allow(dead_code)]
pub async fn run_parallel_tasks<F, T>(tasks: Vec<F>) -> Vec<T>
where
F: std::future::Future<Output = T> + Send + 'static,
T: Send + 'static,
{
let mut handles = Vec::new();
for task in tasks {
handles.push(tokio::spawn(task));
}
let mut results = Vec::new();
for handle in handles {
results.push(handle.await.unwrap());
}
results
}
#[allow(dead_code)]
pub async fn create_sqlite_file_pool() -> Result<(dbnexus::DbPool, TempDir), dbnexus::DbError> {
let temp_dir = tempfile::Builder::new()
.prefix("dbnexus_tracing_test_")
.tempdir()
.expect("Failed to create temp directory");
let db_path = temp_dir.path().join("test.db");
let db_path_str = db_path.to_string_lossy();
std::fs::File::create(&db_path).expect("Failed to create database file");
let perm_content = r#"
roles:
admin:
tables:
- name: "*"
operations:
- select
- insert
- update
- delete
system:
tables:
- name: "*"
operations:
- select
- insert
- update
- delete
"#;
let perm_file = temp_dir.path().join("permissions.yaml");
std::fs::write(&perm_file, perm_content).expect("Failed to write permissions file");
let config = dbnexus::DbConfig {
url: format!("sqlite://{}", db_path_str),
max_connections: 5,
min_connections: 1,
idle_timeout: 300,
acquire_timeout: 5000,
permissions_path: Some(perm_file.to_string_lossy().to_string()),
..Default::default()
};
let pool = dbnexus::DbPool::with_config(config).await?;
Ok((pool, temp_dir))
}
#[allow(dead_code)]
pub async fn create_test_pool() -> Result<(dbnexus::DbPool, Option<TempDir>), dbnexus::DbError> {
let test_db_type = std::env::var("TEST_DB_TYPE").unwrap_or_else(|_| "sqlite".to_string());
eprintln!("DEBUG: TEST_DB_TYPE = {}", test_db_type);
match test_db_type.as_str() {
"postgres" | "mysql" => {
let (config, temp_dir) = get_test_config_with_permissions(true);
eprintln!(
"DEBUG: Using {} database with URL: {}",
test_db_type,
sanitize_url(&config.url)
);
let pool = dbnexus::DbPool::with_config(config).await?;
Ok((pool, temp_dir))
}
_ => {
eprintln!("DEBUG: Using SQLite file database");
let (pool, temp_dir) = create_sqlite_file_pool().await?;
Ok((pool, Some(temp_dir)))
}
}
}
#[allow(dead_code)]
pub async fn create_tracing_test_table(pool: &dbnexus::DbPool) -> (String, TempDir) {
let temp_dir = TempDir::new().expect("Failed to create temp directory");
let table_name = generate_test_table_name("tracing_test");
let session = pool.get_session("admin").await.expect("Failed to get session");
let create_sql = if pool.config().url.contains("mysql") {
format!(
"CREATE TABLE IF NOT EXISTS {} (
id INTEGER AUTO_INCREMENT PRIMARY KEY,
trace_id VARCHAR(64),
span_id VARCHAR(64),
data TEXT,
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
)",
table_name
)
} else if pool.config().url.contains("postgres") {
format!(
"CREATE TABLE IF NOT EXISTS {} (
id INTEGER GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
trace_id TEXT,
span_id TEXT,
data TEXT,
created_at TEXT DEFAULT CURRENT_TIMESTAMP
)",
table_name
)
} else {
format!(
"CREATE TABLE IF NOT EXISTS {} (
id INTEGER PRIMARY KEY,
trace_id TEXT,
span_id TEXT,
data TEXT,
created_at TEXT DEFAULT CURRENT_TIMESTAMP
)",
table_name
)
};
session
.execute_raw_ddl(&create_sql)
.await
.expect("Failed to create test table");
(table_name, temp_dir)
}
#[allow(dead_code)]
pub async fn cleanup_tracing_test_table(pool: &dbnexus::DbPool, table_name: &str) {
let session = pool.get_session("admin").await.expect("Failed to get session");
let _ = session
.execute_raw_ddl(&format!("DROP TABLE IF EXISTS {}", table_name))
.await;
}
#[allow(dead_code)]
pub async fn verify_trace_injection(pool: &dbnexus::DbPool) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let _session = pool
.get_session("admin")
.await
.map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
let trace_id = "0af7651916cd43dd8448eb211c80319c";
let span_id = "b7ad6b7169203331";
let traceparent = format!("00-{}-{}01", trace_id, span_id);
let mut headers = HashMap::new();
headers.insert("traceparent".to_string(), traceparent);
assert!(
headers.contains_key("traceparent"),
"Headers should contain traceparent"
);
let tp = headers.get("traceparent").unwrap();
assert!(tp.starts_with("00-"), "traceparent format should be valid");
Ok(())
}
#[allow(dead_code)]
pub async fn verify_trace_extraction(
traceparent: &str,
) -> Result<Option<String>, Box<dyn std::error::Error + Send + Sync>> {
let parts: Vec<&str> = traceparent.split('-').collect();
if parts.len() < 3 {
return Ok(None);
}
let trace_id = parts[1];
let span_id = &parts[2][0..16];
assert!(trace_id.len() == 32, "trace_id should be 32 characters");
assert!(span_id.len() == 16, "span_id should be 16 characters");
u64::from_str_radix(&trace_id[0..16], 16).map_err(|_| "Invalid trace_id")?;
u64::from_str_radix(&trace_id[16..32], 16).map_err(|_| "Invalid trace_id")?;
u64::from_str_radix(span_id, 16).map_err(|_| "Invalid span_id")?;
Ok(Some(trace_id.to_string()))
}
#[allow(dead_code)]
pub async fn test_pool_with_trace_context(pool: &dbnexus::DbPool) {
let session = pool.get_session("admin").await.expect("Failed to get session");
assert!(!session.role().is_empty(), "Session should have a role");
let status = pool.status();
assert_eq!(
status.total,
status.active + status.idle,
"Total should equal active + idle"
);
}
#[allow(dead_code)]
pub async fn concurrent_trace_injection_test(pool: &dbnexus::DbPool, num_tasks: usize) -> (usize, usize) {
use futures::future::join_all;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
let success_count = Arc::new(AtomicUsize::new(0));
let mut handles = Vec::new();
for i in 0..num_tasks {
let pool = pool.clone();
let success_count = success_count.clone();
let handle = tokio::spawn(async move {
let role = if i % 2 == 0 { "admin" } else { "system" };
if pool.get_session(role).await.is_ok() {
success_count.fetch_add(1, Ordering::SeqCst);
}
});
handles.push(handle);
}
join_all(handles).await;
let success = success_count.load(Ordering::SeqCst);
(success, num_tasks - success)
}
#[allow(dead_code)]
pub async fn verify_db_operation_with_trace(
pool: &dbnexus::DbPool,
table_name: &str,
trace_id: &str,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let session = pool
.get_session("admin")
.await
.map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
let safe_table_name = validate_table_name(table_name)?;
let safe_trace_id = sanitize_sql_string(trace_id);
session
.execute_raw(&format!(
"INSERT INTO {} (trace_id, data) VALUES ('{}', 'test data')",
safe_table_name, safe_trace_id
))
.await
.map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
let _result = session
.execute_raw(&format!(
"SELECT COUNT(*) FROM {} WHERE trace_id = '{}'",
safe_table_name, safe_trace_id
))
.await
.map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
Ok(())
}
pub(crate) fn validate_table_name(table_name: &str) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
if !table_name.chars().all(|c| c.is_alphanumeric() || c == '_') {
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Invalid table name: contains disallowed characters",
)));
}
if table_name.len() > 64 {
return Err(Box::new(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Invalid table name: too long",
)));
}
Ok(table_name.to_string())
}
pub(crate) fn sanitize_sql_string(input: &str) -> String {
input.replace('\'', "''")
}
#[cfg(feature = "test-utils")]
#[allow(dead_code)]
pub fn create_temp_dir() -> TempDir {
tempfile::Builder::new()
.prefix("dbnexus_test_")
.tempdir()
.expect("Failed to create temp directory")
}