use mongreldb_kit::db::complete_sql_output;
use mongreldb_kit::{
CancelOutcome, Column, ColumnType, Database, KitError, QueryId, Schema, SqlOptions, Table,
};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Barrier};
use std::time::Duration;
fn database() -> Arc<Database> {
let directory = tempfile::tempdir().unwrap().keep();
let schema = Schema::new(vec![Table {
id: 1,
name: "items".into(),
columns: vec![Column::new(1, "id", ColumnType::Int64)],
primary_key: vec!["id".into()],
indexes: vec![],
foreign_keys: vec![],
unique_constraints: vec![],
check_constraints: vec![],
}])
.unwrap();
Arc::new(Database::create(&directory, schema).unwrap())
}
#[test]
fn query_handle_supports_queue_cancel_deadline_conflict_and_reuse() {
let db = database();
let barrier = Arc::new(Barrier::new(2));
let hook_barrier = Arc::clone(&barrier);
let entered_once = Arc::new(AtomicBool::new(false));
let hook_entered_once = Arc::clone(&entered_once);
let (entered_tx, entered_rx) = std::sync::mpsc::channel();
db.set_sql_test_hook(Some(Arc::new(move |point| {
if point == mongreldb_query::SqlTestHookPoint::Planning
&& !hook_entered_once.swap(true, Ordering::AcqRel)
{
entered_tx.send(()).unwrap();
hook_barrier.wait();
}
})))
.unwrap();
let first_id: QueryId = "11112222333344445555666677778888".parse().unwrap();
let first = db
.start_sql(
"SELECT 1",
SqlOptions {
query_id: Some(first_id),
timeout: Some(Duration::from_secs(5)),
},
)
.unwrap();
assert_eq!(first.id(), first_id);
entered_rx.recv().unwrap();
let duplicate = db.start_sql(
"SELECT 2",
SqlOptions {
query_id: Some(first_id),
timeout: None,
},
);
assert!(matches!(duplicate, Err(KitError::QueryConflict { .. })));
let cancel_id: QueryId = "aaaabbbbccccddddeeeeffff00001111".parse().unwrap();
let cancelled = db
.start_sql(
"SELECT 2",
SqlOptions {
query_id: Some(cancel_id),
timeout: Some(Duration::from_secs(5)),
},
)
.unwrap();
assert_eq!(cancelled.cancel(), CancelOutcome::Accepted);
let cancelled_error = cancelled.wait().unwrap_err();
assert!(matches!(
&cancelled_error,
KitError::Cancelled {
outcome,
..
} if !outcome.committed
&& outcome.committed_statements == Some(0)
&& outcome.last_commit_epoch.is_none()
&& outcome.first_commit_statement_index.is_none()
&& outcome.last_commit_statement_index.is_none()
&& outcome.completed_statements == 0
&& outcome.statement_index == 0
));
let metadata = cancelled_error.query_metadata().unwrap();
assert_eq!(metadata.cancel_outcome, None);
assert_eq!(
metadata.cancellation_reason.as_deref(),
Some("client_request")
);
assert_eq!(metadata.retryable, Some(false));
assert_eq!(metadata.server_state.as_deref(), Some("cancelled"));
let timed_out = db
.start_sql(
"SELECT 3",
SqlOptions {
query_id: None,
timeout: Some(Duration::from_millis(10)),
},
)
.unwrap();
assert!(matches!(
timed_out.wait(),
Err(KitError::DeadlineExceeded { .. })
));
assert!(matches!(
db.sql_arrow_with_options(
"SELECT 3",
SqlOptions {
query_id: None,
timeout: Some(Duration::from_millis(10)),
},
),
Err(KitError::DeadlineExceeded { .. })
));
barrier.wait();
assert_eq!(first.wait().unwrap()[0].num_rows(), 1);
db.set_sql_test_hook(None).unwrap();
assert_eq!(db.sql_rows("SELECT 4 AS value").unwrap()[0]["value"], 4);
}
#[test]
fn row_serialization_remains_cancellable() {
let db = database();
let barrier = Arc::new(Barrier::new(2));
let hook_barrier = Arc::clone(&barrier);
let (entered_tx, entered_rx) = std::sync::mpsc::channel();
db.set_sql_test_hook(Some(Arc::new(move |point| {
if point == mongreldb_query::SqlTestHookPoint::BeforeSerializationBatch {
entered_tx.send(()).unwrap();
hook_barrier.wait();
}
})))
.unwrap();
let query_id: QueryId = "99990000111122223333444455556666".parse().unwrap();
let worker_db = Arc::clone(&db);
let worker = std::thread::spawn(move || {
worker_db.sql_rows_with_options(
"SELECT 1 AS value",
SqlOptions {
query_id: Some(query_id),
timeout: Some(Duration::from_secs(5)),
},
)
});
entered_rx.recv().unwrap();
assert_eq!(db.cancel_sql(query_id), CancelOutcome::Accepted);
barrier.wait();
assert!(matches!(
worker.join().unwrap(),
Err(KitError::Cancelled { .. })
));
db.set_sql_test_hook(None).unwrap();
assert_eq!(db.sql_rows("SELECT 2 AS value").unwrap()[0]["value"], 2);
}
#[test]
fn accepted_completion_race_never_returns_success() {
let db = database();
for _ in 0..256 {
let handle = db.start_sql("SELECT 1", SqlOptions::default()).unwrap();
let query_id = handle.id();
let output = handle.wait_for_serialization().unwrap();
let barrier = Arc::new(Barrier::new(2));
let worker_barrier = Arc::clone(&barrier);
let worker = std::thread::spawn(move || {
worker_barrier.wait();
complete_sql_output(output)
});
barrier.wait();
let cancel = db.cancel_sql(query_id);
let result = worker.join().unwrap();
match cancel {
CancelOutcome::Accepted | CancelOutcome::AlreadyCancelling => {
assert!(
matches!(result, Err(KitError::Cancelled { .. })),
"{result:?}"
);
}
CancelOutcome::TooLate | CancelOutcome::AlreadyFinished => {
assert!(result.is_ok(), "{result:?}");
}
CancelOutcome::NotFound => panic!("registered query disappeared before completion"),
}
}
}