ai-gateway-go/internal/audit/query.go

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()
}