Build per-tenant ClickHouse write-routing for ingest (Tantivy still deferred)

ingest tags every record with a tenant_id Kafka header (built previously), but nothing consumed it to actually route the write. This closes that for ClickHouse: enterprise/cmd/enterprise-ingest (a second binary, mirroring enterprise-api) reuses ingest/consumer's own flush loop unchanged, with enterprise/internal/chwriter.Registry -- a per-tenant clickhousewriter.Writer registry -- swapped in as the writer. A batch pulled from the single shared Redpanda topic can mix records from many tenants, so WriteBatch groups by TenantID and dispatches each group to its own tenant's connection, fail- closed on an empty or unrecognized tenant_id. ingest/consumer and ingest/clickhousewriter move out of internal/ (same reason api/internal/* moved earlier this phase: enterprise/ can't import anything under another module's internal/). Their New() constructors now take small local Config structs instead of ingest/internal/config types, so enterprise/ doesn't need that import either. Building this surfaced a real bug: tenantprovision.ProvisionClickHouse only granted SELECT on a tenant's ClickHouse user, correct for chrunner's read-only use but not enough for chwriter reusing the same credential to write -- every real per-tenant write would have failed closed with a permission error. Fixed by widening the grant to SELECT, INSERT; no cross-tenant boundary is crossed by also allowing INSERT within a tenant's own database. Helm gates enterprise-ingest's Deployment on the same ingest.requireTenantCredential flag that already gates tag validation -- write-routing is meaningless without tagging already being required, so they're one decision, not two. docker-compose.yml's version is a disclosed, weaker approximation: it can't achieve Helm's genuine -mode=server/-mode=consumer split, so with the enterprise profile active both ingest and enterprise-ingest independently consume every message via different consumer groups -- harmless duplication for local verification only. Not built: Tantivy's independent Redpanda consumer (search/src/consumer.rs) still doesn't read the tenant_id header at all -- every record still lands in the one shared index regardless of tenant. Not run: the live-ClickHouse- gated tests (chwriter's cross-tenant routing test, tenantprovision's INSERT regression test) -- no Docker/database access in this environment; they're correct Go that has never executed, disclosed as such in docs/security/ threat-model.md and docs/phase-4-runbook.md §14.
2026-08-14 19:26:09 -07:00
parent 17fdc212c2
commit 1de77b969f
26 changed files with 1355 additions and 267 deletions
@@ -0,0 +1,180 @@
+// Package consumer reads normalized-on-write LogRecords back off Redpanda
+// and batch-writes them into ClickHouse. Offsets are committed only after
+// a successful ClickHouse write, so a ClickHouse outage causes redelivery
+// on restart rather than silent data loss (at-least-once, not exactly-once
+// — Phase 0 doesn't dedupe on the consumer side).
+//
+// Moved out of internal/ (was ingest/internal/consumer) once
+// enterprise/cmd/enterprise-ingest needed to run this same flush loop
+// against a per-tenant chWriter -- see clickhousewriter's doc comment
+// for why (same Go internal/-visibility reasoning as every other
+// package this phase moved out of internal/ for a cross-module
+// import). Each record's TenantID (Record.TenantID below) is read from
+// the tenant_id Kafka message header grpcserver.TenantIDHeaderKey
+// documents -- empty when no TenantResolver was configured for the
+// PushBatch call that produced it, exactly as before per-tenant ingest
+// credentials existed. What a chWriter implementation *does* with that
+// tag varies: ingest/cmd/ingest's single-tenant clickhousewriter.Writer
+// ignores it (writes everything to its one configured database, per
+// Phase 0-3 behavior, unchanged); enterprise/internal/chwriter.Registry
+// (only ever wired into enterprise/cmd/enterprise-ingest, never this
+// core binary) routes each record to its tenant's dedicated ClickHouse
+// database instead.
+package consumer
+
+import (
+	"context"
+	"log/slog"
+	"time"
+
+	"github.com/segmentio/kafka-go"
+	"google.golang.org/protobuf/proto"
+
+	logsv1 "github.com/sentry/sentry/proto/sentry/logs/v1"
+)
+
+// TenantIDHeaderKey mirrors ingest/internal/grpcserver.TenantIDHeaderKey
+// -- kept as its own constant (not an import of grpcserver, which is
+// the agent-facing *producer* side, a different concern from this
+// package's consumer side) so this package's dependency list stays
+// narrow. Both must name the same literal; a mismatch would silently
+// stop tenant_id from ever reaching a consumer, so grpcserver's own
+// doc comment on TenantIDHeaderKey cross-references this one.
+const TenantIDHeaderKey = "tenant_id"
+
+// Record pairs a parsed LogRecord with the tenant it was tagged with at
+// ingest time (see the package doc comment).
+type Record struct {
+	TenantID string
+	Record   *logsv1.LogRecord
+}
+
+// chWriter is the subset of a ClickHouse writer this package depends
+// on, kept as an interface so the flush loop is unit-testable without a
+// real ClickHouse connection, and so both the single-tenant
+// (clickhousewriter.Writer, adapted) and multi-tenant
+// (enterprise/internal/chwriter.Registry) implementations can share
+// this exact same consumer loop.
+type chWriter interface {
+	WriteBatch(ctx context.Context, records []Record) error
+}
+
+// reader is the subset of *kafka.Reader used here, as an interface so the
+// flush/commit logic can be tested against a fake without a real broker.
+type reader interface {
+	FetchMessage(ctx context.Context) (kafka.Message, error)
+	CommitMessages(ctx context.Context, msgs ...kafka.Message) error
+	Close() error
+}
+
+// Config is deliberately a local type, not ingest/internal/config's
+// RedpandaConfig/BatchConfig -- same "this package must be importable
+// from enterprise/, so it can't depend on ingest/internal/..." reasoning
+// as clickhousewriter.Config.
+type Config struct {
+	Brokers         []string
+	Topic           string
+	ConsumerGroup   string
+	BatchMaxSize    int
+	FlushIntervalMS int
+}
+
+type Consumer struct {
+	logger *slog.Logger
+	reader reader
+	writer chWriter
+	cfg    Config
+}
+
+func New(logger *slog.Logger, cfg Config, w chWriter) *Consumer {
+	r := kafka.NewReader(kafka.ReaderConfig{
+		Brokers: cfg.Brokers,
+		Topic:   cfg.Topic,
+		GroupID: cfg.ConsumerGroup,
+	})
+	return &Consumer{logger: logger, reader: r, writer: w, cfg: cfg}
+}
+
+func (c *Consumer) Run(ctx context.Context) error {
+	defer c.reader.Close()
+
+	flushInterval := time.Duration(c.cfg.FlushIntervalMS) * time.Millisecond
+	ticker := time.NewTicker(flushInterval)
+	defer ticker.Stop()
+
+	msgCh := make(chan kafka.Message)
+	fetchErrCh := make(chan error, 1)
+
+	go func() {
+		for {
+			m, err := c.reader.FetchMessage(ctx)
+			if err != nil {
+				fetchErrCh <- err
+				return
+			}
+			select {
+			case msgCh <- m:
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+
+	var records []Record
+	var pending []kafka.Message
+
+	flush := func() {
+		if len(records) == 0 {
+			return
+		}
+		if err := c.writer.WriteBatch(ctx, records); err != nil {
+			c.logger.Error("clickhouse batch write failed, offsets not committed, will redeliver",
+				"records", len(records), "error", err)
+		} else if err := c.reader.CommitMessages(ctx, pending...); err != nil {
+			c.logger.Error("committing offsets after clickhouse write", "error", err)
+		} else {
+			c.logger.Debug("batch flushed to clickhouse", "records", len(records))
+		}
+		records = records[:0]
+		pending = pending[:0]
+	}
+
+	for {
+		select {
+		case <-ctx.Done():
+			flush()
+			return nil
+		case err := <-fetchErrCh:
+			flush()
+			if ctx.Err() != nil {
+				return nil
+			}
+			return err
+		case <-ticker.C:
+			flush()
+		case m := <-msgCh:
+			var rec logsv1.LogRecord
+			if err := proto.Unmarshal(m.Value, &rec); err != nil {
+				c.logger.Warn("skipping unparseable message", "error", err, "offset", m.Offset)
+				if cerr := c.reader.CommitMessages(ctx, m); cerr != nil {
+					c.logger.Error("committing offset for poison message", "error", cerr)
+				}
+				continue
+			}
+			records = append(records, Record{TenantID: tenantIDFromHeaders(m.Headers), Record: &rec})
+			pending = append(pending, m)
+			if len(records) >= c.cfg.BatchMaxSize {
+				flush()
+			}
+		}
+	}
+}
+
+func tenantIDFromHeaders(headers []kafka.Header) string {
+	for _, h := range headers {
+		if h.Key == TenantIDHeaderKey {
+			return string(h.Value)
+		}
+	}
+	return ""
+}
@@ -0,0 +1,214 @@
+package consumer
+
+import (
+	"context"
+	"errors"
+	"io"
+	"log/slog"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/segmentio/kafka-go"
+	"google.golang.org/protobuf/proto"
+
+	logsv1 "github.com/sentry/sentry/proto/sentry/logs/v1"
+)
+
+type fakeReader struct {
+	msgs chan kafka.Message
+
+	mu        sync.Mutex
+	committed [][]kafka.Message
+}
+
+func newFakeReader() *fakeReader {
+	return &fakeReader{msgs: make(chan kafka.Message, 16)}
+}
+
+func (f *fakeReader) push(m kafka.Message) { f.msgs <- m }
+
+func (f *fakeReader) FetchMessage(ctx context.Context) (kafka.Message, error) {
+	select {
+	case m := <-f.msgs:
+		return m, nil
+	case <-ctx.Done():
+		return kafka.Message{}, ctx.Err()
+	}
+}
+
+func (f *fakeReader) CommitMessages(_ context.Context, msgs ...kafka.Message) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.committed = append(f.committed, msgs)
+	return nil
+}
+
+func (f *fakeReader) Close() error { return nil }
+
+func (f *fakeReader) commitCount() int {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return len(f.committed)
+}
+
+type fakeWriter struct {
+	mu       sync.Mutex
+	batches  [][]Record
+	failNext bool
+}
+
+func (f *fakeWriter) WriteBatch(_ context.Context, records []Record) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	if f.failNext {
+		f.failNext = false
+		return errors.New("simulated clickhouse failure")
+	}
+	batch := make([]Record, len(records))
+	copy(batch, records)
+	f.batches = append(f.batches, batch)
+	return nil
+}
+
+func (f *fakeWriter) batchCount() int {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return len(f.batches)
+}
+
+func newTestConsumer(r reader, w chWriter, cfg Config) *Consumer {
+	return &Consumer{
+		logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
+		reader: r,
+		writer: w,
+		cfg:    cfg,
+	}
+}
+
+func mustMarshal(t *testing.T, rec *logsv1.LogRecord) []byte {
+	t.Helper()
+	b, err := proto.Marshal(rec)
+	if err != nil {
+		t.Fatalf("marshal: %v", err)
+	}
+	return b
+}
+
+func waitFor(t *testing.T, timeout time.Duration, cond func() bool) {
+	t.Helper()
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		if cond() {
+			return
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+	t.Fatal("condition not met before timeout")
+}
+
+func TestConsumerFlushesOnBatchSize(t *testing.T) {
+	fr := newFakeReader()
+	fw := &fakeWriter{}
+	c := newTestConsumer(fr, fw, Config{BatchMaxSize: 2, FlushIntervalMS: 60_000})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	done := make(chan error, 1)
+	go func() { done <- c.Run(ctx) }()
+
+	fr.push(kafka.Message{Value: mustMarshal(t, &logsv1.LogRecord{Message: "a"})})
+	fr.push(kafka.Message{Value: mustMarshal(t, &logsv1.LogRecord{Message: "b"})})
+
+	waitFor(t, time.Second, func() bool { return fw.batchCount() == 1 })
+
+	fw.mu.Lock()
+	if len(fw.batches[0]) != 2 {
+		t.Fatalf("expected batch of 2 records, got %d", len(fw.batches[0]))
+	}
+	fw.mu.Unlock()
+
+	waitFor(t, time.Second, func() bool { return fr.commitCount() == 1 })
+}
+
+func TestConsumerFlushesOnTimeout(t *testing.T) {
+	fr := newFakeReader()
+	fw := &fakeWriter{}
+	c := newTestConsumer(fr, fw, Config{BatchMaxSize: 1000, FlushIntervalMS: 20})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	done := make(chan error, 1)
+	go func() { done <- c.Run(ctx) }()
+
+	fr.push(kafka.Message{Value: mustMarshal(t, &logsv1.LogRecord{Message: "only-one"})})
+
+	waitFor(t, time.Second, func() bool { return fw.batchCount() == 1 })
+
+	fw.mu.Lock()
+	if len(fw.batches[0]) != 1 {
+		t.Fatalf("expected batch of 1 record, got %d", len(fw.batches[0]))
+	}
+	fw.mu.Unlock()
+}
+
+func TestConsumerDoesNotCommitOnWriteFailure(t *testing.T) {
+	fr := newFakeReader()
+	fw := &fakeWriter{failNext: true}
+	c := newTestConsumer(fr, fw, Config{BatchMaxSize: 1, FlushIntervalMS: 60_000})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	done := make(chan error, 1)
+	go func() { done <- c.Run(ctx) }()
+
+	fr.push(kafka.Message{Value: mustMarshal(t, &logsv1.LogRecord{Message: "will-fail"})})
+
+	// Give the flush a moment to run and fail.
+	time.Sleep(100 * time.Millisecond)
+
+	if got := fr.commitCount(); got != 0 {
+		t.Fatalf("expected no commits after a failed clickhouse write, got %d", got)
+	}
+	// The batch was attempted even though writer returned an error.
+	if fw.batchCount() != 0 {
+		t.Fatalf("fakeWriter should not record a failed batch, got %d recorded", fw.batchCount())
+	}
+}
+
+// TestConsumerExtractsTenantIDFromHeader is the read-side half of the
+// producer/consumer tenant_id contract -- ingest/internal/grpcserver
+// attaches this header on the way in; this proves the consumer reads it
+// back correctly (and that a message with no header at all -- the
+// single-tenant/no-resolver case -- gets an empty TenantID, not an
+// error).
+func TestConsumerExtractsTenantIDFromHeader(t *testing.T) {
+	fr := newFakeReader()
+	fw := &fakeWriter{}
+	c := newTestConsumer(fr, fw, Config{BatchMaxSize: 2, FlushIntervalMS: 60_000})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() { _ = c.Run(ctx) }()
+
+	fr.push(kafka.Message{
+		Value:   mustMarshal(t, &logsv1.LogRecord{Message: "tagged"}),
+		Headers: []kafka.Header{{Key: TenantIDHeaderKey, Value: []byte("acme")}},
+	})
+	fr.push(kafka.Message{Value: mustMarshal(t, &logsv1.LogRecord{Message: "untagged"})})
+
+	waitFor(t, time.Second, func() bool { return fw.batchCount() == 1 })
+
+	fw.mu.Lock()
+	defer fw.mu.Unlock()
+	byMessage := map[string]string{}
+	for _, r := range fw.batches[0] {
+		byMessage[r.Record.GetMessage()] = r.TenantID
+	}
+	if byMessage["tagged"] != "acme" {
+		t.Fatalf("TenantID for the tagged message = %q, want acme", byMessage["tagged"])
+	}
+	if byMessage["untagged"] != "" {
+		t.Fatalf("TenantID for the untagged message = %q, want empty", byMessage["untagged"])
+	}
+}