im_server/loadtest/im-pressure/dbverify.go

package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"sort"
	"strings"

	_ "github.com/go-sql-driver/mysql"
)

type dbVerifyResult struct {
	RunID          string   `json:"run_id"`
	TenantCode     string   `json:"tenant_code"`
	TenantID       int64    `json:"tenant_id"`
	Expected       int      `json:"expected"`
	FoundExact     int      `json:"found_exact"`
	Missing        int      `json:"missing"`
	Mismatch       int      `json:"mismatch"`
	DuplicateRows  int      `json:"duplicate_rows"`
	Unexpected     int      `json:"unexpected_rows"`
	Tables         []string `json:"tables"`
	MissingMIDs    []string `json:"missing_client_mids,omitempty"`
	UnexpectedMIDs []string `json:"unexpected_client_mids,omitempty"`
}

type persistedIdentity struct {
	Topic string
	Seq   int64
}

var dbIdentifierRE = regexp.MustCompile(`^[A-Za-z0-9_]+$`)

func runDBVerify(ctx context.Context, coordinatorPath, output, tenantCode, dsnEnv, registryDB string) error {
	raw, err := os.ReadFile(coordinatorPath)
	if err != nil {
		return err
	}
	var coordinator struct {
		Start    StartInfo                    `json:"start"`
		Ledgers  map[string][]AcceptedMessage `json:"ledgers"`
		Observed map[string][]AcceptedMessage `json:"observed"`
	}
	if err = json.Unmarshal(raw, &coordinator); err != nil {
		return fmt.Errorf("decode coordinator: %w", err)
	}
	if coordinator.Start.RunID == "" {
		return fmt.Errorf("coordinator has no run id")
	}
	dsn := strings.TrimSpace(os.Getenv(dsnEnv))
	if dsn == "" {
		return fmt.Errorf("environment variable %s is required", dsnEnv)
	}
	if registryDB == "" {
		registryDB = strings.TrimSpace(os.Getenv("TENANT_REGISTRY_DB"))
	}
	if registryDB == "" {
		registryDB = "meeting_im"
	}
	if !dbIdentifierRE.MatchString(registryDB) {
		return fmt.Errorf("invalid registry database identifier")
	}
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return err
	}
	defer db.Close()
	if err = db.PingContext(ctx); err != nil {
		return err
	}
	var tenantID int64
	query := fmt.Sprintf("SELECT id FROM `%s`.im_tenant WHERE code=?", registryDB)
	if err = db.QueryRowContext(ctx, query, tenantCode).Scan(&tenantID); err != nil {
		return err
	}

	expected := expectedDBMessages(coordinator.Start.RunID, coordinator.Ledgers, coordinator.Observed)
	months := make(map[string]struct{})
	for _, message := range expected {
		if message.Kind == "group" {
			months[message.SentAt.UTC().Format("200601")] = struct{}{}
		}
	}
	tables := []string{fmt.Sprintf("im_%03d_messages", tenantID)}
	for month := range months {
		tables = append(tables, fmt.Sprintf("im_%03d_group_messages_%s", tenantID, month))
	}
	sort.Strings(tables)
	actual := make(map[string][]persistedIdentity)
	for _, table := range tables {
		if !dbIdentifierRE.MatchString(table) {
			return fmt.Errorf("invalid generated table name %q", table)
		}
		rows, queryErr := db.QueryContext(ctx, fmt.Sprintf(
			"SELECT client_msg_id,topic,seqid FROM `%s` WHERE tenant_id=? AND client_msg_id LIKE ?", table),
			tenantID, coordinator.Start.RunID+"-%")
		if queryErr != nil {
			return fmt.Errorf("query %s: %w", table, queryErr)
		}
		for rows.Next() {
			var mid, topic string
			var seq int64
			if err = rows.Scan(&mid, &topic, &seq); err != nil {
				_ = rows.Close()
				return err
			}
			actual[mid] = append(actual[mid], persistedIdentity{Topic: topic, Seq: seq})
		}
		if err = rows.Close(); err != nil {
			return err
		}
	}
	result := dbVerifyResult{RunID: coordinator.Start.RunID, TenantCode: tenantCode, TenantID: tenantID,
		Expected: len(expected), Tables: tables}
	for mid, want := range expected {
		rows := actual[mid]
		if len(rows) == 0 {
			result.Missing++
			result.MissingMIDs = append(result.MissingMIDs, mid)
			continue
		}
		matched := false
		for _, got := range rows {
			if got.Topic == want.Topic && got.Seq == want.Seq {
				matched = true
			}
		}
		if matched {
			result.FoundExact++
		} else {
			result.Mismatch++
		}
		if len(rows) > 1 {
			result.DuplicateRows += len(rows) - 1
		}
	}
	for mid := range actual {
		if _, ok := expected[mid]; !ok {
			result.Unexpected++
			result.UnexpectedMIDs = append(result.UnexpectedMIDs, mid)
		}
	}
	sort.Strings(result.MissingMIDs)
	sort.Strings(result.UnexpectedMIDs)
	resultRaw, _ := json.MarshalIndent(result, "", "  ")
	if output == "" || output == "-" {
		if _, err = fmt.Fprintln(os.Stdout, string(resultRaw)); err != nil {
			return err
		}
	} else {
		if err = writeJSONFile(output, result); err != nil {
			return err
		}
	}
	if result.Missing != 0 || result.Mismatch != 0 || result.DuplicateRows != 0 || result.Unexpected != 0 {
		return fmt.Errorf("database verification failed: missing=%d mismatch=%d duplicates=%d unexpected=%d",
			result.Missing, result.Mismatch, result.DuplicateRows, result.Unexpected)
	}
	return nil
}

func expectedDBMessages(runID string, ledgers, observed map[string][]AcceptedMessage) map[string]AcceptedMessage {
	all := make([]AcceptedMessage, 0)
	for _, messages := range observed {
		all = append(all, messages...)
	}
	for _, messages := range ledgers {
		all = append(all, messages...)
	}
	expected := make(map[string]AcceptedMessage, len(all))
	for _, message := range mergeMessageLedgers(nil, all) {
		if message.RunID == runID && message.ClientMid != "" {
			expected[message.ClientMid] = message
		}
	}
	return expected
}