GraphQL API and utilities for the rpdata project
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

276 lines
5.1 KiB

  1. package mongodb
  2. import (
  3. "context"
  4. "fmt"
  5. "github.com/globalsign/mgo/bson"
  6. "time"
  7. "git.aiterp.net/rpdata/api/internal/config"
  8. "git.aiterp.net/rpdata/api/repositories"
  9. "github.com/globalsign/mgo"
  10. )
  11. var DisableFixes bool
  12. type MongoDB struct {
  13. session *mgo.Session
  14. changes repositories.ChangeRepository
  15. channels *channelRepository
  16. characters repositories.CharacterRepository
  17. tags repositories.TagRepository
  18. logs *logRepository
  19. posts *postRepository
  20. files *fileRepository
  21. stories repositories.StoryRepository
  22. chapters repositories.ChapterRepository
  23. comments repositories.CommentRepository
  24. keys repositories.KeyRepository
  25. users repositories.UserRepository
  26. }
  27. func (m *MongoDB) Changes() repositories.ChangeRepository {
  28. return m.changes
  29. }
  30. func (m *MongoDB) Characters() repositories.CharacterRepository {
  31. return m.characters
  32. }
  33. func (m *MongoDB) Tags() repositories.TagRepository {
  34. return m.tags
  35. }
  36. func (m *MongoDB) Logs() repositories.LogRepository {
  37. return m.logs
  38. }
  39. func (m *MongoDB) Posts() repositories.PostRepository {
  40. return m.posts
  41. }
  42. func (m *MongoDB) Channels() repositories.ChannelRepository {
  43. return m.channels
  44. }
  45. func (m *MongoDB) Files() repositories.FileRepository {
  46. return m.files
  47. }
  48. func (m *MongoDB) Stories() repositories.StoryRepository {
  49. return m.stories
  50. }
  51. func (m *MongoDB) Chapters() repositories.ChapterRepository {
  52. return m.chapters
  53. }
  54. func (m *MongoDB) Comments() repositories.CommentRepository {
  55. return m.comments
  56. }
  57. func (m *MongoDB) Keys() repositories.KeyRepository {
  58. return m.keys
  59. }
  60. func (m *MongoDB) Users() repositories.UserRepository {
  61. return m.users
  62. }
  63. func (m *MongoDB) Close(ctx context.Context) error {
  64. m.session.Close()
  65. return nil
  66. }
  67. // Init initializes the mongodb database
  68. func Init(cfg config.Database) (*MongoDB, error) {
  69. port := cfg.Port
  70. if port <= 0 {
  71. port = 27017
  72. }
  73. session, err := mgo.DialWithInfo(&mgo.DialInfo{
  74. Addrs: []string{fmt.Sprintf("%s:%d", cfg.Host, port)},
  75. Timeout: 30 * time.Second,
  76. Database: cfg.Db,
  77. Username: cfg.Username,
  78. Password: cfg.Password,
  79. Mechanism: cfg.Mechanism,
  80. Source: cfg.Db,
  81. })
  82. if err != nil {
  83. return nil, err
  84. }
  85. db := session.DB(cfg.Db)
  86. characters, err := newCharacterRepository(db, cfg.RestoreIDs)
  87. if err != nil {
  88. session.Close()
  89. return nil, err
  90. }
  91. channels, err := newChannelRepository(db)
  92. if err != nil {
  93. session.Close()
  94. return nil, err
  95. }
  96. changes, err := newChangeRepository(db, cfg.RestoreIDs)
  97. if err != nil {
  98. session.Close()
  99. return nil, err
  100. }
  101. logs, err := newLogRepository(db, cfg.RestoreIDs)
  102. if err != nil {
  103. session.Close()
  104. return nil, err
  105. }
  106. posts, err := newPostRepository(db, cfg.RestoreIDs)
  107. if err != nil {
  108. session.Close()
  109. return nil, err
  110. }
  111. files, err := newFileRepository(db)
  112. if err != nil {
  113. session.Close()
  114. return nil, err
  115. }
  116. stories, err := newStoryRepository(db, cfg.RestoreIDs)
  117. if err != nil {
  118. session.Close()
  119. return nil, err
  120. }
  121. chapters, err := newChapterRepository(db, cfg.RestoreIDs)
  122. if err != nil {
  123. session.Close()
  124. return nil, err
  125. }
  126. comments, err := newCommentRepository(db, cfg.RestoreIDs)
  127. if err != nil {
  128. session.Close()
  129. return nil, err
  130. }
  131. keys, err := newKeyRepository(db, cfg.RestoreIDs)
  132. if err != nil {
  133. session.Close()
  134. return nil, err
  135. }
  136. users, err := newUserRepository(db)
  137. if err != nil {
  138. session.Close()
  139. return nil, err
  140. }
  141. if !DisableFixes {
  142. go posts.fixPositions(logs)
  143. }
  144. return &MongoDB{
  145. session: session,
  146. changes: changes,
  147. characters: characters,
  148. channels: channels,
  149. tags: newTagRepository(db),
  150. stories: stories,
  151. chapters: chapters,
  152. comments: comments,
  153. logs: logs,
  154. posts: posts,
  155. files: files,
  156. keys: keys,
  157. users: users,
  158. }, nil
  159. }
  160. type counter struct {
  161. coll *mgo.Collection
  162. category string
  163. name string
  164. }
  165. func newCounter(db *mgo.Database, category, name string) *counter {
  166. return &counter{
  167. coll: db.C("core.counters"),
  168. category: category,
  169. name: name,
  170. }
  171. }
  172. func (c *counter) WithName(name string) *counter {
  173. return &counter{
  174. coll: c.coll,
  175. category: c.category,
  176. name: name,
  177. }
  178. }
  179. func (c *counter) WithCategory(category string) *counter {
  180. return &counter{
  181. coll: c.coll,
  182. category: category,
  183. name: c.name,
  184. }
  185. }
  186. func (c *counter) With(category, name string) *counter {
  187. return &counter{
  188. coll: c.coll,
  189. category: category,
  190. name: name,
  191. }
  192. }
  193. func (c *counter) Bump(amount int) error {
  194. id := c.category + "." + c.name
  195. err := c.coll.Update(bson.M{
  196. "_id": id,
  197. "value": bson.M{"$lt": amount},
  198. }, bson.M{
  199. "value": amount,
  200. })
  201. if err != nil {
  202. if err == mgo.ErrNotFound {
  203. return c.coll.Insert(bson.M{
  204. "_id": id,
  205. "value": amount,
  206. })
  207. }
  208. return err
  209. }
  210. return nil
  211. }
  212. func (c *counter) Increment(amount int) (int, error) {
  213. type counterDoc struct {
  214. ID string `bson:"_id"`
  215. Value int `bson:"value"`
  216. }
  217. id := c.category + "." + c.name
  218. doc := counterDoc{}
  219. _, err := c.coll.Find(bson.M{"_id": id}).Apply(mgo.Change{
  220. Update: bson.M{"$inc": bson.M{"value": amount}},
  221. Upsert: true,
  222. ReturnNew: true,
  223. }, &doc)
  224. if err != nil {
  225. return -1, err
  226. }
  227. return doc.Value, nil
  228. }