Merge 'Fix scalar API in extensions, add documentation and error handling' from Preston Thorpe

Closes #728
Changes the API to one macro/annotation on the relevant function
```rust
#[scalar(name = "uuid4_str", alias = "gen_random_uuid")]
fn uuid4_str(_args: &[Value]) -> Value {
    let uuid = uuid::Uuid::new_v4().to_string();
    Value::from_text(uuid)
}

register_extension! {
    scalars: { uuid4_str, uuid4 }
}
```
The only downside of this, is that for functions that use their
arguments, because this is not a trait, there is not really a way of
enforcing the function signature like there is with the other way.
Documentation has been added for this in the `scalar` macro, so
hopefully will not be an issue.
Also this PR cleans up the Aggregate API by changing the `args` and
`name` functions to constant associated types, as well as adds some
error handling and documentation.
```rust
impl AggFunc for Median {
    type State = Vec<f64>;
    const NAME: &'static str = "median";
    const ARGS: i32 = 1;

    fn step(state: &mut Self::State, args: &[Value]) {
        if let Some(val) = args.first().and_then(Value::to_float) {
            state.push(val);
        }
    }
//.. etc
```

Closes #735
This commit is contained in:
Pekka Enberg
2025-01-19 09:53:31 +02:00
11 changed files with 706 additions and 604 deletions

View File

@@ -1,7 +1,5 @@
use crate::{function::ExternalFunc, Database};
use limbo_ext::{
ExtensionApi, InitAggFunction, ResultCode, ScalarFunction, RESULT_ERROR, RESULT_OK,
};
use limbo_ext::{ExtensionApi, InitAggFunction, ResultCode, ScalarFunction};
pub use limbo_ext::{FinalizeFunction, StepFunction, Value as ExtValue, ValueType as ExtValueType};
use std::{
ffi::{c_char, c_void, CStr},
@@ -17,10 +15,10 @@ unsafe extern "C" fn register_scalar_function(
let c_str = unsafe { CStr::from_ptr(name) };
let name_str = match c_str.to_str() {
Ok(s) => s.to_string(),
Err(_) => return RESULT_ERROR,
Err(_) => return ResultCode::InvalidArgs,
};
if ctx.is_null() {
return RESULT_ERROR;
return ResultCode::Error;
}
let db = unsafe { &*(ctx as *const Database) };
db.register_scalar_function_impl(&name_str, func)
@@ -37,10 +35,10 @@ unsafe extern "C" fn register_aggregate_function(
let c_str = unsafe { CStr::from_ptr(name) };
let name_str = match c_str.to_str() {
Ok(s) => s.to_string(),
Err(_) => return RESULT_ERROR,
Err(_) => return ResultCode::InvalidArgs,
};
if ctx.is_null() {
return RESULT_ERROR;
return ResultCode::Error;
}
let db = unsafe { &*(ctx as *const Database) };
db.register_aggregate_function_impl(&name_str, args, (init_func, step_func, finalize_func))
@@ -52,7 +50,7 @@ impl Database {
name.to_string(),
Rc::new(ExternalFunc::new_scalar(name.to_string(), func)),
);
RESULT_OK
ResultCode::OK
}
fn register_aggregate_function_impl(
@@ -65,7 +63,7 @@ impl Database {
name.to_string(),
Rc::new(ExternalFunc::new_aggregate(name.to_string(), args, func)),
);
RESULT_OK
ResultCode::OK
}
pub fn build_limbo_ext(&self) -> ExtensionApi {

View File

@@ -22,7 +22,7 @@ use fallible_iterator::FallibleIterator;
#[cfg(not(target_family = "wasm"))]
use libloading::{Library, Symbol};
#[cfg(not(target_family = "wasm"))]
use limbo_ext::{ExtensionApi, ExtensionEntryPoint, RESULT_OK};
use limbo_ext::{ExtensionApi, ExtensionEntryPoint};
use log::trace;
use schema::Schema;
use sqlite3_parser::ast;
@@ -179,7 +179,7 @@ impl Database {
};
let api_ptr: *const ExtensionApi = Box::into_raw(api);
let result_code = unsafe { entry(api_ptr) };
if result_code == RESULT_OK {
if result_code.is_ok() {
self.syms.borrow_mut().extensions.push((lib, api_ptr));
Ok(())
} else {

View File

@@ -155,7 +155,7 @@ impl OwnedValue {
OwnedValue::Blob(std::rc::Rc::new(blob))
}
ExtValueType::Error => {
let Some(err) = v.to_text() else {
let Some(err) = v.to_error() else {
return OwnedValue::Null;
};
OwnedValue::Text(LimboText::new(Rc::new(err)))

View File

@@ -7,7 +7,7 @@ like traditional `sqlite3` extensions, but are able to be written in much more e
## Currently supported features
- [ x ] **Scalar Functions**: Create scalar functions using the `ScalarDerive` derive macro and `Scalar` trait.
- [ x ] **Scalar Functions**: Create scalar functions using the `scalar` macro.
- [ x ] **Aggregate Functions**: Define aggregate functions with `AggregateDerive` macro and `AggFunc` trait.
- [] **Virtual tables**: TODO
---
@@ -37,41 +37,35 @@ Extensions can be registered with the `register_extension!` macro:
```rust
register_extension!{
scalars: { Double },
scalars: { double }, // name of your function, if different from attribute name
aggregates: { Percentile },
}
```
### Scalar Example:
```rust
use limbo_ext::{register_extension, Value, ScalarDerive, Scalar};
use limbo_ext::{register_extension, Value, scalar};
/// Annotate each with the ScalarDerive macro, and implement the Scalar trait on your struct
#[derive(ScalarDerive)]
struct Double;
impl Scalar for Double {
fn name(&self) -> &'static str { "double" }
fn call(&self, args: &[Value]) -> Value {
if let Some(arg) = args.first() {
match arg.value_type() {
ValueType::Float => {
let val = arg.to_float().unwrap();
Value::from_float(val * 2.0)
}
ValueType::Integer => {
let val = arg.to_integer().unwrap();
Value::from_integer(val * 2)
}
/// Annotate each with the scalar macro, specifying the name you would like to call it with
/// and optionally, an alias.. e.g. SELECT double(4); or SELECT twice(4);
#[scalar(name = "double", alias = "twice")]
fn double(&self, args: &[Value]) -> Value {
if let Some(arg) = args.first() {
match arg.value_type() {
ValueType::Float => {
let val = arg.to_float().unwrap();
Value::from_float(val * 2.0)
}
ValueType::Integer => {
let val = arg.to_integer().unwrap();
Value::from_integer(val * 2)
}
} else {
Value::null()
}
} else {
Value::null()
}
/// OPTIONAL: 'alias' if you would like to provide an additional name
fn alias(&self) -> &'static str { "twice" }
}
```
### Aggregates Example:
@@ -88,14 +82,11 @@ impl AggFunc for Percentile {
/// Define the name you wish to call your function by.
/// e.g. SELECT percentile(value, 40);
fn name(&self) -> &'static str {
"percentile"
}
const NAME: &str = "percentile";
/// Define the number of expected arguments for your function.
const ARGS: i32 = 2;
/// Define the number of arguments your function takes
fn args(&self) -> i32 {
2
}
/// Define a function called on each row/value in a relevant group/column
fn step(state: &mut Self::State, args: &[Value]) {
let (values, p_value, error) = state;
@@ -127,7 +118,7 @@ impl AggFunc for Percentile {
let (mut values, p_value, error) = state;
if let Some(error) = error {
return Value::error(error);
return Value::custom_error(error);
}
if values.is_empty() {

View File

@@ -1,9 +1,7 @@
pub use limbo_macros::{register_extension, AggregateDerive, ScalarDerive};
mod types;
pub use limbo_macros::{register_extension, scalar, AggregateDerive};
use std::os::raw::{c_char, c_void};
pub type ResultCode = i32;
pub const RESULT_OK: ResultCode = 0;
pub const RESULT_ERROR: ResultCode = 1;
pub use types::{ResultCode, Value, ValueType};
#[repr(C)]
pub struct ExtensionApi {
@@ -34,10 +32,6 @@ pub type FinalizeFunction = unsafe extern "C" fn(ctx: *mut AggCtx) -> Value;
pub trait Scalar {
fn call(&self, args: &[Value]) -> Value;
fn name(&self) -> &'static str;
fn alias(&self) -> Option<&'static str> {
None
}
}
#[repr(C)]
@@ -47,268 +41,9 @@ pub struct AggCtx {
pub trait AggFunc {
type State: Default;
const NAME: &'static str;
const ARGS: i32;
fn args(&self) -> i32 {
1
}
fn name(&self) -> &'static str;
fn step(state: &mut Self::State, args: &[Value]);
fn finalize(state: Self::State) -> Value;
}
#[repr(C)]
#[derive(PartialEq, Eq, Clone, Copy)]
pub enum ValueType {
Null,
Integer,
Float,
Text,
Blob,
Error,
}
#[repr(C)]
pub struct Value {
value_type: ValueType,
value: *mut c_void,
}
impl std::fmt::Debug for Value {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self.value_type {
ValueType::Null => write!(f, "Value {{ Null }}"),
ValueType::Integer => write!(f, "Value {{ Integer: {} }}", unsafe {
*(self.value as *const i64)
}),
ValueType::Float => write!(f, "Value {{ Float: {} }}", unsafe {
*(self.value as *const f64)
}),
ValueType::Text => write!(f, "Value {{ Text: {:?} }}", unsafe {
&*(self.value as *const TextValue)
}),
ValueType::Blob => write!(f, "Value {{ Blob: {:?} }}", unsafe {
&*(self.value as *const Blob)
}),
ValueType::Error => write!(f, "Value {{ Error: {:?} }}", unsafe {
&*(self.value as *const TextValue)
}),
}
}
}
#[repr(C)]
pub struct TextValue {
text: *const u8,
len: u32,
}
impl std::fmt::Debug for TextValue {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"TextValue {{ text: {:?}, len: {} }}",
self.text, self.len
)
}
}
impl Default for TextValue {
fn default() -> Self {
Self {
text: std::ptr::null(),
len: 0,
}
}
}
impl TextValue {
pub(crate) fn new(text: *const u8, len: usize) -> Self {
Self {
text,
len: len as u32,
}
}
fn as_str(&self) -> &str {
if self.text.is_null() {
return "";
}
unsafe {
std::str::from_utf8_unchecked(std::slice::from_raw_parts(self.text, self.len as usize))
}
}
}
#[repr(C)]
pub struct Blob {
data: *const u8,
size: u64,
}
impl std::fmt::Debug for Blob {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Blob {{ data: {:?}, size: {} }}", self.data, self.size)
}
}
impl Blob {
pub fn new(data: *const u8, size: u64) -> Self {
Self { data, size }
}
}
impl Value {
pub fn null() -> Self {
Self {
value_type: ValueType::Null,
value: std::ptr::null_mut(),
}
}
pub fn value_type(&self) -> ValueType {
self.value_type
}
pub fn to_float(&self) -> Option<f64> {
if self.value.is_null() {
return None;
}
match self.value_type {
ValueType::Float => Some(unsafe { *(self.value as *const f64) }),
ValueType::Integer => Some(unsafe { *(self.value as *const i64) as f64 }),
ValueType::Text => {
let txt = unsafe { &*(self.value as *const TextValue) };
txt.as_str().parse().ok()
}
_ => None,
}
}
pub fn to_text(&self) -> Option<String> {
if self.value_type != ValueType::Text {
return None;
}
if self.value.is_null() {
return None;
}
let txt = unsafe { &*(self.value as *const TextValue) };
Some(String::from(txt.as_str()))
}
pub fn to_blob(&self) -> Option<Vec<u8>> {
if self.value_type != ValueType::Blob {
return None;
}
if self.value.is_null() {
return None;
}
let blob = unsafe { &*(self.value as *const Blob) };
let slice = unsafe { std::slice::from_raw_parts(blob.data, blob.size as usize) };
Some(slice.to_vec())
}
pub fn to_integer(&self) -> Option<i64> {
if self.value.is_null() {
return None;
}
match self.value_type() {
ValueType::Integer => Some(unsafe { *(self.value as *const i64) }),
ValueType::Float => Some(unsafe { *(self.value as *const f64) } as i64),
ValueType::Text => {
let txt = unsafe { &*(self.value as *const TextValue) };
txt.as_str().parse().ok()
}
_ => None,
}
}
pub fn to_error(&self) -> Option<String> {
if self.value_type != ValueType::Error {
return None;
}
if self.value.is_null() {
return None;
}
let txt = unsafe { &*(self.value as *const TextValue) };
Some(String::from(txt.as_str()))
}
pub fn from_integer(value: i64) -> Self {
let boxed = Box::new(value);
Self {
value_type: ValueType::Integer,
value: Box::into_raw(boxed) as *mut c_void,
}
}
pub fn from_float(value: f64) -> Self {
let boxed = Box::new(value);
Self {
value_type: ValueType::Float,
value: Box::into_raw(boxed) as *mut c_void,
}
}
pub fn from_text(s: String) -> Self {
let buffer = s.into_boxed_str();
let ptr = buffer.as_ptr();
let len = buffer.len();
std::mem::forget(buffer);
let text_value = TextValue::new(ptr, len);
let text_box = Box::new(text_value);
Self {
value_type: ValueType::Text,
value: Box::into_raw(text_box) as *mut c_void,
}
}
pub fn error(s: String) -> Self {
let buffer = s.into_boxed_str();
let ptr = buffer.as_ptr();
let len = buffer.len();
std::mem::forget(buffer);
let text_value = TextValue::new(ptr, len);
let text_box = Box::new(text_value);
Self {
value_type: ValueType::Error,
value: Box::into_raw(text_box) as *mut c_void,
}
}
pub fn from_blob(value: Vec<u8>) -> Self {
let boxed = Box::new(Blob::new(value.as_ptr(), value.len() as u64));
std::mem::forget(value);
Self {
value_type: ValueType::Blob,
value: Box::into_raw(boxed) as *mut c_void,
}
}
/// # Safety
/// consumes the value while freeing the underlying memory with null check.
/// however this does assume that the type was properly constructed with
/// the appropriate value_type and value.
pub unsafe fn free(self) {
if self.value.is_null() {
return;
}
match self.value_type {
ValueType::Integer => {
let _ = Box::from_raw(self.value as *mut i64);
}
ValueType::Float => {
let _ = Box::from_raw(self.value as *mut f64);
}
ValueType::Text => {
let _ = Box::from_raw(self.value as *mut TextValue);
}
ValueType::Blob => {
let _ = Box::from_raw(self.value as *mut Blob);
}
ValueType::Error => {
let _ = Box::from_raw(self.value as *mut TextValue);
}
ValueType::Null => {}
}
}
}

View File

@@ -0,0 +1,358 @@
use std::{fmt::Display, os::raw::c_void};
/// Error type is of type ExtError which can be
/// either a user defined error or an error code
#[repr(C)]
pub enum ResultCode {
OK = 0,
Error = 1,
InvalidArgs = 2,
Unknown = 3,
OoM = 4,
Corrupt = 5,
NotFound = 6,
AlreadyExists = 7,
PermissionDenied = 8,
Aborted = 9,
OutOfRange = 10,
Unimplemented = 11,
Internal = 12,
Unavailable = 13,
}
impl ResultCode {
pub fn is_ok(&self) -> bool {
matches!(self, ResultCode::OK)
}
}
impl Display for ResultCode {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ResultCode::OK => write!(f, "OK"),
ResultCode::Error => write!(f, "Error"),
ResultCode::InvalidArgs => write!(f, "InvalidArgs"),
ResultCode::Unknown => write!(f, "Unknown"),
ResultCode::OoM => write!(f, "Out of Memory"),
ResultCode::Corrupt => write!(f, "Corrupt"),
ResultCode::NotFound => write!(f, "Not Found"),
ResultCode::AlreadyExists => write!(f, "Already Exists"),
ResultCode::PermissionDenied => write!(f, "Permission Denied"),
ResultCode::Aborted => write!(f, "Aborted"),
ResultCode::OutOfRange => write!(f, "Out of Range"),
ResultCode::Unimplemented => write!(f, "Unimplemented"),
ResultCode::Internal => write!(f, "Internal Error"),
ResultCode::Unavailable => write!(f, "Unavailable"),
}
}
}
#[repr(C)]
#[derive(PartialEq, Debug, Eq, Clone, Copy)]
pub enum ValueType {
Null,
Integer,
Float,
Text,
Blob,
Error,
}
#[repr(C)]
pub struct Value {
value_type: ValueType,
value: *mut c_void,
}
impl std::fmt::Debug for Value {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if self.value.is_null() {
return write!(f, "{:?}: Null", self.value_type);
}
match self.value_type {
ValueType::Null => write!(f, "Value {{ Null }}"),
ValueType::Integer => write!(f, "Value {{ Integer: {} }}", unsafe {
*(self.value as *const i64)
}),
ValueType::Float => write!(f, "Value {{ Float: {} }}", unsafe {
*(self.value as *const f64)
}),
ValueType::Text => write!(f, "Value {{ Text: {:?} }}", unsafe {
&*(self.value as *const TextValue)
}),
ValueType::Blob => write!(f, "Value {{ Blob: {:?} }}", unsafe {
&*(self.value as *const Blob)
}),
ValueType::Error => write!(f, "Value {{ Error: {:?} }}", unsafe {
&*(self.value as *const TextValue)
}),
}
}
}
#[repr(C)]
pub struct TextValue {
text: *const u8,
len: u32,
}
impl std::fmt::Debug for TextValue {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"TextValue {{ text: {:?}, len: {} }}",
self.text, self.len
)
}
}
impl Default for TextValue {
fn default() -> Self {
Self {
text: std::ptr::null(),
len: 0,
}
}
}
impl TextValue {
pub(crate) fn new(text: *const u8, len: usize) -> Self {
Self {
text,
len: len as u32,
}
}
fn as_str(&self) -> &str {
if self.text.is_null() {
return "";
}
unsafe {
std::str::from_utf8_unchecked(std::slice::from_raw_parts(self.text, self.len as usize))
}
}
}
#[repr(C)]
pub struct Blob {
data: *const u8,
size: u64,
}
impl std::fmt::Debug for Blob {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Blob {{ data: {:?}, size: {} }}", self.data, self.size)
}
}
impl Blob {
pub fn new(data: *const u8, size: u64) -> Self {
Self { data, size }
}
}
impl Value {
/// Creates a new Value with type Null
pub fn null() -> Self {
Self {
value_type: ValueType::Null,
value: std::ptr::null_mut(),
}
}
/// Returns the value type of the Value
pub fn value_type(&self) -> ValueType {
self.value_type
}
/// Returns the float value if the Value is the proper type
pub fn to_float(&self) -> Option<f64> {
if self.value.is_null() {
return None;
}
match self.value_type {
ValueType::Float => Some(unsafe { *(self.value as *const f64) }),
ValueType::Integer => Some(unsafe { *(self.value as *const i64) as f64 }),
ValueType::Text => {
let txt = unsafe { &*(self.value as *const TextValue) };
txt.as_str().parse().ok()
}
_ => None,
}
}
/// Returns the text value if the Value is the proper type
pub fn to_text(&self) -> Option<String> {
if self.value_type != ValueType::Text {
return None;
}
if self.value.is_null() {
return None;
}
let txt = unsafe { &*(self.value as *const TextValue) };
Some(String::from(txt.as_str()))
}
/// Returns the blob value if the Value is the proper type
pub fn to_blob(&self) -> Option<Vec<u8>> {
if self.value_type != ValueType::Blob {
return None;
}
if self.value.is_null() {
return None;
}
let blob = unsafe { &*(self.value as *const Blob) };
let slice = unsafe { std::slice::from_raw_parts(blob.data, blob.size as usize) };
Some(slice.to_vec())
}
/// Returns the integer value if the Value is the proper type
pub fn to_integer(&self) -> Option<i64> {
if self.value.is_null() {
return None;
}
match self.value_type() {
ValueType::Integer => Some(unsafe { *(self.value as *const i64) }),
ValueType::Float => Some(unsafe { *(self.value as *const f64) } as i64),
ValueType::Text => {
let txt = unsafe { &*(self.value as *const TextValue) };
txt.as_str().parse().ok()
}
_ => None,
}
}
/// Returns the error message if the value is an error
pub fn to_error(&self) -> Option<String> {
if self.value_type != ValueType::Error {
return None;
}
if self.value.is_null() {
return None;
}
let err = unsafe { &*(self.value as *const ExtError) };
match &err.error_type {
ErrorType::User => {
if err.message.is_null() {
return None;
}
let txt = unsafe { &*(err.message as *const TextValue) };
Some(txt.as_str().to_string())
}
ErrorType::ErrCode { code } => Some(format!("{}", code)),
}
}
/// Creates a new integer Value from an i64
pub fn from_integer(value: i64) -> Self {
let boxed = Box::new(value);
Self {
value_type: ValueType::Integer,
value: Box::into_raw(boxed) as *mut c_void,
}
}
/// Creates a new float Value from an f64
pub fn from_float(value: f64) -> Self {
let boxed = Box::new(value);
Self {
value_type: ValueType::Float,
value: Box::into_raw(boxed) as *mut c_void,
}
}
/// Creates a new text Value from a String
pub fn from_text(s: String) -> Self {
let buffer = s.into_boxed_str();
let ptr = buffer.as_ptr();
let len = buffer.len();
std::mem::forget(buffer);
let text_value = TextValue::new(ptr, len);
let text_box = Box::new(text_value);
Self {
value_type: ValueType::Text,
value: Box::into_raw(text_box) as *mut c_void,
}
}
/// Creates a new error Value from a ResultCode
pub fn error(err: ResultCode) -> Self {
let error = ExtError {
error_type: ErrorType::ErrCode { code: err },
message: std::ptr::null_mut(),
};
Self {
value_type: ValueType::Error,
value: Box::into_raw(Box::new(error)) as *mut c_void,
}
}
/// Create a new user defined error Value with a message
pub fn custom_error(s: String) -> Self {
let buffer = s.into_boxed_str();
let ptr = buffer.as_ptr();
let len = buffer.len();
std::mem::forget(buffer);
let text_value = TextValue::new(ptr, len);
let text_box = Box::new(text_value);
let error = ExtError {
error_type: ErrorType::User,
message: Box::into_raw(text_box) as *mut c_void,
};
Self {
value_type: ValueType::Error,
value: Box::into_raw(Box::new(error)) as *mut c_void,
}
}
/// Creates a new blob Value from a Vec<u8>
pub fn from_blob(value: Vec<u8>) -> Self {
let boxed = Box::new(Blob::new(value.as_ptr(), value.len() as u64));
std::mem::forget(value);
Self {
value_type: ValueType::Blob,
value: Box::into_raw(boxed) as *mut c_void,
}
}
/// # Safety
/// consumes the value while freeing the underlying memory with null check.
/// however this does assume that the type was properly constructed with
/// the appropriate value_type and value.
pub unsafe fn free(self) {
if self.value.is_null() {
return;
}
match self.value_type {
ValueType::Integer => {
let _ = Box::from_raw(self.value as *mut i64);
}
ValueType::Float => {
let _ = Box::from_raw(self.value as *mut f64);
}
ValueType::Text => {
let _ = Box::from_raw(self.value as *mut TextValue);
}
ValueType::Blob => {
let _ = Box::from_raw(self.value as *mut Blob);
}
ValueType::Error => {
let _ = Box::from_raw(self.value as *mut ExtError);
}
ValueType::Null => {}
}
}
}
#[repr(C)]
pub struct ExtError {
pub error_type: ErrorType,
pub message: *mut std::ffi::c_void,
}
#[repr(C)]
pub enum ErrorType {
User,
/// User type has a user provided message
ErrCode {
code: ResultCode,
},
}

View File

@@ -1,4 +1,4 @@
use limbo_ext::{register_extension, AggFunc, AggregateDerive, Value};
use limbo_ext::{register_extension, AggFunc, AggregateDerive, ResultCode, Value};
register_extension! {
aggregates: { Median, Percentile, PercentileCont, PercentileDisc }
@@ -9,12 +9,8 @@ struct Median;
impl AggFunc for Median {
type State = Vec<f64>;
fn name(&self) -> &'static str {
"median"
}
fn args(&self) -> i32 {
1
}
const NAME: &'static str = "median";
const ARGS: i32 = 1;
fn step(state: &mut Self::State, args: &[Value]) {
if let Some(val) = args.first().and_then(Value::to_float) {
@@ -45,15 +41,10 @@ impl AggFunc for Median {
struct Percentile;
impl AggFunc for Percentile {
type State = (Vec<f64>, Option<f64>, Option<String>);
type State = (Vec<f64>, Option<f64>, Option<()>);
fn name(&self) -> &'static str {
"percentile"
}
fn args(&self) -> i32 {
2
}
const NAME: &'static str = "percentile";
const ARGS: i32 = 2;
fn step(state: &mut Self::State, args: &[Value]) {
let (values, p_value, err_value) = state;
@@ -62,13 +53,13 @@ impl AggFunc for Percentile {
args.get(1).and_then(Value::to_float),
) {
if !(0.0..=100.0).contains(&p) {
err_value.get_or_insert("percentile value out of range".to_string());
err_value.get_or_insert(());
return;
}
if let Some(existing_p) = *p_value {
if (existing_p - p).abs() >= 0.001 {
err_value.get_or_insert("percentile value out of range".to_string());
err_value.get_or_insert(());
return;
}
} else {
@@ -83,8 +74,8 @@ impl AggFunc for Percentile {
if values.is_empty() {
return Value::null();
}
if let Some(err_value) = err_value {
return Value::error(err_value.clone());
if err_value.is_some() {
return Value::error(ResultCode::Error);
}
if values.len() == 1 {
return Value::from_float(values[0]);
@@ -110,15 +101,10 @@ impl AggFunc for Percentile {
struct PercentileCont;
impl AggFunc for PercentileCont {
type State = (Vec<f64>, Option<f64>, Option<String>);
type State = (Vec<f64>, Option<f64>, Option<()>);
fn name(&self) -> &'static str {
"percentile_cont"
}
fn args(&self) -> i32 {
2
}
const NAME: &'static str = "percentile_cont";
const ARGS: i32 = 2;
fn step(state: &mut Self::State, args: &[Value]) {
let (values, p_value, err_state) = state;
@@ -127,13 +113,13 @@ impl AggFunc for PercentileCont {
args.get(1).and_then(Value::to_float),
) {
if !(0.0..=1.0).contains(&p) {
err_state.get_or_insert("percentile value out of range".to_string());
err_state.get_or_insert(());
return;
}
if let Some(existing_p) = *p_value {
if (existing_p - p).abs() >= 0.001 {
err_state.get_or_insert("percentile value out of range".to_string());
err_state.get_or_insert(());
return;
}
} else {
@@ -148,8 +134,8 @@ impl AggFunc for PercentileCont {
if values.is_empty() {
return Value::null();
}
if let Some(err_state) = err_state {
return Value::error(err_state.clone());
if err_state.is_some() {
return Value::error(ResultCode::Error);
}
if values.len() == 1 {
return Value::from_float(values[0]);
@@ -175,15 +161,10 @@ impl AggFunc for PercentileCont {
struct PercentileDisc;
impl AggFunc for PercentileDisc {
type State = (Vec<f64>, Option<f64>, Option<String>);
type State = (Vec<f64>, Option<f64>, Option<()>);
fn name(&self) -> &'static str {
"percentile_disc"
}
fn args(&self) -> i32 {
2
}
const NAME: &'static str = "percentile_disc";
const ARGS: i32 = 2;
fn step(state: &mut Self::State, args: &[Value]) {
Percentile::step(state, args);
@@ -194,8 +175,8 @@ impl AggFunc for PercentileDisc {
if values.is_empty() {
return Value::null();
}
if let Some(err_value) = err_value {
return Value::error(err_value.clone());
if err_value.is_some() {
return Value::error(ResultCode::Error);
}
let p = p_value.unwrap();

View File

@@ -1,20 +1,13 @@
use limbo_ext::{register_extension, Scalar, ScalarDerive, Value, ValueType};
use limbo_ext::{register_extension, scalar, Value, ValueType};
use regex::Regex;
register_extension! {
scalars: { Regexp, RegexpLike, RegexpSubstr }
scalars: { regexp, regexp_like, regexp_substr }
}
#[derive(ScalarDerive)]
struct Regexp;
impl Scalar for Regexp {
fn name(&self) -> &'static str {
"regexp"
}
fn call(&self, args: &[Value]) -> Value {
regex(&args[0], &args[1])
}
#[scalar(name = "regexp")]
fn regexp(args: &[Value]) -> Value {
regex(&args[0], &args[1])
}
fn regex(pattern: &Value, haystack: &Value) -> Value {
@@ -36,44 +29,30 @@ fn regex(pattern: &Value, haystack: &Value) -> Value {
}
}
#[derive(ScalarDerive)]
struct RegexpLike;
impl Scalar for RegexpLike {
fn name(&self) -> &'static str {
"regexp_like"
}
fn call(&self, args: &[Value]) -> Value {
regex(&args[1], &args[0])
}
#[scalar(name = "regexp_like")]
fn regexp_like(args: &[Value]) -> Value {
regex(&args[1], &args[0])
}
#[derive(ScalarDerive)]
struct RegexpSubstr;
impl Scalar for RegexpSubstr {
fn name(&self) -> &'static str {
"regexp_substr"
}
fn call(&self, args: &[Value]) -> Value {
match (args[0].value_type(), args[1].value_type()) {
(ValueType::Text, ValueType::Text) => {
let Some(haystack) = &args[0].to_text() else {
return Value::null();
};
let Some(pattern) = &args[1].to_text() else {
return Value::null();
};
let re = match Regex::new(pattern) {
Ok(re) => re,
Err(_) => return Value::null(),
};
match re.find(haystack) {
Some(mat) => Value::from_text(mat.as_str().to_string()),
None => Value::null(),
}
#[scalar(name = "regexp_substr")]
fn regexp_substr(&self, args: &[Value]) -> Value {
match (args[0].value_type(), args[1].value_type()) {
(ValueType::Text, ValueType::Text) => {
let Some(haystack) = &args[0].to_text() else {
return Value::null();
};
let Some(pattern) = &args[1].to_text() else {
return Value::null();
};
let re = match Regex::new(pattern) {
Ok(re) => re,
Err(_) => return Value::null(),
};
match re.find(haystack) {
Some(mat) => Value::from_text(mat.as_str().to_string()),
None => Value::null(),
}
_ => Value::null(),
}
_ => Value::null(),
}
}

View File

@@ -1,177 +1,126 @@
use limbo_ext::{register_extension, Scalar, ScalarDerive, Value, ValueType};
use limbo_ext::{register_extension, scalar, Value, ValueType};
register_extension! {
scalars: { Uuid4Str, Uuid4Blob, Uuid7Str, Uuid7Blob, ExecTsFromUuid7, UuidStr, UuidBlob }
scalars: { uuid4_str, uuid4_blob, uuid7_str, uuid7, uuid7_ts, uuid_str, uuid_blob }
}
#[derive(ScalarDerive)]
struct Uuid4Str;
impl Scalar for Uuid4Str {
fn name(&self) -> &'static str {
"uuid4_str"
}
fn alias(&self) -> Option<&'static str> {
Some("gen_random_uuid")
}
fn call(&self, _args: &[Value]) -> Value {
let uuid = uuid::Uuid::new_v4().to_string();
Value::from_text(uuid)
}
#[scalar(name = "uuid4_str", alias = "gen_random_uuid")]
fn uuid4_str(_args: &[Value]) -> Value {
let uuid = uuid::Uuid::new_v4().to_string();
Value::from_text(uuid)
}
#[derive(ScalarDerive)]
struct Uuid4Blob;
impl Scalar for Uuid4Blob {
fn name(&self) -> &'static str {
"uuid4"
}
fn call(&self, _args: &[Value]) -> Value {
let uuid = uuid::Uuid::new_v4();
let bytes = uuid.as_bytes();
Value::from_blob(bytes.to_vec())
}
#[scalar(name = "uuid4")]
fn uuid4_blob(_args: &[Value]) -> Value {
let uuid = uuid::Uuid::new_v4();
let bytes = uuid.as_bytes();
Value::from_blob(bytes.to_vec())
}
#[derive(ScalarDerive)]
struct Uuid7Str;
impl Scalar for Uuid7Str {
fn name(&self) -> &'static str {
"uuid7_str"
}
fn call(&self, args: &[Value]) -> Value {
let timestamp = if args.is_empty() {
let ctx = uuid::ContextV7::new();
uuid::Timestamp::now(ctx)
} else {
match args[0].value_type() {
ValueType::Integer => {
let ctx = uuid::ContextV7::new();
let Some(int) = args[0].to_integer() else {
return Value::null();
};
uuid::Timestamp::from_unix(ctx, int as u64, 0)
}
ValueType::Text => {
let Some(text) = args[0].to_text() else {
return Value::null();
};
match text.parse::<i64>() {
Ok(unix) => {
if unix <= 0 {
return Value::null();
}
uuid::Timestamp::from_unix(uuid::ContextV7::new(), unix as u64, 0)
}
Err(_) => return Value::null(),
}
}
_ => return Value::null(),
}
};
let uuid = uuid::Uuid::new_v7(timestamp);
Value::from_text(uuid.to_string())
}
}
#[derive(ScalarDerive)]
struct Uuid7Blob;
impl Scalar for Uuid7Blob {
fn name(&self) -> &'static str {
"uuid7"
}
fn call(&self, args: &[Value]) -> Value {
let timestamp = if args.is_empty() {
let ctx = uuid::ContextV7::new();
uuid::Timestamp::now(ctx)
} else {
match args[0].value_type() {
ValueType::Integer => {
let ctx = uuid::ContextV7::new();
let Some(int) = args[0].to_integer() else {
return Value::null();
};
uuid::Timestamp::from_unix(ctx, int as u64, 0)
}
_ => return Value::null(),
}
};
let uuid = uuid::Uuid::new_v7(timestamp);
let bytes = uuid.as_bytes();
Value::from_blob(bytes.to_vec())
}
}
#[derive(ScalarDerive)]
struct ExecTsFromUuid7;
impl Scalar for ExecTsFromUuid7 {
fn name(&self) -> &'static str {
"uuid7_timestamp_ms"
}
fn call(&self, args: &[Value]) -> Value {
#[scalar(name = "uuid7_str")]
fn uuid7_str(args: &[Value]) -> Value {
let timestamp = if args.is_empty() {
let ctx = uuid::ContextV7::new();
uuid::Timestamp::now(ctx)
} else {
match args[0].value_type() {
ValueType::Blob => {
let Some(blob) = &args[0].to_blob() else {
ValueType::Integer => {
let ctx = uuid::ContextV7::new();
let Some(int) = args[0].to_integer() else {
return Value::null();
};
let uuid = uuid::Uuid::from_slice(blob.as_slice()).unwrap();
let unix = uuid_to_unix(uuid.as_bytes());
Value::from_integer(unix as i64)
uuid::Timestamp::from_unix(ctx, int as u64, 0)
}
ValueType::Text => {
let Some(text) = args[0].to_text() else {
return Value::null();
};
let Ok(uuid) = uuid::Uuid::parse_str(&text) else {
match text.parse::<i64>() {
Ok(unix) => {
if unix <= 0 {
return Value::null();
}
uuid::Timestamp::from_unix(uuid::ContextV7::new(), unix as u64, 0)
}
Err(_) => return Value::null(),
}
}
_ => return Value::null(),
}
};
let uuid = uuid::Uuid::new_v7(timestamp);
Value::from_text(uuid.to_string())
}
#[scalar(name = "uuid7")]
fn uuid7(&self, args: &[Value]) -> Value {
let timestamp = if args.is_empty() {
let ctx = uuid::ContextV7::new();
uuid::Timestamp::now(ctx)
} else {
match args[0].value_type() {
ValueType::Integer => {
let ctx = uuid::ContextV7::new();
let Some(int) = args[0].to_integer() else {
return Value::null();
};
let unix = uuid_to_unix(uuid.as_bytes());
Value::from_integer(unix as i64)
uuid::Timestamp::from_unix(ctx, int as u64, 0)
}
_ => Value::null(),
_ => return Value::null(),
}
};
let uuid = uuid::Uuid::new_v7(timestamp);
let bytes = uuid.as_bytes();
Value::from_blob(bytes.to_vec())
}
#[scalar(name = "uuid7_timestamp_ms")]
fn uuid7_ts(args: &[Value]) -> Value {
match args[0].value_type() {
ValueType::Blob => {
let Some(blob) = &args[0].to_blob() else {
return Value::null();
};
let uuid = uuid::Uuid::from_slice(blob.as_slice()).unwrap();
let unix = uuid_to_unix(uuid.as_bytes());
Value::from_integer(unix as i64)
}
ValueType::Text => {
let Some(text) = args[0].to_text() else {
return Value::null();
};
let Ok(uuid) = uuid::Uuid::parse_str(&text) else {
return Value::null();
};
let unix = uuid_to_unix(uuid.as_bytes());
Value::from_integer(unix as i64)
}
_ => Value::null(),
}
}
#[derive(ScalarDerive)]
struct UuidStr;
impl Scalar for UuidStr {
fn name(&self) -> &'static str {
"uuid_str"
}
fn call(&self, args: &[Value]) -> Value {
let Some(blob) = args[0].to_blob() else {
return Value::null();
};
let parsed = uuid::Uuid::from_slice(blob.as_slice())
.ok()
.map(|u| u.to_string());
match parsed {
Some(s) => Value::from_text(s),
None => Value::null(),
}
#[scalar(name = "uuid_str")]
fn uuid_str(args: &[Value]) -> Value {
let Some(blob) = args[0].to_blob() else {
return Value::null();
};
let parsed = uuid::Uuid::from_slice(blob.as_slice())
.ok()
.map(|u| u.to_string());
match parsed {
Some(s) => Value::from_text(s),
None => Value::null(),
}
}
#[derive(ScalarDerive)]
struct UuidBlob;
impl Scalar for UuidBlob {
fn name(&self) -> &'static str {
"uuid_blob"
}
fn call(&self, args: &[Value]) -> Value {
let Some(text) = args[0].to_text() else {
return Value::null();
};
match uuid::Uuid::parse_str(&text) {
Ok(uuid) => Value::from_blob(uuid.as_bytes().to_vec()),
Err(_) => Value::null(),
}
#[scalar(name = "uuid_blob")]
fn uuid_blob(&self, args: &[Value]) -> Value {
let Some(text) = args[0].to_text() else {
return Value::null();
};
match uuid::Uuid::parse_str(&text) {
Ok(uuid) => Value::from_blob(uuid.as_bytes().to_vec()),
Err(_) => Value::null(),
}
}

View File

@@ -1,5 +1,7 @@
use syn::parse::ParseStream;
use syn::punctuated::Punctuated;
use syn::{Ident, Token};
use syn::token::Eq;
use syn::{Ident, LitStr, Token};
pub(crate) struct RegisterExtensionInput {
pub aggregates: Vec<Ident>,
@@ -44,3 +46,39 @@ impl syn::parse::Parse for RegisterExtensionInput {
})
}
}
pub(crate) struct ScalarInfo {
pub name: String,
pub alias: Option<String>,
}
impl ScalarInfo {
pub fn new(name: String, alias: Option<String>) -> Self {
Self { name, alias }
}
}
impl syn::parse::Parse for ScalarInfo {
fn parse(input: ParseStream) -> syn::parse::Result<Self> {
let mut name = None;
let mut alias = None;
while !input.is_empty() {
if let Ok(ident) = input.parse::<Ident>() {
if ident.to_string().as_str() == "name" {
let _ = input.parse::<Eq>();
name = Some(input.parse::<LitStr>()?);
} else if ident.to_string().as_str() == "alias" {
let _ = input.parse::<Eq>();
alias = Some(input.parse::<LitStr>()?);
}
}
if input.peek(Token![,]) {
input.parse::<Token![,]>()?;
}
}
let Some(name) = name else {
return Err(input.error("Expected name"));
};
Ok(Self::new(name.value(), alias.map(|i| i.value())))
}
}

View File

@@ -1,7 +1,7 @@
mod args;
use args::RegisterExtensionInput;
use args::{RegisterExtensionInput, ScalarInfo};
use quote::{format_ident, quote};
use syn::{parse_macro_input, DeriveInput};
use syn::{parse_macro_input, DeriveInput, ItemFn};
extern crate proc_macro;
use proc_macro::{token_stream::IntoIter, Group, TokenStream, TokenTree};
use std::collections::HashMap;
@@ -138,71 +138,113 @@ fn generate_get_description(
enum_impl.parse().unwrap()
}
#[proc_macro_derive(ScalarDerive)]
pub fn derive_scalar(input: TokenStream) -> TokenStream {
let ast = parse_macro_input!(input as DeriveInput);
let struct_name = &ast.ident;
let register_fn_name = format_ident!("register_{}", struct_name);
let exec_fn_name = format_ident!("{}_exec", struct_name);
let alias_check = quote! {
if let Some(alias) = scalar.alias() {
let alias_c_name = std::ffi::CString::new(alias).unwrap();
/// Declare a scalar function for your extension. This requires the name:
/// #[scalar(name = "example")] of what you wish to call your function with.
/// Your function __must__ use the signature: `fn (args: &[Value]) -> Value`
/// with proper spelling.
/// ```ignore
/// use limbo_ext::{scalar, Value};
/// #[scalar(name = "double", alias = "twice")] // you can provide an <optional> alias
/// fn double(args: &[Value]) -> Value {
/// match arg.value_type() {
/// ValueType::Float => {
/// let val = arg.to_float().unwrap();
/// Value::from_float(val * 2.0)
/// }
/// ValueType::Integer => {
/// let val = arg.to_integer().unwrap();
/// Value::from_integer(val * 2)
/// }
/// }
/// } else {
/// Value::null()
/// }
/// }
/// ```
#[proc_macro_attribute]
pub fn scalar(attr: TokenStream, input: TokenStream) -> TokenStream {
let ast = parse_macro_input!(input as ItemFn);
let fn_name = &ast.sig.ident;
let scalar_info = parse_macro_input!(attr as ScalarInfo);
let name = &scalar_info.name;
let register_fn_name = format_ident!("register_{}", fn_name);
let fn_body = &ast.block;
let alias_check = if let Some(alias) = &scalar_info.alias {
quote! {
let Ok(alias_c_name) = std::ffi::CString::new(#alias) else {
return ::limbo_ext::ResultCode::Error;
};
(api.register_scalar_function)(
api.ctx,
alias_c_name.as_ptr(),
#exec_fn_name,
#fn_name,
);
}
} else {
quote! {}
};
let expanded = quote! {
impl #struct_name {
#[no_mangle]
pub unsafe extern "C" fn #register_fn_name(
api: *const ::limbo_ext::ExtensionApi
) -> ::limbo_ext::ResultCode {
if api.is_null() {
return ::limbo_ext::RESULT_ERROR;
}
let api = unsafe { &*api };
let scalar = #struct_name;
let name = scalar.name();
let c_name = std::ffi::CString::new(name).unwrap();
(api.register_scalar_function)(
api.ctx,
c_name.as_ptr(),
#exec_fn_name,
);
#alias_check
::limbo_ext::RESULT_OK
#[no_mangle]
pub unsafe extern "C" fn #register_fn_name(
api: *const ::limbo_ext::ExtensionApi
) -> ::limbo_ext::ResultCode {
if api.is_null() {
return ::limbo_ext::ResultCode::Error;
}
let api = unsafe { &*api };
let Ok(c_name) = std::ffi::CString::new(#name) else {
return ::limbo_ext::ResultCode::Error;
};
(api.register_scalar_function)(
api.ctx,
c_name.as_ptr(),
#fn_name,
);
#alias_check
::limbo_ext::ResultCode::OK
}
#[no_mangle]
pub unsafe extern "C" fn #exec_fn_name(
pub unsafe extern "C" fn #fn_name(
argc: i32,
argv: *const ::limbo_ext::Value
) -> ::limbo_ext::Value {
let scalar = #struct_name;
let args_slice = if argv.is_null() || argc <= 0 {
let args = if argv.is_null() || argc <= 0 {
&[]
} else {
unsafe { std::slice::from_raw_parts(argv, argc as usize) }
};
scalar.call(args_slice)
#fn_body
}
};
TokenStream::from(expanded)
}
/// Define an aggregate function for your extension by deriving
/// AggregateDerive on a struct that implements the AggFunc trait.
/// ```ignore
/// use limbo_ext::{register_extension, Value, AggregateDerive, AggFunc};
///
///#[derive(AggregateDerive)]
///struct SumPlusOne;
///
///impl AggFunc for SumPlusOne {
/// type State = i64;
/// const NAME: &'static str = "sum_plus_one";
/// const ARGS: i32 = 1;
/// fn step(state: &mut Self::State, args: &[Value]) {
/// let Some(val) = args[0].to_integer() else {
/// return;
/// };
/// *state += val;
/// }
/// fn finalize(state: Self::State) -> Value {
/// Value::from_integer(state + 1)
/// }
///}
/// ```
#[proc_macro_derive(AggregateDerive)]
pub fn derive_agg_func(input: TokenStream) -> TokenStream {
let ast = parse_macro_input!(input as DeriveInput);
@@ -254,21 +296,20 @@ pub fn derive_agg_func(input: TokenStream) -> TokenStream {
api: *const ::limbo_ext::ExtensionApi
) -> ::limbo_ext::ResultCode {
if api.is_null() {
return ::limbo_ext::RESULT_ERROR;
return ::limbo_ext::ResultCode::Error;
}
let api = &*api;
let agg = #struct_name;
let name_str = agg.name();
let name_str = #struct_name::NAME;
let c_name = match std::ffi::CString::new(name_str) {
Ok(cname) => cname,
Err(_) => return ::limbo_ext::RESULT_ERROR,
Err(_) => return ::limbo_ext::ResultCode::Error,
};
(api.register_aggregate_function)(
api.ctx,
c_name.as_ptr(),
agg.args(),
#struct_name::ARGS,
#struct_name::#init_fn_name
as ::limbo_ext::InitAggFunction,
#struct_name::#step_fn_name
@@ -283,6 +324,38 @@ pub fn derive_agg_func(input: TokenStream) -> TokenStream {
TokenStream::from(expanded)
}
/// Register your extension with 'core' by providing the relevant functions
///```ignore
///use limbo_ext::{register_extension, scalar, Value, AggregateDerive, AggFunc};
///
/// register_extension!{ scalars: { return_one }, aggregates: { SumPlusOne } }
///
///#[scalar(name = "one")]
///fn return_one(args: &[Value]) -> Value {
/// return Value::from_integer(1);
///}
///
///#[derive(AggregateDerive)]
///struct SumPlusOne;
///
///impl AggFunc for SumPlusOne {
/// type State = i64;
/// const NAME: &'static str = "sum_plus_one";
/// const ARGS: i32 = 1;
///
/// fn step(state: &mut Self::State, args: &[Value]) {
/// let Some(val) = args[0].to_integer() else {
/// return;
/// };
/// *state += val;
/// }
///
/// fn finalize(state: Self::State) -> Value {
/// Value::from_integer(state + 1)
/// }
///}
///
/// ```
#[proc_macro]
pub fn register_extension(input: TokenStream) -> TokenStream {
let input_ast = parse_macro_input!(input as RegisterExtensionInput);
@@ -297,8 +370,8 @@ pub fn register_extension(input: TokenStream) -> TokenStream {
syn::Ident::new(&format!("register_{}", scalar_ident), scalar_ident.span());
quote! {
{
let result = unsafe { #scalar_ident::#register_fn(api)};
if result != 0 {
let result = unsafe { #register_fn(api)};
if !result.is_ok() {
return result;
}
}
@@ -310,7 +383,7 @@ pub fn register_extension(input: TokenStream) -> TokenStream {
quote! {
{
let result = unsafe{ #agg_ident::#register_fn(api)};
if result != 0 {
if !result.is_ok() {
return result;
}
}
@@ -319,13 +392,13 @@ pub fn register_extension(input: TokenStream) -> TokenStream {
let expanded = quote! {
#[no_mangle]
pub extern "C" fn register_extension(api: &::limbo_ext::ExtensionApi) -> i32 {
pub extern "C" fn register_extension(api: &::limbo_ext::ExtensionApi) -> ::limbo_ext::ResultCode {
let api = unsafe { &*api };
#(#scalar_calls)*
#(#aggregate_calls)*
::limbo_ext::RESULT_OK
::limbo_ext::ResultCode::OK
}
};