emote/sqlitestore: implement index in terms of sqlite

emote/emote.go

@@ -23,23 +23,23 @@ type Emote struct {
 //
 // Parser assumes that emotes are bound by whitespace.
 type Parser struct {
-	m map[string]string
+	m map[string]Emote
 	// TODO(branden): more efficient data structure; trie?
 }
 
 // NewParser creates a Parser for the given list of emotes.
 func NewParser(emotes ...Emote) Parser {
-	m := make(map[string]string, len(emotes))
+	m := make(map[string]Emote, len(emotes))
 	for _, e := range emotes {
-		m[e.Name] = e.ID
+		m[e.Name] = e
 	}
 	return Parser{m}
 }
 
 // Next parses the next emote instance from the message and returns the
 // remainder of the message text following it.
-// If there is no emote in the message the returned emote is the empty string.
-func (p Parser) Next(text string) (name, id, following string) {
+// If there is no emote in the message the returned emote has an empty ID.
+func (p Parser) Next(text string) (emote Emote, following string) {
 	for text != "" {
 		// First trim any existing space.
 		text = strings.TrimSpace(text)
@@ -54,12 +54,12 @@ func (p Parser) Next(text string) (name, id, following string) {
 		} else {
 			following = ""
 		}
-		id = p.m[word]
-		if id != "" {
-			return word, id, following
+		em := p.m[word]
+		if em.ID != "" {
+			return em, following
 		}
 		text = following
 	}
 	// No emote found.
-	return "", "", ""
+	return Emote{}, ""
 }

emote/emote_test.go

@@ -57,16 +57,16 @@ func TestParser(t *testing.T) {
 			var got []string
 			text := c.text
 			for {
-				name, id, rest := p.Next(text)
-				if id == "" {
+				em, rest := p.Next(text)
+				if em.ID == "" {
 					break
 				}
-				if name != id {
+				if em.Name != em.ID {
 					// Not normally the case, but this test is constructed
 					// so that it is.
-					t.Errorf("wrong id %q for name %q", id, name)
+					t.Errorf("wrong id %q for name %q", em.ID, em.Name)
 				}
-				got = append(got, name)
+				got = append(got, em.Name)
 				text = rest
 			}
 			if diff := cmp.Diff(c.want, got); diff != "" {
@@ -103,7 +103,7 @@ func BenchmarkParser(b *testing.B) {
 			for b.Loop() {
 				text := c.text
 				for text != "" {
-					_, _, text = p.Next(text)
+					_, text = p.Next(text)
 				}
 			}
 		})

emote/sqlitestore/metrics.sql

@@ -0,0 +1,12 @@
+SELECT
+	id,
+	name, -- name, source, &c. are constant per id, and sqlite lets us select without aggregate
+	source,
+	link,
+	image,
+	COUNT(*) AS tokens,
+	COUNT(DISTINCT message) AS messages,
+	COUNT(DISTINCT sender) AS senders
+FROM emote
+WHERE channel = :channel AND time BETWEEN :start AND :end
+GROUP BY id

emote/sqlitestore/schema.sql

@@ -0,0 +1,17 @@
+PRAGMA journal_mode = WAL;
+
+CREATE TABLE IF NOT EXISTS emote (
+	channel TEXT NOT NULL,
+	message TEXT NOT NULL,
+	time INTEGER NOT NULL,
+	sender TEXT NOT NULL,
+	id TEXT NOT NULL,
+	-- The following columns could be normalized in a separate table,
+	-- but they're here for simplicity.
+	name TEXT NOT NULL,
+	source TEXT NOT NULL,
+	link TEXT NOT NULL,
+	image TEXT NOT NULL
+) STRICT;
+
+CREATE INDEX IF NOT EXISTS channel_time ON emote (channel, time);

emote/sqlitestore/store.go

@@ -0,0 +1,114 @@
+package sqlitestore
+
+import (
+	"context"
+	_ "embed"
+	"fmt"
+	"time"
+
+	"zombiezen.com/go/sqlite/sqlitex"
+
+	"git.sunturtle.xyz/zephyr/kaiyan/emote"
+)
+
+type Store struct {
+	pool *sqlitex.Pool
+}
+
+//go:embed schema.sql
+var schemaSQL string
+
+//go:embed metrics.sql
+var metricsSQL string
+
+// Open creates a metrics store using db as its underlying storage.
+// It initializes the schema used for Kaiyan.
+func Open(ctx context.Context, db *sqlitex.Pool) (*Store, error) {
+	conn, err := db.Take(ctx)
+	defer db.Put(conn)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't get connection from pool: %w", err)
+	}
+	if err := sqlitex.ExecuteScript(conn, schemaSQL, nil); err != nil {
+		return nil, fmt.Errorf("couldn't run migration: %w", err)
+	}
+	st := Store{db}
+	return &st, nil
+}
+
+// Close closes the database.
+func (db *Store) Close() error {
+	return db.pool.Close()
+}
+
+// Record stores the emotes in a given message.
+func (db *Store) Record(ctx context.Context, channel, message, sender string, tm time.Time, emotes []emote.Emote) (err error) {
+	conn, err := db.pool.Take(ctx)
+	defer db.pool.Put(conn)
+	if err != nil {
+		return fmt.Errorf("couldn't get connection to record emotes: %w", err)
+	}
+	defer sqlitex.Transaction(conn)(&err)
+
+	st, err := conn.Prepare(`INSERT INTO emote (channel, message, time, sender, id, name, source, link, image) VALUES (:channel, :message, :time, :sender, :id, :name, :source, :link, :image)`)
+	if err != nil {
+		return fmt.Errorf("couldn't prepare emote insert: %w", err)
+	}
+	st.SetText(":channel", channel)
+	st.SetText(":message", message)
+	st.SetInt64(":time", tm.UnixNano())
+	st.SetText(":sender", sender)
+	for _, em := range emotes {
+		st.SetText(":id", em.ID)
+		st.SetText(":name", em.Name)
+		st.SetText(":source", em.Source)
+		st.SetText(":link", em.Link)
+		st.SetText(":image", em.Image)
+		_, err := st.Step()
+		if err != nil {
+			return fmt.Errorf("couldn't insert emote: %w", err)
+		}
+		st.Reset() // NOTE(branden): bound parameters are retained
+	}
+	return nil
+}
+
+// Metrics appends emote usage information for a channel in a given time range.
+func (db *Store) Metrics(ctx context.Context, channel string, start, end time.Time, onto []emote.Metric) ([]emote.Metric, error) {
+	conn, err := db.pool.Take(ctx)
+	defer db.pool.Put(conn)
+	if err != nil {
+		return onto, fmt.Errorf("couldn't get connection for metrics: %w", err)
+	}
+	st, err := conn.Prepare(metricsSQL)
+	if err != nil {
+		return onto, fmt.Errorf("couldn't prepare metrics statement: %w", err)
+	}
+	st.SetText(":channel", channel)
+	st.SetInt64(":start", start.UnixNano())
+	st.SetInt64(":end", end.UnixNano())
+	for {
+		ok, err := st.Step()
+		if err != nil {
+			return onto, fmt.Errorf("couldn't step metrics selection: %w", err)
+		}
+		if !ok {
+			break
+		}
+
+		m := emote.Metric{
+			Emote: emote.Emote{
+				ID:     st.ColumnText(0),
+				Name:   st.ColumnText(1),
+				Source: st.ColumnText(2),
+				Link:   st.ColumnText(3),
+				Image:  st.ColumnText(4),
+			},
+			Tokens:   st.ColumnInt64(5),
+			Messages: st.ColumnInt64(6),
+			Users:    st.ColumnInt64(7),
+		}
+		onto = append(onto, m)
+	}
+	return onto, st.Reset()
+}

emote/sqlitestore/store_test.go

@@ -0,0 +1,141 @@
+package sqlitestore_test
+
+import (
+	"fmt"
+	"slices"
+	"strings"
+	"testing"
+	"time"
+
+	"git.sunturtle.xyz/zephyr/kaiyan/emote"
+	"git.sunturtle.xyz/zephyr/kaiyan/emote/sqlitestore"
+	"github.com/google/go-cmp/cmp"
+	"zombiezen.com/go/sqlite"
+	"zombiezen.com/go/sqlite/sqlitex"
+)
+
+func testDB(name string) *sqlitex.Pool {
+	opts := sqlitex.PoolOptions{
+		Flags: sqlite.OpenCreate | sqlite.OpenReadWrite | sqlite.OpenMemory | sqlite.OpenSharedCache | sqlite.OpenURI,
+	}
+	pool, err := sqlitex.NewPool(fmt.Sprintf("file:%s.db?mode=memory&cache=shared", name), opts)
+	if err != nil {
+		panic(err)
+	}
+	return pool
+}
+
+func TestStore(t *testing.T) {
+	ems := []struct {
+		channel, message, sender string
+		ts                       int64
+		emotes                   []emote.Emote
+	}{
+		{
+			channel: "kessoku",
+			message: "1",
+			sender:  "ryo",
+			ts:      1000,
+			emotes: []emote.Emote{
+				{ID: "nijika", Name: "nijika"},
+				{ID: "kita", Name: "kita"},
+				{ID: "nijika", Name: "nijika"},
+				{ID: "kita", Name: "kita"},
+			},
+		},
+		{
+			channel: "kessoku",
+			message: "2",
+			sender:  "seika",
+			ts:      2000,
+			emotes: []emote.Emote{
+				{ID: "nijika", Name: "nijika"},
+				{ID: "kita", Name: "kita"},
+				{ID: "nijika", Name: "nijika"},
+				{ID: "kita", Name: "kita"},
+			},
+		},
+		{
+			channel: "sickhack",
+			message: "3",
+			sender:  "ryo",
+			ts:      1500,
+			emotes: []emote.Emote{
+				{ID: "nijika", Name: "nijika"},
+				{ID: "kita", Name: "kita"},
+				{ID: "nijika", Name: "nijika"},
+				{ID: "kita", Name: "kita"},
+			},
+		},
+	}
+	st, err := sqlitestore.Open(t.Context(), testDB("TestStore"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, m := range ems {
+		err := st.Record(t.Context(), m.channel, m.message, m.sender, time.Unix(m.ts, 0), m.emotes)
+		if err != nil {
+			t.Errorf("couldn't record message %s: %v", m.message, err)
+		}
+	}
+	cases := []struct {
+		name       string
+		channel    string
+		start, end int64
+		want       []emote.Metric
+	}{
+		{
+			name:    "channel",
+			channel: "kessoku",
+			start:   0,
+			end:     1e6,
+			want: []emote.Metric{
+				{
+					Emote:    emote.Emote{ID: "kita", Name: "kita"},
+					Tokens:   4,
+					Messages: 2,
+					Users:    2,
+				},
+				{
+					Emote:    emote.Emote{ID: "nijika", Name: "nijika"},
+					Tokens:   4,
+					Messages: 2,
+					Users:    2,
+				},
+			},
+		},
+		{
+			name:    "time",
+			channel: "kessoku",
+			start:   1250,
+			end:     1e6,
+			want: []emote.Metric{
+				{
+					Emote:    emote.Emote{ID: "kita", Name: "kita"},
+					Tokens:   2,
+					Messages: 1,
+					Users:    1,
+				},
+				{
+					Emote:    emote.Emote{ID: "nijika", Name: "nijika"},
+					Tokens:   2,
+					Messages: 1,
+					Users:    1,
+				},
+			},
+		},
+	}
+	for _, c := range cases {
+		// Don't use Run because we have earlier t.Fails.
+		got, err := st.Metrics(t.Context(), c.channel, time.Unix(c.start, 0), time.Unix(c.end, 0), nil)
+		if err != nil {
+			t.Error(err)
+		}
+		// Sort results by emote ID so we can easily diff.
+		// In particular, we don't expect the store to sort for us.
+		slices.SortFunc(got, func(a, b emote.Metric) int { return strings.Compare(a.Emote.ID, b.Emote.ID) })
+		if diff := cmp.Diff(c.want, got); diff != "" {
+			t.Errorf("wrong results (+got/-want):\n%s", diff)
+		}
+	}
+}

emote/store.go

@@ -0,0 +1,26 @@
+package emote
+
+import (
+	"context"
+	"time"
+)
+
+// Store is a store of emote usage records.
+type Store interface {
+	// Record stores the emotes in a given message.
+	Record(ctx context.Context, channel, message, sender string, tm time.Time, emotes []Emote) error
+	// Metrics appends emote usage information for a channel in a given time range.
+	Metrics(ctx context.Context, channel string, start, end time.Time, onto []Metric) ([]Metric, error)
+}
+
+// Metric is the metrics for a single emote.
+type Metric struct {
+	// Emote is the emote that the metrics describe.
+	Emote Emote
+	// Tokens is the total number of occurrences of the emote.
+	Tokens int64
+	// Messages is the number of unique messages which contained the emote.
+	Messages int64
+	// Users is the number of users who used the emote.
+	Users int64
+}

go.mod

@@ -6,10 +6,21 @@ require (
 	github.com/go-json-experiment/json v0.0.0-20250223041408-d3c622f1b874
 	github.com/google/go-cmp v0.7.0
 	github.com/twmb/franz-go v1.18.1
+	zombiezen.com/go/sqlite v1.4.0
 )
 
 require (
+	github.com/dustin/go-humanize v1.0.1 // indirect
+	github.com/google/uuid v1.6.0 // indirect
 	github.com/klauspost/compress v1.17.11 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/ncruces/go-strftime v0.1.9 // indirect
 	github.com/pierrec/lz4/v4 v4.1.22 // indirect
+	github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
 	github.com/twmb/franz-go/pkg/kmsg v1.9.0 // indirect
+	golang.org/x/sys v0.22.0 // indirect
+	modernc.org/libc v1.55.3 // indirect
+	modernc.org/mathutil v1.6.0 // indirect
+	modernc.org/memory v1.8.0 // indirect
+	modernc.org/sqlite v1.33.1 // indirect
 )

go.sum

@@ -1,14 +1,59 @@
+github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
+github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/go-json-experiment/json v0.0.0-20250223041408-d3c622f1b874 h1:F8d1AJ6M9UQCavhwmO6ZsrYLfG8zVFWfEfMS2MXPkSY=
 github.com/go-json-experiment/json v0.0.0-20250223041408-d3c622f1b874/go.mod h1:TiCD2a1pcmjd7YnhGH0f/zKNcCD06B029pHhzV23c2M=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/klauspost/compress v1.17.11 h1:In6xLpyWOi1+C7tXUUWv2ot1QvBjxevKAaI6IXrJmUc=
 github.com/klauspost/compress v1.17.11/go.mod h1:pMDklpSncoRMuLFrf1W9Ss9KT+0rH90U12bZKk7uwG0=
+github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
+github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4=
+github.com/ncruces/go-strftime v0.1.9/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
 github.com/pierrec/lz4/v4 v4.1.22 h1:cKFw6uJDK+/gfw5BcDL0JL5aBsAFdsIT18eRtLj7VIU=
 github.com/pierrec/lz4/v4 v4.1.22/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4=
+github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
+github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
 github.com/twmb/franz-go v1.18.1 h1:D75xxCDyvTqBSiImFx2lkPduE39jz1vaD7+FNc+vMkc=
 github.com/twmb/franz-go v1.18.1/go.mod h1:Uzo77TarcLTUZeLuGq+9lNpSkfZI+JErv7YJhlDjs9M=
 github.com/twmb/franz-go/pkg/kmsg v1.9.0 h1:JojYUph2TKAau6SBtErXpXGC7E3gg4vGZMv9xFU/B6M=
 github.com/twmb/franz-go/pkg/kmsg v1.9.0/go.mod h1:CMbfazviCyY6HM0SXuG5t9vOwYDHRCSrJJyBAe5paqg=
 golang.org/x/crypto v0.32.0 h1:euUpcYgM8WcP71gNpTqQCn6rC2t6ULUPiOzfWaXVVfc=
 golang.org/x/crypto v0.32.0/go.mod h1:ZnnJkOaASj8g0AjIduWNlq2NRxL0PlBrbKVyZ6V/Ugc=
+golang.org/x/mod v0.16.0 h1:QX4fJ0Rr5cPQCF7O9lh9Se4pmwfwskqZfq5moyldzic=
+golang.org/x/mod v0.16.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.22.0 h1:RI27ohtqKCnwULzJLqkv897zojh5/DwS/ENaMzUOaWI=
+golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
+golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
+golang.org/x/tools v0.19.0 h1:tfGCXNR1OsFG+sVdLAitlpjAvD/I6dHDKnYrpEZUHkw=
+golang.org/x/tools v0.19.0/go.mod h1:qoJWxmGSIBmAeriMx19ogtrEPrGtDbPK634QFIcLAhc=
+modernc.org/cc/v4 v4.21.4 h1:3Be/Rdo1fpr8GrQ7IVw9OHtplU4gWbb+wNgeoBMmGLQ=
+modernc.org/cc/v4 v4.21.4/go.mod h1:HM7VJTZbUCR3rV8EYBi9wxnJ0ZBRiGE5OeGXNA0IsLQ=
+modernc.org/ccgo/v4 v4.19.2 h1:lwQZgvboKD0jBwdaeVCTouxhxAyN6iawF3STraAal8Y=
+modernc.org/ccgo/v4 v4.19.2/go.mod h1:ysS3mxiMV38XGRTTcgo0DQTeTmAO4oCmJl1nX9VFI3s=
+modernc.org/fileutil v1.3.0 h1:gQ5SIzK3H9kdfai/5x41oQiKValumqNTDXMvKo62HvE=
+modernc.org/fileutil v1.3.0/go.mod h1:XatxS8fZi3pS8/hKG2GH/ArUogfxjpEKs3Ku3aK4JyQ=
+modernc.org/gc/v2 v2.4.1 h1:9cNzOqPyMJBvrUipmynX0ZohMhcxPtMccYgGOJdOiBw=
+modernc.org/gc/v2 v2.4.1/go.mod h1:wzN5dK1AzVGoH6XOzc3YZ+ey/jPgYHLuVckd62P0GYU=
+modernc.org/libc v1.55.3 h1:AzcW1mhlPNrRtjS5sS+eW2ISCgSOLLNyFzRh/V3Qj/U=
+modernc.org/libc v1.55.3/go.mod h1:qFXepLhz+JjFThQ4kzwzOjA/y/artDeg+pcYnY+Q83w=
+modernc.org/mathutil v1.6.0 h1:fRe9+AmYlaej+64JsEEhoWuAYBkOtQiMEU7n/XgfYi4=
+modernc.org/mathutil v1.6.0/go.mod h1:Ui5Q9q1TR2gFm0AQRqQUaBWFLAhQpCwNcuhBOSedWPo=
+modernc.org/memory v1.8.0 h1:IqGTL6eFMaDZZhEWwcREgeMXYwmW83LYW8cROZYkg+E=
+modernc.org/memory v1.8.0/go.mod h1:XPZ936zp5OMKGWPqbD3JShgd/ZoQ7899TUuQqxY+peU=
+modernc.org/opt v0.1.3 h1:3XOZf2yznlhC+ibLltsDGzABUGVx8J6pnFMS3E4dcq4=
+modernc.org/opt v0.1.3/go.mod h1:WdSiB5evDcignE70guQKxYUl14mgWtbClRi5wmkkTX0=
+modernc.org/sortutil v1.2.0 h1:jQiD3PfS2REGJNzNCMMaLSp/wdMNieTbKX920Cqdgqc=
+modernc.org/sortutil v1.2.0/go.mod h1:TKU2s7kJMf1AE84OoiGppNHJwvB753OYfNl2WRb++Ss=
+modernc.org/sqlite v1.33.1 h1:trb6Z3YYoeM9eDL1O8do81kP+0ejv+YzgyFo+Gwy0nM=
+modernc.org/sqlite v1.33.1/go.mod h1:pXV2xHxhzXZsgT/RtTFAPY6JJDEvOTcTdwADQCCWD4k=
+modernc.org/strutil v1.2.0 h1:agBi9dp1I+eOnxXeiZawM8F4LawKv4NzGWSaLfyeNZA=
+modernc.org/strutil v1.2.0/go.mod h1:/mdcBmfOibveCTBxUl5B5l6W+TTH1FXPLHZE6bTosX0=
+modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
+modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=
+zombiezen.com/go/sqlite v1.4.0 h1:N1s3RIljwtp4541Y8rM880qgGIgq3fTD2yks1xftnKU=
+zombiezen.com/go/sqlite v1.4.0/go.mod h1:0w9F1DN9IZj9AcLS9YDKMboubCACkwYCGkzoy3eG5ik=

queue/message.go

@@ -2,6 +2,8 @@ package queue
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"sync"
 
 	"github.com/go-json-experiment/json"
@@ -25,7 +27,7 @@ type Producer interface {
 	Produce(ctx context.Context, rec *kgo.Record, promise func(*kgo.Record, error))
 }
 
-const topic = "kaiyan.chat"
+const Topic = "kaiyan.chat"
 
 var recordPool sync.Pool
 
@@ -37,7 +39,7 @@ func Send(ctx context.Context, cl Producer, msg Message, errs chan<- error) {
 	}
 	rec, _ := recordPool.Get().(*kgo.Record)
 	if rec == nil {
-		rec = &kgo.Record{Topic: topic}
+		rec = &kgo.Record{Topic: Topic}
 	}
 	rec.Key = append(rec.Key[:0], msg.ID...)
 	rec.Value = append(rec.Value[:0], b...)
@@ -51,3 +53,23 @@ func Send(ctx context.Context, cl Producer, msg Message, errs chan<- error) {
 		}
 	})
 }
+
+type Consumer interface {
+	PollFetches(ctx context.Context) kgo.Fetches
+}
+
+func Recv(ctx context.Context, cl Consumer, onto []Message) ([]Message, error) {
+	f := cl.PollFetches(ctx)
+	var errs error
+	f.EachError(func(s string, i int32, err error) {
+		errs = errors.Join(errs, fmt.Errorf("partition %d: %w", i, err))
+	})
+	f.EachRecord(func(r *kgo.Record) {
+		var msg Message
+		if err := json.Unmarshal(r.Value, &msg); err != nil {
+			errs = errors.Join(errs, fmt.Errorf("partition %d: %w", r.Partition, err))
+		}
+		onto = append(onto, msg)
+	})
+	return onto, errs
+}

queue/message_test.go

@@ -4,9 +4,10 @@ import (
 	"context"
 	"testing"
 
-	"git.sunturtle.xyz/zephyr/kaiyan/queue"
 	"github.com/go-json-experiment/json"
 	"github.com/twmb/franz-go/pkg/kgo"
+
+	"git.sunturtle.xyz/zephyr/kaiyan/queue"
 )
 
 type spyProducer struct {
@@ -48,3 +49,57 @@ func TestSend(t *testing.T) {
 		t.Errorf("message did not round-trip:\nwant %+v\ngot  %+v", msg, got)
 	}
 }
+
+type spyConsumer struct {
+	get kgo.Fetches
+}
+
+func (c *spyConsumer) PollFetches(ctx context.Context) kgo.Fetches {
+	return c.get
+}
+
+func TestRecv(t *testing.T) {
+	msg := queue.Message{
+		ID:      "bocchi",
+		Channel: "kessoku",
+		Sender:  queue.Sender(make([]byte, 16), "kessoku", "ryō"),
+		Text:    "bocchi the rock!",
+	}
+	val, err := json.Marshal(&msg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	p := spyConsumer{
+		get: kgo.Fetches{
+			{
+				Topics: []kgo.FetchTopic{
+					{
+						Topic: queue.Topic,
+						Partitions: []kgo.FetchPartition{
+							{
+								Partition: 1,
+								Records: []*kgo.Record{
+									{
+										Key:       []byte("bocchi"),
+										Value:     val,
+										Partition: 1,
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+	got, err := queue.Recv(context.Background(), &p, nil)
+	if err != nil {
+		t.Error(err)
+	}
+	if len(got) != 1 {
+		t.Fatalf("wrong number of messages: want 1, got %d", len(got))
+	}
+	if got[0] != msg {
+		t.Errorf("message did not round-trip:\nwant %+v\ngot  %+v", msg, got[0])
+	}
+}