use geo_traits::{Dimensions, GeometryTrait};
use geopackage_core::geometry::encode_gpb;
use geopackage_core::ident::quote;
use geopackage_core::triggers;
use geopackage_core::types::ZmFlag;
use rusqlite::types::Value as SqlValue;
use rusqlite::{Connection, OptionalExtension, Transaction, params_from_iter};
use crate::bulk::{self, BulkIndexOptions};
use crate::index::drop_all_rtree_triggers;
use crate::value::value_to_sql;
use crate::{Error, Layer, Result, Value};
#[derive(Debug, Clone)]
pub struct NewFeature<G> {
pub fid: Option<i64>,
pub geometry: Option<G>,
pub values: Vec<Value>,
}
impl<G> NewFeature<G> {
pub fn new(geometry: G, values: Vec<Value>) -> Self {
Self {
fid: None,
geometry: Some(geometry),
values,
}
}
pub fn attributes(values: Vec<Value>) -> Self {
Self {
fid: None,
geometry: None,
values,
}
}
#[must_use]
pub fn with_fid(mut self, fid: i64) -> Self {
self.fid = Some(fid);
self
}
}
#[derive(Debug)]
struct GeomTarget {
name: String,
quoted_name: String,
srs_id: i32,
z: ZmFlag,
m: ZmFlag,
}
#[derive(Debug, Clone, Copy)]
struct BboxFold {
min_x: f64,
max_x: f64,
min_y: f64,
max_y: f64,
seen: bool,
}
impl BboxFold {
fn new() -> Self {
Self {
min_x: f64::INFINITY,
max_x: f64::NEG_INFINITY,
min_y: f64::INFINITY,
max_y: f64::NEG_INFINITY,
seen: false,
}
}
fn seed(&mut self, existing: Option<[f64; 4]>) {
if let Some([min_x, max_x, min_y, max_y]) = existing {
self.min_x = min_x;
self.max_x = max_x;
self.min_y = min_y;
self.max_y = max_y;
self.seen = true;
}
}
fn add(&mut self, [min_x, max_x, min_y, max_y]: [f64; 4]) {
self.min_x = self.min_x.min(min_x);
self.max_x = self.max_x.max(max_x);
self.min_y = self.min_y.min(min_y);
self.max_y = self.max_y.max(max_y);
self.seen = true;
}
fn bounds(&self) -> Option<[f64; 4]> {
self.seen
.then_some([self.min_x, self.max_x, self.min_y, self.max_y])
}
}
pub struct FeatureWriter<'conn> {
tx: Transaction<'conn>,
table_name: String,
quoted_table: String,
pk_expr: String,
value_columns: Vec<String>,
geometry: Option<GeomTarget>,
bbox: BboxFold,
dirty: bool,
bbox_dirty: bool,
}
impl<'a> Layer<'a> {
pub fn writer(&self) -> Result<FeatureWriter<'a>> {
let conn: &Connection = self.gpkg().connection();
let tx = conn.unchecked_transaction()?;
let existing = read_contents_bbox(&tx, self.table_name())?;
let pk_name = self.primary_key_column();
let pk_expr = match pk_name {
Some(pk) => quote(pk)?,
None => "rowid".to_owned(),
};
let value_columns = self
.value_columns()
.iter()
.filter(|c| Some(c.name.as_str()) != pk_name)
.map(|c| quote(&c.name))
.collect::<std::result::Result<Vec<_>, _>>()?;
let geometry = match self.geometry_column() {
Some(g) => Some(GeomTarget {
name: g.column_name.clone(),
quoted_name: quote(&g.column_name)?,
srs_id: g.srs_id,
z: g.z,
m: g.m,
}),
None => None,
};
let mut bbox = BboxFold::new();
bbox.seed(existing);
Ok(FeatureWriter {
tx,
table_name: self.table_name().to_owned(),
quoted_table: quote(self.table_name())?,
pk_expr,
value_columns,
geometry,
bbox,
dirty: false,
bbox_dirty: false,
})
}
pub fn write_all<G, I>(&self, features: I, batch_size: usize) -> Result<Vec<i64>>
where
G: GeometryTrait<T = f64>,
I: IntoIterator<Item = NewFeature<G>>,
{
self.write_all_with(features, batch_size, BulkIndexOptions::default())
}
pub fn write_all_with<G, I>(
&self,
features: I,
batch_size: usize,
options: BulkIndexOptions,
) -> Result<Vec<i64>>
where
G: GeometryTrait<T = f64>,
I: IntoIterator<Item = NewFeature<G>>,
{
let iter = features.into_iter();
if self.bulk_write_applicable(iter.size_hint().0, options)? {
self.write_all_bulk(iter)
} else {
self.write_all_batched(iter, batch_size)
}
}
fn write_all_batched<G, I>(&self, features: I, batch_size: usize) -> Result<Vec<i64>>
where
G: GeometryTrait<T = f64>,
I: IntoIterator<Item = NewFeature<G>>,
{
let mut fids = Vec::new();
let mut iter = features.into_iter();
let mut batch = self.writer()?;
let mut in_batch = 0usize;
let mut wrote_any = false;
for feature in iter.by_ref() {
let fid = match &feature.geometry {
Some(geometry) => batch.insert(feature.fid, geometry, &feature.values)?,
None => batch.insert_row(feature.fid, &feature.values)?,
};
fids.push(fid);
wrote_any = true;
in_batch += 1;
if batch_size != 0 && in_batch >= batch_size {
batch.commit()?;
batch = self.writer()?;
in_batch = 0;
wrote_any = false;
}
}
if wrote_any || in_batch > 0 {
batch.commit()?;
} else {
drop(batch);
}
Ok(fids)
}
fn bulk_write_applicable(
&self,
size_hint_lower: usize,
options: BulkIndexOptions,
) -> Result<bool> {
if size_hint_lower < options.bulk_threshold {
return Ok(false);
}
let Some(geom) = self.geometry_column() else {
return Ok(false);
};
if self.primary_key_column().is_none() || !self.has_spatial_index()? {
return Ok(false);
}
let rtree = triggers::rtree_table_name(self.table_name(), &geom.column_name);
let count: i64 = self.gpkg().connection().query_row(
&format!("SELECT count(*) FROM {}", quote(&rtree)?),
[],
|r| r.get(0),
)?;
Ok(count == 0)
}
fn write_all_bulk<G, I>(&self, features: I) -> Result<Vec<i64>>
where
G: GeometryTrait<T = f64>,
I: IntoIterator<Item = NewFeature<G>>,
{
let geom = self
.geometry_column()
.ok_or_else(|| Error::NoGeometryColumn {
table_name: self.table_name().to_owned(),
})?;
let pk = self
.primary_key_column()
.ok_or_else(|| Error::NoPrimaryKey {
table_name: self.table_name().to_owned(),
})?;
let table = self.table_name();
let column = &geom.column_name;
let rtree = triggers::rtree_table_name(table, column);
let conn = self.gpkg().connection();
drop_all_rtree_triggers(conn, table, column)?;
let result = (|| -> Result<Vec<i64>> {
let mut fids = Vec::new();
let mut writer = self.writer()?;
for feature in features {
let fid = match &feature.geometry {
Some(geometry) => writer.insert(feature.fid, geometry, &feature.values)?,
None => writer.insert_row(feature.fid, &feature.values)?,
};
fids.push(fid);
}
writer.commit()?;
bulk::fill_index(conn, table, column, pk, &rtree, bulk::no_tamper, |conn| {
for sql in triggers::create_triggers_sql(table, column, pk)? {
conn.execute_batch(&sql)?;
}
Ok(())
})?;
Ok(fids)
})();
if result.is_err() {
restore_index_after_failed_bulk(conn, table, column, pk, &rtree);
}
result
}
}
fn restore_index_after_failed_bulk(
conn: &Connection,
table: &str,
column: &str,
pk: &str,
rtree: &str,
) {
let restore = || -> Result<()> {
let tx = conn.unchecked_transaction()?;
drop_all_rtree_triggers(&tx, table, column)?;
tx.execute_batch(&format!("DROP TABLE IF EXISTS {}", quote(rtree)?))?;
tx.execute_batch(&triggers::create_rtree_table_sql(table, column)?)?;
for sql in triggers::create_triggers_sql(table, column, pk)? {
tx.execute_batch(&sql)?;
}
tx.execute_batch(&triggers::populate_rtree_sql(table, column, pk)?)?;
tx.commit()?;
Ok(())
};
if restore().is_err() {
}
}
impl FeatureWriter<'_> {
pub fn insert<G: GeometryTrait<T = f64>>(
&mut self,
fid: Option<i64>,
geometry: &G,
values: &[Value],
) -> Result<i64> {
self.check_values(values)?;
let (blob, xy) = self.encode_geometry(geometry)?;
let sql = self.insert_sql(fid.is_some(), true);
let mut binds: Vec<SqlValue> = Vec::with_capacity(values.len() + 2);
if let Some(id) = fid {
binds.push(SqlValue::Integer(id));
}
binds.extend(values.iter().map(value_to_sql));
binds.push(SqlValue::Blob(blob));
let assigned = self.exec_insert(&sql, &binds, fid)?;
if let Some(envelope) = xy {
self.bbox.add(envelope);
self.bbox_dirty = true;
}
self.dirty = true;
Ok(assigned)
}
pub fn insert_row(&mut self, fid: Option<i64>, values: &[Value]) -> Result<i64> {
self.check_values(values)?;
let sql = self.insert_sql(fid.is_some(), false);
let mut binds: Vec<SqlValue> = Vec::with_capacity(values.len() + 1);
if let Some(id) = fid {
binds.push(SqlValue::Integer(id));
}
binds.extend(values.iter().map(value_to_sql));
let assigned = self.exec_insert(&sql, &binds, fid)?;
self.dirty = true;
Ok(assigned)
}
pub fn update<G: GeometryTrait<T = f64>>(
&mut self,
fid: i64,
geometry: &G,
values: &[Value],
) -> Result<bool> {
self.check_values(values)?;
let (blob, xy) = self.encode_geometry(geometry)?;
let sql = self.update_sql(true);
let mut binds: Vec<SqlValue> = Vec::with_capacity(values.len() + 2);
binds.extend(values.iter().map(value_to_sql));
binds.push(SqlValue::Blob(blob));
binds.push(SqlValue::Integer(fid));
let matched = self.exec_update(&sql, &binds)?;
if matched {
if let Some(envelope) = xy {
self.bbox.add(envelope);
self.bbox_dirty = true;
}
self.dirty = true;
}
Ok(matched)
}
pub fn update_row(&mut self, fid: i64, values: &[Value]) -> Result<bool> {
self.check_values(values)?;
let sql = self.update_sql(false);
let mut binds: Vec<SqlValue> = Vec::with_capacity(values.len() + 1);
binds.extend(values.iter().map(value_to_sql));
binds.push(SqlValue::Integer(fid));
let matched = self.exec_update(&sql, &binds)?;
if matched {
self.dirty = true;
}
Ok(matched)
}
pub fn delete(&mut self, fid: i64) -> Result<bool> {
let sql = format!(
"DELETE FROM {} WHERE {} = ?1",
self.quoted_table, self.pk_expr
);
let matched = {
let mut stmt = self.tx.prepare_cached(&sql)?;
stmt.execute([fid])? > 0
};
if matched {
self.dirty = true;
}
Ok(matched)
}
pub fn commit(self) -> Result<()> {
let Self {
tx,
table_name,
bbox,
dirty,
bbox_dirty,
..
} = self;
if dirty {
tx.execute(
"UPDATE gpkg_contents \
SET last_change = strftime('%Y-%m-%dT%H:%M:%fZ','now') \
WHERE table_name = ?1",
[&table_name],
)?;
}
if bbox_dirty && let Some([min_x, max_x, min_y, max_y]) = bbox.bounds() {
tx.execute(
"UPDATE gpkg_contents \
SET min_x = ?1, min_y = ?2, max_x = ?3, max_y = ?4 \
WHERE table_name = ?5",
rusqlite::params![min_x, min_y, max_x, max_y, table_name],
)?;
}
tx.commit()?;
Ok(())
}
fn encode_geometry<G: GeometryTrait<T = f64>>(
&self,
geometry: &G,
) -> Result<(Vec<u8>, Option<[f64; 4]>)> {
let geom = self
.geometry
.as_ref()
.ok_or_else(|| Error::NoGeometryColumn {
table_name: self.table_name.clone(),
})?;
let dim = geometry.dim();
let has_z = matches!(dim, Dimensions::Xyz | Dimensions::Xyzm);
let has_m = matches!(dim, Dimensions::Xym | Dimensions::Xyzm);
self.check_zm("z", geom.z, has_z, &geom.name)?;
self.check_zm("m", geom.m, has_m, &geom.name)?;
encode_gpb(geometry, geom.srs_id).map_err(|e| Error::Core(e.into()))
}
fn check_zm(
&self,
dimension: &'static str,
constraint: ZmFlag,
present: bool,
column: &str,
) -> Result<()> {
let ok = match constraint {
ZmFlag::Prohibited => !present,
ZmFlag::Mandatory => present,
ZmFlag::Optional => true,
_ => true,
};
if ok {
return Ok(());
}
Err(Error::ZmViolation {
table_name: self.table_name.clone(),
column: column.to_owned(),
dimension,
constraint,
verb: if present { "carries" } else { "lacks" },
})
}
fn check_values(&self, values: &[Value]) -> Result<()> {
if values.len() == self.value_columns.len() {
return Ok(());
}
Err(Error::ValueCountMismatch {
table_name: self.table_name.clone(),
expected: self.value_columns.len(),
found: values.len(),
})
}
fn insert_sql(&self, with_fid: bool, with_geometry: bool) -> String {
let mut columns: Vec<&str> = Vec::with_capacity(self.value_columns.len() + 2);
if with_fid {
columns.push(&self.pk_expr);
}
for column in &self.value_columns {
columns.push(column);
}
if with_geometry && let Some(geom) = &self.geometry {
columns.push(&geom.quoted_name);
}
if columns.is_empty() {
return format!("INSERT INTO {} DEFAULT VALUES", self.quoted_table);
}
let placeholders = (1..=columns.len())
.map(|i| format!("?{i}"))
.collect::<Vec<_>>()
.join(", ");
format!(
"INSERT INTO {} ({}) VALUES ({placeholders})",
self.quoted_table,
columns.join(", ")
)
}
fn update_sql(&self, with_geometry: bool) -> String {
let mut assignments: Vec<String> = Vec::with_capacity(self.value_columns.len() + 1);
let mut index = 1;
for column in &self.value_columns {
assignments.push(format!("{column} = ?{index}"));
index += 1;
}
if with_geometry && let Some(geom) = &self.geometry {
assignments.push(format!("{} = ?{index}", geom.quoted_name));
index += 1;
}
if assignments.is_empty() {
assignments.push(format!("{pk} = {pk}", pk = self.pk_expr));
}
format!(
"UPDATE {} SET {} WHERE {} = ?{index}",
self.quoted_table,
assignments.join(", "),
self.pk_expr
)
}
fn exec_insert(&self, sql: &str, binds: &[SqlValue], fid: Option<i64>) -> Result<i64> {
let mut stmt = self.tx.prepare_cached(sql)?;
stmt.execute(params_from_iter(binds.iter()))?;
Ok(fid.unwrap_or_else(|| self.tx.last_insert_rowid()))
}
fn exec_update(&self, sql: &str, binds: &[SqlValue]) -> Result<bool> {
let mut stmt = self.tx.prepare_cached(sql)?;
Ok(stmt.execute(params_from_iter(binds.iter()))? > 0)
}
}
fn read_contents_bbox(conn: &Connection, table: &str) -> Result<Option<[f64; 4]>> {
let row = conn
.query_row(
"SELECT min_x, min_y, max_x, max_y FROM gpkg_contents WHERE table_name = ?1",
[table],
|r| {
Ok((
r.get::<_, Option<f64>>(0)?,
r.get::<_, Option<f64>>(1)?,
r.get::<_, Option<f64>>(2)?,
r.get::<_, Option<f64>>(3)?,
))
},
)
.optional()?;
Ok(match row {
Some((Some(min_x), Some(min_y), Some(max_x), Some(max_y))) => {
Some([min_x, max_x, min_y, max_y])
}
_ => None,
})
}