pgrestore

[ManReadApp]

It's probably quite inefficient, but it works.

async fn pg_restore(client: &mut tokio_postgres::Client, path: &Path) {
    let mut builder = vec![];
    let mut queries = vec![];
    let mut copies = vec![];
    let mut copy = false;
    for line in read_to_string(path).unwrap().lines() {
        if line.starts_with("--") || line.is_empty() {
            continue;
        }
        builder.push(line.to_string());
        if line.ends_with("FROM stdin;") {
            copy = true;
        }
        if copy {
            if line == "\\." {
                builder.pop();
                let query = builder.join("\n");
                copies.push(query);
                builder = vec![];
                copy = false;
            }
        } else {
            if line.ends_with(";") && !line.ends_with("\\;") {
                let query = builder.join("\n");
                queries.push(query);
                builder = vec![];
            }
        }
    }
    client.batch_execute(&queries.join("\n")).await.unwrap();
    for copy in copies {
        if let Some((query, lines)) = copy.split_once("\n") {
            let mut stream = stream::iter(
                lines
                    .split("\n")
                    .map(|s| Bytes::from(format!("{s}\n")))
                    .map(Ok::<_, tokio_postgres::Error>),
            );
            let sink: CopyInSink<Bytes> = client.copy_in(query).await.unwrap();
            pin_mut!(sink);
            sink.send_all(&mut stream).await.unwrap();
            let _ = sink.finish().await.unwrap();
        }
    }
}