use crate::config::RedshiftCopyFormat;
use faucet_core::FaucetError;
use faucet_core::util::quote_ident;
use serde_json::Value;
pub(crate) fn qualified_table_ref(schema: Option<&str>, table: &str) -> String {
match schema {
Some(s) => format!("{}.{}", quote_ident(s), quote_ident(table)),
None => quote_ident(table),
}
}
pub(crate) fn sql_string_literal(s: &str) -> String {
format!("'{}'", s.replace('\'', "''"))
}
pub(crate) fn s3_uri(bucket: &str, key: &str) -> String {
format!("s3://{bucket}/{key}")
}
pub(crate) fn copy_statement(
table_ref: &str,
columns: Option<&[String]>,
s3_path: &str,
iam_role: &str,
region: Option<&str>,
format: RedshiftCopyFormat,
) -> String {
let col_clause = match (format, columns) {
(RedshiftCopyFormat::Csv, Some(cols)) if !cols.is_empty() => {
let list = cols
.iter()
.map(|c| quote_ident(c))
.collect::<Vec<_>>()
.join(", ");
format!(" ({list})")
}
_ => String::new(),
};
let format_clause = match format {
RedshiftCopyFormat::Jsonl => "FORMAT AS JSON 'auto'".to_string(),
RedshiftCopyFormat::Csv => "FORMAT AS CSV".to_string(),
};
let mut sql = format!(
"COPY {table_ref}{col_clause} FROM {} IAM_ROLE {} {format_clause}",
sql_string_literal(s3_path),
sql_string_literal(iam_role),
);
if let Some(r) = region.filter(|r| !r.trim().is_empty()) {
sql.push_str(&format!(" REGION {}", sql_string_literal(r)));
}
sql
}
pub(crate) fn insert_statement(table_ref: &str, columns: &[String], num_rows: usize) -> String {
let col_list = columns
.iter()
.map(|c| quote_ident(c))
.collect::<Vec<_>>()
.join(", ");
let num_cols = columns.len();
let mut tuples = Vec::with_capacity(num_rows);
for row in 0..num_rows {
let start = row * num_cols + 1;
let placeholders = (0..num_cols)
.map(|c| format!("${}", start + c))
.collect::<Vec<_>>()
.join(", ");
tuples.push(format!("({placeholders})"));
}
format!(
"INSERT INTO {table_ref} ({col_list}) VALUES {}",
tuples.join(", ")
)
}
pub(crate) fn columns_present<'a>(
records: &[Value],
table_columns: &'a [String],
) -> Vec<&'a String> {
table_columns
.iter()
.filter(|col| {
records
.iter()
.any(|r| r.as_object().is_some_and(|o| o.contains_key(col.as_str())))
})
.collect()
}
pub(crate) fn serialize_jsonl(records: &[Value]) -> Result<Vec<u8>, FaucetError> {
let mut buf = Vec::new();
for record in records {
let line = serde_json::to_vec(record)
.map_err(|e| FaucetError::Sink(format!("redshift: JSON serialization failed: {e}")))?;
buf.extend_from_slice(&line);
buf.push(b'\n');
}
Ok(buf)
}
fn csv_cell(v: Option<&Value>) -> String {
match v {
None | Some(Value::Null) => String::new(),
Some(Value::String(s)) => quote_csv(s),
Some(Value::Bool(b)) => b.to_string(),
Some(Value::Number(n)) => n.to_string(),
Some(other) => quote_csv(&other.to_string()),
}
}
fn quote_csv(s: &str) -> String {
if s.contains([',', '"', '\n', '\r']) {
format!("\"{}\"", s.replace('"', "\"\""))
} else {
s.to_string()
}
}
pub(crate) fn serialize_csv(records: &[Value], columns: &[String]) -> Result<Vec<u8>, FaucetError> {
let mut buf = String::new();
for record in records {
let obj = record.as_object().ok_or_else(|| {
FaucetError::Sink("redshift: CSV requires JSON object records".into())
})?;
let cells: Vec<String> = columns.iter().map(|c| csv_cell(obj.get(c))).collect();
buf.push_str(&cells.join(","));
buf.push('\n');
}
Ok(buf.into_bytes())
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn table_ref_qualified_and_bare() {
assert_eq!(qualified_table_ref(None, "events"), "\"events\"");
assert_eq!(
qualified_table_ref(Some("analytics"), "events"),
"\"analytics\".\"events\""
);
}
#[test]
fn table_ref_escapes_quotes() {
assert_eq!(
qualified_table_ref(Some("we\"ird"), "ta\"ble"),
"\"we\"\"ird\".\"ta\"\"ble\""
);
}
#[test]
fn string_literal_doubles_single_quotes() {
assert_eq!(sql_string_literal("a'b"), "'a''b'");
assert_eq!(sql_string_literal("plain"), "'plain'");
}
#[test]
fn copy_jsonl_omits_columns_and_uses_json_auto() {
let sql = copy_statement(
"\"public\".\"events\"",
Some(&["id".into(), "name".into()]),
"s3://bucket/staging/abc.jsonl",
"arn:aws:iam::123:role/redshift",
None,
RedshiftCopyFormat::Jsonl,
);
assert_eq!(
sql,
"COPY \"public\".\"events\" FROM 's3://bucket/staging/abc.jsonl' \
IAM_ROLE 'arn:aws:iam::123:role/redshift' FORMAT AS JSON 'auto'"
);
}
#[test]
fn copy_csv_includes_ordered_columns() {
let sql = copy_statement(
"\"events\"",
Some(&["id".into(), "name".into()]),
"s3://b/k.csv",
"arn:role",
Some("us-east-1"),
RedshiftCopyFormat::Csv,
);
assert_eq!(
sql,
"COPY \"events\" (\"id\", \"name\") FROM 's3://b/k.csv' IAM_ROLE 'arn:role' \
FORMAT AS CSV REGION 'us-east-1'"
);
}
#[test]
fn copy_escapes_injection_in_role_and_path() {
let sql = copy_statement(
"\"t\"",
None,
"s3://b/k'; DROP TABLE t;--.jsonl",
"arn'evil",
None,
RedshiftCopyFormat::Jsonl,
);
assert!(sql.contains("'s3://b/k''; DROP TABLE t;--.jsonl'"));
assert!(sql.contains("IAM_ROLE 'arn''evil'"));
}
#[test]
fn s3_uri_joins_bucket_and_key() {
assert_eq!(
s3_uri("my-bucket", "prefix/obj.jsonl"),
"s3://my-bucket/prefix/obj.jsonl"
);
}
#[test]
fn insert_statement_numbers_placeholders_per_row() {
let sql = insert_statement("\"t\"", &["a".into(), "b".into()], 2);
assert_eq!(
sql,
"INSERT INTO \"t\" (\"a\", \"b\") VALUES ($1, $2), ($3, $4)"
);
}
#[test]
fn insert_statement_single_row() {
let sql = insert_statement("\"t\"", &["a".into()], 1);
assert_eq!(sql, "INSERT INTO \"t\" (\"a\") VALUES ($1)");
}
#[test]
fn columns_present_is_union_in_table_order() {
let table = vec!["id".to_string(), "name".to_string(), "email".to_string()];
let records = vec![json!({"id": 1, "email": "a@b.c"}), json!({"name": "Bob"})];
let present = columns_present(&records, &table);
assert_eq!(
present.into_iter().cloned().collect::<Vec<_>>(),
vec!["id".to_string(), "name".to_string(), "email".to_string()]
);
}
#[test]
fn columns_present_drops_absent_columns() {
let table = vec!["id".to_string(), "unused".to_string()];
let records = vec![json!({"id": 1})];
let present = columns_present(&records, &table);
assert_eq!(present, vec![&"id".to_string()]);
}
#[test]
fn serialize_jsonl_is_newline_delimited() {
let out = serialize_jsonl(&[json!({"id": 1}), json!({"id": 2})]).unwrap();
let text = String::from_utf8(out).unwrap();
let lines: Vec<&str> = text.trim().split('\n').collect();
assert_eq!(lines.len(), 2);
assert_eq!(
serde_json::from_str::<Value>(lines[0]).unwrap()["id"],
json!(1)
);
}
#[test]
fn serialize_csv_orders_columns_and_quotes() {
let cols = vec!["id".to_string(), "note".to_string()];
let records = vec![
json!({"id": 1, "note": "hello, world"}),
json!({"id": 2, "note": "quote\"inside"}),
json!({"id": 3}), ];
let out = serialize_csv(&records, &cols).unwrap();
let text = String::from_utf8(out).unwrap();
let lines: Vec<&str> = text.trim_end().split('\n').collect();
assert_eq!(lines[0], "1,\"hello, world\"");
assert_eq!(lines[1], "2,\"quote\"\"inside\"");
assert_eq!(lines[2], "3,");
}
#[test]
fn serialize_csv_renders_scalar_types() {
let cols = vec!["b".to_string(), "n".to_string(), "nul".to_string()];
let out = serialize_csv(&[json!({"b": true, "n": 4.5, "nul": null})], &cols).unwrap();
assert_eq!(String::from_utf8(out).unwrap(), "true,4.5,\n");
}
#[test]
fn serialize_csv_rejects_non_object() {
assert!(serialize_csv(&[json!(5)], &["a".to_string()]).is_err());
}
}