sub2api-cn-relay-manager/internal/store/sqlite/user_keys_repo.go

package sqlite

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"strings"
)

type UserKeyRecord struct {
	ID             int64    `json:"-"`
	KeyID          string   `json:"key_id"`
	OwnerSubjectID string   `json:"owner_subject_id"`
	KeyFingerprint string   `json:"key_fingerprint"`
	MaskedPreview  string   `json:"masked_preview"`
	DisplayName    string   `json:"display_name"`
	LogicalGroupID string   `json:"logical_group_id"`
	AllowedModels  []string `json:"allowed_models"`
	AdminStatus    string   `json:"admin_status"`
	QuotaStatus    string   `json:"quota_status"`
	LastUsedAt     string   `json:"last_used_at,omitempty"`
	CreatedAt      string   `json:"created_at"`
	ExpiresAt      string   `json:"expires_at,omitempty"`
	UpdatedAt      string   `json:"updated_at"`
}

type UserKeysRepo struct {
	db execQuerier
}

func newUserKeysRepo(db execQuerier) *UserKeysRepo {
	return &UserKeysRepo{db: db}
}

func (r *UserKeysRepo) Create(ctx context.Context, key UserKeyRecord) (int64, error) {
	modelsJSON, err := json.Marshal(key.AllowedModels)
	if err != nil {
		return 0, fmt.Errorf("marshal allowed_models: %w", err)
	}
	result, err := r.db.ExecContext(ctx, `
		INSERT INTO user_keys (
			key_id, owner_subject_id, key_fingerprint, masked_preview,
			display_name, logical_group_id, allowed_models,
			admin_status, quota_status
		) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
		key.KeyID, key.OwnerSubjectID, key.KeyFingerprint, key.MaskedPreview,
		key.DisplayName, key.LogicalGroupID, string(modelsJSON),
		key.AdminStatus, key.QuotaStatus,
	)
	if err != nil {
		return 0, fmt.Errorf("insert user_key: %w", err)
	}
	return result.LastInsertId()
}

func scanUserKeys(rows *sql.Rows) ([]UserKeyRecord, error) {
	var keys []UserKeyRecord
	for rows.Next() {
		var k UserKeyRecord
		var modelsJSON, lastUsedAt, expiresAt sql.NullString
		err := rows.Scan(
			&k.ID, &k.KeyID, &k.OwnerSubjectID, &k.KeyFingerprint, &k.MaskedPreview,
			&k.DisplayName, &k.LogicalGroupID, &modelsJSON,
			&k.AdminStatus, &k.QuotaStatus, &lastUsedAt, &k.CreatedAt, &expiresAt, &k.UpdatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("scan user_key: %w", err)
		}
		k.LastUsedAt = lastUsedAt.String
		k.ExpiresAt = expiresAt.String
		if modelsJSON.String != "" {
			json.Unmarshal([]byte(modelsJSON.String), &k.AllowedModels)
		}
		keys = append(keys, k)
	}
	return keys, rows.Err()
}

func scanOneUserKey(row *sql.Row) (*UserKeyRecord, error) {
	var k UserKeyRecord
	var modelsJSON, lastUsedAt, expiresAt sql.NullString
	err := row.Scan(
		&k.ID, &k.KeyID, &k.OwnerSubjectID, &k.KeyFingerprint, &k.MaskedPreview,
		&k.DisplayName, &k.LogicalGroupID, &modelsJSON,
		&k.AdminStatus, &k.QuotaStatus, &lastUsedAt, &k.CreatedAt, &expiresAt, &k.UpdatedAt,
	)
	if err != nil {
		return nil, err
	}
	k.LastUsedAt = lastUsedAt.String
	k.ExpiresAt = expiresAt.String
	if modelsJSON.String != "" {
		json.Unmarshal([]byte(modelsJSON.String), &k.AllowedModels)
	}
	return &k, nil
}

func (r *UserKeysRepo) ListByOwner(ctx context.Context, subjectID string) ([]UserKeyRecord, error) {
	rows, err := r.db.QueryContext(ctx, `
		SELECT id, key_id, owner_subject_id, key_fingerprint, masked_preview,
		       display_name, logical_group_id, allowed_models,
		       admin_status, quota_status, last_used_at, created_at, expires_at, updated_at
		FROM user_keys WHERE owner_subject_id = ? ORDER BY created_at DESC`, subjectID)
	if err != nil {
		return nil, fmt.Errorf("list user_keys: %w", err)
	}
	defer rows.Close()
	return scanUserKeys(rows)
}

func (r *UserKeysRepo) ListByFingerprint(ctx context.Context, fingerprint string) ([]UserKeyRecord, error) {
	rows, err := r.db.QueryContext(ctx, `
		SELECT id, key_id, owner_subject_id, key_fingerprint, masked_preview,
		       display_name, logical_group_id, allowed_models,
		       admin_status, quota_status, last_used_at, created_at, expires_at, updated_at
		FROM user_keys WHERE key_fingerprint = ? ORDER BY created_at DESC`, fingerprint)
	if err != nil {
		return nil, fmt.Errorf("list user_keys by fingerprint: %w", err)
	}
	defer rows.Close()
	return scanUserKeys(rows)
}

func (r *UserKeysRepo) GetByID(ctx context.Context, keyID string) (*UserKeyRecord, error) {
	row := r.db.QueryRowContext(ctx, `
		SELECT id, key_id, owner_subject_id, key_fingerprint, masked_preview,
		       display_name, logical_group_id, allowed_models,
		       admin_status, quota_status, last_used_at, created_at, expires_at, updated_at
		FROM user_keys WHERE key_id = ?`, keyID)
	k, err := scanOneUserKey(row)
	if err != nil {
		return nil, fmt.Errorf("get user_key %s: %w", keyID, err)
	}
	return k, nil
}

func (r *UserKeysRepo) UpdateStatus(ctx context.Context, keyID string, adminStatus string) error {
	status := strings.ToLower(strings.TrimSpace(adminStatus))
	valid := map[string]bool{"active": true, "paused": true, "disabled": true, "retired": true}
	if !valid[status] {
		return fmt.Errorf("invalid admin_status: %s", adminStatus)
	}
	result, err := r.db.ExecContext(ctx,
		`UPDATE user_keys SET admin_status = ?, updated_at = strftime('%Y-%m-%dT%H:%M:%SZ','now') WHERE key_id = ?`,
		status, keyID)
	if err != nil {
		return fmt.Errorf("update user_key status: %w", err)
	}
	rowsAffected, rowsErr := result.RowsAffected()
	if rowsErr == nil && rowsAffected == 0 {
		return fmt.Errorf("update user_key status: %w", sql.ErrNoRows)
	}
	return nil
}

func (r *UserKeysRepo) UpdateSecret(ctx context.Context, keyID, fingerprint, maskedPreview, adminStatus string) error {
	keyID = strings.TrimSpace(keyID)
	fingerprint = strings.TrimSpace(fingerprint)
	maskedPreview = strings.TrimSpace(maskedPreview)
	adminStatus = strings.ToLower(strings.TrimSpace(adminStatus))
	if keyID == "" {
		return fmt.Errorf("key_id is required")
	}
	if fingerprint == "" {
		return fmt.Errorf("key_fingerprint is required")
	}
	if maskedPreview == "" {
		return fmt.Errorf("masked_preview is required")
	}
	valid := map[string]bool{"active": true, "paused": true, "disabled": true, "retired": true}
	if !valid[adminStatus] {
		return fmt.Errorf("invalid admin_status: %s", adminStatus)
	}
	result, err := r.db.ExecContext(ctx,
		`UPDATE user_keys
		 SET key_fingerprint = ?, masked_preview = ?, admin_status = ?, updated_at = strftime('%Y-%m-%dT%H:%M:%SZ','now')
		 WHERE key_id = ?`,
		fingerprint, maskedPreview, adminStatus, keyID,
	)
	if err != nil {
		return fmt.Errorf("update user_key secret: %w", err)
	}
	rowsAffected, rowsErr := result.RowsAffected()
	if rowsErr == nil && rowsAffected == 0 {
		return fmt.Errorf("update user_key secret: %w", sql.ErrNoRows)
	}
	return nil
}

func (r *UserKeysRepo) TouchLastUsed(ctx context.Context, keyID string) error {
	result, err := r.db.ExecContext(ctx,
		`UPDATE user_keys SET last_used_at = strftime('%Y-%m-%dT%H:%M:%SZ','now') WHERE key_id = ?`, keyID)
	if err != nil {
		return err
	}
	rowsAffected, rowsErr := result.RowsAffected()
	if rowsErr == nil && rowsAffected == 0 {
		return fmt.Errorf("touch user key last_used_at: %w", sql.ErrNoRows)
	}
	return nil
}