Skip to content

Commit 3a7ad4d

Browse files
badmonster0badmonster0
committed
Add date/time types support in Rust core.
1 parent 7c590c2 commit 3a7ad4d

File tree

7 files changed

+188
-12
lines changed

7 files changed

+188
-12
lines changed

‎Cargo.toml‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ name = "cocoindex_engine"
1111
crate-type = ["cdylib"]
1212

1313
[dependencies]
14-
pyo3 = { version = "0.23.5" }
14+
pyo3 = { version = "0.23.5", features = ["chrono"] }
1515
anyhow = { version = "1.0.97", features = ["std"] }
1616
async-trait = "0.1.88"
1717
axum = "0.7.9"
@@ -25,7 +25,12 @@ log = "0.4.26"
2525
regex = "1.11.1"
2626
serde = { version = "1.0.219", features = ["derive"] }
2727
serde_json = "1.0.140"
28-
sqlx = { version = "0.8.3", features = ["chrono", "postgres", "runtime-tokio", "uuid"] }
28+
sqlx = { version = "0.8.3", features = [
29+
"chrono",
30+
"postgres",
31+
"runtime-tokio",
32+
"uuid",
33+
] }
2934
tokio = { version = "1.44.1", features = [
3035
"macros",
3136
"rt-multi-thread",

‎src/base/json_schema.rs‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,16 +49,33 @@ impl ToJsonSchema for schema::BasicValueType {
4949
max_items: Some(2),
5050
..Default::default()
5151
}));
52-
schema
53-
.metadata
54-
.get_or_insert_with(Default::default)
55-
.description =
52+
schema.metadata.get_or_insert_default().description =
5653
Some("A range, start pos (inclusive), end pos (exclusive).".to_string());
5754
}
5855
schema::BasicValueType::Uuid => {
5956
schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String)));
6057
schema.format = Some("uuid".to_string());
6158
}
59+
schema::BasicValueType::Date => {
60+
schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String)));
61+
schema.format = Some("date".to_string());
62+
}
63+
schema::BasicValueType::Time => {
64+
schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String)));
65+
schema.format = Some("time".to_string());
66+
}
67+
schema::BasicValueType::LocalDateTime => {
68+
schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String)));
69+
schema.format = Some("date-time".to_string());
70+
schema.metadata.get_or_insert_default().description =
71+
Some("Date time without timezone offset, YYYY-MM-DDThh:mm:ss".to_string());
72+
}
73+
schema::BasicValueType::OffsetDateTime => {
74+
schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String)));
75+
schema.format = Some("date-time".to_string());
76+
schema.metadata.get_or_insert_default().description =
77+
Some("Date time with timezone offset in RFC3339".to_string());
78+
}
6279
schema::BasicValueType::Json => {
6380
// Can be any value. No type constraint.
6481
}

‎src/base/schema.rs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,18 @@ pub enum BasicValueType {
3838
/// A UUID.
3939
Uuid,
4040

41+
/// Date (without time within the current day).
42+
Date,
43+
44+
/// Time of the day.
45+
Time,
46+
47+
/// Local date and time, without timezone.
48+
LocalDateTime,
49+
50+
/// Date and time with timezone.
51+
OffsetDateTime,
52+
4153
/// A JSON value.
4254
Json,
4355

@@ -56,6 +68,10 @@ impl std::fmt::Display for BasicValueType {
5668
BasicValueType::Float64 => write!(f, "float64"),
5769
BasicValueType::Range => write!(f, "range"),
5870
BasicValueType::Uuid => write!(f, "uuid"),
71+
BasicValueType::Date => write!(f, "date"),
72+
BasicValueType::Time => write!(f, "time"),
73+
BasicValueType::LocalDateTime => write!(f, "local_datetime"),
74+
BasicValueType::OffsetDateTime => write!(f, "offset_datetime"),
5975
BasicValueType::Json => write!(f, "json"),
6076
BasicValueType::Vector(s) => write!(
6177
f,

‎src/base/value.rs‎

Lines changed: 97 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ pub enum KeyValue {
7777
Int64(i64),
7878
Range(RangeValue),
7979
Uuid(uuid::Uuid),
80+
Date(chrono::NaiveDate),
8081
Struct(Vec<KeyValue>),
8182
}
8283

@@ -122,6 +123,18 @@ impl From<RangeValue> for KeyValue {
122123
}
123124
}
124125

126+
impl From<uuid::Uuid> for KeyValue {
127+
fn from(value: uuid::Uuid) -> Self {
128+
KeyValue::Uuid(value)
129+
}
130+
}
131+
132+
impl From<chrono::NaiveDate> for KeyValue {
133+
fn from(value: chrono::NaiveDate) -> Self {
134+
KeyValue::Date(value)
135+
}
136+
}
137+
125138
impl From<Vec<KeyValue>> for KeyValue {
126139
fn from(value: Vec<KeyValue>) -> Self {
127140
KeyValue::Struct(value)
@@ -143,6 +156,7 @@ impl std::fmt::Display for KeyValue {
143156
KeyValue::Int64(v) => write!(f, "{}", v),
144157
KeyValue::Range(v) => write!(f, "[{}, {})", v.start, v.end),
145158
KeyValue::Uuid(v) => write!(f, "{}", v),
159+
KeyValue::Date(v) => write!(f, "{}", v),
146160
KeyValue::Struct(v) => {
147161
write!(
148162
f,
@@ -172,17 +186,19 @@ impl KeyValue {
172186
KeyValue::Bytes(Arc::from(BASE64_STANDARD.decode(v)?))
173187
}
174188
BasicValueType::Str { .. } => KeyValue::Str(Arc::from(v)),
175-
BasicValueType::Bool => KeyValue::Bool(v.parse::<bool>()?),
176-
BasicValueType::Int64 => KeyValue::Int64(v.parse::<i64>()?),
189+
BasicValueType::Bool => KeyValue::Bool(v.parse()?),
190+
BasicValueType::Int64 => KeyValue::Int64(v.parse()?),
177191
BasicValueType::Range => {
178192
let v2 = values_iter
179193
.next()
180194
.ok_or_else(|| api_error!("Key parts less than expected"))?;
181195
KeyValue::Range(RangeValue {
182-
start: v.parse::<usize>()?,
183-
end: v2.parse::<usize>()?,
196+
start: v.parse()?,
197+
end: v2.parse()?,
184198
})
185199
}
200+
BasicValueType::Uuid => KeyValue::Uuid(v.parse()?),
201+
BasicValueType::Date => KeyValue::Date(v.parse()?),
186202
schema => api_bail!("Invalid key type {schema}"),
187203
}
188204
}
@@ -208,6 +224,7 @@ impl KeyValue {
208224
output.push(v.end.to_string());
209225
}
210226
KeyValue::Uuid(v) => output.push(v.to_string()),
227+
KeyValue::Date(v) => output.push(v.to_string()),
211228
KeyValue::Struct(v) => {
212229
for part in v {
213230
part.parts_to_strs(output);
@@ -239,6 +256,7 @@ impl KeyValue {
239256
KeyValue::Int64(_) => "int64",
240257
KeyValue::Range { .. } => "range",
241258
KeyValue::Uuid(_) => "uuid",
259+
KeyValue::Date(_) => "date",
242260
KeyValue::Struct(_) => "struct",
243261
}
244262
}
@@ -278,6 +296,20 @@ impl KeyValue {
278296
}
279297
}
280298

299+
pub fn uuid_value(&self) -> Result<uuid::Uuid> {
300+
match self {
301+
KeyValue::Uuid(v) => Ok(*v),
302+
_ => anyhow::bail!("expected uuid value, but got {}", self.kind_str()),
303+
}
304+
}
305+
306+
pub fn date_value(&self) -> Result<chrono::NaiveDate> {
307+
match self {
308+
KeyValue::Date(v) => Ok(*v),
309+
_ => anyhow::bail!("expected date value, but got {}", self.kind_str()),
310+
}
311+
}
312+
281313
pub fn struct_value(&self) -> Result<&Vec<KeyValue>> {
282314
match self {
283315
KeyValue::Struct(v) => Ok(v),
@@ -304,6 +336,10 @@ pub enum BasicValue {
304336
Float64(f64),
305337
Range(RangeValue),
306338
Uuid(uuid::Uuid),
339+
Date(chrono::NaiveDate),
340+
Time(chrono::NaiveTime),
341+
LocalDateTime(chrono::NaiveDateTime),
342+
OffsetDateTime(chrono::DateTime<chrono::FixedOffset>),
307343
Json(Arc<serde_json::Value>),
308344
Vector(Arc<[BasicValue]>),
309345
}
@@ -356,6 +392,36 @@ impl From<f64> for BasicValue {
356392
}
357393
}
358394

395+
impl From<uuid::Uuid> for BasicValue {
396+
fn from(value: uuid::Uuid) -> Self {
397+
BasicValue::Uuid(value)
398+
}
399+
}
400+
401+
impl From<chrono::NaiveDate> for BasicValue {
402+
fn from(value: chrono::NaiveDate) -> Self {
403+
BasicValue::Date(value)
404+
}
405+
}
406+
407+
impl From<chrono::NaiveTime> for BasicValue {
408+
fn from(value: chrono::NaiveTime) -> Self {
409+
BasicValue::Time(value)
410+
}
411+
}
412+
413+
impl From<chrono::NaiveDateTime> for BasicValue {
414+
fn from(value: chrono::NaiveDateTime) -> Self {
415+
BasicValue::LocalDateTime(value)
416+
}
417+
}
418+
419+
impl From<chrono::DateTime<chrono::FixedOffset>> for BasicValue {
420+
fn from(value: chrono::DateTime<chrono::FixedOffset>) -> Self {
421+
BasicValue::OffsetDateTime(value)
422+
}
423+
}
424+
359425
impl From<serde_json::Value> for BasicValue {
360426
fn from(value: serde_json::Value) -> Self {
361427
BasicValue::Json(Arc::from(value))
@@ -379,8 +445,12 @@ impl BasicValue {
379445
BasicValue::Int64(v) => KeyValue::Int64(v),
380446
BasicValue::Range(v) => KeyValue::Range(v),
381447
BasicValue::Uuid(v) => KeyValue::Uuid(v),
448+
BasicValue::Date(v) => KeyValue::Date(v),
382449
BasicValue::Float32(_)
383450
| BasicValue::Float64(_)
451+
| BasicValue::Time(_)
452+
| BasicValue::LocalDateTime(_)
453+
| BasicValue::OffsetDateTime(_)
384454
| BasicValue::Json(_)
385455
| BasicValue::Vector(_) => api_bail!("invalid key value type"),
386456
};
@@ -395,8 +465,12 @@ impl BasicValue {
395465
BasicValue::Int64(v) => KeyValue::Int64(*v),
396466
BasicValue::Range(v) => KeyValue::Range(*v),
397467
BasicValue::Uuid(v) => KeyValue::Uuid(*v),
468+
BasicValue::Date(v) => KeyValue::Date(*v),
398469
BasicValue::Float32(_)
399470
| BasicValue::Float64(_)
471+
| BasicValue::Time(_)
472+
| BasicValue::LocalDateTime(_)
473+
| BasicValue::OffsetDateTime(_)
400474
| BasicValue::Json(_)
401475
| BasicValue::Vector(_) => api_bail!("invalid key value type"),
402476
};
@@ -413,6 +487,10 @@ impl BasicValue {
413487
BasicValue::Float64(_) => "float64",
414488
BasicValue::Range(_) => "range",
415489
BasicValue::Uuid(_) => "uuid",
490+
BasicValue::Date(_) => "date",
491+
BasicValue::Time(_) => "time",
492+
BasicValue::LocalDateTime(_) => "local_datetime",
493+
BasicValue::OffsetDateTime(_) => "offset_datetime",
416494
BasicValue::Json(_) => "json",
417495
BasicValue::Vector(_) => "vector",
418496
}
@@ -445,6 +523,7 @@ impl From<KeyValue> for Value {
445523
KeyValue::Int64(v) => Value::Basic(BasicValue::Int64(v)),
446524
KeyValue::Range(v) => Value::Basic(BasicValue::Range(v)),
447525
KeyValue::Uuid(v) => Value::Basic(BasicValue::Uuid(v)),
526+
KeyValue::Date(v) => Value::Basic(BasicValue::Date(v)),
448527
KeyValue::Struct(v) => Value::Struct(FieldValues {
449528
fields: v.into_iter().map(Value::from).collect(),
450529
}),
@@ -744,6 +823,12 @@ impl serde::Serialize for BasicValue {
744823
BasicValue::Float64(v) => serializer.serialize_f64(*v),
745824
BasicValue::Range(v) => v.serialize(serializer),
746825
BasicValue::Uuid(v) => serializer.serialize_str(&v.to_string()),
826+
BasicValue::Date(v) => serializer.serialize_str(&v.to_string()),
827+
BasicValue::Time(v) => serializer.serialize_str(&v.to_string()),
828+
BasicValue::LocalDateTime(v) => serializer.serialize_str(&v.to_string()),
829+
BasicValue::OffsetDateTime(v) => {
830+
serializer.serialize_str(&v.to_rfc3339_opts(chrono::SecondsFormat::AutoSi, true))
831+
}
747832
BasicValue::Json(v) => v.serialize(serializer),
748833
BasicValue::Vector(v) => v.serialize(serializer),
749834
}
@@ -774,8 +859,14 @@ impl BasicValue {
774859
.ok_or_else(|| anyhow::anyhow!("invalid fp64 value {v}"))?,
775860
),
776861
(v, BasicValueType::Range) => BasicValue::Range(serde_json::from_value(v)?),
777-
(serde_json::Value::String(v), BasicValueType::Uuid) => {
778-
BasicValue::Uuid(uuid::Uuid::parse_str(v.as_str())?)
862+
(serde_json::Value::String(v), BasicValueType::Uuid) => BasicValue::Uuid(v.parse()?),
863+
(serde_json::Value::String(v), BasicValueType::Date) => BasicValue::Date(v.parse()?),
864+
(serde_json::Value::String(v), BasicValueType::Time) => BasicValue::Time(v.parse()?),
865+
(serde_json::Value::String(v), BasicValueType::LocalDateTime) => {
866+
BasicValue::LocalDateTime(v.parse()?)
867+
}
868+
(serde_json::Value::String(v), BasicValueType::OffsetDateTime) => {
869+
BasicValue::OffsetDateTime(chrono::DateTime::parse_from_rfc3339(&v)?)
779870
}
780871
(v, BasicValueType::Json) => BasicValue::Json(Arc::from(v)),
781872
(

‎src/execution/query.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,10 @@ impl SimpleSemanticsQueryHandler {
8888
| value::BasicValue::Bool(_)
8989
| value::BasicValue::Range(_)
9090
| value::BasicValue::Uuid(_)
91+
| value::BasicValue::Date(_)
92+
| value::BasicValue::Time(_)
93+
| value::BasicValue::LocalDateTime(_)
94+
| value::BasicValue::OffsetDateTime(_)
9195
| value::BasicValue::Json(_)
9296
| value::BasicValue::Vector(_) => {
9397
bail!("Query results is not a vector of number")

‎src/ops/storages/postgres.rs‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,9 @@ fn bind_key_field<'arg>(
8282
KeyValue::Uuid(v) => {
8383
builder.push_bind(v);
8484
}
85+
KeyValue::Date(v) => {
86+
builder.push_bind(v);
87+
}
8588
KeyValue::Struct(fields) => {
8689
builder.push_bind(sqlx::types::Json(fields));
8790
}
@@ -123,6 +126,18 @@ fn bind_value_field<'arg>(
123126
BasicValue::Uuid(v) => {
124127
builder.push_bind(v);
125128
}
129+
BasicValue::Date(v) => {
130+
builder.push_bind(v);
131+
}
132+
BasicValue::Time(v) => {
133+
builder.push_bind(v);
134+
}
135+
BasicValue::LocalDateTime(v) => {
136+
builder.push_bind(v);
137+
}
138+
BasicValue::OffsetDateTime(v) => {
139+
builder.push_bind(v);
140+
}
126141
BasicValue::Json(v) => {
127142
builder.push_bind(sqlx::types::Json(&**v));
128143
}
@@ -196,6 +211,18 @@ fn from_pg_value(row: &PgRow, field_idx: usize, typ: &ValueType) -> Result<Value
196211
BasicValueType::Uuid => row
197212
.try_get::<Option<Uuid>, _>(field_idx)?
198213
.map(BasicValue::Uuid),
214+
BasicValueType::Date => row
215+
.try_get::<Option<chrono::NaiveDate>, _>(field_idx)?
216+
.map(BasicValue::Date),
217+
BasicValueType::Time => row
218+
.try_get::<Option<chrono::NaiveTime>, _>(field_idx)?
219+
.map(BasicValue::Time),
220+
BasicValueType::LocalDateTime => row
221+
.try_get::<Option<chrono::NaiveDateTime>, _>(field_idx)?
222+
.map(BasicValue::LocalDateTime),
223+
BasicValueType::OffsetDateTime => row
224+
.try_get::<Option<chrono::DateTime<chrono::FixedOffset>>, _>(field_idx)?
225+
.map(BasicValue::OffsetDateTime),
199226
BasicValueType::Json => row
200227
.try_get::<Option<serde_json::Value>, _>(field_idx)?
201228
.map(|v| BasicValue::Json(Arc::from(v))),
@@ -666,6 +693,10 @@ fn to_column_type_sql(column_type: &ValueType) -> Cow<'static, str> {
666693
BasicValueType::Float64 => "double precision".into(),
667694
BasicValueType::Range => "int8range".into(),
668695
BasicValueType::Uuid => "uuid".into(),
696+
BasicValueType::Date => "date".into(),
697+
BasicValueType::Time => "time".into(),
698+
BasicValueType::LocalDateTime => "timestamp".into(),
699+
BasicValueType::OffsetDateTime => "timestamp with time zone".into(),
669700
BasicValueType::Json => "jsonb".into(),
670701
BasicValueType::Vector(vec_schema) => {
671702
if convertible_to_pgvector(vec_schema) {

0 commit comments

Comments
 (0)