package pg import ( "context" "database/sql" "fmt" "strings" "time" "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/store" ) type PGUsageEventStore struct { db *sql.DB } func NewPGUsageEventStore(db *sql.DB) *PGUsageEventStore { return &PGUsageEventStore{db: db} } const usageEventFieldCount = 31 const usageRollupFieldCount = 24 func (s *PGUsageEventStore) InsertEvent(ctx context.Context, event *store.UsageEvent) error { if event == nil { return nil } return s.InsertEvents(ctx, []store.UsageEvent{*event}) } func (s *PGUsageEventStore) InsertEvents(ctx context.Context, events []store.UsageEvent) error { if len(events) == 0 { return nil } for i := range events { prepareUsageEvent(ctx, &events[i]) } vals := make([]string, len(events)) args := make([]any, 0, len(events)*usageEventFieldCount) for i, event := range events { base := i * usageEventFieldCount placeholders := make([]string, usageEventFieldCount) for j := range usageEventFieldCount { placeholders[j] = fmt.Sprintf("$%d", base+j+1) } vals[i] = "(" + strings.Join(placeholders, ", ") + ")" args = append(args, event.ID, event.TenantID, event.EventTime, event.BucketHour, event.EventType, event.ResourceType, event.ResourceName, event.ResourceID, event.Source, nilUUID(event.AgentID), nilUUID(event.TeamID), nilUUID(event.TraceID), nilUUID(event.SpanID), event.RunID, event.SessionKey, event.Channel, event.Provider, event.Model, event.Status, event.InputTokens, event.OutputTokens, event.TotalTokens, event.CacheReadTokens, event.CacheCreateTokens, event.ThinkingTokens, event.CostUSD, event.DurationMS, event.CallCount, event.ErrorCount, jsonOrNull(event.Metadata), event.CreatedAt, ) } query := `INSERT INTO usage_events ( id, tenant_id, event_time, bucket_hour, event_type, resource_type, resource_name, resource_id, source, agent_id, team_id, trace_id, span_id, run_id, session_key, channel, provider, model, status, input_tokens, output_tokens, total_tokens, cache_read_tokens, cache_create_tokens, thinking_tokens, cost_usd, duration_ms, call_count, error_count, metadata, created_at ) VALUES ` + strings.Join(vals, ", ") + ` ON CONFLICT DO NOTHING` _, err := s.db.ExecContext(ctx, query, args...) return err } func (s *PGUsageEventStore) RefreshEventRollupHour(ctx context.Context, bucketHour time.Time) error { start := bucketHour.UTC().Truncate(time.Hour) end := start.Add(time.Hour) rows, err := s.db.QueryContext(ctx, `SELECT tenant_id, bucket_hour, event_type, resource_type, resource_name, source, agent_id, channel, provider, model, status, COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), COALESCE(SUM(total_tokens), 0), COALESCE(SUM(cache_read_tokens), 0), COALESCE(SUM(cache_create_tokens), 0), COALESCE(SUM(thinking_tokens), 0), COALESCE(SUM(cost_usd), 0), CASE WHEN COALESCE(SUM(call_count), 0) > 0 THEN COALESCE(SUM(duration_ms * call_count), 0) / SUM(call_count) ELSE 0 END, COALESCE(SUM(call_count), 0), COALESCE(SUM(error_count), 0) FROM usage_events WHERE event_time >= $1 AND event_time < $2 GROUP BY tenant_id, bucket_hour, event_type, resource_type, resource_name, source, agent_id, channel, provider, model, status`, start, end) if err != nil { return fmt.Errorf("aggregate usage event rollup: %w", err) } defer rows.Close() now := time.Now().UTC() var rollups []store.UsageEventRollup for rows.Next() { rollup := store.UsageEventRollup{ID: uuid.New(), CreatedAt: now, UpdatedAt: now} if err := rows.Scan( &rollup.TenantID, &rollup.BucketHour, &rollup.EventType, &rollup.ResourceType, &rollup.ResourceName, &rollup.Source, &rollup.AgentID, &rollup.Channel, &rollup.Provider, &rollup.Model, &rollup.Status, &rollup.InputTokens, &rollup.OutputTokens, &rollup.TotalTokens, &rollup.CacheReadTokens, &rollup.CacheCreateTokens, &rollup.ThinkingTokens, &rollup.CostUSD, &rollup.DurationMS, &rollup.CallCount, &rollup.ErrorCount, ); err != nil { return fmt.Errorf("scan usage event rollup: %w", err) } rollups = append(rollups, rollup) } if err := rows.Err(); err != nil { return err } return s.upsertEventRollups(ctx, rollups) } func (s *PGUsageEventStore) GetLatestEventRollupBucket(ctx context.Context) (*time.Time, error) { var t sql.NullTime err := s.db.QueryRowContext(ctx, `SELECT MAX(bucket_hour) FROM usage_event_rollups`).Scan(&t) if err != nil { return nil, fmt.Errorf("get latest event rollup bucket: %w", err) } if !t.Valid { return nil, nil } return &t.Time, nil } func (s *PGUsageEventStore) upsertEventRollups(ctx context.Context, rollups []store.UsageEventRollup) error { if len(rollups) == 0 { return nil } vals := make([]string, len(rollups)) args := make([]any, 0, len(rollups)*usageRollupFieldCount) for i, rollup := range rollups { base := i * usageRollupFieldCount placeholders := make([]string, usageRollupFieldCount) for j := range usageRollupFieldCount { placeholders[j] = fmt.Sprintf("$%d", base+j+1) } vals[i] = "(" + strings.Join(placeholders, ", ") + ")" args = append(args, rollup.ID, rollup.TenantID, rollup.BucketHour, rollup.EventType, rollup.ResourceType, rollup.ResourceName, rollup.Source, nilUUID(rollup.AgentID), rollup.Channel, rollup.Provider, rollup.Model, rollup.Status, rollup.InputTokens, rollup.OutputTokens, rollup.TotalTokens, rollup.CacheReadTokens, rollup.CacheCreateTokens, rollup.ThinkingTokens, rollup.CostUSD, rollup.DurationMS, rollup.CallCount, rollup.ErrorCount, rollup.CreatedAt, rollup.UpdatedAt, ) } query := `INSERT INTO usage_event_rollups ( id, tenant_id, bucket_hour, event_type, resource_type, resource_name, source, agent_id, channel, provider, model, status, input_tokens, output_tokens, total_tokens, cache_read_tokens, cache_create_tokens, thinking_tokens, cost_usd, duration_ms, call_count, error_count, created_at, updated_at ) VALUES ` + strings.Join(vals, ", ") + ` ON CONFLICT ( tenant_id, bucket_hour, event_type, resource_type, resource_name, source, COALESCE(agent_id, '00000000-0000-0000-0000-000000000000'::uuid), channel, provider, model, status ) DO UPDATE SET input_tokens = EXCLUDED.input_tokens, output_tokens = EXCLUDED.output_tokens, total_tokens = EXCLUDED.total_tokens, cache_read_tokens = EXCLUDED.cache_read_tokens, cache_create_tokens = EXCLUDED.cache_create_tokens, thinking_tokens = EXCLUDED.thinking_tokens, cost_usd = EXCLUDED.cost_usd, duration_ms = EXCLUDED.duration_ms, call_count = EXCLUDED.call_count, error_count = EXCLUDED.error_count, updated_at = EXCLUDED.updated_at` _, err := s.db.ExecContext(ctx, query, args...) return err } func (s *PGUsageEventStore) GetEventTimeSeries(ctx context.Context, q store.UsageEventQuery) ([]store.UsageEventTimeSeries, error) { bucketExpr := "bucket_hour" if q.GroupBy == "day" { bucketExpr = "date_trunc('day', bucket_hour)" } where, args := buildUsageEventWhere(ctx, q, "bucket_hour") query := fmt.Sprintf(`SELECT %s AS bucket_time, COALESCE(SUM(call_count), 0), COALESCE(SUM(error_count), 0), COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), COALESCE(SUM(total_tokens), 0), COALESCE(SUM(cache_read_tokens), 0), COALESCE(SUM(cache_create_tokens), 0), COALESCE(SUM(thinking_tokens), 0), COALESCE(SUM(cost_usd), 0), CASE WHEN COALESCE(SUM(call_count), 0) > 0 THEN COALESCE(SUM(duration_ms * call_count), 0) / SUM(call_count) ELSE 0 END FROM usage_event_rollups %s GROUP BY bucket_time ORDER BY bucket_time`, bucketExpr, where) rows, err := s.db.QueryContext(ctx, query, args...) if err != nil { return nil, fmt.Errorf("get usage event timeseries: %w", err) } defer rows.Close() var result []store.UsageEventTimeSeries for rows.Next() { var point store.UsageEventTimeSeries if err := rows.Scan( &point.BucketTime, &point.Calls, &point.Errors, &point.InputTokens, &point.OutputTokens, &point.TotalTokens, &point.CacheReadTokens, &point.CacheCreateTokens, &point.ThinkingTokens, &point.CostUSD, &point.AvgDurationMS, ); err != nil { return nil, fmt.Errorf("scan usage event timeseries: %w", err) } result = append(result, point) } return result, rows.Err() } func (s *PGUsageEventStore) GetEventBreakdown(ctx context.Context, q store.UsageEventQuery) ([]store.UsageEventBreakdown, error) { groupCol := usageEventGroupColumn(q.GroupBy) where, args := buildUsageEventWhere(ctx, q, "bucket_hour") if where == "" { where = " WHERE 1=1" } limit := q.Limit if limit <= 0 || limit > 100 { limit = 25 } args = append(args, limit) limitPlaceholder := fmt.Sprintf("$%d", len(args)) query := fmt.Sprintf(`SELECT %s AS key, MIN(event_type), MIN(resource_type), MIN(resource_name), MIN(source), COALESCE(SUM(call_count), 0), COALESCE(SUM(error_count), 0), COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), COALESCE(SUM(total_tokens), 0), COALESCE(SUM(cache_read_tokens), 0), COALESCE(SUM(cache_create_tokens), 0), COALESCE(SUM(thinking_tokens), 0), COALESCE(SUM(cost_usd), 0), CASE WHEN COALESCE(SUM(call_count), 0) > 0 THEN COALESCE(SUM(duration_ms * call_count), 0) / SUM(call_count) ELSE 0 END FROM usage_event_rollups %s GROUP BY %s ORDER BY SUM(call_count) DESC, key ASC LIMIT %s`, groupCol, where, groupCol, limitPlaceholder) rows, err := s.db.QueryContext(ctx, query, args...) if err != nil { return nil, fmt.Errorf("get usage event breakdown: %w", err) } defer rows.Close() var result []store.UsageEventBreakdown for rows.Next() { var row store.UsageEventBreakdown if err := rows.Scan( &row.Key, &row.EventType, &row.ResourceType, &row.ResourceName, &row.Source, &row.Calls, &row.Errors, &row.InputTokens, &row.OutputTokens, &row.TotalTokens, &row.CacheReadTokens, &row.CacheCreateTokens, &row.ThinkingTokens, &row.CostUSD, &row.AvgDurationMS, ); err != nil { return nil, fmt.Errorf("scan usage event breakdown: %w", err) } result = append(result, row) } return result, rows.Err() } func (s *PGUsageEventStore) GetEventSummary(ctx context.Context, q store.UsageEventQuery) (*store.UsageEventSummary, error) { where, args := buildUsageEventWhere(ctx, q, "bucket_hour") query := `SELECT COALESCE(SUM(call_count), 0), COALESCE(SUM(error_count), 0), COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), COALESCE(SUM(total_tokens), 0), COALESCE(SUM(cache_read_tokens), 0), COALESCE(SUM(cache_create_tokens), 0), COALESCE(SUM(thinking_tokens), 0), COALESCE(SUM(cost_usd), 0), CASE WHEN COALESCE(SUM(call_count), 0) > 0 THEN COALESCE(SUM(duration_ms * call_count), 0) / SUM(call_count) ELSE 0 END FROM usage_event_rollups` + where var summary store.UsageEventSummary if err := s.db.QueryRowContext(ctx, query, args...).Scan( &summary.Calls, &summary.Errors, &summary.InputTokens, &summary.OutputTokens, &summary.TotalTokens, &summary.CacheReadTokens, &summary.CacheCreateTokens, &summary.ThinkingTokens, &summary.CostUSD, &summary.AvgDurationMS, ); err != nil { return nil, fmt.Errorf("get usage event summary: %w", err) } return &summary, nil } func prepareUsageEvent(ctx context.Context, event *store.UsageEvent) { if event.ID == uuid.Nil { event.ID = uuid.New() } if event.TenantID == uuid.Nil { event.TenantID = store.TenantIDFromContext(ctx) } if event.TenantID == uuid.Nil { event.TenantID = store.MasterTenantID } if event.EventTime.IsZero() { event.EventTime = time.Now().UTC() } event.EventTime = event.EventTime.UTC() if event.BucketHour.IsZero() { event.BucketHour = event.EventTime.Truncate(time.Hour) } if event.CallCount <= 0 { event.CallCount = 1 } if event.Status == "" { event.Status = "completed" } if event.CreatedAt.IsZero() { event.CreatedAt = time.Now().UTC() } } func buildUsageEventWhere(ctx context.Context, q store.UsageEventQuery, timeColumn string) (string, []any) { var conds []string var args []any idx := 1 if !store.IsCrossTenant(ctx) { if tenantID := store.TenantIDFromContext(ctx); tenantID != uuid.Nil { conds = append(conds, fmt.Sprintf("tenant_id = $%d", idx)) args = append(args, tenantID) idx++ } } add := func(col string, value any) { conds = append(conds, fmt.Sprintf("%s = $%d", col, idx)) args = append(args, value) idx++ } if !q.From.IsZero() { conds = append(conds, fmt.Sprintf("%s >= $%d", timeColumn, idx)) args = append(args, q.From.UTC()) idx++ } if !q.To.IsZero() { conds = append(conds, fmt.Sprintf("%s < $%d", timeColumn, idx)) args = append(args, q.To.UTC()) idx++ } if q.AgentID != nil { add("agent_id", *q.AgentID) } if q.Channel != "" { add("channel", q.Channel) } if q.EventType != "" { add("event_type", q.EventType) } if q.ResourceType != "" { add("resource_type", q.ResourceType) } if q.ResourceName != "" { add("resource_name", q.ResourceName) } if q.Provider != "" { add("provider", q.Provider) } if q.Model != "" { add("model", q.Model) } if q.Status != "" { add("status", q.Status) } if q.Source != "" { add("source", q.Source) } if len(conds) == 0 { return "", nil } return " WHERE " + strings.Join(conds, " AND "), args } func usageEventGroupColumn(groupBy string) string { switch groupBy { case "event_type": return "event_type" case "resource_type": return "resource_type" case "source": return "source" case "status": return "status" case "agent": return "COALESCE(agent_id::TEXT, '')" case "channel": return "channel" case "provider": return "provider" case "model": return "model" default: return "resource_name" } }