feat: binary copy osmpbf

jjcfrancisco · Sep 11, 2024 · 20eda1f · 20eda1f
1 parent 059394b
commit 20eda1f
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 21 deletions.
diff --git a/src/format/common.rs b/src/format/common.rs
@@ -6,6 +6,7 @@ use std::path::Path;
 use crate::pg::binary_copy::Wkb;
 
 // Struct to hold column name and data type
+#[derive(Debug)]
 pub struct NewTableTypes {
     pub column_name: String,
     pub data_type: Type,
@@ -49,6 +50,7 @@ pub enum AcceptedTypes {
     Double(Option<f32>),
     Text(Option<String>),
     Bool(Option<bool>),
+    Array(Option<Vec<String>>),
     Geometry(Option<Wkb>),
 }
 

diff --git a/src/format/osmpbf.rs b/src/format/osmpbf.rs
@@ -1,5 +1,9 @@
+use crate::format::common::{AcceptedTypes, NewTableTypes, Row, Rows};
+use crate::pg::binary_copy::Wkb;
 use crate::utils::cli::Cli;
 use crate::Result;
+use postgres::types::Type;
+use wkb::geom_to_wkb;
 
 use geo::{Geometry, LineString, Point, Polygon};
 use osmpbf::{Element, ElementReader};
@@ -8,7 +12,7 @@ use std::collections::HashMap;
 #[derive(Debug)]
 #[allow(dead_code)]
 struct OsmPbf {
-    tags: Vec<(String, String)>,
+    tags: Vec<String>,
     geometry: geo::Geometry<f64>,
 }
 
@@ -22,7 +26,7 @@ fn build_nodes(file_path: &str) -> Result<HashMap<i64, OsmPbf>> {
                 OsmPbf {
                     tags: node
                         .tags()
-                        .map(|(key, value)| (key.to_string(), value.to_string()))
+                        .map(|(key, value)| format!("{}={}", key, value))
                         .collect(),
                     geometry: Geometry::Point(Point::new(node.lon(), node.lat())),
                 },
@@ -34,7 +38,7 @@ fn build_nodes(file_path: &str) -> Result<HashMap<i64, OsmPbf>> {
                 OsmPbf {
                     tags: dense_node
                         .tags()
-                        .map(|(key, value)| (key.to_string(), value.to_string()))
+                        .map(|(key, value)| format!("{}={}", key, value))
                         .collect(),
                     geometry: Geometry::Point(Point::new(dense_node.lon(), dense_node.lat())),
                 },
@@ -46,7 +50,7 @@ fn build_nodes(file_path: &str) -> Result<HashMap<i64, OsmPbf>> {
     Ok(nodes)
 }
 
-fn build_polygon(way: &osmpbf::Way, nodes: &HashMap<i64, OsmPbf>) -> OsmPbf {
+fn build_line(way: &osmpbf::Way, nodes: &HashMap<i64, OsmPbf>) -> OsmPbf {
     let mut points: Vec<(f64, f64)> = Vec::new();
     way.refs().for_each(|node_id| {
         if let Some(node) = nodes.get(&node_id) {
@@ -55,17 +59,16 @@ fn build_polygon(way: &osmpbf::Way, nodes: &HashMap<i64, OsmPbf>) -> OsmPbf {
             }
         }
     });
-
     OsmPbf {
         tags: way
             .tags()
-            .map(|(key, value)| (key.to_string(), value.to_string()))
+            .map(|(key, value)| format!("{}={}", key, value))
             .collect(),
-        geometry: Geometry::Polygon(Polygon::new(LineString::from(points), vec![])),
+        geometry: Geometry::LineString(LineString::from(points)),
     }
 }
 
-fn build_line(way: &osmpbf::Way, nodes: &HashMap<i64, OsmPbf>) -> OsmPbf {
+fn build_polygon(way: &osmpbf::Way, nodes: &HashMap<i64, OsmPbf>) -> OsmPbf {
     let mut points: Vec<(f64, f64)> = Vec::new();
     way.refs().for_each(|node_id| {
         if let Some(node) = nodes.get(&node_id) {
@@ -74,12 +77,13 @@ fn build_line(way: &osmpbf::Way, nodes: &HashMap<i64, OsmPbf>) -> OsmPbf {
             }
         }
     });
+
     OsmPbf {
         tags: way
             .tags()
-            .map(|(key, value)| (key.to_string(), value.to_string()))
+            .map(|(key, value)| format!("{}={}", key, value))
             .collect(),
-        geometry: Geometry::LineString(LineString::from(points)),
+        geometry: Geometry::Polygon(Polygon::new(LineString::from(points), vec![])),
     }
 }
 
@@ -109,11 +113,39 @@ fn build_polys_and_lines(file_path: &str, nodes: &HashMap<i64, OsmPbf>) -> Resul
     Ok(all)
 }
 
-pub fn read_osmpbf(args: &Cli) -> Result<()> {
+fn osmpbf_to_rows(all: &Vec<OsmPbf>) -> Result<Rows> {
+    let mut rows = Rows::new();
+
+    for item in all.into_iter() {
+        let mut row = Row::new();
+        // Tags
+        row.add(AcceptedTypes::Array(Some(item.tags.clone())));
+        // Geom
+        let wkb =
+            geom_to_wkb(&item.geometry.clone()).expect("❌ Could not convert geometry to WKB");
+        row.add(AcceptedTypes::Geometry(Some(Wkb { geometry: wkb })));
+        rows.add(row);
+    }
+
+    Ok(rows)
+}
+
+pub fn determine_data_types() -> Result<Vec<NewTableTypes>> {
+    let mut data_types: Vec<NewTableTypes> = Vec::new();
+    data_types.push(NewTableTypes {
+        column_name: "tags".to_string(),
+        data_type: Type::TEXT_ARRAY,
+    });
+
+    Ok(data_types)
+}
+
+pub fn read_osmpbf(args: &Cli) -> Result<Rows> {
     let nodes = build_nodes(&args.input)?;
     let all = build_polys_and_lines(&args.input, &nodes)?;
+    let rows = osmpbf_to_rows(&all)?;
 
-    Ok(())
+    Ok(rows)
 }
 
 #[cfg(test)]

diff --git a/src/pg/binary_copy.rs b/src/pg/binary_copy.rs
@@ -73,7 +73,7 @@ pub fn insert_rows(
     let mut writer = BinaryCopyInWriter::new(writer, &types);
 
     // Use to test if types are correct
-    if std::env::var("DEBUG").is_ok() {
+    if cfg!(debug_assertions) {
         println!("DEBUG || {:?}", types);
     }
 
@@ -99,6 +99,9 @@ pub fn insert_rows(
                 AcceptedTypes::Bool(value) => {
                     tosql.push(value);
                 }
+                AcceptedTypes::Array(value) => {
+                    tosql.push(value);
+                }
                 AcceptedTypes::Geometry(value) => {
                     tosql.push(value);
                 }

diff --git a/src/pg/crud.rs b/src/pg/crud.rs
@@ -55,14 +55,17 @@ pub fn create_table(
             Type::BOOL => {
                 query.push_str(&format!("{} BOOL,", column.column_name));
             }
+            Type::TEXT_ARRAY => {
+                query.push_str(&format!("{} TEXT[],", column.column_name));
+            }
             _ => println!("❌ Type currently not supported"),
         }
     }
     query.push_str(&format!("geom Geometry(Geometry, {})", srid));
     query.push_str(");");
 
     // Debugging
-    if std::env::var("DEBUG").is_ok() {
+    if cfg!(debug_assertions) {
         println!("DEBUG || {}", query);
     }
 

diff --git a/src/utils/cli.rs b/src/utils/cli.rs
@@ -1,6 +1,6 @@
 use crate::format::common::{determine_file_type, FileType};
-use crate::format::{geojson, osmpbf};
 use crate::format::shapefile;
+use crate::format::{geojson, osmpbf};
 use crate::pg::binary_copy::{infer_geom_type, insert_rows};
 use crate::pg::crud::{check_table_exists, create_schema, create_table, drop_table, get_stmt};
 use crate::utils::validate::validate_args;
@@ -33,7 +33,7 @@ pub struct Cli {
     pub srid: Option<i32>,
 
     /// Mode: overwrite or fail. Optional.
-    #[arg(short, long)]
+    #[arg(short, long, default_value = "fail")]
     pub mode: Option<String>,
 
     /// Reproject: reproject to 4326 or 3857. Optional.
@@ -60,11 +60,7 @@ pub fn run() -> Result<()> {
             geojson::read_geojson(&args)?,
             geojson::determine_data_types(&args.input)?,
         ),
-        FileType::Osmpbf => {
-            // Unimplemented
-            osmpbf::read_osmpbf(&args)?;
-            unimplemented!()
-        }
+        FileType::Osmpbf => (osmpbf::read_osmpbf(&args)?, osmpbf::determine_data_types()?),
     };
 
     // If mode not present, check if table exists