github.com/qichengzx/mattermost-server@v4.5.1-0.20180604164826-2c75247c97d0+incompatible/cmd/mattermost/commands/sampledata.go (about) 1 // Copyright (c) 2016-present Mattermost, Inc. All Rights Reserved. 2 // See License.txt for license information. 3 4 package commands 5 6 import ( 7 "encoding/json" 8 "errors" 9 "fmt" 10 "io/ioutil" 11 "math/rand" 12 "os" 13 "path" 14 "sort" 15 "strings" 16 "time" 17 18 "github.com/icrowley/fake" 19 "github.com/mattermost/mattermost-server/app" 20 "github.com/spf13/cobra" 21 ) 22 23 var SampleDataCmd = &cobra.Command{ 24 Use: "sampledata", 25 Short: "Generate sample data", 26 RunE: sampleDataCmdF, 27 } 28 29 func init() { 30 SampleDataCmd.Flags().Int64P("seed", "s", 1, "Seed used for generating the random data (Different seeds generate different data).") 31 SampleDataCmd.Flags().IntP("teams", "t", 2, "The number of sample teams.") 32 SampleDataCmd.Flags().Int("channels-per-team", 10, "The number of sample channels per team.") 33 SampleDataCmd.Flags().IntP("users", "u", 15, "The number of sample users.") 34 SampleDataCmd.Flags().Int("team-memberships", 2, "The number of sample team memberships per user.") 35 SampleDataCmd.Flags().Int("channel-memberships", 5, "The number of sample channel memberships per user in a team.") 36 SampleDataCmd.Flags().Int("posts-per-channel", 100, "The number of sample post per channel.") 37 SampleDataCmd.Flags().Int("direct-channels", 30, "The number of sample direct message channels.") 38 SampleDataCmd.Flags().Int("posts-per-direct-channel", 15, "The number of sample posts per direct message channel.") 39 SampleDataCmd.Flags().Int("group-channels", 15, "The number of sample group message channels.") 40 SampleDataCmd.Flags().Int("posts-per-group-channel", 30, "The number of sample posts per group message channel.") 41 SampleDataCmd.Flags().IntP("workers", "w", 2, "How many workers to run during the import.") 42 SampleDataCmd.Flags().String("profile-images", "", "Optional. Path to folder with images to randomly pick as user profile image.") 43 SampleDataCmd.Flags().StringP("bulk", "b", "", "Optional. Path to write a JSONL bulk file instead of loading into the database.") 44 RootCmd.AddCommand(SampleDataCmd) 45 } 46 47 func sliceIncludes(vs []string, t string) bool { 48 for _, v := range vs { 49 if v == t { 50 return true 51 } 52 } 53 return false 54 } 55 56 func randomPastTime(seconds int) int64 { 57 now := time.Now() 58 today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.FixedZone("UTC", 0)) 59 return today.Unix() - int64(rand.Intn(seconds*1000)) 60 } 61 62 func randomEmoji() string { 63 emojis := []string{"+1", "-1", "heart", "blush"} 64 return emojis[rand.Intn(len(emojis))] 65 } 66 67 func randomReaction(users []string, parentCreateAt int64) app.ReactionImportData { 68 user := users[rand.Intn(len(users))] 69 emoji := randomEmoji() 70 date := parentCreateAt + int64(rand.Intn(100000)) 71 return app.ReactionImportData{ 72 User: &user, 73 EmojiName: &emoji, 74 CreateAt: &date, 75 } 76 } 77 78 func randomReply(users []string, parentCreateAt int64) app.ReplyImportData { 79 user := users[rand.Intn(len(users))] 80 message := randomMessage(users) 81 date := parentCreateAt + int64(rand.Intn(100000)) 82 return app.ReplyImportData{ 83 User: &user, 84 Message: &message, 85 CreateAt: &date, 86 } 87 } 88 89 func randomMessage(users []string) string { 90 var message string 91 switch rand.Intn(30) { 92 case 0: 93 mention := users[rand.Intn(len(users))] 94 message = "@" + mention + " " + fake.Sentence() 95 case 1: 96 switch rand.Intn(2) { 97 case 0: 98 mattermostVideos := []string{"Q4MgnxbpZas", "BFo7E9-Kc_E", "LsMLR-BHsKg", "MRmGDhlMhNA", "mUOPxT7VgWc"} 99 message = "https://www.youtube.com/watch?v=" + mattermostVideos[rand.Intn(len(mattermostVideos))] 100 case 1: 101 mattermostTweets := []string{"943119062334353408", "949370809528832005", "948539688171819009", "939122439115681792", "938061722027425797"} 102 message = "https://twitter.com/mattermosthq/status/" + mattermostTweets[rand.Intn(len(mattermostTweets))] 103 } 104 case 2: 105 message = "" 106 if rand.Intn(2) == 0 { 107 message += fake.Sentence() 108 } 109 for i := 0; i < rand.Intn(4)+1; i++ { 110 message += "\n * " + fake.Word() 111 } 112 default: 113 if rand.Intn(2) == 0 { 114 message = fake.Sentence() 115 } else { 116 message = fake.Paragraph() 117 } 118 if rand.Intn(3) == 0 { 119 message += "\n" + fake.Sentence() 120 } 121 if rand.Intn(3) == 0 { 122 message += "\n" + fake.Sentence() 123 } 124 if rand.Intn(3) == 0 { 125 message += "\n" + fake.Sentence() 126 } 127 } 128 return message 129 } 130 131 func sampleDataCmdF(command *cobra.Command, args []string) error { 132 a, err := InitDBCommandContextCobra(command) 133 if err != nil { 134 return err 135 } 136 defer a.Shutdown() 137 138 seed, err := command.Flags().GetInt64("seed") 139 if err != nil { 140 return errors.New("Invalid seed parameter") 141 } 142 bulk, err := command.Flags().GetString("bulk") 143 if err != nil { 144 return errors.New("Invalid bulk parameter") 145 } 146 teams, err := command.Flags().GetInt("teams") 147 if err != nil || teams < 0 { 148 return errors.New("Invalid teams parameter") 149 } 150 channelsPerTeam, err := command.Flags().GetInt("channels-per-team") 151 if err != nil || channelsPerTeam < 0 { 152 return errors.New("Invalid channels-per-team parameter") 153 } 154 users, err := command.Flags().GetInt("users") 155 if err != nil || users < 0 { 156 return errors.New("Invalid users parameter") 157 } 158 teamMemberships, err := command.Flags().GetInt("team-memberships") 159 if err != nil || teamMemberships < 0 { 160 return errors.New("Invalid team-memberships parameter") 161 } 162 channelMemberships, err := command.Flags().GetInt("channel-memberships") 163 if err != nil || channelMemberships < 0 { 164 return errors.New("Invalid channel-memberships parameter") 165 } 166 postsPerChannel, err := command.Flags().GetInt("posts-per-channel") 167 if err != nil || postsPerChannel < 0 { 168 return errors.New("Invalid posts-per-channel parameter") 169 } 170 directChannels, err := command.Flags().GetInt("direct-channels") 171 if err != nil || directChannels < 0 { 172 return errors.New("Invalid direct-channels parameter") 173 } 174 postsPerDirectChannel, err := command.Flags().GetInt("posts-per-direct-channel") 175 if err != nil || postsPerDirectChannel < 0 { 176 return errors.New("Invalid posts-per-direct-channel parameter") 177 } 178 groupChannels, err := command.Flags().GetInt("group-channels") 179 if err != nil || groupChannels < 0 { 180 return errors.New("Invalid group-channels parameter") 181 } 182 postsPerGroupChannel, err := command.Flags().GetInt("posts-per-group-channel") 183 if err != nil || postsPerGroupChannel < 0 { 184 return errors.New("Invalid posts-per-group-channel parameter") 185 } 186 workers, err := command.Flags().GetInt("workers") 187 if err != nil { 188 return errors.New("Invalid workers parameter") 189 } 190 profileImagesPath, err := command.Flags().GetString("profile-images") 191 if err != nil { 192 return errors.New("Invalid profile-images parameter") 193 } 194 profileImages := []string{} 195 if profileImagesPath != "" { 196 profileImagesStat, err := os.Stat(profileImagesPath) 197 if os.IsNotExist(err) { 198 return errors.New("Profile images folder doesn't exists.") 199 } 200 if !profileImagesStat.IsDir() { 201 return errors.New("profile-images parameters must be a folder path.") 202 } 203 profileImagesFiles, err := ioutil.ReadDir(profileImagesPath) 204 if err != nil { 205 return errors.New("Invalid profile-images parameter") 206 } 207 for _, profileImage := range profileImagesFiles { 208 profileImages = append(profileImages, path.Join(profileImagesPath, profileImage.Name())) 209 } 210 sort.Strings(profileImages) 211 } 212 213 if workers < 1 { 214 return errors.New("You must have at least one worker.") 215 } 216 if teamMemberships > teams { 217 return errors.New("You can't have more team memberships than teams.") 218 } 219 if channelMemberships > channelsPerTeam { 220 return errors.New("You can't have more channel memberships than channels per team.") 221 } 222 223 var bulkFile *os.File 224 switch bulk { 225 case "": 226 bulkFile, err = ioutil.TempFile("", ".mattermost-sample-data-") 227 defer os.Remove(bulkFile.Name()) 228 if err != nil { 229 return errors.New("Unable to open temporary file.") 230 } 231 case "-": 232 bulkFile = os.Stdout 233 default: 234 bulkFile, err = os.OpenFile(bulk, os.O_RDWR|os.O_CREATE, 0755) 235 if err != nil { 236 return errors.New("Unable to write into the \"" + bulk + "\" file.") 237 } 238 } 239 240 encoder := json.NewEncoder(bulkFile) 241 version := 1 242 encoder.Encode(app.LineImportData{Type: "version", Version: &version}) 243 244 fake.Seed(seed) 245 rand.Seed(seed) 246 247 teamsAndChannels := make(map[string][]string) 248 for i := 0; i < teams; i++ { 249 teamLine := createTeam(i) 250 teamsAndChannels[*teamLine.Team.Name] = []string{} 251 encoder.Encode(teamLine) 252 } 253 254 teamsList := []string{} 255 for teamName := range teamsAndChannels { 256 teamsList = append(teamsList, teamName) 257 } 258 sort.Strings(teamsList) 259 260 for _, teamName := range teamsList { 261 for i := 0; i < channelsPerTeam; i++ { 262 channelLine := createChannel(i, teamName) 263 teamsAndChannels[teamName] = append(teamsAndChannels[teamName], *channelLine.Channel.Name) 264 encoder.Encode(channelLine) 265 } 266 } 267 268 allUsers := []string{} 269 for i := 0; i < users; i++ { 270 userLine := createUser(i, teamMemberships, channelMemberships, teamsAndChannels, profileImages) 271 encoder.Encode(userLine) 272 allUsers = append(allUsers, *userLine.User.Username) 273 } 274 275 for team, channels := range teamsAndChannels { 276 for _, channel := range channels { 277 for i := 0; i < postsPerChannel; i++ { 278 postLine := createPost(team, channel, allUsers) 279 encoder.Encode(postLine) 280 } 281 } 282 } 283 284 for i := 0; i < directChannels; i++ { 285 user1 := allUsers[rand.Intn(len(allUsers))] 286 user2 := allUsers[rand.Intn(len(allUsers))] 287 channelLine := createDirectChannel([]string{user1, user2}) 288 encoder.Encode(channelLine) 289 for j := 0; j < postsPerDirectChannel; j++ { 290 postLine := createDirectPost([]string{user1, user2}) 291 encoder.Encode(postLine) 292 } 293 } 294 295 for i := 0; i < groupChannels; i++ { 296 users := []string{} 297 totalUsers := 3 + rand.Intn(3) 298 for len(users) < totalUsers { 299 user := allUsers[rand.Intn(len(allUsers))] 300 if !sliceIncludes(users, user) { 301 users = append(users, user) 302 } 303 } 304 channelLine := createDirectChannel(users) 305 encoder.Encode(channelLine) 306 for j := 0; j < postsPerGroupChannel; j++ { 307 postLine := createDirectPost(users) 308 encoder.Encode(postLine) 309 } 310 } 311 312 if bulk == "" { 313 _, err := bulkFile.Seek(0, 0) 314 if err != nil { 315 return errors.New("Unable to read correctly the temporary file.") 316 } 317 importErr, lineNumber := a.BulkImport(bulkFile, false, workers) 318 if importErr != nil { 319 return fmt.Errorf("%s: %s, %s (line: %d)", importErr.Where, importErr.Message, importErr.DetailedError, lineNumber) 320 } 321 } else if bulk != "-" { 322 err := bulkFile.Close() 323 if err != nil { 324 return errors.New("Unable to close correctly the output file") 325 } 326 } 327 328 return nil 329 } 330 331 func createUser(idx int, teamMemberships int, channelMemberships int, teamsAndChannels map[string][]string, profileImages []string) app.LineImportData { 332 password := fmt.Sprintf("user-%d", idx) 333 email := fmt.Sprintf("user-%d@sample.mattermost.com", idx) 334 firstName := fake.FirstName() 335 lastName := fake.LastName() 336 username := fmt.Sprintf("%s.%s", strings.ToLower(firstName), strings.ToLower(lastName)) 337 if idx == 0 { 338 username = "sysadmin" 339 password = "sysadmin" 340 email = "sysadmin@sample.mattermost.com" 341 } else if idx == 1 { 342 username = "user-1" 343 } 344 position := fake.JobTitle() 345 roles := "system_user" 346 if idx%5 == 0 { 347 roles = "system_admin system_user" 348 } 349 350 // The 75% of the users have custom profile image 351 var profileImage *string = nil 352 if rand.Intn(4) != 0 { 353 profileImageSelector := rand.Int() 354 if len(profileImages) > 0 { 355 profileImage = &profileImages[profileImageSelector%len(profileImages)] 356 } 357 } 358 359 useMilitaryTime := "false" 360 if idx != 0 && rand.Intn(2) == 0 { 361 useMilitaryTime = "true" 362 } 363 364 collapsePreviews := "false" 365 if idx != 0 && rand.Intn(2) == 0 { 366 collapsePreviews = "true" 367 } 368 369 messageDisplay := "clean" 370 if idx != 0 && rand.Intn(2) == 0 { 371 messageDisplay = "compact" 372 } 373 374 channelDisplayMode := "full" 375 if idx != 0 && rand.Intn(2) == 0 { 376 channelDisplayMode = "centered" 377 } 378 379 // Some users has nickname 380 nickname := "" 381 if rand.Intn(5) == 0 { 382 nickname = fake.Company() 383 } 384 385 // Half of users skip tutorial 386 tutorialStep := "999" 387 switch rand.Intn(6) { 388 case 1: 389 tutorialStep = "1" 390 case 2: 391 tutorialStep = "2" 392 case 3: 393 tutorialStep = "3" 394 } 395 396 teams := []app.UserTeamImportData{} 397 possibleTeams := []string{} 398 for teamName := range teamsAndChannels { 399 possibleTeams = append(possibleTeams, teamName) 400 } 401 sort.Strings(possibleTeams) 402 for x := 0; x < teamMemberships; x++ { 403 if len(possibleTeams) == 0 { 404 break 405 } 406 position := rand.Intn(len(possibleTeams)) 407 team := possibleTeams[position] 408 possibleTeams = append(possibleTeams[:position], possibleTeams[position+1:]...) 409 if teamChannels, err := teamsAndChannels[team]; err { 410 teams = append(teams, createTeamMembership(channelMemberships, teamChannels, &team)) 411 } 412 } 413 414 user := app.UserImportData{ 415 ProfileImage: profileImage, 416 Username: &username, 417 Email: &email, 418 Password: &password, 419 Nickname: &nickname, 420 FirstName: &firstName, 421 LastName: &lastName, 422 Position: &position, 423 Roles: &roles, 424 Teams: &teams, 425 UseMilitaryTime: &useMilitaryTime, 426 CollapsePreviews: &collapsePreviews, 427 MessageDisplay: &messageDisplay, 428 ChannelDisplayMode: &channelDisplayMode, 429 TutorialStep: &tutorialStep, 430 } 431 return app.LineImportData{ 432 Type: "user", 433 User: &user, 434 } 435 } 436 437 func createTeamMembership(numOfchannels int, teamChannels []string, teamName *string) app.UserTeamImportData { 438 roles := "team_user" 439 if rand.Intn(5) == 0 { 440 roles = "team_user team_admin" 441 } 442 channels := []app.UserChannelImportData{} 443 teamChannelsCopy := append([]string(nil), teamChannels...) 444 for x := 0; x < numOfchannels; x++ { 445 if len(teamChannelsCopy) == 0 { 446 break 447 } 448 position := rand.Intn(len(teamChannelsCopy)) 449 channelName := teamChannelsCopy[position] 450 teamChannelsCopy = append(teamChannelsCopy[:position], teamChannelsCopy[position+1:]...) 451 channels = append(channels, createChannelMembership(channelName)) 452 } 453 454 return app.UserTeamImportData{ 455 Name: teamName, 456 Roles: &roles, 457 Channels: &channels, 458 } 459 } 460 461 func createChannelMembership(channelName string) app.UserChannelImportData { 462 roles := "channel_user" 463 if rand.Intn(5) == 0 { 464 roles = "channel_user channel_admin" 465 } 466 favorite := rand.Intn(5) == 0 467 468 return app.UserChannelImportData{ 469 Name: &channelName, 470 Roles: &roles, 471 Favorite: &favorite, 472 } 473 } 474 475 func createTeam(idx int) app.LineImportData { 476 displayName := fake.Word() 477 name := fmt.Sprintf("%s-%d", fake.Word(), idx) 478 allowOpenInvite := rand.Intn(2) == 0 479 480 description := fake.Paragraph() 481 if len(description) > 255 { 482 description = description[0:255] 483 } 484 485 teamType := "O" 486 if rand.Intn(2) == 0 { 487 teamType = "I" 488 } 489 490 team := app.TeamImportData{ 491 DisplayName: &displayName, 492 Name: &name, 493 AllowOpenInvite: &allowOpenInvite, 494 Description: &description, 495 Type: &teamType, 496 } 497 return app.LineImportData{ 498 Type: "team", 499 Team: &team, 500 } 501 } 502 503 func createChannel(idx int, teamName string) app.LineImportData { 504 displayName := fake.Word() 505 name := fmt.Sprintf("%s-%d", fake.Word(), idx) 506 header := fake.Paragraph() 507 purpose := fake.Paragraph() 508 509 if len(purpose) > 250 { 510 purpose = purpose[0:250] 511 } 512 513 channelType := "P" 514 if rand.Intn(2) == 0 { 515 channelType = "O" 516 } 517 518 channel := app.ChannelImportData{ 519 Team: &teamName, 520 Name: &name, 521 DisplayName: &displayName, 522 Type: &channelType, 523 Header: &header, 524 Purpose: &purpose, 525 } 526 return app.LineImportData{ 527 Type: "channel", 528 Channel: &channel, 529 } 530 } 531 532 func createPost(team string, channel string, allUsers []string) app.LineImportData { 533 message := randomMessage(allUsers) 534 create_at := randomPastTime(50000) 535 user := allUsers[rand.Intn(len(allUsers))] 536 537 // Some messages are flagged by an user 538 flagged_by := []string{} 539 if rand.Intn(10) == 0 { 540 flagged_by = append(flagged_by, allUsers[rand.Intn(len(allUsers))]) 541 } 542 543 reactions := []app.ReactionImportData{} 544 if rand.Intn(10) == 0 { 545 for { 546 reactions = append(reactions, randomReaction(allUsers, create_at)) 547 if rand.Intn(3) == 0 { 548 break 549 } 550 } 551 } 552 553 replies := []app.ReplyImportData{} 554 if rand.Intn(10) == 0 { 555 for { 556 replies = append(replies, randomReply(allUsers, create_at)) 557 if rand.Intn(4) == 0 { 558 break 559 } 560 } 561 } 562 563 post := app.PostImportData{ 564 Team: &team, 565 Channel: &channel, 566 User: &user, 567 Message: &message, 568 CreateAt: &create_at, 569 FlaggedBy: &flagged_by, 570 Reactions: &reactions, 571 Replies: &replies, 572 } 573 return app.LineImportData{ 574 Type: "post", 575 Post: &post, 576 } 577 } 578 579 func createDirectChannel(members []string) app.LineImportData { 580 header := fake.Sentence() 581 582 channel := app.DirectChannelImportData{ 583 Members: &members, 584 Header: &header, 585 } 586 return app.LineImportData{ 587 Type: "direct_channel", 588 DirectChannel: &channel, 589 } 590 } 591 592 func createDirectPost(members []string) app.LineImportData { 593 message := randomMessage(members) 594 create_at := randomPastTime(50000) 595 user := members[rand.Intn(len(members))] 596 597 // Some messages are flagged by an user 598 flagged_by := []string{} 599 if rand.Intn(10) == 0 { 600 flagged_by = append(flagged_by, members[rand.Intn(len(members))]) 601 } 602 603 reactions := []app.ReactionImportData{} 604 if rand.Intn(10) == 0 { 605 for { 606 reactions = append(reactions, randomReaction(members, create_at)) 607 if rand.Intn(3) == 0 { 608 break 609 } 610 } 611 } 612 613 replies := []app.ReplyImportData{} 614 if rand.Intn(10) == 0 { 615 for { 616 replies = append(replies, randomReply(members, create_at)) 617 if rand.Intn(4) == 0 { 618 break 619 } 620 } 621 } 622 623 post := app.DirectPostImportData{ 624 ChannelMembers: &members, 625 User: &user, 626 Message: &message, 627 CreateAt: &create_at, 628 FlaggedBy: &flagged_by, 629 Reactions: &reactions, 630 Replies: &replies, 631 } 632 return app.LineImportData{ 633 Type: "direct_post", 634 DirectPost: &post, 635 } 636 }