Merge pull request #153 from dead10ck/sql

Implement SQL queries

Author: jdroenner

src/dataset.rs

@@ -1,18 +1,25 @@
-use std::{ffi::CString, ffi::NulError, path::Path, ptr};
+use ptr::null_mut;
+use std::convert::TryInto;
+use std::{
+    ffi::NulError,
+    ffi::{CStr, CString},
+    path::Path,
+    ptr,
+};
 
+use crate::errors::*;
 use crate::utils::{_last_cpl_err, _last_null_pointer_err, _string};
+use crate::vector::sql;
+use crate::vector::Geometry;
 use crate::{
     gdal_major_object::MajorObject, raster::RasterBand, spatial_ref::SpatialRef, vector::Layer,
     Driver, Metadata,
 };
+use gdal_sys::OGRGeometryH;
 use gdal_sys::{
     self, CPLErr, GDALAccess, GDALDatasetH, GDALMajorObjectH, OGRErr, OGRLayerH, OGRwkbGeometryType,
 };
 use libc::{c_double, c_int, c_uint};
-use ptr::null_mut;
-
-use crate::errors::*;
-use std::convert::TryInto;
 
 use bitflags::bitflags;
 
@@ -535,6 +542,107 @@ impl Dataset {
         }
         Ok(Transaction::new(self))
     }
+
+    /// Execute a SQL query against the Dataset. It is equivalent to calling
+    /// [`GDALDatasetExecuteSQL`](https://gdal.org/api/raster_c_api.html#_CPPv421GDALDatasetExecuteSQL12GDALDatasetHPKc12OGRGeometryHPKc).
+    /// Returns a [`sql::ResultSet`], which can be treated just as any other [`Layer`].
+    ///
+    /// Queries such as `ALTER TABLE`, `CREATE INDEX`, etc. have no [`sql::ResultSet`], and return
+    /// `None`, which is distinct from an empty [`sql::ResultSet`].
+    ///
+    /// # Arguments
+    ///
+    /// * `query`: The SQL query
+    /// * `spatial_filter`: Limit results of the query to features that intersect the given
+    ///   [`Geometry`]
+    /// * `dialect`: The dialect of SQL to use. See
+    ///   <https://gdal.org/user/ogr_sql_sqlite_dialect.html>
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use gdal::Dataset;
+    /// # use std::path::Path;
+    /// use gdal::vector::sql;
+    ///
+    /// let ds = Dataset::open(Path::new("fixtures/roads.geojson")).unwrap();
+    /// let query = "SELECT kind, is_bridge, highway FROM roads WHERE highway = 'pedestrian'";
+    /// let result_set = ds.execute_sql(query, None, sql::Dialect::DEFAULT).unwrap().unwrap();
+    ///
+    /// assert_eq!(10, result_set.feature_count());
+    ///
+    /// for feature in result_set.features() {
+    ///     let highway = feature
+    ///         .field("highway")
+    ///         .unwrap()
+    ///         .unwrap()
+    ///         .into_string()
+    ///         .unwrap();
+    ///
+    ///     assert_eq!("pedestrian", highway);
+    /// }
+    /// ```
+    pub fn execute_sql<S: AsRef<str>>(
+        &self,
+        query: S,
+        spatial_filter: Option<&Geometry>,
+        dialect: sql::Dialect,
+    ) -> Result<Option<sql::ResultSet>> {
+        let query = CString::new(query.as_ref())?;
+
+        let dialect_c_str = match dialect {
+            sql::Dialect::DEFAULT => None,
+            sql::Dialect::OGR => Some(unsafe { CStr::from_bytes_with_nul_unchecked(sql::OGRSQL) }),
+            sql::Dialect::SQLITE => {
+                Some(unsafe { CStr::from_bytes_with_nul_unchecked(sql::SQLITE) })
+            }
+        };
+
+        self._execute_sql(query, spatial_filter, dialect_c_str)
+    }
+
+    fn _execute_sql(
+        &self,
+        query: CString,
+        spatial_filter: Option<&Geometry>,
+        dialect_c_str: Option<&CStr>,
+    ) -> Result<Option<sql::ResultSet>> {
+        let mut filter_geom: OGRGeometryH = std::ptr::null_mut();
+
+        let dialect_ptr = match dialect_c_str {
+            None => std::ptr::null(),
+            Some(ref d) => d.as_ptr(),
+        };
+
+        if let Some(spatial_filter) = spatial_filter {
+            filter_geom = unsafe { spatial_filter.c_geometry() };
+        }
+
+        let c_dataset = unsafe { self.c_dataset() };
+
+        unsafe { gdal_sys::CPLErrorReset() };
+
+        let c_layer = unsafe {
+            gdal_sys::GDALDatasetExecuteSQL(c_dataset, query.as_ptr(), filter_geom, dialect_ptr)
+        };
+
+        let cpl_err = unsafe { gdal_sys::CPLGetLastErrorType() };
+
+        if cpl_err != CPLErr::CE_None {
+            return Err(_last_cpl_err(cpl_err));
+        }
+
+        if c_layer.is_null() {
+            return Ok(None);
+        }
+
+        let layer = unsafe { Layer::from_c_layer(self, c_layer) };
+
+        Ok(Some(sql::ResultSet {
+            layer,
+            dataset: c_dataset,
+        }))
+    }
 }
 
 pub struct LayerIterator<'a> {

src/errors.rs

@@ -9,6 +9,8 @@ pub type Result<T> = std::result::Result<T, GdalError>;
 pub enum GdalError {
     #[error("FfiNulError")]
     FfiNulError(#[from] std::ffi::NulError),
+    #[error("FfiIntoStringError")]
+    FfiIntoStringError(#[from] std::ffi::IntoStringError),
     #[error("StrUtf8Error")]
     StrUtf8Error(#[from] std::str::Utf8Error),
     #[cfg(feature = "ndarray")]

src/vector/defn.rs

@@ -11,6 +11,7 @@ use crate::errors::*;
 /// Layer definition
 ///
 /// Defines the fields available for features in a layer.
+#[derive(Debug)]
 pub struct Defn {
     c_defn: OGRFeatureDefnH,
 }

src/vector/feature.rs

@@ -15,6 +15,7 @@ use crate::errors::*;
 use std::slice;
 
 /// OGR Feature
+#[derive(Debug)]
 pub struct Feature<'a> {
     _defn: &'a Defn,
     c_feature: OGRFeatureH,
@@ -72,12 +73,16 @@ impl<'a> Feature<'a> {
     /// If the field has an unsupported type, returns a [`GdalError::UnhandledFieldType`].
     ///
     /// If the field is null, returns `None`.
-    pub fn field(&self, name: &str) -> Result<Option<FieldValue>> {
-        let c_name = CString::new(name)?;
+    pub fn field<S: AsRef<str>>(&self, name: S) -> Result<Option<FieldValue>> {
+        let c_name = CString::new(name.as_ref())?;
+        self._field(c_name)
+    }
+
+    fn _field(&self, c_name: CString) -> Result<Option<FieldValue>> {
         let field_id = unsafe { gdal_sys::OGR_F_GetFieldIndex(self.c_feature, c_name.as_ptr()) };
         if field_id == -1 {
             Err(GdalError::InvalidFieldName {
-                field_name: name.to_string(),
+                field_name: c_name.into_string()?,
                 method_name: "OGR_F_GetFieldIndex",
             })
         } else {

src/vector/layer.rs

@@ -25,6 +25,7 @@ use crate::errors::*;
 ///     // do something with each feature
 /// }
 /// ```
+#[derive(Debug)]
 pub struct Layer<'a> {
     c_layer: OGRLayerH,
     defn: Defn,

src/vector/mod.rs

@@ -22,6 +22,7 @@ mod geo_to_gdal;
 mod geometry;
 mod layer;
 mod ops;
+pub mod sql;
 
 pub use defn::{Defn, Field, FieldIterator};
 pub use feature::{Feature, FieldValue, FieldValueIterator};

src/vector/sql.rs

@@ -0,0 +1,46 @@
+use std::ops::Deref;
+
+use gdal_sys::GDALDatasetH;
+
+use crate::vector::Layer;
+
+/// The result of a SQL query executed by
+/// [`Dataset::execute_sql()`](crate::Dataset::execute_sql()). It is just a thin wrapper around a
+/// [`Layer`], and you can treat it as such.
+#[derive(Debug)]
+pub struct ResultSet<'a> {
+    pub(crate) layer: Layer<'a>,
+    pub(crate) dataset: GDALDatasetH,
+}
+
+impl<'a> Deref for ResultSet<'a> {
+    type Target = Layer<'a>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.layer
+    }
+}
+
+impl<'a> Drop for ResultSet<'a> {
+    fn drop(&mut self) {
+        unsafe { gdal_sys::GDALDatasetReleaseResultSet(self.dataset, self.layer.c_layer()) };
+    }
+}
+
+/// Represents valid SQL dialects to use in SQL queries. See
+/// <https://gdal.org/user/ogr_sql_sqlite_dialect.html>
+pub enum Dialect {
+    /// Use the default dialect. This is OGR SQL unless the underlying driver has a native dialect,
+    /// such as MySQL, Postgres, Oracle, etc.
+    DEFAULT,
+
+    /// Explicitly choose OGR SQL regardless of if the underlying driver has a native dialect.
+    OGR,
+
+    /// SQLite dialect. If the data set is not actually a SQLite database, then a virtual SQLite
+    /// table is created to execute the query.
+    SQLITE,
+}
+
+pub(crate) const OGRSQL: &[u8] = b"OGRSQL\0";
+pub(crate) const SQLITE: &[u8] = b"SQLITE\0";

src/vector/vector_tests/mod.rs

@@ -3,13 +3,14 @@ use super::{
 };
 use crate::spatial_ref::SpatialRef;
 use crate::{assert_almost_eq, Dataset, Driver};
-use std::path::Path;
 
 mod convert_geo;
+mod sql;
 
+#[macro_export]
 macro_rules! fixture {
     ($name:expr) => {
-        Path::new(file!())
+        std::path::Path::new(file!())
             .parent()
             .unwrap()
             .parent()

src/vector/vector_tests/sql.rs

@@ -0,0 +1,133 @@
+use std::collections::HashSet;
+
+use crate::{
+    fixture,
+    vector::{sql, Geometry},
+    Dataset,
+};
+
+#[test]
+fn test_sql() {
+    let ds = Dataset::open(fixture!("roads.geojson")).unwrap();
+    let query = "SELECT kind, is_bridge, highway FROM roads WHERE highway = 'pedestrian'";
+    let result_set = ds
+        .execute_sql(query, None, sql::Dialect::DEFAULT)
+        .unwrap()
+        .unwrap();
+
+    let field_names: HashSet<_> = result_set
+        .defn()
+        .fields()
+        .map(|field| field.name())
+        .collect();
+
+    let mut correct_field_names = HashSet::new();
+    correct_field_names.insert("kind".into());
+    correct_field_names.insert("is_bridge".into());
+    correct_field_names.insert("highway".into());
+
+    assert_eq!(correct_field_names, field_names);
+    assert_eq!(10, result_set.feature_count());
+
+    for feature in result_set.features() {
+        let highway = feature
+            .field("highway")
+            .unwrap()
+            .unwrap()
+            .into_string()
+            .unwrap();
+
+        assert_eq!("pedestrian", highway);
+    }
+}
+
+#[test]
+fn test_sql_with_spatial_filter() {
+    let query = "SELECT * FROM roads WHERE highway = 'pedestrian'";
+    let ds = Dataset::open(fixture!("roads.geojson")).unwrap();
+    let bbox = Geometry::bbox(26.1017, 44.4297, 26.1025, 44.4303).unwrap();
+    let result_set = ds
+        .execute_sql(query, Some(&bbox), sql::Dialect::DEFAULT)
+        .unwrap()
+        .unwrap();
+
+    assert_eq!(2, result_set.feature_count());
+    let mut correct_fids = HashSet::new();
+    correct_fids.insert(252725993);
+    correct_fids.insert(23489656);
+
+    let mut fids = HashSet::new();
+    for feature in result_set.features() {
+        let highway = feature
+            .field("highway")
+            .unwrap()
+            .unwrap()
+            .into_string()
+            .unwrap();
+
+        assert_eq!("pedestrian", highway);
+        fids.insert(feature.fid().unwrap());
+    }
+
+    assert_eq!(correct_fids, fids);
+}
+
+#[test]
+fn test_sql_with_dialect() {
+    let query = "SELECT * FROM roads WHERE highway = 'pedestrian' and NumPoints(GEOMETRY) = 3";
+    let ds = Dataset::open(fixture!("roads.geojson")).unwrap();
+    let bbox = Geometry::bbox(26.1017, 44.4297, 26.1025, 44.4303).unwrap();
+    let result_set = ds
+        .execute_sql(query, Some(&bbox), sql::Dialect::SQLITE)
+        .unwrap()
+        .unwrap();
+
+    assert_eq!(1, result_set.feature_count());
+    let mut features: Vec<_> = result_set.features().collect();
+    let feature = features.pop().unwrap();
+    let highway = feature
+        .field("highway")
+        .unwrap()
+        .unwrap()
+        .into_string()
+        .unwrap();
+
+    assert_eq!("pedestrian", highway);
+}
+
+#[test]
+fn test_sql_empty_result() {
+    let ds = Dataset::open(fixture!("roads.geojson")).unwrap();
+    let query = "SELECT kind, is_bridge, highway FROM roads WHERE highway = 'jazz hands 👐'";
+    let result_set = ds
+        .execute_sql(query, None, sql::Dialect::DEFAULT)
+        .unwrap()
+        .unwrap();
+    assert_eq!(0, result_set.feature_count());
+    assert_eq!(0, result_set.features().count());
+}
+
+#[test]
+fn test_sql_no_result() {
+    let ds = Dataset::open(fixture!("roads.geojson")).unwrap();
+    let query = "ALTER TABLE roads ADD COLUMN fun integer";
+    let result_set = ds.execute_sql(query, None, sql::Dialect::DEFAULT).unwrap();
+    assert!(result_set.is_none());
+}
+
+#[test]
+fn test_sql_bad_query() {
+    let ds = Dataset::open(fixture!("roads.geojson")).unwrap();
+
+    let query = "SELECT nope FROM roads";
+    let result_set = ds.execute_sql(query, None, sql::Dialect::DEFAULT);
+    assert!(result_set.is_err());
+
+    let query = "SELECT nope FROM";
+    let result_set = ds.execute_sql(query, None, sql::Dialect::DEFAULT);
+    assert!(result_set.is_err());
+
+    let query = "SELECT ninetynineredballoons(highway) FROM roads";
+    let result_set = ds.execute_sql(query, None, sql::Dialect::DEFAULT);
+    assert!(result_set.is_err());
+}