user_setting.go 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  1. package sqlite
  2. import (
  3. "context"
  4. "database/sql"
  5. "errors"
  6. "strings"
  7. "google.golang.org/protobuf/encoding/protojson"
  8. storepb "github.com/usememos/memos/proto/gen/store"
  9. "github.com/usememos/memos/store"
  10. )
  11. func (d *DB) UpsertUserSetting(ctx context.Context, upsert *store.UserSetting) (*store.UserSetting, error) {
  12. stmt := `
  13. INSERT INTO user_setting (
  14. user_id, key, value
  15. )
  16. VALUES (?, ?, ?)
  17. ON CONFLICT(user_id, key) DO UPDATE
  18. SET value = EXCLUDED.value
  19. `
  20. if _, err := d.db.ExecContext(ctx, stmt, upsert.UserID, upsert.Key, upsert.Value); err != nil {
  21. return nil, err
  22. }
  23. return upsert, nil
  24. }
  25. func (d *DB) ListUserSettings(ctx context.Context, find *store.FindUserSetting) ([]*store.UserSetting, error) {
  26. where, args := []string{"1 = 1"}, []any{}
  27. if v := find.Key; v != "" {
  28. where, args = append(where, "key = ?"), append(args, v)
  29. }
  30. if v := find.UserID; v != nil {
  31. where, args = append(where, "user_id = ?"), append(args, *find.UserID)
  32. }
  33. query := `
  34. SELECT
  35. user_id,
  36. key,
  37. value
  38. FROM user_setting
  39. WHERE ` + strings.Join(where, " AND ")
  40. rows, err := d.db.QueryContext(ctx, query, args...)
  41. if err != nil {
  42. return nil, err
  43. }
  44. defer rows.Close()
  45. userSettingList := make([]*store.UserSetting, 0)
  46. for rows.Next() {
  47. var userSetting store.UserSetting
  48. if err := rows.Scan(
  49. &userSetting.UserID,
  50. &userSetting.Key,
  51. &userSetting.Value,
  52. ); err != nil {
  53. return nil, err
  54. }
  55. userSettingList = append(userSettingList, &userSetting)
  56. }
  57. if err := rows.Err(); err != nil {
  58. return nil, err
  59. }
  60. return userSettingList, nil
  61. }
  62. func (d *DB) UpsertUserSettingV1(ctx context.Context, upsert *storepb.UserSetting) (*storepb.UserSetting, error) {
  63. stmt := `
  64. INSERT INTO user_setting (
  65. user_id, key, value
  66. )
  67. VALUES (?, ?, ?)
  68. ON CONFLICT(user_id, key) DO UPDATE
  69. SET value = EXCLUDED.value
  70. `
  71. var valueString string
  72. if upsert.Key == storepb.UserSettingKey_USER_SETTING_ACCESS_TOKENS {
  73. valueBytes, err := protojson.Marshal(upsert.GetAccessTokens())
  74. if err != nil {
  75. return nil, err
  76. }
  77. valueString = string(valueBytes)
  78. } else {
  79. return nil, errors.New("invalid user setting key")
  80. }
  81. if _, err := d.db.ExecContext(ctx, stmt, upsert.UserId, upsert.Key.String(), valueString); err != nil {
  82. return nil, err
  83. }
  84. return upsert, nil
  85. }
  86. func (d *DB) ListUserSettingsV1(ctx context.Context, find *store.FindUserSettingV1) ([]*storepb.UserSetting, error) {
  87. where, args := []string{"1 = 1"}, []any{}
  88. if v := find.Key; v != storepb.UserSettingKey_USER_SETTING_KEY_UNSPECIFIED {
  89. where, args = append(where, "key = ?"), append(args, v.String())
  90. }
  91. if v := find.UserID; v != nil {
  92. where, args = append(where, "user_id = ?"), append(args, *find.UserID)
  93. }
  94. query := `
  95. SELECT
  96. user_id,
  97. key,
  98. value
  99. FROM user_setting
  100. WHERE ` + strings.Join(where, " AND ")
  101. rows, err := d.db.QueryContext(ctx, query, args...)
  102. if err != nil {
  103. return nil, err
  104. }
  105. defer rows.Close()
  106. userSettingList := make([]*storepb.UserSetting, 0)
  107. for rows.Next() {
  108. userSetting := &storepb.UserSetting{}
  109. var keyString, valueString string
  110. if err := rows.Scan(
  111. &userSetting.UserId,
  112. &keyString,
  113. &valueString,
  114. ); err != nil {
  115. return nil, err
  116. }
  117. userSetting.Key = storepb.UserSettingKey(storepb.UserSettingKey_value[keyString])
  118. if userSetting.Key == storepb.UserSettingKey_USER_SETTING_ACCESS_TOKENS {
  119. accessTokensUserSetting := &storepb.AccessTokensUserSetting{}
  120. if err := protojson.Unmarshal([]byte(valueString), accessTokensUserSetting); err != nil {
  121. return nil, err
  122. }
  123. userSetting.Value = &storepb.UserSetting_AccessTokens{
  124. AccessTokens: accessTokensUserSetting,
  125. }
  126. } else {
  127. // Skip unknown user setting v1 key.
  128. continue
  129. }
  130. userSettingList = append(userSettingList, userSetting)
  131. }
  132. if err := rows.Err(); err != nil {
  133. return nil, err
  134. }
  135. return userSettingList, nil
  136. }
  137. func vacuumUserSetting(ctx context.Context, tx *sql.Tx) error {
  138. stmt := `
  139. DELETE FROM
  140. user_setting
  141. WHERE
  142. user_id NOT IN (
  143. SELECT
  144. id
  145. FROM
  146. user
  147. )`
  148. _, err := tx.ExecContext(ctx, stmt)
  149. if err != nil {
  150. return err
  151. }
  152. return nil
  153. }