From 99ad7ff766ea44576e2fad4a62e11eb10e7aa14a Mon Sep 17 00:00:00 2001
From: Nuno Cruces <ncruces@users.noreply.github.com>
Date: Fri, 26 Jan 2024 23:52:45 +0000
Subject: [PATCH] Collation callback.

---
 conn.go      |  1 +
 func.go      | 28 +++++++++++++++++++++++++---
 func_test.go | 13 ++++++++-----
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/conn.go b/conn.go
index a23dd60..410e1d1 100644
--- a/conn.go
+++ b/conn.go
@@ -21,6 +21,7 @@ type Conn struct {
 	interrupt context.Context
 	pending   *Stmt
 	log       func(code xErrorCode, msg string)
+	collation func(name string)
 	arena     arena
 
 	handle uint32
diff --git a/func.go b/func.go
index 7c69963..85650d4 100644
--- a/func.go
+++ b/func.go
@@ -8,8 +8,25 @@ import (
 	"github.com/tetratelabs/wazero/api"
 )
 
-// AnyCollationNeeded registers a fake collating function
-// for any unknown collating sequence.
+// CollationNeeded registers a callback to be invoked
+// whenever an unknown collation sequence is required.
+//
+// https://sqlite.org/c3ref/collation_needed.html
+func (c *Conn) CollationNeeded(cb func(name string)) error {
+	var enable uint64
+	if cb != nil {
+		enable = 1
+	}
+	r := c.call("sqlite3_collation_needed_go", uint64(c.handle), enable)
+	if err := c.error(r); err != nil {
+		return err
+	}
+	c.collation = cb
+	return nil
+}
+
+// AnyCollationNeeded uses [Conn.CollationNeeded] to register
+// a fake collating function for any unknown collating sequence.
 // The fake collating function works like BINARY.
 //
 // This can be used to load schemas that contain
@@ -106,7 +123,12 @@ func destroyCallback(ctx context.Context, mod api.Module, pApp uint32) {
 	util.DelHandle(ctx, pApp)
 }
 
-func collationCallback(ctx context.Context, mod api.Module, pArg, pDB, eTextRep, zName uint32) {}
+func collationCallback(ctx context.Context, mod api.Module, pArg, pDB, eTextRep, zName uint32) {
+	if c, ok := ctx.Value(connKey{}).(*Conn); ok && c.handle == pDB && c.collation != nil {
+		name := util.ReadString(mod, zName, _MAX_NAME)
+		c.collation(name)
+	}
+}
 
 func compareCallback(ctx context.Context, mod api.Module, pApp, nKey1, pKey1, nKey2, pKey2 uint32) uint32 {
 	fn := util.GetHandle(ctx, pApp).(func(a, b []byte) int)
diff --git a/func_test.go b/func_test.go
index c27eabd..baf4050 100644
--- a/func_test.go
+++ b/func_test.go
@@ -6,11 +6,9 @@ import (
 	"log"
 	"regexp"
 
-	"golang.org/x/text/collate"
-	"golang.org/x/text/language"
-
 	"github.com/ncruces/go-sqlite3"
 	_ "github.com/ncruces/go-sqlite3/embed"
+	"github.com/ncruces/go-sqlite3/ext/unicode"
 )
 
 func ExampleConn_CreateCollation() {
@@ -30,12 +28,17 @@ func ExampleConn_CreateCollation() {
 		log.Fatal(err)
 	}
 
-	err = db.CreateCollation("french", collate.New(language.French).Compare)
+	err = db.CollationNeeded(func(name string) {
+		err := unicode.RegisterCollation(db, name, name)
+		if err != nil {
+			log.Fatal(err)
+		}
+	})
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	stmt, _, err := db.Prepare(`SELECT word FROM words ORDER BY word COLLATE french`)
+	stmt, _, err := db.Prepare(`SELECT word FROM words ORDER BY word COLLATE fr_FR`)
 	if err != nil {
 		log.Fatal(err)
 	}