Skip to content

Proof of concept for passing query time context to UDF functions #418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Proof of concept for passing query time context to UDF functions #418

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions connection.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,18 @@ type Conn struct {
conn mapping.Connection
closed bool
tx bool
state connState
}

type connState struct {
currCtx context.Context
}

func (conn *Conn) setCurrCtx(ctx context.Context) func() {
conn.state.currCtx = ctx
return func() {
// conn.state.currCtx = nil
}
}

// CheckNamedValue implements the driver.NamedValueChecker interface.
Expand All @@ -30,6 +42,9 @@ func (conn *Conn) CheckNamedValue(nv *driver.NamedValue) error {
// ExecContext executes a query that doesn't return rows, such as an INSERT or UPDATE.
// It implements the driver.ExecerContext interface.
func (conn *Conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
cleanupCtx := conn.setCurrCtx(ctx)
defer cleanupCtx()

prepared, err := conn.prepareStmts(ctx, query)
if err != nil {
return nil, err
Expand All @@ -53,6 +68,9 @@ func (conn *Conn) ExecContext(ctx context.Context, query string, args []driver.N
// QueryContext executes a query that may return rows, such as a SELECT.
// It implements the driver.QueryerContext interface.
func (conn *Conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
cleanupCtx := conn.setCurrCtx(ctx)
defer cleanupCtx()

prepared, err := conn.prepareStmts(ctx, query)
if err != nil {
return nil, err
Expand Down
66 changes: 37 additions & 29 deletions scalar_udf.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ typedef void (*scalar_udf_delete_callback_t)(void *);
import "C"

import (
"context"
"database/sql"
"database/sql/driver"
"runtime"
Expand Down Expand Up @@ -47,7 +48,7 @@ type ScalarFuncConfig struct {
type ScalarFuncExecutor struct {
// RowExecutor accepts a row-based execution function.
// []driver.Value contains the row values, and it returns the row execution result, or error.
RowExecutor func(values []driver.Value) (any, error)
RowExecutor func(ctx context.Context, values []driver.Value) (any, error)
}

// ScalarFunc is the user-defined scalar function interface.
Expand All @@ -64,15 +65,15 @@ type ScalarFunc interface {
// name is the function name, and f is the scalar function's interface ScalarFunc.
// RegisterScalarUDF takes ownership of f, so you must pass it as a pointer.
func RegisterScalarUDF(c *sql.Conn, name string, f ScalarFunc) error {
function, err := createScalarFunc(name, f)
if err != nil {
return getError(errAPI, err)
}
defer mapping.DestroyScalarFunction(&function)

// Register the function on the underlying driver connection exposed by c.Raw.
err = c.Raw(func(driverConn any) error {
err := c.Raw(func(driverConn any) error {
conn := driverConn.(*Conn)
function, err := createScalarFunc(name, f, &conn.state)
if err != nil {
return getError(errAPI, err)
}
defer mapping.DestroyScalarFunction(&function)

state := mapping.RegisterScalarFunction(conn.conn, function)
if state == mapping.StateError {
return getError(errAPI, errScalarUDFCreate)
Expand All @@ -91,25 +92,25 @@ func RegisterScalarUDF(c *sql.Conn, name string, f ScalarFunc) error {
func RegisterScalarUDFSet(c *sql.Conn, name string, functions ...ScalarFunc) error {
set := mapping.CreateScalarFunctionSet(name)

// Create each function and add it to the set.
for i, f := range functions {
function, err := createScalarFunc(name, f)
if err != nil {
mapping.DestroyScalarFunctionSet(&set)
return getError(errAPI, err)
}

state := mapping.AddScalarFunctionToSet(set, function)
mapping.DestroyScalarFunction(&function)
if state == mapping.StateError {
mapping.DestroyScalarFunctionSet(&set)
return getError(errAPI, addIndexToError(errScalarUDFAddToSet, i))
}
}

// Register the function set on the underlying driver connection exposed by c.Raw.
err := c.Raw(func(driverConn any) error {
conn := driverConn.(*Conn)
connState := &conn.state
// Create each function and add it to the set.
for i, f := range functions {
function, err := createScalarFunc(name, f, connState)
if err != nil {
mapping.DestroyScalarFunctionSet(&set)
return getError(errAPI, err)
}

state := mapping.AddScalarFunctionToSet(set, function)
mapping.DestroyScalarFunction(&function)
if state == mapping.StateError {
mapping.DestroyScalarFunctionSet(&set)
return getError(errAPI, addIndexToError(errScalarUDFAddToSet, i))
}
}
state := mapping.RegisterScalarFunctionSet(conn.conn, set)
mapping.DestroyScalarFunctionSet(&set)
if state == mapping.StateError {
Expand All @@ -127,7 +128,9 @@ func scalar_udf_callback(functionInfoPtr unsafe.Pointer, inputPtr unsafe.Pointer
output := mapping.Vector{Ptr: outputPtr}

extraInfo := mapping.ScalarFunctionGetExtraInfo(functionInfo)
function := getPinned[ScalarFunc](extraInfo)
functionPinnedValue := getPinned[scalarFuncPinnedValue](extraInfo)
function := functionPinnedValue.f
ctx := functionPinnedValue.connState.currCtx

// Initialize the input chunk.
var inputChunk DataChunk
Expand Down Expand Up @@ -177,7 +180,7 @@ func scalar_udf_callback(functionInfoPtr unsafe.Pointer, inputPtr unsafe.Pointer

// Execute the function.
var val any
if val, err = executor.RowExecutor(values); err != nil {
if val, err = executor.RowExecutor(ctx, values); err != nil {
mapping.ScalarFunctionSetError(functionInfo, getError(errAPI, err).Error())
return
}
Expand Down Expand Up @@ -238,7 +241,12 @@ func registerResultParams(config ScalarFuncConfig, f mapping.ScalarFunction) err
return nil
}

func createScalarFunc(name string, f ScalarFunc) (mapping.ScalarFunction, error) {
type scalarFuncPinnedValue struct {
f ScalarFunc
connState *connState
}

func createScalarFunc(name string, f ScalarFunc, connState *connState) (mapping.ScalarFunction, error) {
if name == "" {
return mapping.ScalarFunction{}, errScalarUDFNoName
}
Expand Down Expand Up @@ -274,9 +282,9 @@ func createScalarFunc(name string, f ScalarFunc) (mapping.ScalarFunction, error)
mapping.ScalarFunctionSetFunction(function, callbackPtr)

// Pin the ScalarFunc f.
value := pinnedValue[ScalarFunc]{
value := pinnedValue[scalarFuncPinnedValue]{
pinner: &runtime.Pinner{},
value: f,
value: scalarFuncPinnedValue{f: f, connState: connState},
}
h := cgo.NewHandle(value)
value.pinner.Pin(&h)
Expand Down
65 changes: 59 additions & 6 deletions scalar_udf_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,25 +25,26 @@ type (
errResultNilSUDF struct{}
errResultAnySUDF struct{}
errExecSUDF struct{}
contextSUDF struct{}
)

func simpleSum(values []driver.Value) (any, error) {
func simpleSum(ctx context.Context, values []driver.Value) (any, error) {
if values[0] == nil || values[1] == nil {
return nil, nil
}
val := values[0].(int32) + values[1].(int32)
return val, nil
}

func constantOne([]driver.Value) (any, error) {
func constantOne(ctx context.Context, _ []driver.Value) (any, error) {
return int32(1), nil
}

func identity(values []driver.Value) (any, error) {
func identity(ctx context.Context, values []driver.Value) (any, error) {
return values[0], nil
}

func variadicSum(values []driver.Value) (any, error) {
func variadicSum(ctx context.Context, values []driver.Value) (any, error) {
sum := int32(0)
for _, val := range values {
if val == nil {
Expand All @@ -54,7 +55,7 @@ func variadicSum(values []driver.Value) (any, error) {
return sum, nil
}

func nilCount(values []driver.Value) (any, error) {
func nilCount(ctx context.Context, values []driver.Value) (any, error) {
count := int32(0)
for _, val := range values {
if val == nil {
Expand All @@ -64,10 +65,19 @@ func nilCount(values []driver.Value) (any, error) {
return count, nil
}

func constantError([]driver.Value) (any, error) {
func constantError(ctx context.Context, _ []driver.Value) (any, error) {
return nil, errors.New("test invalid execution")
}

func contextUDF(ctx context.Context, values []driver.Value) (any, error) {
key, isString := values[0].(string)
if !isString {
return nil, errors.New("expected string key")
}

return ctx.Value(key), nil
}

func (*simpleSUDF) Config() ScalarFuncConfig {
return ScalarFuncConfig{[]TypeInfo{currentInfo, currentInfo}, currentInfo, nil, false, false}
}
Expand Down Expand Up @@ -168,6 +178,20 @@ func (*errExecSUDF) Executor() ScalarFuncExecutor {
return ScalarFuncExecutor{constantError}
}

func (*contextSUDF) Config() ScalarFuncConfig {
return ScalarFuncConfig{
[]TypeInfo{currentInfo},
currentInfo,
nil,
false,
false,
}
}

func (*contextSUDF) Executor() ScalarFuncExecutor {
return ScalarFuncExecutor{contextUDF}
}

func TestSimpleScalarUDF(t *testing.T) {
db := openDbWrapper(t, ``)
defer closeDbWrapper(t, db)
Expand Down Expand Up @@ -431,3 +455,32 @@ func TestErrScalarUDFClosedConn(t *testing.T) {
err = RegisterScalarUDF(conn, "closed_con", errClosedConUDF)
require.ErrorContains(t, err, sql.ErrConnDone.Error())
}

func TestContextScalarUDF(t *testing.T) {
db := openDbWrapper(t, ``)
defer closeDbWrapper(t, db)

ctx := context.WithValue(context.Background(), "test_key", "test_value")
conn := openConnWrapper(t, db, ctx)
defer closeConnWrapper(t, conn)

var err error
currentInfo, err = NewTypeInfo(TYPE_VARCHAR)
require.NoError(t, err)

var udf *contextSUDF
err = RegisterScalarUDF(conn, "context_test", udf)
require.NoError(t, err)

var result string
row := db.QueryRowContext(ctx, `SELECT context_test('test_key') AS result`)
require.NoError(t, row.Scan(&result))
require.Equal(t, "test_value", result)

// Test with non-existent key
var nilResult *string
row = db.QueryRowContext(ctx, `SELECT context_test('non_existent') AS result`)
require.NoError(t, row.Scan(&nilResult))
require.Nil(t, nilResult)
}

11 changes: 11 additions & 0 deletions statement.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,6 +304,9 @@ func (s *Stmt) Exec(args []driver.Value) (driver.Result, error) {
// ExecContext executes a query that doesn't return rows, such as an INSERT or UPDATE.
// It implements the driver.StmtExecContext interface.
func (s *Stmt) ExecContext(ctx context.Context, nargs []driver.NamedValue) (driver.Result, error) {
cleanupCtx := s.conn.setCurrCtx(ctx)
defer cleanupCtx()

res, err := s.execute(ctx, nargs)
if err != nil {
return nil, err
Expand All @@ -328,6 +331,9 @@ func (s *Stmt) ExecBound(ctx context.Context) (driver.Result, error) {
return nil, errNotBound
}

cleanupCtx := s.conn.setCurrCtx(ctx)
defer cleanupCtx()

res, err := s.executeBound(ctx)
if err != nil {
return nil, err
Expand All @@ -346,6 +352,9 @@ func (s *Stmt) Query(args []driver.Value) (driver.Rows, error) {
// QueryContext executes a query that may return rows, such as a SELECT.
// It implements the driver.StmtQueryContext interface.
func (s *Stmt) QueryContext(ctx context.Context, nargs []driver.NamedValue) (driver.Rows, error) {
cleanupCtx := s.conn.setCurrCtx(ctx)
defer cleanupCtx()

res, err := s.execute(ctx, nargs)
if err != nil {
return nil, err
Expand All @@ -367,6 +376,8 @@ func (s *Stmt) QueryBound(ctx context.Context) (driver.Rows, error) {
if !s.bound {
return nil, errNotBound
}
cleanupCtx := s.conn.setCurrCtx(ctx)
defer cleanupCtx()

res, err := s.executeBound(ctx)
if err != nil {
Expand Down
Loading