memos/store/user_setting.go

package store

import (
	"context"

	"github.com/pkg/errors"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/timestamppb"

	storepb "github.com/usememos/memos/proto/gen/store"
)

type UserSetting struct {
	UserID int32
	Key    storepb.UserSetting_Key
	Value  string
}

type FindUserSetting struct {
	UserID *int32
	Key    storepb.UserSetting_Key
}

type DeleteUserSetting struct {
	UserID *int32
	Key    storepb.UserSetting_Key
}

// RefreshTokenQueryResult contains the result of querying a refresh token.
type RefreshTokenQueryResult struct {
	UserID       int32
	RefreshToken *storepb.RefreshTokensUserSetting_RefreshToken
}

// PATQueryResult contains the result of querying a PAT by hash.
type PATQueryResult struct {
	UserID int32
	User   *User
	PAT    *storepb.PersonalAccessTokensUserSetting_PersonalAccessToken
}

func (s *Store) UpsertUserSetting(ctx context.Context, upsert *storepb.UserSetting) (*storepb.UserSetting, error) {
	userSettingRaw, err := convertUserSettingToRaw(upsert)
	if err != nil {
		return nil, err
	}
	userSettingRaw, err = s.driver.UpsertUserSetting(ctx, userSettingRaw)
	if err != nil {
		return nil, err
	}

	userSetting, err := convertUserSettingFromRaw(userSettingRaw)
	if err != nil {
		return nil, err
	}
	if userSetting == nil {
		return nil, errors.New("unexpected nil user setting")
	}
	s.userSettingCache.Set(ctx, getUserSettingCacheKey(userSetting.UserId, userSetting.Key.String()), userSetting)
	return userSetting, nil
}

func (s *Store) ListUserSettings(ctx context.Context, find *FindUserSetting) ([]*storepb.UserSetting, error) {
	userSettingRawList, err := s.driver.ListUserSettings(ctx, find)
	if err != nil {
		return nil, err
	}

	userSettings := []*storepb.UserSetting{}
	for _, userSettingRaw := range userSettingRawList {
		userSetting, err := convertUserSettingFromRaw(userSettingRaw)
		if err != nil {
			return nil, err
		}
		if userSetting == nil {
			continue
		}
		s.userSettingCache.Set(ctx, getUserSettingCacheKey(userSetting.UserId, userSetting.Key.String()), userSetting)
		userSettings = append(userSettings, userSetting)
	}
	return userSettings, nil
}

func (s *Store) GetUserSetting(ctx context.Context, find *FindUserSetting) (*storepb.UserSetting, error) {
	if find.UserID != nil {
		if cache, ok := s.userSettingCache.Get(ctx, getUserSettingCacheKey(*find.UserID, find.Key.String())); ok {
			userSetting, ok := cache.(*storepb.UserSetting)
			if ok {
				return userSetting, nil
			}
		}
	}

	list, err := s.ListUserSettings(ctx, find)
	if err != nil {
		return nil, err
	}
	if len(list) == 0 {
		return nil, nil
	}
	if len(list) > 1 {
		return nil, errors.Errorf("expected 1 user setting, but got %d", len(list))
	}

	userSetting := list[0]
	s.userSettingCache.Set(ctx, getUserSettingCacheKey(userSetting.UserId, userSetting.Key.String()), userSetting)
	return userSetting, nil
}

func (s *Store) DeleteUserSettings(ctx context.Context, delete *DeleteUserSetting) error {
	existing, err := s.ListUserSettings(ctx, &FindUserSetting{
		UserID: delete.UserID,
		Key:    delete.Key,
	})
	if err != nil {
		return err
	}
	if err := s.driver.DeleteUserSettings(ctx, delete); err != nil {
		return err
	}
	for _, setting := range existing {
		s.userSettingCache.Delete(ctx, getUserSettingCacheKey(setting.UserId, setting.Key.String()))
	}
	return nil
}

// GetUserByPATHash finds a user by PAT hash.
func (s *Store) GetUserByPATHash(ctx context.Context, tokenHash string) (*PATQueryResult, error) {
	result, err := s.driver.GetUserByPATHash(ctx, tokenHash)
	if err != nil {
		return nil, err
	}

	// Fetch user info
	user, err := s.GetUser(ctx, &FindUser{ID: &result.UserID})
	if err != nil {
		return nil, err
	}
	if user == nil {
		return nil, errors.New("user not found for PAT")
	}
	result.User = user

	return result, nil
}

// GetUserRefreshTokens returns the refresh tokens of the user.
func (s *Store) GetUserRefreshTokens(ctx context.Context, userID int32) ([]*storepb.RefreshTokensUserSetting_RefreshToken, error) {
	userSetting, err := s.GetUserSetting(ctx, &FindUserSetting{
		UserID: &userID,
		Key:    storepb.UserSetting_REFRESH_TOKENS,
	})
	if err != nil {
		return nil, err
	}
	if userSetting == nil {
		return []*storepb.RefreshTokensUserSetting_RefreshToken{}, nil
	}
	return userSetting.GetRefreshTokens().RefreshTokens, nil
}

// AddUserRefreshToken adds a new refresh token for the user.
func (s *Store) AddUserRefreshToken(ctx context.Context, userID int32, token *storepb.RefreshTokensUserSetting_RefreshToken) error {
	tokens, err := s.GetUserRefreshTokens(ctx, userID)
	if err != nil {
		return err
	}

	tokens = append(tokens, token)

	_, err = s.UpsertUserSetting(ctx, &storepb.UserSetting{
		UserId: userID,
		Key:    storepb.UserSetting_REFRESH_TOKENS,
		Value: &storepb.UserSetting_RefreshTokens{
			RefreshTokens: &storepb.RefreshTokensUserSetting{
				RefreshTokens: tokens,
			},
		},
	})
	return err
}

// RemoveUserRefreshToken removes a refresh token from the user.
func (s *Store) RemoveUserRefreshToken(ctx context.Context, userID int32, tokenID string) error {
	existingTokens, err := s.GetUserRefreshTokens(ctx, userID)
	if err != nil {
		return err
	}

	newTokens := make([]*storepb.RefreshTokensUserSetting_RefreshToken, 0, len(existingTokens))
	for _, token := range existingTokens {
		if token.TokenId != tokenID {
			newTokens = append(newTokens, token)
		}
	}

	_, err = s.UpsertUserSetting(ctx, &storepb.UserSetting{
		UserId: userID,
		Key:    storepb.UserSetting_REFRESH_TOKENS,
		Value: &storepb.UserSetting_RefreshTokens{
			RefreshTokens: &storepb.RefreshTokensUserSetting{
				RefreshTokens: newTokens,
			},
		},
	})
	return err
}

// GetUserRefreshTokenByID returns a specific refresh token.
func (s *Store) GetUserRefreshTokenByID(ctx context.Context, userID int32, tokenID string) (*storepb.RefreshTokensUserSetting_RefreshToken, error) {
	tokens, err := s.GetUserRefreshTokens(ctx, userID)
	if err != nil {
		return nil, err
	}
	for _, token := range tokens {
		if token.TokenId == tokenID {
			return token, nil
		}
	}
	return nil, nil
}

// GetUserPersonalAccessTokens returns the PATs of the user.
func (s *Store) GetUserPersonalAccessTokens(ctx context.Context, userID int32) ([]*storepb.PersonalAccessTokensUserSetting_PersonalAccessToken, error) {
	userSetting, err := s.GetUserSetting(ctx, &FindUserSetting{
		UserID: &userID,
		Key:    storepb.UserSetting_PERSONAL_ACCESS_TOKENS,
	})
	if err != nil {
		return nil, err
	}
	if userSetting == nil {
		return []*storepb.PersonalAccessTokensUserSetting_PersonalAccessToken{}, nil
	}
	return userSetting.GetPersonalAccessTokens().Tokens, nil
}

// AddUserPersonalAccessToken adds a new PAT for the user.
func (s *Store) AddUserPersonalAccessToken(ctx context.Context, userID int32, token *storepb.PersonalAccessTokensUserSetting_PersonalAccessToken) error {
	s.patMu.Lock()
	defer s.patMu.Unlock()

	tokens, err := s.GetUserPersonalAccessTokens(ctx, userID)
	if err != nil {
		return err
	}

	tokens = append(tokens, token)

	_, err = s.UpsertUserSetting(ctx, &storepb.UserSetting{
		UserId: userID,
		Key:    storepb.UserSetting_PERSONAL_ACCESS_TOKENS,
		Value: &storepb.UserSetting_PersonalAccessTokens{
			PersonalAccessTokens: &storepb.PersonalAccessTokensUserSetting{
				Tokens: tokens,
			},
		},
	})
	return err
}

// RemoveUserPersonalAccessToken removes a PAT from the user.
func (s *Store) RemoveUserPersonalAccessToken(ctx context.Context, userID int32, tokenID string) error {
	s.patMu.Lock()
	defer s.patMu.Unlock()

	existingTokens, err := s.GetUserPersonalAccessTokens(ctx, userID)
	if err != nil {
		return err
	}

	newTokens := make([]*storepb.PersonalAccessTokensUserSetting_PersonalAccessToken, 0, len(existingTokens))
	for _, token := range existingTokens {
		if token.TokenId != tokenID {
			newTokens = append(newTokens, token)
		}
	}

	_, err = s.UpsertUserSetting(ctx, &storepb.UserSetting{
		UserId: userID,
		Key:    storepb.UserSetting_PERSONAL_ACCESS_TOKENS,
		Value: &storepb.UserSetting_PersonalAccessTokens{
			PersonalAccessTokens: &storepb.PersonalAccessTokensUserSetting{
				Tokens: newTokens,
			},
		},
	})
	return err
}

// UpdatePATLastUsed updates the last_used_at timestamp of a PAT.
func (s *Store) UpdatePATLastUsed(ctx context.Context, userID int32, tokenID string, lastUsed *timestamppb.Timestamp) error {
	s.patMu.Lock()
	defer s.patMu.Unlock()

	tokens, err := s.GetUserPersonalAccessTokens(ctx, userID)
	if err != nil {
		return err
	}

	for i, token := range tokens {
		if token.TokenId == tokenID {
			// Concurrent requests can finish out of order. Never let an older usage
			// timestamp overwrite a newer one.
			if lastUsed != nil && token.LastUsedAt != nil && !token.LastUsedAt.AsTime().Before(lastUsed.AsTime()) {
				return nil
			}

			updatedToken, ok := proto.Clone(token).(*storepb.PersonalAccessTokensUserSetting_PersonalAccessToken)
			if !ok {
				return errors.Errorf("failed to clone personal access token")
			}
			updatedToken.LastUsedAt = lastUsed
			updatedTokens := make([]*storepb.PersonalAccessTokensUserSetting_PersonalAccessToken, len(tokens))
			copy(updatedTokens, tokens)
			updatedTokens[i] = updatedToken

			_, err = s.UpsertUserSetting(ctx, &storepb.UserSetting{
				UserId: userID,
				Key:    storepb.UserSetting_PERSONAL_ACCESS_TOKENS,
				Value: &storepb.UserSetting_PersonalAccessTokens{
					PersonalAccessTokens: &storepb.PersonalAccessTokensUserSetting{
						Tokens: updatedTokens,
					},
				},
			})
			return err
		}
	}

	return nil
}

// GetUserWebhooks returns the webhooks of the user.
func (s *Store) GetUserWebhooks(ctx context.Context, userID int32) ([]*storepb.WebhooksUserSetting_Webhook, error) {
	userSetting, err := s.GetUserSetting(ctx, &FindUserSetting{
		UserID: &userID,
		Key:    storepb.UserSetting_WEBHOOKS,
	})
	if err != nil {
		return nil, err
	}
	if userSetting == nil {
		return []*storepb.WebhooksUserSetting_Webhook{}, nil
	}

	webhooksUserSetting := userSetting.GetWebhooks()
	return webhooksUserSetting.Webhooks, nil
}

// AddUserWebhook adds a new webhook for the user.
func (s *Store) AddUserWebhook(ctx context.Context, userID int32, webhook *storepb.WebhooksUserSetting_Webhook) error {
	existingWebhooks, err := s.GetUserWebhooks(ctx, userID)
	if err != nil {
		return err
	}

	// Check if webhook already exists, update if it does
	var updatedWebhooks []*storepb.WebhooksUserSetting_Webhook
	webhookExists := false
	for _, existing := range existingWebhooks {
		if existing.Id == webhook.Id {
			updatedWebhooks = append(updatedWebhooks, webhook)
			webhookExists = true
		} else {
			updatedWebhooks = append(updatedWebhooks, existing)
		}
	}

	// If webhook doesn't exist, add it
	if !webhookExists {
		updatedWebhooks = append(updatedWebhooks, webhook)
	}

	_, err = s.UpsertUserSetting(ctx, &storepb.UserSetting{
		UserId: userID,
		Key:    storepb.UserSetting_WEBHOOKS,
		Value: &storepb.UserSetting_Webhooks{
			Webhooks: &storepb.WebhooksUserSetting{
				Webhooks: updatedWebhooks,
			},
		},
	})

	return err
}

// RemoveUserWebhook removes the webhook of the user.
func (s *Store) RemoveUserWebhook(ctx context.Context, userID int32, webhookID string) error {
	oldWebhooks, err := s.GetUserWebhooks(ctx, userID)
	if err != nil {
		return err
	}

	newWebhooks := make([]*storepb.WebhooksUserSetting_Webhook, 0, len(oldWebhooks))
	for _, webhook := range oldWebhooks {
		if webhookID != webhook.Id {
			newWebhooks = append(newWebhooks, webhook)
		}
	}

	_, err = s.UpsertUserSetting(ctx, &storepb.UserSetting{
		UserId: userID,
		Key:    storepb.UserSetting_WEBHOOKS,
		Value: &storepb.UserSetting_Webhooks{
			Webhooks: &storepb.WebhooksUserSetting{
				Webhooks: newWebhooks,
			},
		},
	})

	return err
}

// UpdateUserWebhook updates an existing webhook for the user.
func (s *Store) UpdateUserWebhook(ctx context.Context, userID int32, webhook *storepb.WebhooksUserSetting_Webhook) error {
	webhooks, err := s.GetUserWebhooks(ctx, userID)
	if err != nil {
		return err
	}

	for i, existing := range webhooks {
		if existing.Id == webhook.Id {
			webhooks[i] = webhook
			break
		}
	}

	_, err = s.UpsertUserSetting(ctx, &storepb.UserSetting{
		UserId: userID,
		Key:    storepb.UserSetting_WEBHOOKS,
		Value: &storepb.UserSetting_Webhooks{
			Webhooks: &storepb.WebhooksUserSetting{
				Webhooks: webhooks,
			},
		},
	})

	return err
}

// GetUserMemoViews returns the memo views of the user.
func (s *Store) GetUserMemoViews(ctx context.Context, userID int32) ([]*storepb.MemoViewsUserSetting_MemoView, error) {
	userSetting, err := s.GetUserSetting(ctx, &FindUserSetting{
		UserID: &userID,
		Key:    storepb.UserSetting_MEMO_VIEWS,
	})
	if err != nil {
		return nil, errors.Wrap(err, "get memo views user setting")
	}
	if userSetting == nil {
		return []*storepb.MemoViewsUserSetting_MemoView{}, nil
	}

	memoViews := userSetting.GetMemoViews().GetMemoViews()
	clonedMemoViews := make([]*storepb.MemoViewsUserSetting_MemoView, len(memoViews))
	for i, memoView := range memoViews {
		if memoView != nil {
			clonedMemoView, ok := proto.Clone(memoView).(*storepb.MemoViewsUserSetting_MemoView)
			if !ok {
				return nil, errors.New("failed to clone memo view")
			}
			clonedMemoViews[i] = clonedMemoView
		}
	}
	return clonedMemoViews, nil
}

// AddUserMemoView appends a new memo view for the user.
func (s *Store) AddUserMemoView(ctx context.Context, userID int32, memoView *storepb.MemoViewsUserSetting_MemoView) error {
	s.memoViewMu.Lock()
	defer s.memoViewMu.Unlock()

	existing, err := s.GetUserMemoViews(ctx, userID)
	if err != nil {
		return errors.Wrap(err, "get existing memo views")
	}

	// Build a fresh slice so the cached setting is never mutated in place.
	memoViews := make([]*storepb.MemoViewsUserSetting_MemoView, 0, len(existing)+1)
	memoViews = append(memoViews, existing...)
	memoViews = append(memoViews, memoView)

	if err := s.upsertUserMemoViews(ctx, userID, memoViews); err != nil {
		return errors.Wrap(err, "add memo view")
	}
	return nil
}

// UpdateUserMemoView applies the non-nil field updates to the memo view carrying the same ID.
// It returns nil when no matching memo view is found.
func (s *Store) UpdateUserMemoView(
	ctx context.Context,
	userID int32,
	memoViewID string,
	title *string,
	filter *string,
) (*storepb.MemoViewsUserSetting_MemoView, error) {
	s.memoViewMu.Lock()
	defer s.memoViewMu.Unlock()

	existing, err := s.GetUserMemoViews(ctx, userID)
	if err != nil {
		return nil, errors.Wrap(err, "get existing memo views")
	}

	var updatedMemoView *storepb.MemoViewsUserSetting_MemoView
	memoViews := make([]*storepb.MemoViewsUserSetting_MemoView, 0, len(existing))
	for _, item := range existing {
		if item.GetId() != memoViewID {
			memoViews = append(memoViews, item)
			continue
		}

		updatedMemoView = &storepb.MemoViewsUserSetting_MemoView{
			Id:     item.GetId(),
			Title:  item.GetTitle(),
			Filter: item.GetFilter(),
		}
		if title != nil {
			updatedMemoView.Title = *title
		}
		if filter != nil {
			updatedMemoView.Filter = *filter
		}
		memoViews = append(memoViews, updatedMemoView)
	}
	if updatedMemoView == nil {
		return nil, nil
	}

	if err := s.upsertUserMemoViews(ctx, userID, memoViews); err != nil {
		return nil, errors.Wrap(err, "update memo view")
	}
	return updatedMemoView, nil
}

// RemoveUserMemoView removes the memo view of the user.
// It reports whether a matching memo view was found.
func (s *Store) RemoveUserMemoView(ctx context.Context, userID int32, memoViewID string) (bool, error) {
	s.memoViewMu.Lock()
	defer s.memoViewMu.Unlock()

	existing, err := s.GetUserMemoViews(ctx, userID)
	if err != nil {
		return false, errors.Wrap(err, "get existing memo views")
	}

	found := false
	memoViews := make([]*storepb.MemoViewsUserSetting_MemoView, 0, len(existing))
	for _, item := range existing {
		if item.GetId() == memoViewID {
			found = true
			continue
		}
		memoViews = append(memoViews, item)
	}
	if !found {
		return false, nil
	}

	if err := s.upsertUserMemoViews(ctx, userID, memoViews); err != nil {
		return false, errors.Wrap(err, "remove memo view")
	}
	return true, nil
}

func (s *Store) upsertUserMemoViews(ctx context.Context, userID int32, memoViews []*storepb.MemoViewsUserSetting_MemoView) error {
	_, err := s.UpsertUserSetting(ctx, &storepb.UserSetting{
		UserId: userID,
		Key:    storepb.UserSetting_MEMO_VIEWS,
		Value: &storepb.UserSetting_MemoViews{
			MemoViews: &storepb.MemoViewsUserSetting{
				MemoViews: memoViews,
			},
		},
	})
	return errors.Wrap(err, "upsert memo views user setting")
}

func convertUserSettingFromRaw(raw *UserSetting) (*storepb.UserSetting, error) {
	userSetting := &storepb.UserSetting{
		UserId: raw.UserID,
		Key:    raw.Key,
	}

	switch raw.Key {
	case storepb.UserSetting_MEMO_VIEWS:
		memoViewsUserSetting := &storepb.MemoViewsUserSetting{}
		if err := protojsonUnmarshaler.Unmarshal([]byte(raw.Value), memoViewsUserSetting); err != nil {
			return nil, err
		}
		userSetting.Value = &storepb.UserSetting_MemoViews{MemoViews: memoViewsUserSetting}
	case storepb.UserSetting_GENERAL:
		generalUserSetting := &storepb.GeneralUserSetting{}
		if err := protojsonUnmarshaler.Unmarshal([]byte(raw.Value), generalUserSetting); err != nil {
			return nil, err
		}
		userSetting.Value = &storepb.UserSetting_General{General: generalUserSetting}
	case storepb.UserSetting_TAGS:
		tagsUserSetting := &storepb.TagsUserSetting{}
		if err := protojsonUnmarshaler.Unmarshal([]byte(raw.Value), tagsUserSetting); err != nil {
			return nil, errors.Wrap(err, "unmarshal tags user setting")
		}
		userSetting.Value = &storepb.UserSetting_Tags{Tags: tagsUserSetting}
	case storepb.UserSetting_REFRESH_TOKENS:
		refreshTokensUserSetting := &storepb.RefreshTokensUserSetting{}
		if err := protojsonUnmarshaler.Unmarshal([]byte(raw.Value), refreshTokensUserSetting); err != nil {
			return nil, err
		}
		userSetting.Value = &storepb.UserSetting_RefreshTokens{RefreshTokens: refreshTokensUserSetting}
	case storepb.UserSetting_PERSONAL_ACCESS_TOKENS:
		patsUserSetting := &storepb.PersonalAccessTokensUserSetting{}
		if err := protojsonUnmarshaler.Unmarshal([]byte(raw.Value), patsUserSetting); err != nil {
			return nil, err
		}
		userSetting.Value = &storepb.UserSetting_PersonalAccessTokens{PersonalAccessTokens: patsUserSetting}
	case storepb.UserSetting_WEBHOOKS:
		webhooksUserSetting := &storepb.WebhooksUserSetting{}
		if err := protojsonUnmarshaler.Unmarshal([]byte(raw.Value), webhooksUserSetting); err != nil {
			return nil, err
		}
		userSetting.Value = &storepb.UserSetting_Webhooks{Webhooks: webhooksUserSetting}
	default:
		return nil, nil
	}
	return userSetting, nil
}

func convertUserSettingToRaw(userSetting *storepb.UserSetting) (*UserSetting, error) {
	raw := &UserSetting{
		UserID: userSetting.UserId,
		Key:    userSetting.Key,
	}

	switch userSetting.Key {
	case storepb.UserSetting_MEMO_VIEWS:
		memoViewsUserSetting := userSetting.GetMemoViews()
		value, err := protojson.Marshal(memoViewsUserSetting)
		if err != nil {
			return nil, err
		}
		raw.Value = string(value)
	case storepb.UserSetting_GENERAL:
		generalUserSetting := userSetting.GetGeneral()
		value, err := protojson.Marshal(generalUserSetting)
		if err != nil {
			return nil, err
		}
		raw.Value = string(value)
	case storepb.UserSetting_TAGS:
		tagsUserSetting := userSetting.GetTags()
		value, err := protojson.Marshal(tagsUserSetting)
		if err != nil {
			return nil, errors.Wrap(err, "marshal tags user setting")
		}
		raw.Value = string(value)
	case storepb.UserSetting_REFRESH_TOKENS:
		refreshTokensUserSetting := userSetting.GetRefreshTokens()
		value, err := protojson.Marshal(refreshTokensUserSetting)
		if err != nil {
			return nil, err
		}
		raw.Value = string(value)
	case storepb.UserSetting_PERSONAL_ACCESS_TOKENS:
		patsUserSetting := userSetting.GetPersonalAccessTokens()
		value, err := protojson.Marshal(patsUserSetting)
		if err != nil {
			return nil, err
		}
		raw.Value = string(value)
	case storepb.UserSetting_WEBHOOKS:
		webhooksUserSetting := userSetting.GetWebhooks()
		value, err := protojson.Marshal(webhooksUserSetting)
		if err != nil {
			return nil, err
		}
		raw.Value = string(value)
	default:
		return nil, errors.Errorf("unsupported user setting key: %v", userSetting.Key)
	}
	return raw, nil
}