cairnobs/enterprise/internal/chwriter/chwriter.go

// Package chwriter is enterprise/internal/chrunner's write-side
// counterpart -- the tenant-scoped implementation ingest/consumer.
// Consumer needs to route each batch's records into their own tenant's
// dedicated ClickHouse database, instead of the one shared table
// ingest/cmd/ingest's single-tenant mode always writes to. Requires
// importing ingest/clickhousewriter and ingest/consumer directly (see
// enterprise/go.mod's replace directive) -- same allowed
// "enterprise -> core" import direction chrunner uses for
// api/querylang/executor, just against a different core module.
//
// Design mirrors chrunner.Registry closely: one fully separate
// *clickhousewriter.Writer (and the driver.Conn under it) per tenant,
// built once at construction from an immutable map -- never a shared
// pool with session-level USE, for the same concurrency reasons
// chrunner's doc comment explains. The one real difference:
// chrunner.RunSQL resolves exactly one tenant per call from ctx (a
// single request always belongs to one identity); WriteBatch resolves
// per *record*, since one Kafka batch pulled off the shared
// sentry.logs.raw topic can freely mix records from many different
// tenants -- see ingest/internal/grpcserver's doc comment for why
// there's one shared topic, not topic-per-tenant.
package chwriter

import (
	"context"
	"fmt"

	"github.com/sentry/sentry/ingest/clickhousewriter"
	"github.com/sentry/sentry/ingest/consumer"
	logsv1 "github.com/sentry/sentry/proto/sentry/logs/v1"
)

// DataSource mirrors chrunner.DataSource -- deliberately not
// enterprise/internal/rbacstore.DataSource itself, so this package
// doesn't need to import rbacstore just to describe "an address and a
// credential." Callers (enterprise-ingest's main.go) adapt rbacstore
// rows into this.
type DataSource struct {
	TenantID string
	Database string
	Username string
	Password string
}

// Registry implements ingest/consumer's chWriter interface
// (WriteBatch(ctx, []consumer.Record) error) by routing each record to
// its tenant's dedicated connection. Immutable after New returns -- see
// this file's doc comment.
type Registry struct {
	writers map[string]*clickhousewriter.Writer
	closers []func()
}

// New opens one real ClickHouse connection per DataSource (same native
// address for all of them, different per-tenant credentials -- tenants
// sharing a physical ClickHouse server today, same as chrunner). Fails
// closed: if any one tenant's connection can't be opened, the whole
// Registry fails to construct rather than silently running with a
// partial tenant set.
func New(ctx context.Context, addr string, sources []DataSource) (*Registry, error) {
	reg := &Registry{writers: make(map[string]*clickhousewriter.Writer, len(sources))}
	for _, src := range sources {
		w, err := clickhousewriter.New(ctx, clickhousewriter.Config{
			Addr: addr, Database: src.Database, Username: src.Username, Password: src.Password,
		})
		if err != nil {
			reg.Close()
			return nil, fmt.Errorf("chwriter: opening connection for tenant %q: %w", src.TenantID, err)
		}
		reg.writers[src.TenantID] = w
		reg.closers = append(reg.closers, func() { _ = w.Close() })
	}
	return reg, nil
}

// Close releases every underlying connection -- call once at process
// shutdown, same lifecycle as chrunner.Registry.Close.
func (r *Registry) Close() {
	for _, c := range r.closers {
		c()
	}
}

// WriteBatch implements ingest/consumer's chWriter interface. Groups
// records by TenantID and writes each tenant's group through its own
// dedicated connection -- fails the *whole* call (matching
// ingest/consumer's existing all-or-nothing batch contract: a failed
// WriteBatch means no offsets are committed and the entire batch is
// redelivered, never partial credit) if any record's tenant is empty
// (no TenantResolver was configured for the PushBatch call that
// produced it -- a multi-tenant deployment must never silently write an
// untagged record somewhere) or unrecognized (not yet provisioned, or
// provisioning failed). Fail closed, same reasoning
// chrunner.Registry.RunSQL's doc comment gives for the read side.
//
// A permanently-unprovisioned or permanently-mistagged tenant would
// stall this consumer's offset progress entirely (every redelivery of
// that batch fails the same way) -- a real, disclosed limitation of
// reusing ingest/consumer's existing all-or-nothing contract rather
// than building new partial-batch-success semantics nothing else in
// this codebase has either. See /docs/phase-4-runbook.md.
func (r *Registry) WriteBatch(ctx context.Context, records []consumer.Record) error {
	byTenant := make(map[string][]consumer.Record, len(records))
	for _, rec := range records {
		byTenant[rec.TenantID] = append(byTenant[rec.TenantID], rec)
	}

	for tenantID, group := range byTenant {
		if tenantID == "" {
			return fmt.Errorf("chwriter: %d record(s) in this batch have no tenant_id, refusing to write any of it", len(group))
		}
		writer, ok := r.writers[tenantID]
		if !ok {
			return fmt.Errorf("chwriter: tenant %q has no provisioned ClickHouse connection, refusing to write %d record(s)", tenantID, len(group))
		}
		plain := make([]*logsv1.LogRecord, len(group))
		for i, rec := range group {
			plain[i] = rec.Record
		}
		if err := writer.WriteBatch(ctx, plain); err != nil {
			return fmt.Errorf("chwriter: writing batch for tenant %q: %w", tenantID, err)
		}
	}
	return nil
}