package usage import ( "context" "database/sql" "time" ) type Repository struct { db *sql.DB } func NewRepository(db *sql.DB) *Repository { return &Repository{db: db} } func (r *Repository) RunMigrations(ctx context.Context) error { migrations := []string{ `CREATE TABLE IF NOT EXISTS usage_metrics ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), user_id UUID NOT NULL, date DATE NOT NULL, tier VARCHAR(50) NOT NULL DEFAULT 'free', api_requests INT DEFAULT 0, llm_requests INT DEFAULT 0, llm_tokens INT DEFAULT 0, search_requests INT DEFAULT 0, storage_used BIGINT DEFAULT 0, created_at TIMESTAMPTZ DEFAULT NOW(), updated_at TIMESTAMPTZ DEFAULT NOW(), UNIQUE(user_id, date) )`, `CREATE INDEX IF NOT EXISTS idx_usage_user_date ON usage_metrics(user_id, date DESC)`, `CREATE INDEX IF NOT EXISTS idx_usage_date ON usage_metrics(date)`, } for _, m := range migrations { if _, err := r.db.ExecContext(ctx, m); err != nil { return err } } return nil } func (r *Repository) GetTodayUsage(ctx context.Context, userID string) (*UsageMetric, error) { today := time.Now().Format("2006-01-02") return r.GetUsageByDate(ctx, userID, today) } func (r *Repository) GetUsageByDate(ctx context.Context, userID, date string) (*UsageMetric, error) { query := ` SELECT id, user_id, date, tier, api_requests, llm_requests, llm_tokens, search_requests, storage_used, created_at, updated_at FROM usage_metrics WHERE user_id = $1 AND date = $2 ` var m UsageMetric var dateVal time.Time err := r.db.QueryRowContext(ctx, query, userID, date).Scan( &m.ID, &m.UserID, &dateVal, &m.Tier, &m.APIRequests, &m.LLMRequests, &m.LLMTokens, &m.SearchReqs, &m.StorageUsed, &m.CreatedAt, &m.UpdatedAt, ) if err == sql.ErrNoRows { return nil, nil } if err != nil { return nil, err } m.Date = dateVal.Format("2006-01-02") return &m, nil } func (r *Repository) IncrementAPIRequests(ctx context.Context, userID, tier string) error { today := time.Now().Format("2006-01-02") query := ` INSERT INTO usage_metrics (user_id, date, tier, api_requests) VALUES ($1, $2, $3, 1) ON CONFLICT (user_id, date) DO UPDATE SET api_requests = usage_metrics.api_requests + 1, updated_at = NOW() ` _, err := r.db.ExecContext(ctx, query, userID, today, tier) return err } func (r *Repository) IncrementLLMUsage(ctx context.Context, userID, tier string, tokens int) error { today := time.Now().Format("2006-01-02") query := ` INSERT INTO usage_metrics (user_id, date, tier, llm_requests, llm_tokens) VALUES ($1, $2, $3, 1, $4) ON CONFLICT (user_id, date) DO UPDATE SET llm_requests = usage_metrics.llm_requests + 1, llm_tokens = usage_metrics.llm_tokens + $4, updated_at = NOW() ` _, err := r.db.ExecContext(ctx, query, userID, today, tier, tokens) return err } func (r *Repository) IncrementSearchRequests(ctx context.Context, userID, tier string) error { today := time.Now().Format("2006-01-02") query := ` INSERT INTO usage_metrics (user_id, date, tier, search_requests) VALUES ($1, $2, $3, 1) ON CONFLICT (user_id, date) DO UPDATE SET search_requests = usage_metrics.search_requests + 1, updated_at = NOW() ` _, err := r.db.ExecContext(ctx, query, userID, today, tier) return err } func (r *Repository) UpdateStorageUsed(ctx context.Context, userID, tier string, bytes int64) error { today := time.Now().Format("2006-01-02") query := ` INSERT INTO usage_metrics (user_id, date, tier, storage_used) VALUES ($1, $2, $3, $4) ON CONFLICT (user_id, date) DO UPDATE SET storage_used = $4, updated_at = NOW() ` _, err := r.db.ExecContext(ctx, query, userID, today, tier, bytes) return err } func (r *Repository) GetUsageHistory(ctx context.Context, userID string, days int) ([]*UsageMetric, error) { query := ` SELECT id, user_id, date, tier, api_requests, llm_requests, llm_tokens, search_requests, storage_used, created_at, updated_at FROM usage_metrics WHERE user_id = $1 AND date >= CURRENT_DATE - $2::int ORDER BY date DESC ` rows, err := r.db.QueryContext(ctx, query, userID, days) if err != nil { return nil, err } defer rows.Close() var metrics []*UsageMetric for rows.Next() { var m UsageMetric var dateVal time.Time if err := rows.Scan( &m.ID, &m.UserID, &dateVal, &m.Tier, &m.APIRequests, &m.LLMRequests, &m.LLMTokens, &m.SearchReqs, &m.StorageUsed, &m.CreatedAt, &m.UpdatedAt, ); err != nil { return nil, err } m.Date = dateVal.Format("2006-01-02") metrics = append(metrics, &m) } return metrics, nil } func (r *Repository) CheckLLMLimits(ctx context.Context, userID, tier string) (bool, string) { usage, err := r.GetTodayUsage(ctx, userID) if err != nil { return true, "" } limits := GetLimits(tier) if usage != nil { if usage.LLMRequests >= limits.LLMRequestsPerDay { return false, "Daily LLM request limit exceeded" } if usage.LLMTokens >= limits.LLMTokensPerDay { return false, "Daily LLM token limit exceeded" } } return true, "" }