From f38097f6d54d6a490ca00878a493737770b1238d Mon Sep 17 00:00:00 2001 From: Frank Jimenez Date: Sun, 7 Jul 2024 22:21:56 +0100 Subject: [PATCH] organise modules --- justfile | 8 ++ src/file_types/common.rs | 72 +++++++++++ src/file_types/geo.rs | 48 +++++++ src/file_types/mod.rs | 3 + src/{utils/shp.rs => file_types/shapefile.rs} | 26 ++-- src/main.rs | 1 + src/pg/binary_copy.rs | 21 ++-- src/pg/crud.rs | 8 +- src/utils/cli.rs | 10 +- src/utils/mod.rs | 118 ------------------ 10 files changed, 166 insertions(+), 149 deletions(-) create mode 100644 src/file_types/common.rs create mode 100644 src/file_types/geo.rs create mode 100644 src/file_types/mod.rs rename src/{utils/shp.rs => file_types/shapefile.rs} (80%) diff --git a/justfile b/justfile index f7a12f9..04f7b01 100644 --- a/justfile +++ b/justfile @@ -1,4 +1,12 @@ @try: cargo run -- --input ./examples/shapefile/small-example.shp \ --uri postgresql://pio:password@localhost:25432/popgis \ + --schema test \ + --table waters + +@build-and-run: + cargo build --release + cd ./target/release/ && ./popgis --input ~/Downloads/water-polygons-split-4326/water_polygons.shp \ + --uri postgresql://pio:password@localhost:25432/popgis \ + --schema osm \ --table waters diff --git a/src/file_types/common.rs b/src/file_types/common.rs new file mode 100644 index 0000000..c7996df --- /dev/null +++ b/src/file_types/common.rs @@ -0,0 +1,72 @@ +use crate::Result; + +use postgres::types::Type; +use std::path::Path; + +use crate::pg::binary_copy::Wkb; + +// Struct to hold column name and data type +pub struct NewTableTypes { + pub column_name: String, + pub data_type: Type, +} + +#[derive(Debug)] +pub struct Row { + pub columns: Vec, +} + +#[derive(Debug)] +pub struct Rows { + pub row: Vec, +} + +impl Row { + pub fn new() -> Self { + Row { columns: Vec::new() } + } + pub fn add(&mut self, column: AcceptedTypes) { + self.columns.push(column); + } +} + +impl Rows { + pub fn new() -> Self { + Rows { row: Vec::new() } + } + pub fn add(&mut self, row: Row) { + self.row.push(row); + } +} + +// Enum to hold accepted data types +#[derive(Debug)] +pub enum AcceptedTypes { + Int(Option), + Float(Option), + Double(Option), + Text(Option), + Bool(Option), + Geometry(Option), +} + +// Create enum of supported file types +pub enum FileType { + Shapefile, + GeoJson, +} + +pub fn determine_file_type(input_file: &str) -> Result { + let file_extension = Path::new(input_file) + .extension() + .expect("No file extension found"); + let file_extension_str = file_extension + .to_str() + .expect("Could not convert file extension to string"); + match file_extension_str { + "shp" => Ok(FileType::Shapefile), + "json" => Ok(FileType::GeoJson), + _ => Err("Unsupported file type".into()), + } +} + diff --git a/src/file_types/geo.rs b/src/file_types/geo.rs new file mode 100644 index 0000000..a9d102c --- /dev/null +++ b/src/file_types/geo.rs @@ -0,0 +1,48 @@ +use crate::Result; +use geo::Coord; +use shapefile::Shape; + +pub fn to_geo(shape: &Shape) -> Result> { + match shape { + Shape::Point(p) => Ok(geo::Point::new(p.x, p.y).into()), + Shape::Polyline(p) => { + let mut coords: Vec = Vec::new(); + for part in p.parts().iter() { + for point in part.iter() { + coords.push(Coord::from((point.x, point.y))); + } + } + Ok(geo::LineString::new(coords).into()) + } + Shape::Polygon(p) => { + let mut outer_placeholder: Vec<(f64, f64)> = Vec::new(); + let mut inner_rings: Vec = Vec::new(); + + for ring_type in p.rings() { + match ring_type { + //Gather all outer rings + shapefile::PolygonRing::Outer(out) => { + out.iter().for_each(|p| outer_placeholder.push((p.x, p.y))) + } + //Gather all inner rings + shapefile::PolygonRing::Inner(inn) => { + let mut inner_ring: Vec<(f64, f64)> = Vec::new(); + inn.iter().for_each(|p| inner_ring.push((p.x, p.y))); + let ls = geo::LineString::from(inner_ring); + inner_rings.push(ls); + } + } + } + + let outer_ring = geo::LineString::from(outer_placeholder); + if inner_rings.is_empty() { + let poly = geo::Polygon::new(outer_ring, vec![]); + Ok(geo::Geometry::from(poly)) + } else { + let poly = geo::Polygon::new(outer_ring, inner_rings); + Ok(geo::Geometry::from(poly)) + } + } + _ => Err("Unsupported shape type".into()), + } +} diff --git a/src/file_types/mod.rs b/src/file_types/mod.rs new file mode 100644 index 0000000..e9389da --- /dev/null +++ b/src/file_types/mod.rs @@ -0,0 +1,3 @@ +pub mod common; +mod geo; +pub mod shapefile; diff --git a/src/utils/shp.rs b/src/file_types/shapefile.rs similarity index 80% rename from src/utils/shp.rs rename to src/file_types/shapefile.rs index 3db03a7..367d797 100644 --- a/src/utils/shp.rs +++ b/src/file_types/shapefile.rs @@ -3,8 +3,8 @@ use postgres::types::Type; use shapefile::dbase::FieldValue; use crate::pg::binary_copy::Wkb; -use crate::utils::to_geo; -use crate::utils::{AcceptedTypes, NewTableTypes, Row, Rows}; +use crate::file_types::geo::to_geo; +use crate::file_types::common::{AcceptedTypes, NewTableTypes, Row, Rows}; use wkb::geom_to_wkb; pub fn determine_data_types(file_path: &str) -> Result> { @@ -67,30 +67,22 @@ pub fn read_shapefile(file_path: &str) -> Result { for (_, data_type) in record.into_iter() { match data_type { FieldValue::Numeric(value) => { - if let Some(value) = value { - row.add(AcceptedTypes::Float(value)); - } + row.add(AcceptedTypes::Float(value)); } FieldValue::Float(value) => { - if let Some(value) = value { - row.add(AcceptedTypes::Double(value as f64)); - } + row.add(AcceptedTypes::Double(value)); } FieldValue::Double(value) => { - row.add(AcceptedTypes::Double(value)); + row.add(AcceptedTypes::Float(Some(value))); } FieldValue::Integer(value) => { - row.add(AcceptedTypes::Int(value as i64)); + row.add(AcceptedTypes::Int(Some(value))); } FieldValue::Character(value) => { - if let Some(value) = value { - row.add(AcceptedTypes::Text(value)); - } + row.add(AcceptedTypes::Text(value)); } FieldValue::Logical(value) => { - if let Some(value) = value { - row.add(AcceptedTypes::Bool(value)); - } + row.add(AcceptedTypes::Bool(value)); } _ => println!("Type currently not supported"), } @@ -98,7 +90,7 @@ pub fn read_shapefile(file_path: &str) -> Result { let geom = to_geo(&shape)?; let wkb = geom_to_wkb(&geom).expect("Failed to insert node into database"); - row.add(AcceptedTypes::Geometry(Wkb { geometry: wkb })); + row.add(AcceptedTypes::Geometry(Some(Wkb { geometry: wkb }))); rows.add(row); } diff --git a/src/main.rs b/src/main.rs index aa8d473..15e3b5f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ pub type Error = Box; mod utils; mod pg; +mod file_types; use utils::cli::run; diff --git a/src/pg/binary_copy.rs b/src/pg/binary_copy.rs index 02412a9..979e894 100644 --- a/src/pg/binary_copy.rs +++ b/src/pg/binary_copy.rs @@ -9,7 +9,7 @@ use postgres::binary_copy::BinaryCopyInWriter; use postgres::CopyInWriter; use crate::pg::crud::create_connection; -use crate::utils::{AcceptedTypes, NewTableTypes, Rows}; +use crate::file_types::common::{AcceptedTypes, NewTableTypes, Rows}; #[derive(Debug)] pub struct Wkb { @@ -72,36 +72,37 @@ pub fn insert_rows<'a>( let mut writer = BinaryCopyInWriter::new(writer, &types); - println!("{:?}", types); + // Use to test if types are correct + // println!("{:?}", types); for row in rows.row.iter() { // Transform row into vector of ToSql - let mut vec: Vec<&(dyn ToSql + Sync)> = Vec::new(); + let mut tosql: Vec<&(dyn ToSql + Sync)> = Vec::new(); for column in row.columns.iter() { match column { AcceptedTypes::Int(value) => { - vec.push(value); + tosql.push(value); } AcceptedTypes::Float(value) => { - vec.push(value); + tosql.push(value); } AcceptedTypes::Double(value) => { - vec.push(value); + tosql.push(value); } AcceptedTypes::Text(value) => { - vec.push(value); + tosql.push(value); } AcceptedTypes::Bool(value) => { - vec.push(value); + tosql.push(value); } AcceptedTypes::Geometry(value) => { - vec.push(value); + tosql.push(value); } } } // Convert the vector to a slice of references - let vec_slice: &[&(dyn ToSql + Sync)] = &vec; + let vec_slice: &[&(dyn ToSql + Sync)] = &tosql; // Write row to database writer diff --git a/src/pg/crud.rs b/src/pg/crud.rs index 054144f..6f698b3 100644 --- a/src/pg/crud.rs +++ b/src/pg/crud.rs @@ -4,13 +4,19 @@ use postgres::Statement; use postgres::{Client, NoTls}; -use crate::utils::NewTableTypes; +use crate::file_types::common::NewTableTypes; pub fn create_connection(uri: &str) -> Result { let client = Client::connect(uri, NoTls)?; Ok(client) } +pub fn create_schema(schema_name: &str, uri: &str) -> Result<()> { + let mut client = create_connection(uri)?; + client.batch_execute(&format!("CREATE SCHEMA IF NOT EXISTS {}", schema_name))?; + Ok(()) +} + pub fn create_table( table_name: &str, schema_name: &Option, diff --git a/src/utils/cli.rs b/src/utils/cli.rs index c96a1ee..926d3bf 100644 --- a/src/utils/cli.rs +++ b/src/utils/cli.rs @@ -1,8 +1,8 @@ use crate::Result; -use crate::utils::{determine_file_type, FileType}; +use crate::file_types::common::{FileType, determine_file_type}; use crate::utils::validate::validate_args; -use crate::utils::shp::{read_shapefile, determine_data_types}; -use crate::pg::crud::create_table; +use crate::file_types::shapefile::{read_shapefile, determine_data_types}; +use crate::pg::crud::{create_table, create_schema}; use crate::pg::binary_copy::{infer_geom_type, insert_rows}; use clap::Parser; @@ -37,6 +37,10 @@ pub fn run() -> Result<()> { } }; let config = determine_data_types(&args.input)?; + // If schema present, create schema + if let Some(schema) = &args.schema { + create_schema(&schema, &args.uri)?; + } let stmt = create_table(&args.table, &args.schema, &config, &args.uri)?; let geom_type = infer_geom_type(stmt)?; insert_rows(&rows, &config, geom_type, &args.uri, &args.schema, &args.table)?; diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 0f89e4c..6cd948c 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,120 +1,2 @@ -use crate::Result; -use std::path::Path; -use postgres::types::Type; -use geo::Coord; -use shapefile::Shape; -use crate::pg::binary_copy::Wkb; - pub mod cli; -pub mod shp; mod validate; - -// Struct to hold column name and data type -pub struct NewTableTypes { - pub column_name: String, - pub data_type: Type, -} - -#[derive(Debug)] -pub struct Row { - pub columns: Vec, -} - -#[derive(Debug)] -pub struct Rows { - pub row: Vec, -} - -impl Row { - pub fn new() -> Self { - Row { columns: Vec::new() } - } - pub fn add(&mut self, column: AcceptedTypes) { - self.columns.push(column); - } -} - -impl Rows { - pub fn new() -> Self { - Rows { row: Vec::new() } - } - pub fn add(&mut self, row: Row) { - self.row.push(row); - } -} - -// Enum to hold accepted data types -#[derive(Debug)] -pub enum AcceptedTypes { - Int(i64), - Float(f64), - Double(f64), - Text(String), - Bool(bool), - Geometry(Wkb), -} - -// Create enum of supported file types -pub enum FileType { - Shapefile, - GeoJson, -} - -fn determine_file_type(input_file: &str) -> Result { - let file_extension = Path::new(input_file) - .extension() - .expect("No file extension found"); - let file_extension_str = file_extension - .to_str() - .expect("Could not convert file extension to string"); - match file_extension_str { - "shp" => Ok(FileType::Shapefile), - "json" => Ok(FileType::GeoJson), - _ => Err("Unsupported file type".into()), - } -} - -pub fn to_geo(shape: &Shape) -> Result> { - match shape { - Shape::Point(p) => Ok(geo::Point::new(p.x, p.y).into()), - Shape::Polyline(p) => { - let mut coords: Vec = Vec::new(); - for part in p.parts().iter() { - for point in part.iter() { - coords.push(Coord::from((point.x, point.y))); - } - } - Ok(geo::LineString::new(coords).into()) - } - Shape::Polygon(p) => { - let mut outer_placeholder: Vec<(f64, f64)> = Vec::new(); - let mut inner_rings: Vec = Vec::new(); - - for ring_type in p.rings() { - match ring_type { - //Gather all outer rings - shapefile::PolygonRing::Outer(out) => { - out.iter().for_each(|p| outer_placeholder.push((p.x, p.y))) - } - //Gather all inner rings - shapefile::PolygonRing::Inner(inn) => { - let mut inner_ring: Vec<(f64, f64)> = Vec::new(); - inn.iter().for_each(|p| inner_ring.push((p.x, p.y))); - let ls = geo::LineString::from(inner_ring); - inner_rings.push(ls); - } - } - } - - let outer_ring = geo::LineString::from(outer_placeholder); - if inner_rings.is_empty() { - let poly = geo::Polygon::new(outer_ring, vec![]); - Ok(geo::Geometry::from(poly)) - } else { - let poly = geo::Polygon::new(outer_ring, inner_rings); - Ok(geo::Geometry::from(poly)) - } - } - _ => Err("Unsupported shape type".into()), - } -}