use std::sync::OnceLock;
use faucet_common_redshift::RedshiftConnection;
use faucet_core::Sink;
use faucet_sink_redshift::{
RedshiftCopyFormat, RedshiftSink, RedshiftSinkConfig, RedshiftWriteStrategy,
};
use serde_json::{Value, json};
use sqlx::Row;
use testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner};
use testcontainers_modules::postgres::Postgres;
fn serial() -> &'static tokio::sync::Mutex<()> {
static SERIAL: OnceLock<tokio::sync::Mutex<()>> = OnceLock::new();
SERIAL.get_or_init(|| tokio::sync::Mutex::new(()))
}
async fn start_postgres() -> (ContainerAsync<Postgres>, u16) {
let image = Postgres::default().with_tag("16-alpine");
let container: ContainerAsync<Postgres> =
image.start().await.expect("postgres container start");
let port = container
.get_host_port_ipv4(5432)
.await
.expect("postgres port");
(container, port)
}
fn redshift_conn(port: u16) -> RedshiftConnection {
let mut conn = RedshiftConnection::new("127.0.0.1", "postgres", "postgres", "postgres");
conn.port = port;
conn.tls = false;
conn
}
async fn seed_pool(port: u16) -> sqlx::PgPool {
let url = format!("postgres://postgres:postgres@127.0.0.1:{port}/postgres");
sqlx::PgPool::connect(&url)
.await
.expect("seed pool connect")
}
fn insert_config(port: u16, table: &str, batch_size: usize) -> RedshiftSinkConfig {
RedshiftSinkConfig {
connection: redshift_conn(port),
table_name: table.into(),
schema: None,
write_strategy: RedshiftWriteStrategy::Insert,
copy_format: RedshiftCopyFormat::Jsonl,
staging_bucket: None,
staging_prefix: String::new(),
iam_role: None,
region: None,
endpoint_url: None,
batch_size,
max_connections: 2,
}
}
async fn install_insert_counter(pool: &sqlx::PgPool, table: &str) {
sqlx::query("CREATE TABLE insert_calls (calls BIGINT NOT NULL)")
.execute(pool)
.await
.expect("create counter table");
sqlx::query("INSERT INTO insert_calls (calls) VALUES (0)")
.execute(pool)
.await
.expect("seed counter");
sqlx::query(
"CREATE OR REPLACE FUNCTION bump_insert_calls() RETURNS TRIGGER AS $$ \
BEGIN UPDATE insert_calls SET calls = calls + 1; RETURN NULL; END; \
$$ LANGUAGE plpgsql",
)
.execute(pool)
.await
.expect("create trigger fn");
sqlx::query(&format!(
"CREATE TRIGGER count_inserts AFTER INSERT ON \"{table}\" \
FOR EACH STATEMENT EXECUTE FUNCTION bump_insert_calls()"
))
.execute(pool)
.await
.expect("attach trigger");
}
async fn insert_call_count(pool: &sqlx::PgPool) -> i64 {
sqlx::query_scalar("SELECT calls FROM insert_calls")
.fetch_one(pool)
.await
.expect("query counter")
}
async fn row_count(pool: &sqlx::PgPool, table: &str) -> i64 {
sqlx::query_scalar(&format!("SELECT COUNT(*)::BIGINT FROM {table}"))
.fetch_one(pool)
.await
.expect("count")
}
#[tokio::test(flavor = "multi_thread")]
async fn insert_writes_typed_rows() {
let _guard = serial().lock().await;
let (_container, port) = start_postgres().await;
let pool = seed_pool(port).await;
sqlx::query(
"CREATE TABLE events (\
id BIGINT, name TEXT, amount DOUBLE PRECISION, active BOOLEAN, note TEXT)",
)
.execute(&pool)
.await
.expect("create table");
let sink = RedshiftSink::new(insert_config(port, "events", 1000))
.await
.expect("sink builds");
let records = vec![
json!({"id": 1, "name": "alice", "amount": 1.5, "active": true, "note": null}),
json!({"id": 2, "name": "bob", "amount": 2.5, "active": false}),
];
let written = sink.write_batch(&records).await.expect("insert runs");
assert_eq!(written, 2);
assert_eq!(row_count(&pool, "events").await, 2);
let row = sqlx::query("SELECT id, name, amount, active, note FROM events WHERE id = 1")
.fetch_one(&pool)
.await
.expect("read back row 1");
assert_eq!(row.get::<i64, _>("id"), 1);
assert_eq!(row.get::<String, _>("name"), "alice");
assert!((row.get::<f64, _>("amount") - 1.5).abs() < 1e-9);
assert!(row.get::<bool, _>("active"));
assert_eq!(row.get::<Option<String>, _>("note"), None);
let note2: Option<String> = sqlx::query_scalar("SELECT note FROM events WHERE id = 2")
.fetch_one(&pool)
.await
.expect("row 2 note");
assert_eq!(note2, None, "missing key binds SQL NULL");
pool.close().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn write_batch_re_chunks_by_batch_size() {
let _guard = serial().lock().await;
let (_container, port) = start_postgres().await;
let pool = seed_pool(port).await;
sqlx::query("CREATE TABLE events (id BIGINT, name TEXT)")
.execute(&pool)
.await
.expect("create table");
install_insert_counter(&pool, "events").await;
let sink = RedshiftSink::new(insert_config(port, "events", 2))
.await
.expect("sink builds");
let records: Vec<Value> = (1..=5).map(|i| json!({"id": i, "name": "r"})).collect();
let written = sink.write_batch(&records).await.expect("write");
assert_eq!(written, 5);
assert_eq!(row_count(&pool, "events").await, 5);
assert_eq!(
insert_call_count(&pool).await,
3,
"5 rows at batch_size 2 → 3 INSERT statements (2 + 2 + 1)"
);
sink.flush().await.expect("flush");
pool.close().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn batch_size_zero_writes_single_statement() {
let _guard = serial().lock().await;
let (_container, port) = start_postgres().await;
let pool = seed_pool(port).await;
sqlx::query("CREATE TABLE events (id BIGINT, name TEXT)")
.execute(&pool)
.await
.expect("create table");
install_insert_counter(&pool, "events").await;
let sink = RedshiftSink::new(insert_config(port, "events", 0))
.await
.expect("sink builds");
let records: Vec<Value> = (1..=4).map(|i| json!({"id": i, "name": "r"})).collect();
assert_eq!(sink.write_batch(&records).await.expect("write"), 4);
assert_eq!(row_count(&pool, "events").await, 4);
assert_eq!(
insert_call_count(&pool).await,
1,
"batch_size 0 drains the slice in one INSERT statement"
);
pool.close().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn insert_with_no_matching_columns_is_a_noop() {
let _guard = serial().lock().await;
let (_container, port) = start_postgres().await;
let pool = seed_pool(port).await;
sqlx::query("CREATE TABLE events (id BIGINT)")
.execute(&pool)
.await
.expect("create table");
let sink = RedshiftSink::new(insert_config(port, "events", 1000))
.await
.expect("sink builds");
let written = sink
.write_batch(&[json!({"unknown": 1})])
.await
.expect("write");
assert_eq!(written, 0);
assert_eq!(row_count(&pool, "events").await, 0);
pool.close().await;
}
#[tokio::test(flavor = "multi_thread")]
async fn insert_into_missing_table_errors() {
let _guard = serial().lock().await;
let (_container, port) = start_postgres().await;
let sink = RedshiftSink::new(insert_config(port, "does_not_exist", 1000))
.await
.expect("sink builds");
let err = sink
.write_batch(&[json!({"id": 1})])
.await
.expect_err("missing table must error");
assert!(
matches!(err, faucet_core::FaucetError::Sink(_)),
"got {err:?}"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn empty_write_and_write_modes() {
let _guard = serial().lock().await;
let (_container, port) = start_postgres().await;
let sink = RedshiftSink::new(insert_config(port, "events", 1000))
.await
.expect("sink builds");
assert_eq!(sink.write_batch(&[]).await.expect("empty write"), 0);
assert_eq!(
sink.supported_write_modes(),
[faucet_core::WriteMode::Append].as_slice()
);
}
#[tokio::test(flavor = "multi_thread")]
async fn check_probe_passes() {
let _guard = serial().lock().await;
let (_container, port) = start_postgres().await;
let sink = RedshiftSink::new(insert_config(port, "events", 1000))
.await
.expect("sink builds");
let ctx = faucet_core::check::CheckContext {
timeout: std::time::Duration::from_secs(10),
};
let report = sink.check(&ctx).await.expect("check runs");
assert!(
report
.probes
.iter()
.all(|p| matches!(p.status, faucet_core::check::ProbeStatus::Pass)),
"all probes should pass against a reachable database: {report:?}"
);
}