From f351f3acde14c839b04789cb25ffae9bbd15b5ae Mon Sep 17 00:00:00 2001 From: LJ Date: Mon, 24 Mar 2025 17:19:00 -0700 Subject: [PATCH] Add uuid type support in core engine. --- Cargo.toml | 3 ++- src/base/json_schema.rs | 10 +++++++--- src/base/schema.rs | 4 ++++ src/base/value.rs | 11 +++++++++++ src/execution/query.rs | 1 + src/ops/storages/postgres.rs | 17 ++++++++++++++--- 6 files changed, 39 insertions(+), 7 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2f59b9af3..ed2690e1e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,7 +25,7 @@ log = "0.4.26" regex = "1.11.1" serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.140" -sqlx = { version = "0.8.3", features = ["chrono", "postgres", "runtime-tokio"] } +sqlx = { version = "0.8.3", features = ["chrono", "postgres", "runtime-tokio", "uuid"] } tokio = { version = "1.44.1", features = [ "macros", "rt-multi-thread", @@ -89,3 +89,4 @@ rustls = { version = "0.23.25" } http-body-util = "0.1.3" yaml-rust2 = "0.10.0" urlencoding = "2.1.3" +uuid = "1.16.0" diff --git a/src/base/json_schema.rs b/src/base/json_schema.rs index 4a7ee1253..5a0840dfd 100644 --- a/src/base/json_schema.rs +++ b/src/base/json_schema.rs @@ -26,9 +26,6 @@ impl ToJsonSchema for schema::BasicValueType { schema::BasicValueType::Float32 | schema::BasicValueType::Float64 => { schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::Number))); } - schema::BasicValueType::Json => { - // Can be any value. No type constraint. - } schema::BasicValueType::Range => { schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::Array))); schema.array = Some(Box::new(ArrayValidation { @@ -51,6 +48,13 @@ impl ToJsonSchema for schema::BasicValueType { .description = Some("A range, start pos (inclusive), end pos (exclusive).".to_string()); } + schema::BasicValueType::Uuid => { + schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String))); + schema.format = Some("uuid".to_string()); + } + schema::BasicValueType::Json => { + // Can be any value. No type constraint. + } schema::BasicValueType::Vector(s) => { schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::Array))); schema.array = Some(Box::new(ArrayValidation { diff --git a/src/base/schema.rs b/src/base/schema.rs index 4c54ca38c..de10e6ab8 100644 --- a/src/base/schema.rs +++ b/src/base/schema.rs @@ -35,6 +35,9 @@ pub enum BasicValueType { /// A range, with a start offset and a length. Range, + /// A UUID. + Uuid, + /// A JSON value. Json, @@ -52,6 +55,7 @@ impl std::fmt::Display for BasicValueType { BasicValueType::Float32 => write!(f, "float32"), BasicValueType::Float64 => write!(f, "float64"), BasicValueType::Range => write!(f, "range"), + BasicValueType::Uuid => write!(f, "uuid"), BasicValueType::Json => write!(f, "json"), BasicValueType::Vector(s) => write!( f, diff --git a/src/base/value.rs b/src/base/value.rs index a465bd07b..905da0bd8 100644 --- a/src/base/value.rs +++ b/src/base/value.rs @@ -9,6 +9,7 @@ use serde::{ Deserialize, Serialize, }; use std::{collections::BTreeMap, ops::Deref, sync::Arc}; +use uuid::Uuid; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] pub struct RangeValue { @@ -76,6 +77,7 @@ pub enum KeyValue { Bool(bool), Int64(i64), Range(RangeValue), + Uuid(Uuid), Struct(Vec), } @@ -141,6 +143,7 @@ impl std::fmt::Display for KeyValue { KeyValue::Bool(v) => write!(f, "{}", v), KeyValue::Int64(v) => write!(f, "{}", v), KeyValue::Range(v) => write!(f, "[{}, {})", v.start, v.end), + KeyValue::Uuid(v) => write!(f, "{}", v), KeyValue::Struct(v) => { write!( f, @@ -205,6 +208,7 @@ impl KeyValue { output.push(v.start.to_string()); output.push(v.end.to_string()); } + KeyValue::Uuid(v) => output.push(v.to_string()), KeyValue::Struct(v) => { for part in v { part.parts_to_strs(output); @@ -235,6 +239,7 @@ impl KeyValue { KeyValue::Bool(_) => "bool", KeyValue::Int64(_) => "int64", KeyValue::Range { .. } => "range", + KeyValue::Uuid(_) => "uuid", KeyValue::Struct(_) => "struct", } } @@ -299,6 +304,7 @@ pub enum BasicValue { Float32(f32), Float64(f64), Range(RangeValue), + Uuid(Uuid), Json(Arc), Vector(Arc<[BasicValue]>), } @@ -373,6 +379,7 @@ impl BasicValue { BasicValue::Bool(v) => KeyValue::Bool(v), BasicValue::Int64(v) => KeyValue::Int64(v), BasicValue::Range(v) => KeyValue::Range(v), + BasicValue::Uuid(v) => KeyValue::Uuid(v), BasicValue::Float32(_) | BasicValue::Float64(_) | BasicValue::Json(_) @@ -388,6 +395,7 @@ impl BasicValue { BasicValue::Bool(v) => KeyValue::Bool(*v), BasicValue::Int64(v) => KeyValue::Int64(*v), BasicValue::Range(v) => KeyValue::Range(*v), + BasicValue::Uuid(v) => KeyValue::Uuid(*v), BasicValue::Float32(_) | BasicValue::Float64(_) | BasicValue::Json(_) @@ -405,6 +413,7 @@ impl BasicValue { BasicValue::Float32(_) => "float32", BasicValue::Float64(_) => "float64", BasicValue::Range(_) => "range", + BasicValue::Uuid(_) => "uuid", BasicValue::Json(_) => "json", BasicValue::Vector(_) => "vector", } @@ -436,6 +445,7 @@ impl From for Value { KeyValue::Bool(v) => Value::Basic(BasicValue::Bool(v)), KeyValue::Int64(v) => Value::Basic(BasicValue::Int64(v)), KeyValue::Range(v) => Value::Basic(BasicValue::Range(v)), + KeyValue::Uuid(v) => Value::Basic(BasicValue::Uuid(v)), KeyValue::Struct(v) => Value::Struct(FieldValues { fields: v.into_iter().map(Value::from).collect(), }), @@ -734,6 +744,7 @@ impl serde::Serialize for BasicValue { BasicValue::Float32(v) => serializer.serialize_f32(*v), BasicValue::Float64(v) => serializer.serialize_f64(*v), BasicValue::Range(v) => v.serialize(serializer), + BasicValue::Uuid(v) => serializer.serialize_str(&v.to_string()), BasicValue::Json(v) => v.serialize(serializer), BasicValue::Vector(v) => v.serialize(serializer), } diff --git a/src/execution/query.rs b/src/execution/query.rs index ebd180cd9..82031c93c 100644 --- a/src/execution/query.rs +++ b/src/execution/query.rs @@ -87,6 +87,7 @@ impl SimpleSemanticsQueryHandler { | value::BasicValue::Str(_) | value::BasicValue::Bool(_) | value::BasicValue::Range(_) + | value::BasicValue::Uuid(_) | value::BasicValue::Json(_) | value::BasicValue::Vector(_) => { bail!("Query results is not a vector of number") diff --git a/src/ops/storages/postgres.rs b/src/ops/storages/postgres.rs index dcedb5cc1..3b63edfe7 100644 --- a/src/ops/storages/postgres.rs +++ b/src/ops/storages/postgres.rs @@ -21,6 +21,7 @@ use serde::Serialize; use sqlx::postgres::types::PgRange; use sqlx::postgres::PgRow; use sqlx::{PgPool, Row}; +use uuid::Uuid; #[derive(Debug, Deserialize)] pub struct Spec { @@ -78,6 +79,9 @@ fn bind_key_field<'arg>( end: Bound::Excluded(v.end as i64), }); } + KeyValue::Uuid(v) => { + builder.push_bind(v); + } KeyValue::Struct(fields) => { builder.push_bind(sqlx::types::Json(fields)); } @@ -110,15 +114,18 @@ fn bind_value_field<'arg>( BasicValue::Float64(v) => { builder.push_bind(v); } - BasicValue::Json(v) => { - builder.push_bind(sqlx::types::Json(&**v)); - } BasicValue::Range(v) => { builder.push_bind(PgRange { start: Bound::Included(v.start as i64), end: Bound::Excluded(v.end as i64), }); } + BasicValue::Uuid(v) => { + builder.push_bind(v); + } + BasicValue::Json(v) => { + builder.push_bind(sqlx::types::Json(&**v)); + } BasicValue::Vector(v) => match &field_schema.value_type.typ { ValueType::Basic(BasicValueType::Vector(vs)) if convertible_to_pgvector(vs) => { let vec = v @@ -186,6 +193,9 @@ fn from_pg_value(row: &PgRow, field_idx: usize, typ: &ValueType) -> Result anyhow::bail!("invalid range value"), }) .transpose()?, + BasicValueType::Uuid => row + .try_get::, _>(field_idx)? + .map(BasicValue::Uuid), BasicValueType::Json => row .try_get::, _>(field_idx)? .map(|v| BasicValue::Json(Arc::from(v))), @@ -655,6 +665,7 @@ fn to_column_type_sql(column_type: &ValueType) -> Cow<'static, str> { BasicValueType::Float32 => "real".into(), BasicValueType::Float64 => "double precision".into(), BasicValueType::Range => "int8range".into(), + BasicValueType::Uuid => "uuid".into(), BasicValueType::Json => "jsonb".into(), BasicValueType::Vector(vec_schema) => { if convertible_to_pgvector(vec_schema) {