Add Dataset::maybe_run_in_batch for faster DS write when possible

Author: Atreyagaurav

CHANGES.md

@@ -28,6 +28,7 @@
   - Add `Defn::field_index` and `Feature::field_index` ([#581](https://github.com/georust/gdal/pull/581))
   - Add `Defn::geometry_field_index` and `Feature::geometry_field_index` ([#594](https://github.com/georust/gdal/pull/594))
   - Add `Dataset::has_capability` for dataset capability check ([#581](https://github.com/georust/gdal/pull/585))
+  - Add `Dataset::maybe_run_in_batch` which runs a user function inside a transaction if the dataset supports it, as an optimization for I/O; if the function fails, the transaction may or may not be rolled back ([#584](https://github.com/georust/gdal/pull/584))
 
 ### Fixed
 

src/test_utils.rs

@@ -136,6 +136,41 @@ pub fn open_gpkg_for_update(path: &Path) -> (TempPath, Dataset) {
     (temp_path, ds)
 }
 
+/// Copies the given file to a temporary file and opens it for writing. When the returned
+/// `TempPath` is dropped, the file is deleted.
+pub fn open_dataset_for_update(path: &Path) -> (TempPath, Dataset) {
+    use std::fs;
+    use std::io::Write;
+
+    let input_data = fs::read(path).unwrap();
+    let (mut file, temp_path) = tempfile::Builder::new()
+        // using the whole filename as suffix should be fine (can't
+        // use .extension() for .shp.zip and such)
+        .suffix(
+            path.file_name()
+                .unwrap_or_default()
+                .to_string_lossy()
+                .as_ref(),
+        )
+        .tempfile()
+        .unwrap()
+        .into_parts();
+    file.write_all(&input_data).unwrap();
+    // Close the temporary file so that Dataset can open it safely even if the filesystem uses
+    // exclusive locking (Windows?).
+    drop(file);
+
+    let ds = Dataset::open_ex(
+        &temp_path,
+        DatasetOptions {
+            open_flags: GDALAccess::GA_Update.into(),
+            ..DatasetOptions::default()
+        },
+    )
+    .unwrap();
+    (temp_path, ds)
+}
+
 /// Assert numerical difference between two expressions is less than
 /// 64-bit machine epsilon or a specified epsilon.
 ///

src/vector/transaction.rs

@@ -177,11 +177,52 @@ impl Dataset {
         }
         Ok(Transaction::new(self))
     }
+
+    /// Optionally start a transaction before running `func` for performance
+    ///
+    /// This uses transaction if the dataset supports it, otherwise it
+    /// runs the `func` function as it is on the dataset. If the
+    /// closure results in error, and it is a transaction, it is
+    /// rolled back. If the rollback fails, it not reported.
+    pub fn maybe_run_in_batch(
+        &mut self,
+        mut func: impl FnMut(&mut Dataset) -> Result<()>,
+    ) -> Result<()> {
+        let force = 0; // since this is for speed
+        let rv = unsafe { gdal_sys::GDALDatasetStartTransaction(self.c_dataset(), force) };
+        if rv == OGRErr::OGRERR_UNSUPPORTED_OPERATION {
+            func(self)
+        } else if rv == OGRErr::OGRERR_NONE {
+            let res = func(self);
+            if res.is_ok() {
+                let rv = unsafe { gdal_sys::GDALDatasetCommitTransaction(self.c_dataset()) };
+                if rv != OGRErr::OGRERR_NONE {
+                    Err(GdalError::OgrError {
+                        err: rv,
+                        method_name: "GDALDatasetCommitTransaction",
+                    })
+                } else {
+                    res
+                }
+            } else {
+                let _ = unsafe { gdal_sys::GDALDatasetRollbackTransaction(self.c_dataset()) };
+                // ignore rollback error because it's not guaranteed to be
+                // supported, and the `func` failure is more important.
+                res
+            }
+        } else {
+            // transaction supported but failed
+            Err(GdalError::OgrError {
+                err: rv,
+                method_name: "GDALDatasetStartTransaction",
+            })
+        }
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::test_utils::{fixture, open_gpkg_for_update};
+    use crate::test_utils::{fixture, open_dataset_for_update, open_gpkg_for_update};
     use crate::vector::{Geometry, LayerAccess};
     use crate::Dataset;
 
@@ -241,4 +282,30 @@ mod tests {
         let mut ds = Dataset::open(fixture("roads.geojson")).unwrap();
         assert!(ds.start_transaction().is_err());
     }
+
+    #[test]
+    fn test_maybe_run_in_batch() {
+        let (_temp_path, mut ds) = open_gpkg_for_update(&fixture("poly.gpkg"));
+        let orig_feature_count = ds.layer(0).unwrap().feature_count();
+
+        let res = ds.maybe_run_in_batch(|d| {
+            let mut layer = d.layer(0).unwrap();
+            layer.create_feature(polygon())
+        });
+        assert!(res.is_ok());
+        assert_eq!(ds.layer(0).unwrap().feature_count(), orig_feature_count + 1);
+    }
+
+    #[test]
+    fn test_maybe_transaction_unsupported() {
+        let (_temp_path, mut ds) = open_dataset_for_update(&fixture("roads.geojson"));
+        let orig_feature_count = ds.layer(0).unwrap().feature_count();
+
+        let res = ds.maybe_run_in_batch(|d| {
+            let mut layer = d.layer(0).unwrap();
+            layer.create_feature(polygon())
+        });
+        assert!(res.is_ok());
+        assert_eq!(ds.layer(0).unwrap().feature_count(), orig_feature_count + 1);
+    }
 }