byteor_adapters/

postgres.rs

use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use bytes::Bytes;
use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod};
use futures_util::SinkExt;
use serde_json::{Map, Value};
use tokio::pin;
use tokio::runtime::Runtime;
use tokio_postgres::types::Json;
use tokio_postgres::{NoTls, SimpleQueryMessage};

use crate::{Adapter, AdapterError, EgressAdapter, IngressAdapter};

pub(crate) struct PostgresFlushResult {
    pub rows_affected: u64,
}

#[derive(Debug)]
struct PostgresIngressStatus {
    query_state: String,
    row_format: String,
    last_query_rows: u64,
    pending_rows: u64,
}

impl PostgresIngressStatus {
    fn new(row_format: PostgresRowFormat) -> Self {
        Self {
            query_state: "idle".to_string(),
            row_format: match row_format {
                PostgresRowFormat::Json => "json".to_string(),
                PostgresRowFormat::Columns => "columns".to_string(),
            },
            last_query_rows: 0,
            pending_rows: 0,
        }
    }

    fn status_fields(&self) -> serde_json::Map<String, serde_json::Value> {
        let mut details = serde_json::Map::new();
        details.insert(
            "query_state".to_string(),
            serde_json::Value::String(self.query_state.clone()),
        );
        details.insert(
            "row_format".to_string(),
            serde_json::Value::String(self.row_format.clone()),
        );
        details.insert(
            "last_query_rows".to_string(),
            serde_json::Value::from(self.last_query_rows),
        );
        details.insert(
            "pending_rows".to_string(),
            serde_json::Value::from(self.pending_rows),
        );
        details
    }
}

#[derive(Debug)]
struct PostgresEgressStatus {
    write_state: String,
    write_mode: String,
    buffered_rows: u64,
    last_flush_rows: u64,
    last_rows_affected: Option<u64>,
}

impl PostgresEgressStatus {
    fn new(mode: PostgresWriteMode) -> Self {
        Self {
            write_state: "idle".to_string(),
            write_mode: match mode {
                PostgresWriteMode::Insert => "insert".to_string(),
                PostgresWriteMode::Upsert => "upsert".to_string(),
                PostgresWriteMode::Copy => "copy".to_string(),
            },
            buffered_rows: 0,
            last_flush_rows: 0,
            last_rows_affected: None,
        }
    }

    fn status_fields(&self) -> serde_json::Map<String, serde_json::Value> {
        let mut details = serde_json::Map::new();
        details.insert(
            "write_state".to_string(),
            serde_json::Value::String(self.write_state.clone()),
        );
        details.insert(
            "write_mode".to_string(),
            serde_json::Value::String(self.write_mode.clone()),
        );
        details.insert(
            "buffered_rows".to_string(),
            serde_json::Value::from(self.buffered_rows),
        );
        details.insert(
            "last_flush_rows".to_string(),
            serde_json::Value::from(self.last_flush_rows),
        );
        if let Some(rows_affected) = self.last_rows_affected {
            details.insert(
                "last_rows_affected".to_string(),
                serde_json::Value::from(rows_affected),
            );
        }
        details
    }
}

/// Result row serialization mode for PostgreSQL ingress.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PostgresRowFormat {
    /// Emit each row as a JSON object keyed by column name.
    Json,
    /// Emit each row as `{"columns": [...], "values": [...]}` using string values.
    Columns,
}

impl PostgresRowFormat {
    fn parse(value: Option<&str>) -> Result<Self, AdapterError> {
        match value.map(str::trim).filter(|value| !value.is_empty()) {
            None | Some("json") => Ok(Self::Json),
            Some("columns") => Ok(Self::Columns),
            Some(other) => Err(AdapterError::Config {
                detail: format!("unsupported postgres row_format: {other}"),
            }),
        }
    }
}

/// Persistence mode for PostgreSQL egress.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PostgresWriteMode {
    /// Insert rows into the target table.
    Insert,
    /// Upsert rows into the target table using a conflict key.
    Upsert,
    /// Bulk load rows into the target table using `COPY FROM STDIN`.
    Copy,
}

impl PostgresWriteMode {
    fn parse(value: &str) -> Result<Self, AdapterError> {
        match value.trim() {
            "insert" => Ok(Self::Insert),
            "upsert" => Ok(Self::Upsert),
            "copy" => Ok(Self::Copy),
            other => Err(AdapterError::Config {
                detail: format!("unsupported postgres write mode: {other}"),
            }),
        }
    }
}

/// Internal PostgreSQL ingress client contract used by the adapter and crate-local tests.
pub(crate) trait PostgresIngressClient: Send {
    fn read_next(&mut self, timeout: Duration) -> Result<Option<Vec<u8>>, AdapterError>;
}

/// Internal PostgreSQL egress client contract used by the adapter and crate-local tests.
pub(crate) trait PostgresEgressClient: Send {
    fn flush_rows(&mut self, rows: &[Map<String, Value>]) -> Result<PostgresFlushResult, AdapterError>;
}

struct AsyncPostgresIngressClient {
    runtime: Runtime,
    pool: Pool,
    query: String,
    row_format: PostgresRowFormat,
    max_row_bytes: usize,
    pending: VecDeque<Vec<u8>>,
    status: Arc<Mutex<PostgresIngressStatus>>,
}

impl AsyncPostgresIngressClient {
    fn connect(
        dsn: &str,
        query: &str,
        row_format: PostgresRowFormat,
        max_row_bytes: usize,
        status: Arc<Mutex<PostgresIngressStatus>>,
    ) -> Result<Self, AdapterError> {
        Ok(Self {
            runtime: build_runtime("postgres_ingress_runtime_build")?,
            pool: build_pool(dsn)?,
            query: query.trim().to_string(),
            row_format,
            max_row_bytes,
            pending: VecDeque::new(),
            status,
        })
    }

    async fn fetch_pending_json(
        pool: Pool,
        query: String,
        max_row_bytes: usize,
    ) -> Result<Vec<Vec<u8>>, AdapterError> {
        let client = pool.get().await.map_err(|err| AdapterError::Postgres {
            op: "postgres_pool_get",
            detail: err.to_string(),
        })?;
        let wrapped_query = format!("SELECT row_to_json(t)::text FROM ({query}) AS t");
        let rows = client
            .query(wrapped_query.as_str(), &[])
            .await
            .map_err(|err| AdapterError::Postgres {
                op: "postgres_query_rows",
                detail: err.to_string(),
            })?;

        rows.into_iter()
            .map(|row| {
                let payload: String = row.get(0);
                enforce_payload_limit(
                    "postgres_ingress_row_to_json",
                    payload.as_bytes(),
                    max_row_bytes,
                )
            })
            .collect()
    }

    async fn fetch_pending_columns(
        pool: Pool,
        query: String,
        max_row_bytes: usize,
    ) -> Result<Vec<Vec<u8>>, AdapterError> {
        let client = pool.get().await.map_err(|err| AdapterError::Postgres {
            op: "postgres_pool_get",
            detail: err.to_string(),
        })?;
        let messages =
            client
                .simple_query(query.as_str())
                .await
                .map_err(|err| AdapterError::Postgres {
                    op: "postgres_simple_query",
                    detail: err.to_string(),
                })?;

        let mut payloads = Vec::new();
        for message in messages {
            let SimpleQueryMessage::Row(row) = message else {
                continue;
            };

            let columns = row
                .columns()
                .iter()
                .map(|column| Value::String(column.name().to_string()))
                .collect::<Vec<_>>();
            let values = row
                .columns()
                .iter()
                .enumerate()
                .map(|(index, _)| match row.get(index) {
                    Some(value) => Value::String(value.to_string()),
                    None => Value::Null,
                })
                .collect::<Vec<_>>();
            let payload = serde_json::to_vec(&serde_json::json!({
                "columns": columns,
                "values": values,
            }))
            .map_err(|err| AdapterError::Postgres {
                op: "postgres_columns_serialize",
                detail: err.to_string(),
            })?;
            payloads.push(enforce_payload_limit(
                "postgres_columns_serialize",
                &payload,
                max_row_bytes,
            )?);
        }

        Ok(payloads)
    }
}

impl PostgresIngressClient for AsyncPostgresIngressClient {
    fn read_next(&mut self, timeout: Duration) -> Result<Option<Vec<u8>>, AdapterError> {
        if let Some(payload) = self.pending.pop_front() {
            self.status.lock().unwrap().pending_rows = self.pending.len() as u64;
            return Ok(Some(payload));
        }

        self.status.lock().unwrap().query_state = "querying".to_string();

        let pool = self.pool.clone();
        let query = self.query.clone();
        let max_row_bytes = self.max_row_bytes;
        let fetched = match self.row_format {
            PostgresRowFormat::Json => {
                self.runtime
                    .block_on(Self::fetch_pending_json(pool, query, max_row_bytes))?
            }
            PostgresRowFormat::Columns => {
                self.runtime
                    .block_on(Self::fetch_pending_columns(pool, query, max_row_bytes))?
            }
        };

        if fetched.is_empty() {
            let mut status = self.status.lock().unwrap();
            status.query_state = "idle".to_string();
            status.last_query_rows = 0;
            status.pending_rows = 0;
            std::thread::sleep(timeout);
            return Ok(None);
        }

        {
            let mut status = self.status.lock().unwrap();
            status.query_state = "rows_buffered".to_string();
            status.last_query_rows = fetched.len() as u64;
            status.pending_rows = fetched.len().saturating_sub(1) as u64;
        }
        self.pending.extend(fetched);
        Ok(self.pending.pop_front())
    }
}

struct AsyncPostgresEgressClient {
    runtime: Runtime,
    pool: Pool,
    table: String,
    mode: PostgresWriteMode,
    conflict_key: Option<String>,
}

impl AsyncPostgresEgressClient {
    fn connect(
        dsn: &str,
        table: &str,
        mode: PostgresWriteMode,
        conflict_key: Option<&str>,
    ) -> Result<Self, AdapterError> {
        let table = quote_relation(table)?;
        let conflict_key = conflict_key.map(quote_identifier).transpose()?;
        if matches!(mode, PostgresWriteMode::Upsert) && conflict_key.is_none() {
            return Err(AdapterError::Config {
                detail: "postgres upsert mode requires a conflict_key".into(),
            });
        }

        Ok(Self {
            runtime: build_runtime("postgres_egress_runtime_build")?,
            pool: build_pool(dsn)?,
            table,
            mode,
            conflict_key,
        })
    }

    async fn flush_insert_or_upsert(
        pool: Pool,
        table: String,
        mode: PostgresWriteMode,
        conflict_key: Option<String>,
        rows: Vec<Map<String, Value>>,
    ) -> Result<PostgresFlushResult, AdapterError> {
        let columns = union_keys(&rows)?;
        if columns.is_empty() {
            return Err(AdapterError::Config {
                detail: "postgres egress row payload must contain at least one column".into(),
            });
        }

        if let Some(conflict_key) = conflict_key.as_deref() {
            let conflict_field = conflict_key.trim_matches('"');
            for row in &rows {
                if !row.contains_key(conflict_field) {
                    return Err(AdapterError::Config {
                        detail: format!(
                            "postgres upsert payload missing conflict_key field: {conflict_field}"
                        ),
                    });
                }
            }
        }

        let rows_json = Value::Array(rows.into_iter().map(Value::Object).collect());
        let sql = build_insert_sql(&table, &columns, mode, conflict_key.as_deref())?;
        let client = pool.get().await.map_err(|err| AdapterError::Postgres {
            op: "postgres_pool_get",
            detail: err.to_string(),
        })?;
        let statement =
            client
                .prepare_cached(sql.as_str())
                .await
                .map_err(|err| AdapterError::Postgres {
                    op: "postgres_prepare_cached",
                    detail: err.to_string(),
                })?;
        client
            .execute(&statement, &[&Json(rows_json)])
            .await
            .map_err(|err| AdapterError::Postgres {
                op: "postgres_execute",
                detail: err.to_string(),
            })
            .map(|rows_affected| PostgresFlushResult { rows_affected })
    }

    async fn flush_copy(
        pool: Pool,
        table: String,
        rows: Vec<Map<String, Value>>,
    ) -> Result<PostgresFlushResult, AdapterError> {
        let columns = union_keys(&rows)?;
        if columns.is_empty() {
            return Err(AdapterError::Config {
                detail: "postgres copy payload must contain at least one column".into(),
            });
        }

        let copy_sql = format!(
            "COPY {} ({}) FROM STDIN WITH (FORMAT csv, NULL '\\N')",
            table,
            columns.join(", ")
        );
        let csv = render_copy_csv(&rows, &columns)?;
        let client = pool.get().await.map_err(|err| AdapterError::Postgres {
            op: "postgres_pool_get",
            detail: err.to_string(),
        })?;
        let sink =
            client
                .copy_in(copy_sql.as_str())
                .await
                .map_err(|err| AdapterError::Postgres {
                    op: "postgres_copy_in",
                    detail: err.to_string(),
                })?;
        pin!(sink);
        sink.send(Bytes::from(csv))
            .await
            .map_err(|err| AdapterError::Postgres {
                op: "postgres_copy_send",
                detail: err.to_string(),
            })?;
        let rows_affected = sink.as_mut().finish().await.map_err(|err| AdapterError::Postgres {
            op: "postgres_copy_close",
            detail: err.to_string(),
        })?;
        Ok(PostgresFlushResult { rows_affected })
    }
}

impl PostgresEgressClient for AsyncPostgresEgressClient {
    fn flush_rows(&mut self, rows: &[Map<String, Value>]) -> Result<PostgresFlushResult, AdapterError> {
        let rows = rows.to_vec();
        let pool = self.pool.clone();
        let table = self.table.clone();
        match self.mode {
            PostgresWriteMode::Insert | PostgresWriteMode::Upsert => {
                self.runtime.block_on(Self::flush_insert_or_upsert(
                    pool,
                    table,
                    self.mode,
                    self.conflict_key.clone(),
                    rows,
                ))
            }
            PostgresWriteMode::Copy => self.runtime.block_on(Self::flush_copy(pool, table, rows)),
        }
    }
}

fn build_runtime(op: &'static str) -> Result<Runtime, AdapterError> {
    tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .map_err(|err| AdapterError::Postgres {
            op,
            detail: err.to_string(),
        })
}

fn build_pool(dsn: &str) -> Result<Pool, AdapterError> {
    let config = dsn
        .parse::<tokio_postgres::Config>()
        .map_err(|err| AdapterError::Config {
            detail: format!("invalid postgres dsn: {err}"),
        })?;
    let manager = Manager::from_config(
        config,
        NoTls,
        ManagerConfig {
            recycling_method: RecyclingMethod::Fast,
        },
    );
    Pool::builder(manager)
        .max_size(16)
        .build()
        .map_err(|err| AdapterError::Postgres {
            op: "postgres_pool_build",
            detail: err.to_string(),
        })
}

fn enforce_payload_limit(
    op: &'static str,
    payload: &[u8],
    max_row_bytes: usize,
) -> Result<Vec<u8>, AdapterError> {
    if payload.len() > max_row_bytes {
        return Err(AdapterError::Postgres {
            op,
            detail: format!(
                "postgres row too large len={} max={max_row_bytes}",
                payload.len()
            ),
        });
    }
    Ok(payload.to_vec())
}

fn quote_relation(value: &str) -> Result<String, AdapterError> {
    let segments = value
        .split('.')
        .map(quote_identifier)
        .collect::<Result<Vec<_>, _>>()?;
    if segments.is_empty() {
        return Err(AdapterError::Config {
            detail: "postgres table must not be empty".into(),
        });
    }
    Ok(segments.join("."))
}

fn quote_identifier(value: &str) -> Result<String, AdapterError> {
    let value = value.trim();
    let mut chars = value.chars();
    let Some(first) = chars.next() else {
        return Err(AdapterError::Config {
            detail: "postgres identifier must not be empty".into(),
        });
    };
    if !(first == '_' || first.is_ascii_alphabetic())
        || !chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric())
    {
        return Err(AdapterError::Config {
            detail: format!("unsupported postgres identifier: {value}"),
        });
    }
    Ok(format!("\"{value}\""))
}

fn union_keys(rows: &[Map<String, Value>]) -> Result<Vec<String>, AdapterError> {
    let mut keys = rows
        .iter()
        .flat_map(|row| row.keys())
        .map(|key| quote_identifier(key))
        .collect::<Result<Vec<_>, _>>()?;
    keys.sort();
    keys.dedup();
    Ok(keys)
}

fn build_insert_sql(
    table: &str,
    columns: &[String],
    mode: PostgresWriteMode,
    conflict_key: Option<&str>,
) -> Result<String, AdapterError> {
    let columns_csv = columns.join(", ");
    let mut sql = format!(
        "INSERT INTO {table} ({columns_csv}) SELECT {columns_csv} FROM jsonb_populate_recordset(NULL::{table}, $1::jsonb)"
    );
    if matches!(mode, PostgresWriteMode::Upsert) {
        let conflict_key = conflict_key.ok_or_else(|| AdapterError::Config {
            detail: "postgres upsert mode requires a conflict_key".into(),
        })?;
        let updates = columns
            .iter()
            .filter(|column| column.as_str() != conflict_key)
            .map(|column| format!("{column} = EXCLUDED.{column}"))
            .collect::<Vec<_>>();
        if updates.is_empty() {
            sql.push_str(&format!(" ON CONFLICT ({conflict_key}) DO NOTHING"));
        } else {
            sql.push_str(&format!(
                " ON CONFLICT ({conflict_key}) DO UPDATE SET {}",
                updates.join(", ")
            ));
        }
    }
    Ok(sql)
}

fn render_copy_csv(
    rows: &[Map<String, Value>],
    columns: &[String],
) -> Result<Vec<u8>, AdapterError> {
    let mut csv = String::new();
    for row in rows {
        for (index, column) in columns.iter().enumerate() {
            if index > 0 {
                csv.push(',');
            }
            let key = column.trim_matches('"');
            push_csv_field(&mut csv, row.get(key))?;
        }
        csv.push('\n');
    }
    Ok(csv.into_bytes())
}

fn push_csv_field(out: &mut String, value: Option<&Value>) -> Result<(), AdapterError> {
    match value {
        None | Some(Value::Null) => out.push_str("\\N"),
        Some(Value::String(text)) => push_csv_string(out, text),
        Some(Value::Bool(flag)) => push_csv_string(out, if *flag { "true" } else { "false" }),
        Some(Value::Number(number)) => push_csv_string(out, &number.to_string()),
        Some(other) => {
            let text = serde_json::to_string(other).map_err(|err| AdapterError::Postgres {
                op: "postgres_copy_serialize",
                detail: err.to_string(),
            })?;
            push_csv_string(out, &text);
        }
    }
    Ok(())
}

fn push_csv_string(out: &mut String, value: &str) {
    let must_quote = value == "\\N"
        || value.contains(',')
        || value.contains('"')
        || value.contains('\n')
        || value.contains('\r');
    if must_quote {
        out.push('"');
        for ch in value.chars() {
            if ch == '"' {
                out.push('"');
            }
            out.push(ch);
        }
        out.push('"');
    } else {
        out.push_str(value);
    }
}

fn parse_rows_payload(payload: &[u8]) -> Result<Vec<Map<String, Value>>, AdapterError> {
    let value: Value = serde_json::from_slice(payload).map_err(|err| AdapterError::Postgres {
        op: "postgres_payload_parse",
        detail: err.to_string(),
    })?;
    match value {
        Value::Object(object) => Ok(vec![normalize_row_object(object)?]),
        Value::Array(items) => items
            .into_iter()
            .map(|item| match item {
                Value::Object(object) => normalize_row_object(object),
                _ => Err(AdapterError::Config {
                    detail: "postgres payload arrays must contain JSON objects".into(),
                }),
            })
            .collect(),
        _ => Err(AdapterError::Config {
            detail: "postgres payload must be a JSON object or array of objects".into(),
        }),
    }
}

fn normalize_row_object(
    mut object: Map<String, Value>,
) -> Result<Map<String, Value>, AdapterError> {
    if let (Some(columns), Some(values)) = (object.remove("columns"), object.remove("values")) {
        return row_from_columns_payload(columns, values);
    }
    Ok(object)
}

fn row_from_columns_payload(
    columns: Value,
    values: Value,
) -> Result<Map<String, Value>, AdapterError> {
    let Value::Array(columns) = columns else {
        return Err(AdapterError::Config {
            detail: "postgres columns payload must provide an array of column names".into(),
        });
    };
    let Value::Array(values) = values else {
        return Err(AdapterError::Config {
            detail: "postgres columns payload must provide an array of values".into(),
        });
    };
    if columns.len() != values.len() {
        return Err(AdapterError::Config {
            detail: "postgres columns payload column/value lengths do not match".into(),
        });
    }

    let mut row = Map::new();
    for (column, value) in columns.into_iter().zip(values) {
        let Value::String(column_name) = column else {
            return Err(AdapterError::Config {
                detail: "postgres columns payload column names must be strings".into(),
            });
        };
        row.insert(column_name, value);
    }
    Ok(row)
}

/// Minimal blocking PostgreSQL ingress adapter.
///
/// The adapter executes the configured query whenever its internal queue is empty. Query authors
/// should provide incremental or tailing SQL if they want exactly-once style progression.
pub struct PostgresIngress {
    name: String,
    client: Box<dyn PostgresIngressClient>,
    poll_interval: Duration,
    status: Arc<Mutex<PostgresIngressStatus>>,
}

impl PostgresIngress {
    /// Create a blocking PostgreSQL ingress adapter.
    pub fn new(
        name: impl Into<String>,
        dsn: &str,
        query: &str,
        poll_interval: Duration,
        max_row_bytes: usize,
        row_format: Option<&str>,
    ) -> Result<Self, AdapterError> {
        let row_format = PostgresRowFormat::parse(row_format)?;
        let status = Arc::new(Mutex::new(PostgresIngressStatus::new(row_format)));
        let client = AsyncPostgresIngressClient::connect(
            dsn,
            query,
            row_format,
            max_row_bytes,
            Arc::clone(&status),
        )?;
        Ok(Self {
            name: name.into(),
            client: Box::new(client),
            poll_interval,
            status,
        })
    }

    #[cfg(test)]
    /// Build a PostgreSQL ingress adapter around a crate-local fake client.
    pub(crate) fn with_client(
        name: impl Into<String>,
        client: impl PostgresIngressClient + 'static,
        poll_interval: Duration,
    ) -> Self {
        Self {
            name: name.into(),
            client: Box::new(client),
            poll_interval,
            status: Arc::new(Mutex::new(PostgresIngressStatus::new(PostgresRowFormat::Json))),
        }
    }
}

impl Adapter for PostgresIngress {
    fn name(&self) -> &str {
        &self.name
    }

    fn transport_kind(&self) -> &str {
        "postgres"
    }

    fn status_fields(&self) -> serde_json::Map<String, serde_json::Value> {
        self.status.lock().unwrap().status_fields()
    }
}

impl IngressAdapter for PostgresIngress {
    fn read_next(&mut self, out: &mut [u8]) -> Result<Option<usize>, AdapterError> {
        let Some(payload) = self.client.read_next(self.poll_interval)? else {
            return Ok(None);
        };
        if payload.len() > out.len() {
            return Err(AdapterError::Postgres {
                op: "postgres_read_next",
                detail: format!("payload too large len={} buf={}", payload.len(), out.len()),
            });
        }
        out[..payload.len()].copy_from_slice(&payload);
        Ok(Some(payload.len()))
    }
}

/// Minimal blocking PostgreSQL egress adapter.
///
/// Input messages must be JSON objects, arrays of JSON objects, or ingress-style
/// `{"columns": [...], "values": [...]}` envelopes.
pub struct PostgresEgress {
    name: String,
    client: Box<dyn PostgresEgressClient>,
    buffered_rows: Vec<Map<String, Value>>,
    status: Arc<Mutex<PostgresEgressStatus>>,
}

impl PostgresEgress {
    /// Create a blocking PostgreSQL egress adapter.
    pub fn new(
        name: impl Into<String>,
        dsn: &str,
        table: &str,
        mode: &str,
        conflict_key: Option<&str>,
    ) -> Result<Self, AdapterError> {
        let mode = PostgresWriteMode::parse(mode)?;
        let client = AsyncPostgresEgressClient::connect(
            dsn,
            table,
            mode,
            conflict_key,
        )?;
        Ok(Self {
            name: name.into(),
            client: Box::new(client),
            buffered_rows: Vec::new(),
            status: Arc::new(Mutex::new(PostgresEgressStatus::new(mode))),
        })
    }

    #[cfg(test)]
    /// Build a PostgreSQL egress adapter around a crate-local fake client.
    pub(crate) fn with_client(
        name: impl Into<String>,
        client: impl PostgresEgressClient + 'static,
    ) -> Self {
        Self {
            name: name.into(),
            client: Box::new(client),
            buffered_rows: Vec::new(),
            status: Arc::new(Mutex::new(PostgresEgressStatus::new(PostgresWriteMode::Insert))),
        }
    }
}

impl Adapter for PostgresEgress {
    fn name(&self) -> &str {
        &self.name
    }

    fn transport_kind(&self) -> &str {
        "postgres"
    }

    fn status_fields(&self) -> serde_json::Map<String, serde_json::Value> {
        self.status.lock().unwrap().status_fields()
    }
}

impl EgressAdapter for PostgresEgress {
    fn write_msg(&mut self, msg: &[u8]) -> Result<(), AdapterError> {
        self.buffered_rows.extend(parse_rows_payload(msg)?);
        let mut status = self.status.lock().unwrap();
        status.write_state = "buffering".to_string();
        status.buffered_rows = self.buffered_rows.len() as u64;
        Ok(())
    }

    fn flush(&mut self) -> Result<(), AdapterError> {
        if self.buffered_rows.is_empty() {
            return Ok(());
        }
        {
            let mut status = self.status.lock().unwrap();
            status.write_state = "flushing".to_string();
            status.last_flush_rows = self.buffered_rows.len() as u64;
        }
        let result = self.client.flush_rows(&self.buffered_rows)?;
        self.buffered_rows.clear();
        let mut status = self.status.lock().unwrap();
        status.write_state = "idle".to_string();
        status.buffered_rows = 0;
        status.last_rows_affected = Some(result.rows_affected);
        Ok(())
    }
}