Add timestamp wrapper type (#105)

* Add timestamp wrapper type * fix build * Fix tests * Address comments * cargo fmt * add query by duration Change-Id: Ia1df252ab09ec93680d9f1e6a3dba196c6dfe478 Co-authored-by: wangfenjin <wangfenj@gmail.com>
duckdb · Dec 12, 2022 · e666e08 · e666e08
1 parent d2f611f
commit e666e08
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 6 deletions.
diff --git a/src/appender.rs b/src/appender.rs
@@ -6,7 +6,7 @@ use std::iter::IntoIterator;
 use std::os::raw::c_char;
 
 use crate::error::result_from_duckdb_appender;
-use crate::types::{ToSql, ToSqlOutput};
+use crate::types::{TimeUnit, ToSql, ToSqlOutput};
 use crate::Error;
 
 /// Appender for fast import data
@@ -108,6 +108,15 @@ impl Appender<'_> {
             ValueRef::Text(s) => unsafe {
                 ffi::duckdb_append_varchar_length(ptr, s.as_ptr() as *const c_char, s.len() as u64)
             },
+            ValueRef::Timestamp(u, i) => unsafe {
+                let micros = match u {
+                    TimeUnit::Second => i * 1_000_000,
+                    TimeUnit::Millisecond => i * 1_000,
+                    TimeUnit::Microsecond => i,
+                    TimeUnit::Nanosecond => i / 1_000,
+                };
+                ffi::duckdb_append_timestamp(ptr, ffi::duckdb_timestamp { micros })
+            },
             ValueRef::Blob(b) => unsafe { ffi::duckdb_append_blob(ptr, b.as_ptr() as *const c_void, b.len() as u64) },
             _ => unreachable!("not supported"),
         };
@@ -215,8 +224,25 @@ mod test {
             app.append_row(["2022-04-09 15:56:37.544"])?;
         }
 
-        let val = db.query_row("SELECT x FROM foo", [], |row| <(String,)>::try_from(row))?;
-        assert_eq!(val, ("2022-04-09 15:56:37.544".to_string(),));
+        let val = db.query_row("SELECT x FROM foo", [], |row| <(i64,)>::try_from(row))?;
+        assert_eq!(val, (1649519797544000,));
+        Ok(())
+    }
+
+    #[test]
+    fn test_append_timestamp() -> Result<()> {
+        use std::time::Duration;
+        let db = Connection::open_in_memory()?;
+        db.execute_batch("CREATE TABLE foo(x TIMESTAMP)")?;
+
+        let d = Duration::from_secs(1);
+        {
+            let mut app = db.appender("foo")?;
+            app.append_row([d])?;
+        }
+
+        let val = db.query_row("SELECT x FROM foo where x=?", [d], |row| <(i32,)>::try_from(row))?;
+        assert_eq!(val, (d.as_micros() as i32,));
         Ok(())
     }
 }
diff --git a/src/statement.rs b/src/statement.rs
@@ -7,7 +7,7 @@ use super::ffi;
 use super::{AndThenRows, Connection, Error, MappedRows, Params, RawStatement, Result, Row, Rows, ValueRef};
 use crate::arrow_batch::Arrow;
 use crate::error::result_from_duckdb_prepare;
-use crate::types::{ToSql, ToSqlOutput};
+use crate::types::{TimeUnit, ToSql, ToSqlOutput};
 
 use arrow::array::StructArray;
 use arrow::datatypes::DataType;
@@ -434,7 +434,16 @@ impl Statement<'_> {
             ValueRef::Blob(b) => unsafe {
                 ffi::duckdb_bind_blob(ptr, col as u64, b.as_ptr() as *const c_void, b.len() as u64)
             },
-            _ => unreachable!("not supported"),
+            ValueRef::Timestamp(u, i) => unsafe {
+                let micros = match u {
+                    TimeUnit::Second => i * 1_000_000,
+                    TimeUnit::Millisecond => i * 1_000,
+                    TimeUnit::Microsecond => i,
+                    TimeUnit::Nanosecond => i / 1_000,
+                };
+                ffi::duckdb_bind_timestamp(ptr, col as u64, ffi::duckdb_timestamp { micros })
+            },
+            _ => unreachable!("not supported: {}", value.data_type()),
         };
         result_from_duckdb_prepare(rc, ptr)
     }

diff --git a/src/types/to_sql.rs b/src/types/to_sql.rs
@@ -1,4 +1,4 @@
-use super::{Null, Value, ValueRef};
+use super::{Null, TimeUnit, Value, ValueRef};
 use crate::Result;
 use std::borrow::Cow;
 
@@ -202,6 +202,15 @@ impl<T: ToSql> ToSql for Option<T> {
     }
 }
 
+impl ToSql for std::time::Duration {
+    fn to_sql(&self) -> crate::Result<ToSqlOutput<'_>> {
+        Ok(ToSqlOutput::Owned(Value::Timestamp(
+            TimeUnit::Microsecond,
+            self.as_micros() as i64,
+        )))
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::ToSql;