fix build

1 files changed, 187 insertions, 0 deletions

diff --git a/tracecallback.go b/tracecallback.go
index 6a59282..01923f2 100644
--- a/tracecallback.go
+++ b/tracecallback.go
@@ -3,6 +3,7 @@
 //
 // Use of this source code is governed by an MIT-style
 // license that can be found in the LICENSE file.
+// +build trace
 
 package sqlite3
 
@@ -13,11 +14,23 @@ package sqlite3
 #include <sqlite3.h>
 #endif
 #include <stdlib.h>
+
+void stepTrampoline(sqlite3_context*, int, sqlite3_value**);
+void doneTrampoline(sqlite3_context*);
+
+int
+_sqlite3_trace_v2(sqlite3* db, unsigned mask, int(*xc)(unsigned,void*,void*,void*), void *ctx) {
+  return sqlite3_trace_v2(db, mask, xc, ctx);
+}
+
+void traceCallbackTrampoline(unsigned traceEventCode, void *ctx, void *p, void *x);
 */
 import "C"
 
 import (
+	"errors"
 	"fmt"
+	"reflect"
 	"strings"
 	"sync"
 	"unsafe"
@@ -232,3 +245,177 @@ func popTraceMapping(connHandle uintptr) (TraceConfig, bool) {
 	}
 	return entryCopy.config, found
 }
+
+// RegisterAggregator makes a Go type available as a SQLite aggregation function.
+//
+// Because aggregation is incremental, it's implemented in Go with a
+// type that has 2 methods: func Step(values) accumulates one row of
+// data into the accumulator, and func Done() ret finalizes and
+// returns the aggregate value. "values" and "ret" may be any type
+// supported by RegisterFunc.
+//
+// RegisterAggregator takes as implementation a constructor function
+// that constructs an instance of the aggregator type each time an
+// aggregation begins. The constructor must return a pointer to a
+// type, or an interface that implements Step() and Done().
+//
+// The constructor function and the Step/Done methods may optionally
+// return an error in addition to their other return values.
+//
+// See _example/go_custom_funcs for a detailed example.
+func (c *SQLiteConn) RegisterAggregator(name string, impl interface{}, pure bool) error {
+	var ai aggInfo
+	ai.constructor = reflect.ValueOf(impl)
+	t := ai.constructor.Type()
+	if t.Kind() != reflect.Func {
+		return errors.New("non-function passed to RegisterAggregator")
+	}
+	if t.NumOut() != 1 && t.NumOut() != 2 {
+		return errors.New("SQLite aggregator constructors must return 1 or 2 values")
+	}
+	if t.NumOut() == 2 && !t.Out(1).Implements(reflect.TypeOf((*error)(nil)).Elem()) {
+		return errors.New("Second return value of SQLite function must be error")
+	}
+	if t.NumIn() != 0 {
+		return errors.New("SQLite aggregator constructors must not have arguments")
+	}
+
+	agg := t.Out(0)
+	switch agg.Kind() {
+	case reflect.Ptr, reflect.Interface:
+	default:
+		return errors.New("SQlite aggregator constructor must return a pointer object")
+	}
+	stepFn, found := agg.MethodByName("Step")
+	if !found {
+		return errors.New("SQlite aggregator doesn't have a Step() function")
+	}
+	step := stepFn.Type
+	if step.NumOut() != 0 && step.NumOut() != 1 {
+		return errors.New("SQlite aggregator Step() function must return 0 or 1 values")
+	}
+	if step.NumOut() == 1 && !step.Out(0).Implements(reflect.TypeOf((*error)(nil)).Elem()) {
+		return errors.New("type of SQlite aggregator Step() return value must be error")
+	}
+
+	stepNArgs := step.NumIn()
+	start := 0
+	if agg.Kind() == reflect.Ptr {
+		// Skip over the method receiver
+		stepNArgs--
+		start++
+	}
+	if step.IsVariadic() {
+		stepNArgs--
+	}
+	for i := start; i < start+stepNArgs; i++ {
+		conv, err := callbackArg(step.In(i))
+		if err != nil {
+			return err
+		}
+		ai.stepArgConverters = append(ai.stepArgConverters, conv)
+	}
+	if step.IsVariadic() {
+		conv, err := callbackArg(t.In(start + stepNArgs).Elem())
+		if err != nil {
+			return err
+		}
+		ai.stepVariadicConverter = conv
+		// Pass -1 to sqlite so that it allows any number of
+		// arguments. The call helper verifies that the minimum number
+		// of arguments is present for variadic functions.
+		stepNArgs = -1
+	}
+
+	doneFn, found := agg.MethodByName("Done")
+	if !found {
+		return errors.New("SQlite aggregator doesn't have a Done() function")
+	}
+	done := doneFn.Type
+	doneNArgs := done.NumIn()
+	if agg.Kind() == reflect.Ptr {
+		// Skip over the method receiver
+		doneNArgs--
+	}
+	if doneNArgs != 0 {
+		return errors.New("SQlite aggregator Done() function must have no arguments")
+	}
+	if done.NumOut() != 1 && done.NumOut() != 2 {
+		return errors.New("SQLite aggregator Done() function must return 1 or 2 values")
+	}
+	if done.NumOut() == 2 && !done.Out(1).Implements(reflect.TypeOf((*error)(nil)).Elem()) {
+		return errors.New("second return value of SQLite aggregator Done() function must be error")
+	}
+
+	conv, err := callbackRet(done.Out(0))
+	if err != nil {
+		return err
+	}
+	ai.doneRetConverter = conv
+	ai.active = make(map[int64]reflect.Value)
+	ai.next = 1
+
+	// ai must outlast the database connection, or we'll have dangling pointers.
+	c.aggregators = append(c.aggregators, &ai)
+
+	cname := C.CString(name)
+	defer C.free(unsafe.Pointer(cname))
+	opts := C.SQLITE_UTF8
+	if pure {
+		opts |= C.SQLITE_DETERMINISTIC
+	}
+	rv := C._sqlite3_create_function(c.db, cname, C.int(stepNArgs), C.int(opts), C.uintptr_t(newHandle(c, &ai)), nil, (*[0]byte)(unsafe.Pointer(C.stepTrampoline)), (*[0]byte)(unsafe.Pointer(C.doneTrampoline)))
+	if rv != C.SQLITE_OK {
+		return c.lastError()
+	}
+	return nil
+}
+
+// SetTrace installs or removes the trace callback for the given database connection.
+// It's not named 'RegisterTrace' because only one callback can be kept and called.
+// Calling SetTrace a second time on same database connection
+// overrides (cancels) any prior callback and all its settings:
+// event mask, etc.
+func (c *SQLiteConn) SetTrace(requested *TraceConfig) error {
+	connHandle := uintptr(unsafe.Pointer(c.db))
+
+	_, _ = popTraceMapping(connHandle)
+
+	if requested == nil {
+		// The traceMap entry was deleted already by popTraceMapping():
+		// can disable all events now, no need to watch for TraceClose.
+		err := c.setSQLiteTrace(0)
+		return err
+	}
+
+	reqCopy := *requested
+
+	// Disable potentially expensive operations
+	// if their result will not be used. We are doing this
+	// just in case the caller provided nonsensical input.
+	if reqCopy.EventMask&TraceStmt == 0 {
+		reqCopy.WantExpandedSQL = false
+	}
+
+	addTraceMapping(connHandle, reqCopy)
+
+	// The callback trampoline function does cleanup on Close event,
+	// regardless of the presence or absence of the user callback.
+	// Therefore it needs the Close event to be selected:
+	actualEventMask := reqCopy.EventMask | TraceClose
+	err := c.setSQLiteTrace(actualEventMask)
+	return err
+}
+
+func (c *SQLiteConn) setSQLiteTrace(sqliteEventMask uint) error {
+	rv := C._sqlite3_trace_v2(c.db,
+		C.uint(sqliteEventMask),
+		(*[0]byte)(unsafe.Pointer(C.traceCallbackTrampoline)),
+		unsafe.Pointer(c.db)) // Fourth arg is same as first: we are
+	// passing the database connection handle as callback context.
+
+	if rv != C.SQLITE_OK {
+		return c.lastError()
+	}
+	return nil
+}