Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

305 строки
8.0 KiB

  1. package helpers
  2. import (
  3. "fmt"
  4. "os"
  5. "github.com/QuantumNous/new-api/common"
  6. "github.com/QuantumNous/new-api/model"
  7. "github.com/glebarez/sqlite"
  8. "gorm.io/gorm"
  9. )
  10. const (
  11. CNMergeTargetUserID = 11
  12. CNSyncedCopyOwnerUserID = 12
  13. CNTakenUsernameUserID = 13
  14. CNSoftAffOwnerUserID = 14
  15. OVCreateUserID = 10000001
  16. OVMergeUserID = 10000002
  17. OVHardSyncedUserID = 10000003
  18. OVUsernameConflictUserID = 10000004
  19. OVSoftAffUserID = 10000005
  20. OVSyncedCopyUserID = 10000100
  21. UsernameCNMergeTarget = "cn-merge-target"
  22. UsernameCNSyncedCopyOwner = "cn-synced-copy-owner"
  23. UsernameCNTakenUsername = "taken-username"
  24. UsernameCNSoftAffOwner = "aff-owner"
  25. UsernameOVCreate = "ov-create"
  26. UsernameOVMerge = "ov-merge"
  27. UsernameOVHardSynced = "ov-hard-synced"
  28. UsernameOVUsernameConflict = "taken-username"
  29. UsernameOVSoftAff = "ov-aff-soft"
  30. UsernameOVSyncedCopy = "ov-existing-synced-copy"
  31. EmailCNMergeTarget = "merge@example.com"
  32. EmailCNSyncedCopyOwner = "synced@example.com"
  33. EmailCNTakenUsername = "taken@example.com"
  34. EmailCNSoftAffOwner = "soft-owner@example.com"
  35. EmailOVCreate = "create@example.com"
  36. EmailOVMerge = "merge@example.com"
  37. EmailOVHardSynced = "synced@example.com"
  38. EmailOVUsernameConflict = "username-conflict@example.com"
  39. EmailOVSoftAff = "soft@example.com"
  40. EmailOVSyncedCopy = "copy@example.com"
  41. AffCodeCNMergeTarget = "CNM1"
  42. AffCodeCNSyncedCopyOwner = "CNH3"
  43. AffCodeCNTakenUsername = "CNT1"
  44. AffCodeCNSoftAffOwner = "SOFT1"
  45. AffCodeOVCreate = "OVC1"
  46. AffCodeOVMerge = "OVM2"
  47. AffCodeOVHardSynced = "OVH3"
  48. AffCodeOVUsername = "OVU4"
  49. AffCodeOVSyncedCopy = "OVS1"
  50. )
  51. func ResetDB(path string) error {
  52. if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
  53. return err
  54. }
  55. return nil
  56. }
  57. func OpenSQLite(path string) (*gorm.DB, error) {
  58. db, err := gorm.Open(sqlite.Open(path), &gorm.Config{})
  59. if err != nil {
  60. return nil, err
  61. }
  62. sqlDB, err := db.DB()
  63. if err != nil {
  64. return nil, err
  65. }
  66. sqlDB.SetMaxOpenConns(1)
  67. return db, nil
  68. }
  69. func CloseSQLite(db *gorm.DB) error {
  70. if db == nil {
  71. return nil
  72. }
  73. sqlDB, err := db.DB()
  74. if err != nil {
  75. return err
  76. }
  77. return sqlDB.Close()
  78. }
  79. func SeedCNData(dbPath string) error {
  80. db, err := OpenSQLite(dbPath)
  81. if err != nil {
  82. return fmt.Errorf("open CN sqlite: %w", err)
  83. }
  84. defer func() {
  85. _ = CloseSQLite(db)
  86. }()
  87. passwordHash, err := common.Password2Hash("fixture-password")
  88. if err != nil {
  89. return fmt.Errorf("hash fixture password: %w", err)
  90. }
  91. now := common.GetTimestamp()
  92. users := []*model.User{
  93. {
  94. Id: CNMergeTargetUserID,
  95. Username: UsernameCNMergeTarget,
  96. Password: passwordHash,
  97. DisplayName: "CN Merge Target",
  98. Role: common.RoleCommonUser,
  99. Status: common.UserStatusEnabled,
  100. Email: EmailCNMergeTarget,
  101. Group: "default",
  102. Quota: 500,
  103. AffCode: AffCodeCNMergeTarget,
  104. Source: common.UserSourceLocal,
  105. CreatedAt: now,
  106. },
  107. {
  108. Id: CNSyncedCopyOwnerUserID,
  109. Username: UsernameCNSyncedCopyOwner,
  110. Password: passwordHash,
  111. DisplayName: "CN Synced Copy Owner",
  112. Role: common.RoleCommonUser,
  113. Status: common.UserStatusEnabled,
  114. Email: EmailCNSyncedCopyOwner,
  115. Group: "default",
  116. Quota: 300,
  117. AffCode: AffCodeCNSyncedCopyOwner,
  118. Source: common.UserSourceLocal,
  119. CreatedAt: now,
  120. },
  121. {
  122. Id: CNTakenUsernameUserID,
  123. Username: UsernameCNTakenUsername,
  124. Password: passwordHash,
  125. DisplayName: "CN Username Conflict",
  126. Role: common.RoleCommonUser,
  127. Status: common.UserStatusEnabled,
  128. Email: EmailCNTakenUsername,
  129. Group: "default",
  130. Quota: 50,
  131. AffCode: AffCodeCNTakenUsername,
  132. Source: common.UserSourceLocal,
  133. CreatedAt: now,
  134. },
  135. {
  136. Id: CNSoftAffOwnerUserID,
  137. Username: UsernameCNSoftAffOwner,
  138. Password: passwordHash,
  139. DisplayName: "CN Soft Aff Owner",
  140. Role: common.RoleCommonUser,
  141. Status: common.UserStatusEnabled,
  142. Email: EmailCNSoftAffOwner,
  143. Group: "default",
  144. Quota: 80,
  145. AffCode: AffCodeCNSoftAffOwner,
  146. Source: common.UserSourceLocal,
  147. CreatedAt: now,
  148. },
  149. }
  150. for _, user := range users {
  151. if err := db.Create(user).Error; err != nil {
  152. return fmt.Errorf("seed CN user %s: %w", user.Username, err)
  153. }
  154. }
  155. return nil
  156. }
  157. func SeedOVData(dbPath string) error {
  158. db, err := OpenSQLite(dbPath)
  159. if err != nil {
  160. return fmt.Errorf("open OV sqlite: %w", err)
  161. }
  162. defer func() {
  163. _ = CloseSQLite(db)
  164. }()
  165. passwordHash, err := common.Password2Hash("fixture-password")
  166. if err != nil {
  167. return fmt.Errorf("hash fixture password: %w", err)
  168. }
  169. now := common.GetTimestamp()
  170. users := []*model.User{
  171. {
  172. Id: OVCreateUserID,
  173. Username: UsernameOVCreate,
  174. Password: passwordHash,
  175. DisplayName: "OV Create",
  176. Role: common.RoleCommonUser,
  177. Status: common.UserStatusEnabled,
  178. Email: EmailOVCreate,
  179. Group: "default",
  180. Quota: 100,
  181. AffCode: AffCodeOVCreate,
  182. Source: common.UserSourceLocal,
  183. CreatedAt: now,
  184. },
  185. {
  186. Id: OVMergeUserID,
  187. Username: UsernameOVMerge,
  188. Password: passwordHash,
  189. DisplayName: "OV Merge",
  190. Role: common.RoleCommonUser,
  191. Status: common.UserStatusEnabled,
  192. Email: EmailOVMerge,
  193. Group: "default",
  194. Quota: 200,
  195. AffCode: AffCodeOVMerge,
  196. Source: common.UserSourceLocal,
  197. CreatedAt: now,
  198. },
  199. {
  200. Id: OVHardSyncedUserID,
  201. Username: UsernameOVHardSynced,
  202. Password: passwordHash,
  203. DisplayName: "OV Hard Synced",
  204. Role: common.RoleCommonUser,
  205. Status: common.UserStatusEnabled,
  206. Email: EmailOVHardSynced,
  207. Group: "default",
  208. Quota: 300,
  209. AffCode: AffCodeOVHardSynced,
  210. Source: common.UserSourceLocal,
  211. CreatedAt: now,
  212. },
  213. {
  214. Id: OVUsernameConflictUserID,
  215. Username: UsernameOVUsernameConflict,
  216. Password: passwordHash,
  217. DisplayName: "OV Username Conflict",
  218. Role: common.RoleCommonUser,
  219. Status: common.UserStatusEnabled,
  220. Email: EmailOVUsernameConflict,
  221. Group: "default",
  222. Quota: 120,
  223. AffCode: AffCodeOVUsername,
  224. Source: common.UserSourceLocal,
  225. CreatedAt: now,
  226. },
  227. {
  228. Id: OVSoftAffUserID,
  229. Username: UsernameOVSoftAff,
  230. Password: passwordHash,
  231. DisplayName: "OV Soft Aff",
  232. Role: common.RoleCommonUser,
  233. Status: common.UserStatusEnabled,
  234. Email: EmailOVSoftAff,
  235. Group: "default",
  236. Quota: 150,
  237. AffCode: AffCodeCNSoftAffOwner,
  238. Source: common.UserSourceLocal,
  239. CreatedAt: now,
  240. },
  241. }
  242. for _, user := range users {
  243. if err := db.Create(user).Error; err != nil {
  244. return fmt.Errorf("seed OV user %s: %w", user.Username, err)
  245. }
  246. }
  247. return SeedOVSyncedCopy(dbPath)
  248. }
  249. func SeedOVSyncedCopy(dbPath string) error {
  250. db, err := OpenSQLite(dbPath)
  251. if err != nil {
  252. return fmt.Errorf("open OV sqlite for synced copy: %w", err)
  253. }
  254. defer func() {
  255. _ = CloseSQLite(db)
  256. }()
  257. passwordHash, err := common.Password2Hash("fixture-password")
  258. if err != nil {
  259. return fmt.Errorf("hash fixture password for synced copy: %w", err)
  260. }
  261. now := common.GetTimestamp()
  262. user := &model.User{
  263. Id: OVSyncedCopyUserID,
  264. Username: UsernameOVSyncedCopy,
  265. Password: passwordHash,
  266. DisplayName: "OV Existing Synced Copy",
  267. Role: common.RoleCommonUser,
  268. Status: common.UserStatusEnabled,
  269. Email: EmailOVSyncedCopy,
  270. Group: "default",
  271. Quota: 0,
  272. AffCode: AffCodeOVSyncedCopy,
  273. Source: common.UserSourceSynced,
  274. RemoteUserId: CNSyncedCopyOwnerUserID,
  275. SyncedQuota: 300,
  276. LastSyncAt: now,
  277. CreatedAt: now,
  278. }
  279. if err := db.Create(user).Error; err != nil {
  280. return fmt.Errorf("seed OV synced copy %s: %w", user.Username, err)
  281. }
  282. return nil
  283. }