pipu-music/services/api/internal/auth/pgstore.go

package auth

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
)

// querier — общее у пула и транзакции.
type querier interface {
	Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
	Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}

// PgStore — Store поверх PostgreSQL.
type PgStore struct {
	pool *pgxpool.Pool
	q    querier
}

// NewPgStore создаёт Store на пуле соединений.
func NewPgStore(pool *pgxpool.Pool) *PgStore { return &PgStore{pool: pool, q: pool} }

// WithTx реализует Store. Вложенный вызов выполняется в уже открытой транзакции.
func (s *PgStore) WithTx(ctx context.Context, fn func(Store) error) error {
	if s.pool == nil {
		return fn(s)
	}
	return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { return fn(&PgStore{q: tx}) })
}

const userColumns = `id, username, display_name, password_hash, role, locale, ui_prefs, must_change_password,
	must_change_credentials, is_blocked, token_version, failed_login_count, locked_until, created_at, last_login_at`

func scanUser(row pgx.Row) (User, error) {
	var (
		u     User
		role  string
		prefs []byte
	)
	err := row.Scan(&u.ID, &u.Username, &u.DisplayName, &u.PasswordHash, &role, &u.Locale, &prefs,
		&u.MustChangePassword, &u.MustChangeCredentials, &u.IsBlocked, &u.TokenVersion, &u.FailedLoginCount,
		&u.LockedUntil, &u.CreatedAt, &u.LastLoginAt)
	if errors.Is(err, pgx.ErrNoRows) {
		return User{}, ErrNotFound
	}
	if err != nil {
		return User{}, fmt.Errorf("auth: scan user: %w", err)
	}
	u.Role = Role(role)
	if err := json.Unmarshal(prefs, &u.UIPrefs); err != nil {
		return User{}, fmt.Errorf("auth: decode ui_prefs: %w", err)
	}
	return u, nil
}

// CountAdmins реализует Store. Заблокированные админы не считаются: войти они не могут.
func (s *PgStore) CountAdmins(ctx context.Context) (int, error) {
	var n int
	err := s.q.QueryRow(ctx, `SELECT count(*) FROM users WHERE role = 'admin' AND NOT is_blocked`).Scan(&n)
	if err != nil {
		return 0, fmt.Errorf("auth: count admins: %w", err)
	}
	return n, nil
}

// CreateUser реализует Store.
func (s *PgStore) CreateUser(ctx context.Context, u User, createdBy *uuid.UUID) error {
	prefs, err := json.Marshal(u.UIPrefs)
	if err != nil {
		return err
	}
	_, err = s.q.Exec(ctx, `
		INSERT INTO users (id, username, display_name, password_hash, role, locale, ui_prefs,
			must_change_password, must_change_credentials, created_by, created_at, updated_at)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $11)`,
		u.ID, u.Username, u.DisplayName, u.PasswordHash, string(u.Role), u.Locale, prefs,
		u.MustChangePassword, u.MustChangeCredentials, createdBy, u.CreatedAt)
	if isUniqueViolation(err) {
		return ErrUsernameTaken
	}
	if err != nil {
		return fmt.Errorf("auth: create user: %w", err)
	}
	return nil
}

func isUniqueViolation(err error) bool {
	var pgErr *pgconn.PgError
	return errors.As(err, &pgErr) && pgErr.Code == "23505"
}

// GetUserByID реализует Store.
func (s *PgStore) GetUserByID(ctx context.Context, id uuid.UUID) (User, error) {
	return scanUser(s.q.QueryRow(ctx, `SELECT `+userColumns+` FROM users WHERE id = $1`, id))
}

// GetUserByUsername реализует Store (username — citext, регистр не важен).
func (s *PgStore) GetUserByUsername(ctx context.Context, username string) (User, error) {
	return scanUser(s.q.QueryRow(ctx, `SELECT `+userColumns+` FROM users WHERE username = $1`, username))
}

// RecordLoginFailure реализует Store.
func (s *PgStore) RecordLoginFailure(ctx context.Context, id uuid.UUID, failedCount int, lockedUntil *time.Time) error {
	_, err := s.q.Exec(ctx, `UPDATE users SET failed_login_count = $2, locked_until = $3 WHERE id = $1`,
		id, failedCount, lockedUntil)
	if err != nil {
		return fmt.Errorf("auth: record login failure: %w", err)
	}
	return nil
}

// RecordLoginSuccess реализует Store.
func (s *PgStore) RecordLoginSuccess(ctx context.Context, id uuid.UUID, at time.Time, newHash *string) error {
	_, err := s.q.Exec(ctx, `
		UPDATE users SET failed_login_count = 0, locked_until = NULL, last_login_at = $2,
			password_hash = COALESCE($3, password_hash)
		WHERE id = $1`, id, at, newHash)
	if err != nil {
		return fmt.Errorf("auth: record login success: %w", err)
	}
	return nil
}

// SetPassword реализует Store.
func (s *PgStore) SetPassword(ctx context.Context, id uuid.UUID, hash string) error {
	_, err := s.q.Exec(ctx, `
		UPDATE users SET password_hash = $2, must_change_password = false,
			token_version = token_version + 1, updated_at = now()
		WHERE id = $1`, id, hash)
	if err != nil {
		return fmt.Errorf("auth: set password: %w", err)
	}
	return nil
}

// SetCredentials реализует Store.
func (s *PgStore) SetCredentials(ctx context.Context, id uuid.UUID, username, hash string) error {
	_, err := s.q.Exec(ctx, `
		UPDATE users SET username = $2, password_hash = $3, must_change_credentials = false,
			must_change_password = false, token_version = token_version + 1, updated_at = now()
		WHERE id = $1`, id, username, hash)
	if isUniqueViolation(err) {
		return ErrUsernameTaken
	}
	if err != nil {
		return fmt.Errorf("auth: set credentials: %w", err)
	}
	return nil
}

// UpdateProfile реализует Store.
func (s *PgStore) UpdateProfile(ctx context.Context, id uuid.UUID, upd ProfileUpdate) (User, error) {
	var (
		setName bool
		name    *string
		prefs   []byte
		err     error
	)
	if upd.DisplayName != nil {
		setName, name = true, *upd.DisplayName
	}
	if upd.UIPrefs != nil {
		if prefs, err = json.Marshal(upd.UIPrefs); err != nil {
			return User{}, err
		}
	}
	return scanUser(s.q.QueryRow(ctx, `
		UPDATE users SET
			display_name = CASE WHEN $2 THEN $3 ELSE display_name END,
			locale = COALESCE($4, locale),
			ui_prefs = COALESCE($5::jsonb, ui_prefs),
			updated_at = now()
		WHERE id = $1
		RETURNING `+userColumns, id, setName, name, upd.Locale, prefs))
}

// InsertRefreshToken реализует Store.
func (s *PgStore) InsertRefreshToken(ctx context.Context, t RefreshToken) error {
	_, err := s.q.Exec(ctx, `
		INSERT INTO refresh_tokens (id, user_id, family_id, family_started_at, token_hash, device_name,
			client_type, created_at, expires_at)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`,
		t.ID, t.UserID, t.FamilyID, t.FamilyStartedAt, t.TokenHash, t.DeviceName, string(t.ClientType),
		t.CreatedAt, t.ExpiresAt)
	if err != nil {
		return fmt.Errorf("auth: insert refresh token: %w", err)
	}
	return nil
}

// GetRefreshTokenForUpdate реализует Store.
func (s *PgStore) GetRefreshTokenForUpdate(ctx context.Context, hash []byte) (RefreshToken, error) {
	var (
		t  RefreshToken
		ct string
	)
	err := s.q.QueryRow(ctx, `
		SELECT id, user_id, family_id, family_started_at, token_hash, device_name, client_type,
			created_at, expires_at, used_at, revoked_at
		FROM refresh_tokens WHERE token_hash = $1 FOR UPDATE`, hash).
		Scan(&t.ID, &t.UserID, &t.FamilyID, &t.FamilyStartedAt, &t.TokenHash, &t.DeviceName, &ct,
			&t.CreatedAt, &t.ExpiresAt, &t.UsedAt, &t.RevokedAt)
	if errors.Is(err, pgx.ErrNoRows) {
		return RefreshToken{}, ErrNotFound
	}
	if err != nil {
		return RefreshToken{}, fmt.Errorf("auth: get refresh token: %w", err)
	}
	t.ClientType = ClientType(ct)
	return t, nil
}

// MarkRefreshTokenUsed реализует Store.
func (s *PgStore) MarkRefreshTokenUsed(ctx context.Context, id uuid.UUID, at time.Time) error {
	_, err := s.q.Exec(ctx, `UPDATE refresh_tokens SET used_at = $2 WHERE id = $1`, id, at)
	if err != nil {
		return fmt.Errorf("auth: mark refresh token used: %w", err)
	}
	return nil
}

// RevokeFamily реализует Store.
func (s *PgStore) RevokeFamily(ctx context.Context, familyID uuid.UUID, at time.Time) error {
	_, err := s.q.Exec(ctx, `UPDATE refresh_tokens SET revoked_at = $2 WHERE family_id = $1 AND revoked_at IS NULL`,
		familyID, at)
	if err != nil {
		return fmt.Errorf("auth: revoke family: %w", err)
	}
	return nil
}

// RevokeUserSessions реализует Store.
func (s *PgStore) RevokeUserSessions(ctx context.Context, userID uuid.UUID, at time.Time) error {
	_, err := s.q.Exec(ctx, `UPDATE refresh_tokens SET revoked_at = $2 WHERE user_id = $1 AND revoked_at IS NULL`,
		userID, at)
	if err != nil {
		return fmt.Errorf("auth: revoke user sessions: %w", err)
	}
	return nil
}

// ListSessions реализует Store.
func (s *PgStore) ListSessions(ctx context.Context, userID uuid.UUID, now time.Time) ([]Session, error) {
	rows, err := s.q.Query(ctx, `
		SELECT family_id, device_name, client_type, family_started_at, created_at
		FROM refresh_tokens
		WHERE user_id = $1 AND used_at IS NULL AND revoked_at IS NULL AND expires_at > $2
		ORDER BY created_at DESC`, userID, now)
	if err != nil {
		return nil, fmt.Errorf("auth: list sessions: %w", err)
	}
	out, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (Session, error) {
		var (
			ss Session
			ct string
		)
		err := row.Scan(&ss.ID, &ss.DeviceName, &ct, &ss.CreatedAt, &ss.LastUsedAt)
		ss.ClientType = ClientType(ct)
		return ss, err
	})
	if err != nil {
		return nil, fmt.Errorf("auth: list sessions: %w", err)
	}
	return out, nil
}

// InsertAudit реализует Store.
func (s *PgStore) InsertAudit(ctx context.Context, id uuid.UUID, e AuditEntry) error {
	var payload []byte
	if e.Payload != nil {
		var err error
		if payload, err = json.Marshal(e.Payload); err != nil {
			return err
		}
	}
	_, err := s.q.Exec(ctx, `
		INSERT INTO audit_log (id, actor_id, action, target_type, target_id, payload)
		VALUES ($1, $2, $3, NULLIF($4, ''), $5, $6)`, id, e.ActorID, e.Action, e.TargetType, e.TargetID, payload)
	if err != nil {
		return fmt.Errorf("auth: insert audit: %w", err)
	}
	return nil
}

// GetSettingBool реализует Store: отсутствующий ключ — false.
func (s *PgStore) GetSettingBool(ctx context.Context, key string) (bool, error) {
	var v bool
	err := s.q.QueryRow(ctx, `SELECT value::boolean FROM app_settings WHERE key = $1`, key).Scan(&v)
	if errors.Is(err, pgx.ErrNoRows) {
		return false, nil
	}
	if err != nil {
		return false, fmt.Errorf("auth: get setting %s: %w", key, err)
	}
	return v, nil
}

// SetSetting реализует Store.
func (s *PgStore) SetSetting(ctx context.Context, key string, value any) error {
	raw, err := json.Marshal(value)
	if err != nil {
		return err
	}
	_, err = s.q.Exec(ctx, `
		INSERT INTO app_settings (key, value, updated_at) VALUES ($1, $2, now())
		ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = now()`, key, raw)
	if err != nil {
		return fmt.Errorf("auth: set setting %s: %w", key, err)
	}
	return nil
}

var _ Store = (*PgStore)(nil)