Add Archive::read_table_rows

.github/workflows/test.yaml

@@ -30,12 +30,12 @@ jobs:
           key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
           restore-keys: ${{ runner.os }}-cargo-
 
-      - name: Build
-        run: cargo build
-
-      - name: Run tests
+      - name: Run tests without features
         run: cargo test
 
+      - name: Run tests with features
+        run: cargo test --all-features
+
       - name: Lint
         run: |
           cargo fmt --all -- --check

Cargo.toml

@@ -11,14 +11,20 @@ keywords = ["postgresql", "pgdump"]
 categories = ["database"]
 exclude = [".github"]
 
+[features]
+tabledata = ["dep:pg_query", "dep:csv", "dep:serde"]
+
 [lib]
 name = "pgarchive"
 path = "src/lib.rs"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
 chrono = "0.4.30"
+csv = { version = "1.3.0", optional = true }
 flate2 = "1.0.27"
+pg_query = { version = "0.8.2", optional = true }
+serde = { version = "1.0.188", optional = true }
 thiserror = "1.0.49"
 
 [dev-dependencies]

src/archive.rs

@@ -226,6 +226,16 @@ impl Archive {
             )),
         }
     }
+
+    #[cfg(feature = "tabledata")]
+    pub fn read_table_rows(
+        &self,
+        f: &mut File,
+        table: &str,
+    ) -> Result<csv::Reader<Box<dyn io::Read>>, ArchiveError> {
+        use crate::tabledata::table_data_reader;
+        table_data_reader(self, f, table)
+    }
 }
 
 #[cfg(test)]

src/lib.rs

@@ -18,9 +18,10 @@
 //! ```
 mod archive;
 mod io;
+#[cfg(feature = "tabledata")]
+mod tabledata;
 mod toc;
 mod types;
-
 pub use archive::Archive;
 pub use toc::{TocEntry, ID};
 pub use types::{ArchiveError, CompressionMethod, Section, Version};

src/tabledata.rs

@@ -0,0 +1,69 @@
+use crate::archive::Archive;
+use crate::types::{ArchiveError, Section};
+use pg_query::{NodeEnum, NodeRef};
+use std::fs::File;
+use std::io::Read;
+
+#[cfg(feature = "tabledata")]
+pub fn table_data_reader(
+    archive: &Archive,
+    file: &mut File,
+    table: &str,
+) -> Result<csv::Reader<Box<dyn Read>>, ArchiveError> {
+    let create_entry = archive
+        .find_toc_entry(Section::PreData, "TABLE", table)
+        .ok_or(ArchiveError::NoDataPresent)?;
+    let columns = table_column_names(&create_entry.defn).or(Err(ArchiveError::InvalidData(
+        "invalid CREATE TABLE statement".into(),
+    )))?;
+
+    let data_entry = archive
+        .find_toc_entry(Section::Data, "TABLE DATA", table)
+        .ok_or(ArchiveError::NoDataPresent)
+        .unwrap();
+    let data = archive.read_data(file, data_entry).unwrap();
+    let mut rdr = csv::ReaderBuilder::new()
+        .delimiter(b'\t')
+        .quoting(false)
+        .flexible(false)
+        .from_reader(data);
+    rdr.set_headers(columns.into());
+    Ok(rdr)
+}
+
+#[cfg(feature = "tabledata")]
+fn table_column_names(create_stmt: &str) -> Result<Vec<String>, pg_query::Error> {
+    let result = pg_query::parse(create_stmt)?;
+    let stmt = result.protobuf.nodes()[0].0;
+    match stmt {
+        NodeRef::CreateStmt(table_info) => Ok(table_info
+            .table_elts
+            .iter()
+            .filter_map(|e| match &e.node {
+                Some(NodeEnum::ColumnDef(cd)) => Some(cd.as_ref().colname.clone()),
+                _ => None,
+            })
+            .collect()),
+        _ => Err(pg_query::Error::Parse("invalid statement type".into())),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_table_column_names() {
+        assert!(table_column_names(
+            "CREATE DATABASE pizza WITH TEMPLATE = template0 ENCODING = 'UTF8' LOCALE = 'C';"
+        )
+        .is_err());
+
+        let columns = table_column_names(
+            "CREATE TABLE public.pizza (pizza_id integer NOT NULL, name text NOT NULL);",
+        );
+        assert!(columns.is_ok());
+        let columns = columns.unwrap();
+        assert_eq!(columns, vec!["pizza_id", "name"]);
+    }
+}

tests/data_test.rs

@@ -20,3 +20,26 @@ fn test_table_data() -> Result<(), pgarchive::ArchiveError> {
     );
     Ok(())
 }
+
+#[cfg(feature = "tabledata")]
+#[test]
+fn test_table_rows() -> Result<(), Box<dyn std::error::Error>> {
+    use csv::StringRecord;
+
+    let cargo_path = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests");
+    let mut f = File::open(cargo_path.join("test.pgdump"))?;
+    let archive = pgarchive::Archive::parse(&mut f)?;
+    let mut reader = archive.read_table_rows(&mut f, "pizza")?;
+    let rows: Vec<StringRecord> = reader
+        .records()
+        .into_iter()
+        .filter(|r| r.is_ok())
+        .map(|r| r.unwrap())
+        .collect();
+    assert_eq!(rows.len(), 5);
+    assert_eq!(
+        rows.first().unwrap().iter().collect::<Vec<&str>>(),
+        vec!["1", "The Classic"]
+    );
+    Ok(())
+}