diff --git a/internal/db/db.go b/internal/db/db.go index 824ae05..e133fc0 100644 --- a/internal/db/db.go +++ b/internal/db/db.go @@ -7,6 +7,7 @@ import ( "fmt" "regexp" "sync" + "uuid" "github.com/mattn/go-sqlite3" ) @@ -55,6 +56,13 @@ func Connect(ctx context.Context, cfg Config) (*SQLite, error) { driverOnce.Do(func() { sql.Register("odidere_sqlite", &sqlite3.SQLiteDriver{ ConnectHook: func(conn *sqlite3.SQLiteConn) error { + if err := conn.RegisterFunc( + "uuid4", + func() string { return uuid.New().String() }, + false, + ); err != nil { + return fmt.Errorf("failed to register uuid4: %w", err) + } if err := conn.RegisterFunc( "regexp", func(re, s string) (bool, error) { diff --git a/internal/db/queries.go b/internal/db/queries.go index 9dc22a0..acdd1f8 100644 --- a/internal/db/queries.go +++ b/internal/db/queries.go @@ -313,6 +313,8 @@ type InsertMessageParams struct { Content string // ConversationId identifies the conversation that owns the new row. ConversationId string + // Id is the new row's UUID. Nil generates a version 4 UUID. + Id *uuid.UUID // Metadata is inference JSON converted to JSONB for storage. // Nil stores SQL NULL; malformed JSON causes insertion to fail. Metadata jsontext.Value @@ -324,6 +326,12 @@ type InsertMessageParams struct { // InsertMessage creates a new message. func (q *Queries) InsertMessage(ctx context.Context, arg InsertMessageParams) (*Message, error) { + var id uuid.UUID + if arg.Id == nil { + id = uuid.NewV4() + } else { + id = *arg.Id + } var metadataArg any if arg.Metadata != nil { metadataArg = string(arg.Metadata) @@ -331,7 +339,7 @@ func (q *Queries) InsertMessage(ctx context.Context, arg InsertMessageParams) (* row := q.db.QueryRowContext( ctx, insertMessage, - uuid.New().String(), + id.String(), arg.Content, arg.ConversationId, metadataArg, diff --git a/internal/db/queries_test.go b/internal/db/queries_test.go index 7e3bb03..279a3ae 100644 --- a/internal/db/queries_test.go +++ b/internal/db/queries_test.go @@ -8,6 +8,7 @@ import ( "log/slog" "reflect" "testing" + "uuid" "github.com/mattn/go-sqlite3" ) @@ -344,6 +345,53 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { } }) } + suppliedID := uuid.NewV4() + for _, tt := range []struct { + name string + id *uuid.UUID + wantErr bool + }{ + {name: "generated ID"}, + {name: "supplied v4 ID", id: &suppliedID}, + {name: "duplicate ID", id: &suppliedID, wantErr: true}, + } { + t.Run(tt.name, func(t *testing.T) { + seq, err := qw.NextSequence(ctx, c.Id) + if err != nil { + t.Fatal(err) + } + inserted, err := qw.InsertMessage(ctx, InsertMessageParams{ + Content: `{"role":"assistant","content":"answer"}`, + ConversationId: c.Id, + Id: tt.id, + Role: MessageRoleAssistant, + Sequence: seq, + }) + if tt.wantErr { + if err == nil { + t.Fatal("duplicate ID was accepted") + } + return + } + if err != nil { + t.Fatal(err) + } + if _, err := uuid.Parse(inserted.Id); err != nil { + t.Fatal(err) + } + if tt.id != nil && inserted.Id != tt.id.String() { + t.Errorf("inserted ID = %s, want %s", inserted.Id, tt.id.String()) + } + if tt.id == nil { + if err := qw.DeleteMessageId(ctx, inserted.Id); err != nil { + t.Fatal(err) + } + } + }) + } + if err := qw.DeleteMessageId(ctx, suppliedID.String()); err != nil { + t.Fatal(err) + } tx, err := db.Writer.BeginTx(ctx, nil) if err != nil { t.Fatal(err)