Skip to content
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
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -103,3 +103,4 @@ dmypy.json

.cargo/
.claude/prep/
.pi-subagents
18 changes: 18 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

19 changes: 17 additions & 2 deletions nodedb-test-support/src/native_harness/frames.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,8 @@ use tokio::net::TcpStream;
use nodedb_types::protocol::request_fields::RequestFields;
use nodedb_types::protocol::text_fields::TextFields;
use nodedb_types::protocol::{
FRAME_HEADER_LEN, HELLO_ACK_MAGIC, HELLO_ERROR_MAGIC_U32, HelloAckFrame, HelloErrorFrame,
HelloFrame, NativeRequest, NativeResponse, OpCode,
AuthMethod, FRAME_HEADER_LEN, HELLO_ACK_MAGIC, HELLO_ERROR_MAGIC_U32, HelloAckFrame,
HelloErrorFrame, HelloFrame, NativeRequest, NativeResponse, OpCode,
};

/// Perform the handshake with a custom `HelloFrame`.
Expand Down Expand Up @@ -127,6 +127,21 @@ pub async fn send_request(
sonic_rs::from_slice(&response_payload).expect("json decode NativeResponse")
}

/// Authenticate a fresh native connection with an API key. The JSON Auth
/// request also selects JSON framing for the rest of the session.
pub async fn send_api_key_auth(stream: &mut TcpStream, seq: u64, token: String) -> NativeResponse {
send_request(
stream,
seq,
OpCode::Auth,
TextFields {
auth: Some(AuthMethod::ApiKey { token }),
..Default::default()
},
)
.await
}

/// Send a `SHOW`/SQL statement over an established JSON-encoding session and
/// decode the `NativeResponse`. Assumes the session's first frame already
/// selected JSON (see `json_request_gets_json_response`) — callers that open
Expand Down
4 changes: 3 additions & 1 deletion nodedb-test-support/src/native_harness/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,5 +9,7 @@
mod frames;
mod server;

pub use frames::{do_handshake, read_frame, send_request, send_sql, write_frame};
pub use frames::{
do_handshake, read_frame, send_api_key_auth, send_request, send_sql, write_frame,
};
pub use server::NativeTestServer;
20 changes: 12 additions & 8 deletions nodedb-test-support/src/native_harness/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,8 @@ use nodedb::wal::WalManager;
/// A running native-protocol test server.
pub struct NativeTestServer {
pub addr: std::net::SocketAddr,
/// Shared Control-Plane state, exposed so tests can read metrics (e.g.
/// `shared.system_metrics.active_txn_overlays`) that the Data-Plane
/// core also updates via the same `Arc<SystemMetrics>`.
/// Shared Control-Plane state, exposed so tests can inspect the same
/// metrics and authorization stores used by the running server.
pub shared: Arc<SharedState>,
pub(super) shutdown_bus: nodedb::control::shutdown::ShutdownBus,
pub(super) poller_shutdown_tx: tokio::sync::watch::Sender<bool>,
Expand All @@ -35,6 +34,15 @@ impl NativeTestServer {
/// Spawn a single-core NodeDB server with the native listener bound to
/// an ephemeral `127.0.0.1` port (trust-mode auth).
pub async fn start() -> Self {
Self::start_with_auth_mode(AuthMode::Trust).await
}

/// Spawn a single-core server that requires an explicit native Auth frame.
pub async fn start_authenticated() -> Self {
Self::start_with_auth_mode(AuthMode::Password).await
}

async fn start_with_auth_mode(auth_mode: AuthMode) -> Self {
let dir = tempfile::tempdir().expect("tempdir");
let wal_path = dir.path().join("test.wal");
let wal = Arc::new(WalManager::open_for_testing(&wal_path).expect("open wal"));
Expand Down Expand Up @@ -69,7 +77,6 @@ impl NativeTestServer {
let core_dir = dir.path().to_path_buf();
let event_producer = event_producers.into_iter().next().expect("event producer");
let core_array_catalog = shared.array_catalog.clone();
let core_metrics = shared.system_metrics.clone();
let (core_stop_tx, core_stop_rx) = std::sync::mpsc::channel::<()>();
let _core_handle = tokio::task::spawn_blocking(move || {
let mut core = CoreLoop::open_with_array_catalog(
Expand All @@ -82,9 +89,6 @@ impl NativeTestServer {
)
.expect("open core");
core.set_event_producer(event_producer);
if let Some(m) = core_metrics {
core.set_metrics(m);
}
while matches!(
core_stop_rx.try_recv(),
Err(std::sync::mpsc::TryRecvError::Empty)
Expand Down Expand Up @@ -136,7 +140,7 @@ impl NativeTestServer {
listener
.run(nodedb::control::server::listener::ListenerRunParams {
state: shared_listener,
auth_mode: AuthMode::Trust,
auth_mode,
tls_acceptor: None,
conn_semaphore: Arc::new(tokio::sync::Semaphore::new(128)),
startup_gate: test_startup_gate,
Expand Down
16 changes: 13 additions & 3 deletions nodedb-test-support/src/pgwire_harness/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -123,14 +123,24 @@ impl TestServer {
&self,
user: &str,
password: &str,
) -> Result<(tokio_postgres::Client, tokio::task::JoinHandle<()>), String> {
self.connect_as_database(user, password, "nodedb").await
}

/// Open a second pgwire connection under a user-selected database.
pub async fn connect_as_database(
&self,
user: &str,
password: &str,
database: &str,
) -> Result<(tokio_postgres::Client, tokio::task::JoinHandle<()>), String> {
let conn_str = format!(
"host=127.0.0.1 port={} user={} password={} dbname=nodedb",
self.pg_port, user, password
"host=127.0.0.1 port={} user={} password={} dbname={}",
self.pg_port, user, password, database
);
let (client, connection) = tokio_postgres::connect(&conn_str, tokio_postgres::NoTls)
.await
.map_err(|e| pg_error_detail(&e))?;
.map_err(|error| pg_error_detail(&error))?;
let handle = tokio::spawn(async move {
let _ = connection.await;
});
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ impl Scheduler {
database_id: DatabaseId::DEFAULT,
vshard_id: VShardId::new(self.vshard_id),
plan,
// no-determinism: scheduler deadline controls waiting, not ordered state.
deadline: Instant::now()
+ Duration::from_millis(
self.config.epoch_duration_ms * u64::from(self.config.txn_deadline_multiplier),
Expand Down
Loading