mirrors/limbo
1
0
Fork 0
mirror of https://github.com/tursodatabase/limbo.git synced 2025-08-04 01:58:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Fix scalar API in extensions, add some error handling

Browse source
This commit is contained in:
PThorpe92 2025-01-18 15:19:23 -05:00
parent 4a41736f89
commit 956320b7d0
No known key found for this signature in database
GPG key ID: 66DB3FBACBDD05CC
11 changed files with 609 additions and 604 deletions
Download patch file Download diff file Expand all files Collapse all files

57
extensions/core/README.md
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() {

275
extensions/core/src/lib.rs
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 => {}
}
}
}

339
extensions/core/src/types.rs Normal file
View file

@ -0,0 +1,339 @@
use std::{fmt::Display, os::raw::c_void};
#[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, 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 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)),
}
}
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(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,
}
}
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,
}
}
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,
ErrCode { code: ResultCode },
}