package audit import ( "context" "encoding/json" "fmt" "strconv" "strings" "time" "github.com/jackc/pgx/v5/pgxpool" ) type QueryService struct{ pool *pgxpool.Pool } func NewQueryService(pool *pgxpool.Pool) *QueryService { return &QueryService{pool: pool} } type EventView struct { ID string `json:"id"` TenantID *string `json:"tenant_id"` RequestID string `json:"request_id"` APIKeyID *string `json:"api_key_id"` APIKeyName *string `json:"api_key_name"` ProviderCode *string `json:"provider_code"` Model *string `json:"model"` Protocol string `json:"protocol"` StatusCode *int `json:"status_code"` PromptTokens *int64 `json:"prompt_tokens"` CompletionTokens *int64 `json:"completion_tokens"` CostMicrounits *int64 `json:"cost_microunits"` LatencyMS *int `json:"latency_ms"` Labels json.RawMessage `json:"labels"` RecordedAt time.Time `json:"recorded_at"` } type EventFilter struct { From, To time.Time Before *time.Time APIKeyID string Provider string Model string StatusCode *int Limit int } func (s *QueryService) ListEvents(ctx context.Context, filter EventFilter) ([]EventView, error) { if s == nil || s.pool == nil { return nil, fmt.Errorf("audit store unavailable") } where := []string{"a.recorded_at >= $1", "a.recorded_at < $2"} args := []any{filter.From, filter.To} add := func(condition string, value any) { args = append(args, value) where = append(where, fmt.Sprintf(condition, len(args))) } if filter.Before != nil { add("a.recorded_at < $%d", *filter.Before) } if filter.APIKeyID != "" { add("a.api_key_id = $%d", filter.APIKeyID) } if filter.Provider != "" { add("a.provider_code = $%d", filter.Provider) } if filter.Model != "" { add("a.model = $%d", filter.Model) } if filter.StatusCode != nil { add("a.status_code = $%d", *filter.StatusCode) } args = append(args, filter.Limit) query := `SELECT a.id::text,a.tenant_id::text,a.request_id,a.api_key_id::text,k.name, a.provider_code,a.model,a.protocol,a.status_code,a.prompt_tokens,a.completion_tokens,a.cost_microunits,a.latency_ms,a.labels,a.recorded_at FROM gateway.audit_events a LEFT JOIN gateway.api_keys k ON k.id=a.api_key_id WHERE ` + strings.Join(where, " AND ") + ` ORDER BY a.recorded_at DESC,a.id DESC LIMIT $` + strconv.Itoa(len(args)) rows, err := s.pool.Query(ctx, query, args...) if err != nil { return nil, fmt.Errorf("query audit events: %w", err) } defer rows.Close() items := make([]EventView, 0) for rows.Next() { var item EventView if err := rows.Scan(&item.ID, &item.TenantID, &item.RequestID, &item.APIKeyID, &item.APIKeyName, &item.ProviderCode, &item.Model, &item.Protocol, &item.StatusCode, &item.PromptTokens, &item.CompletionTokens, &item.CostMicrounits, &item.LatencyMS, &item.Labels, &item.RecordedAt); err != nil { return nil, fmt.Errorf("scan audit event: %w", err) } items = append(items, item) } return items, rows.Err() } type DailyUsageView struct { Date time.Time `json:"date"` APIKeyID string `json:"api_key_id"` APIKeyName string `json:"api_key_name"` ProviderCode string `json:"provider_code"` Model string `json:"model"` Currency string `json:"currency"` Requests int64 `json:"requests"` FailedRequests int64 `json:"failed_requests"` PromptTokens int64 `json:"prompt_tokens"` CompletionTokens int64 `json:"completion_tokens"` CostMicrounits int64 `json:"cost_microunits"` } type UsageFilter struct { From, To time.Time APIKeyID string Provider string Model string } func (s *QueryService) ListDailyUsage(ctx context.Context, filter UsageFilter) ([]DailyUsageView, error) { if s == nil || s.pool == nil { return nil, fmt.Errorf("usage store unavailable") } where := []string{"u.usage_date >= $1", "u.usage_date <= $2"} args := []any{filter.From, filter.To} add := func(column string, value any) { args = append(args, value) where = append(where, column+" = $"+strconv.Itoa(len(args))) } if filter.APIKeyID != "" { add("u.api_key_id", filter.APIKeyID) } if filter.Provider != "" { add("u.provider_code", filter.Provider) } if filter.Model != "" { add("u.model", filter.Model) } rows, err := s.pool.Query(ctx, `SELECT u.usage_date,u.api_key_id::text,k.name,u.provider_code,u.model,u.currency, u.requests,u.failed_requests,u.prompt_tokens,u.completion_tokens,u.cost_microunits FROM gateway.usage_daily u JOIN gateway.api_keys k ON k.id=u.api_key_id WHERE `+ strings.Join(where, " AND ")+` ORDER BY u.usage_date DESC,k.name,u.provider_code,u.model`, args...) if err != nil { return nil, fmt.Errorf("query daily usage: %w", err) } defer rows.Close() items := make([]DailyUsageView, 0) for rows.Next() { var item DailyUsageView if err := rows.Scan(&item.Date, &item.APIKeyID, &item.APIKeyName, &item.ProviderCode, &item.Model, &item.Currency, &item.Requests, &item.FailedRequests, &item.PromptTokens, &item.CompletionTokens, &item.CostMicrounits); err != nil { return nil, fmt.Errorf("scan daily usage: %w", err) } items = append(items, item) } return items, rows.Err() }