mirror of https://github.com/stashapp/stash.git
712 lines
15 KiB
Go
712 lines
15 KiB
Go
//go:build tools
|
|
// +build tools
|
|
|
|
package main
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"fmt"
|
|
"log"
|
|
"math"
|
|
"math/rand"
|
|
"os"
|
|
"path"
|
|
"strconv"
|
|
"time"
|
|
|
|
"gopkg.in/yaml.v2"
|
|
|
|
"github.com/stashapp/stash/pkg/file"
|
|
"github.com/stashapp/stash/pkg/fsutil"
|
|
"github.com/stashapp/stash/pkg/hash/md5"
|
|
"github.com/stashapp/stash/pkg/models"
|
|
"github.com/stashapp/stash/pkg/sliceutil"
|
|
"github.com/stashapp/stash/pkg/sqlite"
|
|
"github.com/stashapp/stash/pkg/txn"
|
|
)
|
|
|
|
const batchSize = 50000
|
|
|
|
// create an example database by generating a number of scenes, markers,
|
|
// performers, studios, galleries, chapters and tags, and associating between them all
|
|
|
|
type config struct {
|
|
Database string `yaml:"database"`
|
|
Scenes int `yaml:"scenes"`
|
|
Markers int `yaml:"markers"`
|
|
Images int `yaml:"images"`
|
|
Galleries int `yaml:"galleries"`
|
|
Chapters int `yaml:"chapters"`
|
|
Performers int `yaml:"performers"`
|
|
Studios int `yaml:"studios"`
|
|
Tags int `yaml:"tags"`
|
|
Naming namingConfig `yaml:"naming"`
|
|
}
|
|
|
|
var (
|
|
repo models.Repository
|
|
c *config
|
|
db *sqlite.Database
|
|
folderID file.FolderID
|
|
)
|
|
|
|
func main() {
|
|
rand.Seed(time.Now().UnixNano())
|
|
|
|
var err error
|
|
c, err = loadConfig()
|
|
if err != nil {
|
|
log.Fatalf("couldn't load configuration: %v", err)
|
|
}
|
|
|
|
initNaming(*c)
|
|
|
|
db = sqlite.NewDatabase()
|
|
repo = db.TxnRepository()
|
|
|
|
logf("Initializing database...")
|
|
if err = db.Open(c.Database); err != nil {
|
|
log.Fatalf("couldn't initialize database: %v", err)
|
|
}
|
|
logf("Populating database...")
|
|
populateDB()
|
|
}
|
|
|
|
func loadConfig() (*config, error) {
|
|
ret := &config{}
|
|
|
|
file, err := os.Open("config.yml")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer file.Close()
|
|
|
|
parser := yaml.NewDecoder(file)
|
|
parser.SetStrict(true)
|
|
err = parser.Decode(&ret)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return ret, nil
|
|
}
|
|
|
|
func populateDB() {
|
|
makeTags(c.Tags)
|
|
makeStudios(c.Studios)
|
|
makePerformers(c.Performers)
|
|
makeScenes(c.Scenes)
|
|
makeImages(c.Images)
|
|
makeGalleries(c.Galleries)
|
|
makeChapters(c.Chapters)
|
|
makeMarkers(c.Markers)
|
|
}
|
|
|
|
func withTxn(f func(ctx context.Context) error) error {
|
|
return txn.WithTxn(context.Background(), db, f)
|
|
}
|
|
|
|
func retry(attempts int, fn func() error) error {
|
|
var err error
|
|
for tries := 0; tries < attempts; tries++ {
|
|
err = fn()
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
}
|
|
|
|
return err
|
|
}
|
|
|
|
func getOrCreateFolder(ctx context.Context, p string) (*file.Folder, error) {
|
|
ret, err := repo.Folder.FindByPath(ctx, p)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if ret != nil {
|
|
return ret, nil
|
|
}
|
|
|
|
var parentID *file.FolderID
|
|
|
|
if p != "." {
|
|
parent := path.Dir(p)
|
|
parentFolder, err := getOrCreateFolder(ctx, parent)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
parentID = &parentFolder.ID
|
|
}
|
|
|
|
f := file.Folder{
|
|
Path: p,
|
|
ParentFolderID: parentID,
|
|
}
|
|
|
|
if err := repo.Folder.Create(ctx, &f); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
ret = &f
|
|
return ret, nil
|
|
}
|
|
|
|
func makeTags(n int) {
|
|
logf("creating %d tags...", n)
|
|
for i := 0; i < n; i++ {
|
|
if err := retry(100, func() error {
|
|
return withTxn(func(ctx context.Context) error {
|
|
name := names[c.Naming.Tags].generateName(1)
|
|
tag := models.Tag{
|
|
Name: name,
|
|
}
|
|
|
|
created, err := repo.Tag.Create(ctx, tag)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if rand.Intn(100) > 5 {
|
|
t, _, err := repo.Tag.Query(ctx, nil, getRandomFilter(1))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if len(t) > 0 && t[0].ID != created.ID {
|
|
if err := repo.Tag.UpdateParentTags(ctx, created.ID, []int{t[0].ID}); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
})
|
|
}); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func makeStudios(n int) {
|
|
logf("creating %d studios...", n)
|
|
for i := 0; i < n; i++ {
|
|
if err := retry(100, func() error {
|
|
return withTxn(func(ctx context.Context) error {
|
|
name := names[c.Naming.Tags].generateName(rand.Intn(5) + 1)
|
|
studio := models.Studio{
|
|
Name: sql.NullString{String: name, Valid: true},
|
|
Checksum: md5.FromString(name),
|
|
}
|
|
|
|
if rand.Intn(100) > 5 {
|
|
ss, _, err := repo.Studio.Query(ctx, nil, getRandomFilter(1))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if len(ss) > 0 {
|
|
studio.ParentID = sql.NullInt64{
|
|
Int64: int64(ss[0].ID),
|
|
Valid: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
_, err := repo.Studio.Create(ctx, studio)
|
|
return err
|
|
})
|
|
}); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func makePerformers(n int) {
|
|
logf("creating %d performers...", n)
|
|
for i := 0; i < n; i++ {
|
|
if err := retry(100, func() error {
|
|
return withTxn(func(ctx context.Context) error {
|
|
name := generatePerformerName()
|
|
performer := &models.Performer{
|
|
Name: name,
|
|
Checksum: md5.FromString(name),
|
|
}
|
|
|
|
// TODO - set tags
|
|
|
|
err := repo.Performer.Create(ctx, performer)
|
|
if err != nil {
|
|
err = fmt.Errorf("error creating performer with name: %s: %s", performer.Name, err.Error())
|
|
}
|
|
return err
|
|
})
|
|
}); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func generateBaseFile(parentFolderID file.FolderID, path string) *file.BaseFile {
|
|
return &file.BaseFile{
|
|
Basename: path,
|
|
ParentFolderID: parentFolderID,
|
|
Fingerprints: []file.Fingerprint{
|
|
file.Fingerprint{
|
|
Type: "md5",
|
|
Fingerprint: md5.FromString(path),
|
|
},
|
|
file.Fingerprint{
|
|
Type: "oshash",
|
|
Fingerprint: md5.FromString(path),
|
|
},
|
|
},
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
}
|
|
}
|
|
|
|
func generateVideoFile(parentFolderID file.FolderID, path string) file.File {
|
|
w, h := getResolution()
|
|
|
|
return &file.VideoFile{
|
|
BaseFile: generateBaseFile(parentFolderID, path),
|
|
Duration: rand.Float64() * 14400,
|
|
Height: h,
|
|
Width: w,
|
|
}
|
|
}
|
|
|
|
func makeVideoFile(ctx context.Context, path string) (file.File, error) {
|
|
folderPath := fsutil.GetIntraDir(path, 2, 2)
|
|
parentFolder, err := getOrCreateFolder(ctx, folderPath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
f := generateVideoFile(parentFolder.ID, path)
|
|
|
|
if err := repo.File.Create(ctx, f); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return f, nil
|
|
}
|
|
|
|
func logf(f string, args ...interface{}) {
|
|
log.Printf(f+"\n", args...)
|
|
}
|
|
|
|
func makeScenes(n int) {
|
|
logf("creating %d scenes...", n)
|
|
for i := 0; i < n; {
|
|
// do in batches of 1000
|
|
batch := i + batchSize
|
|
|
|
if err := withTxn(func(ctx context.Context) error {
|
|
for ; i < batch && i < n; i++ {
|
|
scene := generateScene(i)
|
|
scene.StudioID = getRandomStudioID(ctx)
|
|
makeSceneRelationships(ctx, &scene)
|
|
|
|
path := md5.FromString("scene/" + strconv.Itoa(i))
|
|
f, err := makeVideoFile(ctx, path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := repo.Scene.Create(ctx, &scene, []file.ID{f.Base().ID}); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}); err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
logf("... created %d scenes", i)
|
|
}
|
|
}
|
|
|
|
func getResolution() (int, int) {
|
|
res := models.AllResolutionEnum[rand.Intn(len(models.AllResolutionEnum))]
|
|
h := res.GetMaxResolution()
|
|
var w int
|
|
if h == 240 || h == 480 || rand.Intn(10) == 9 {
|
|
w = h * 4 / 3
|
|
} else {
|
|
w = h * 16 / 9
|
|
}
|
|
|
|
if rand.Intn(10) == 9 {
|
|
return h, w
|
|
}
|
|
|
|
return w, h
|
|
}
|
|
|
|
func getBool() {
|
|
return rand.Intn(2) == 0
|
|
}
|
|
|
|
func getDate() time.Time {
|
|
s := rand.Int63n(time.Now().Unix())
|
|
|
|
return time.Unix(s, 0)
|
|
}
|
|
|
|
func generateScene(i int) models.Scene {
|
|
return models.Scene{
|
|
Title: names[c.Naming.Scenes].generateName(rand.Intn(7) + 1),
|
|
Date: &models.Date{
|
|
Time: getDate(),
|
|
},
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
}
|
|
}
|
|
|
|
func generateImageFile(parentFolderID file.FolderID, path string) file.File {
|
|
w, h := getResolution()
|
|
|
|
return &file.ImageFile{
|
|
BaseFile: generateBaseFile(parentFolderID, path),
|
|
Height: h,
|
|
Width: w,
|
|
Clip: getBool(),
|
|
}
|
|
}
|
|
|
|
func makeImageFile(ctx context.Context, path string) (file.File, error) {
|
|
folderPath := fsutil.GetIntraDir(path, 2, 2)
|
|
parentFolder, err := getOrCreateFolder(ctx, folderPath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
f := generateImageFile(parentFolder.ID, path)
|
|
|
|
if err := repo.File.Create(ctx, f); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return f, nil
|
|
}
|
|
|
|
func makeImages(n int) {
|
|
logf("creating %d images...", n)
|
|
for i := 0; i < n; {
|
|
// do in batches of 1000
|
|
batch := i + batchSize
|
|
if err := withTxn(func(ctx context.Context) error {
|
|
for ; i < batch && i < n; i++ {
|
|
image := generateImage(i)
|
|
image.StudioID = getRandomStudioID(ctx)
|
|
makeImageRelationships(ctx, &image)
|
|
|
|
path := md5.FromString("image/" + strconv.Itoa(i))
|
|
f, err := makeImageFile(ctx, path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := repo.Image.Create(ctx, &models.ImageCreateInput{
|
|
Image: &image,
|
|
FileIDs: []file.ID{f.Base().ID},
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
logf("... created %d images", i)
|
|
|
|
return nil
|
|
}); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func generateImage(i int) models.Image {
|
|
return models.Image{
|
|
Title: names[c.Naming.Images].generateName(rand.Intn(7) + 1),
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
}
|
|
}
|
|
|
|
func makeGalleries(n int) {
|
|
logf("creating %d galleries...", n)
|
|
for i := 0; i < n; {
|
|
// do in batches of 1000
|
|
batch := i + batchSize
|
|
|
|
if err := withTxn(func(ctx context.Context) error {
|
|
for ; i < batch && i < n; i++ {
|
|
gallery := generateGallery(i)
|
|
gallery.StudioID = getRandomStudioID(ctx)
|
|
gallery.TagIDs = models.NewRelatedIDs(getRandomTags(ctx, 0, 15))
|
|
gallery.PerformerIDs = models.NewRelatedIDs(getRandomPerformers(ctx))
|
|
|
|
path := md5.FromString("gallery/" + strconv.Itoa(i))
|
|
f, err := makeZipFile(ctx, path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := repo.Gallery.Create(ctx, &gallery, []file.ID{f.Base().ID}); err != nil {
|
|
return err
|
|
}
|
|
|
|
makeGalleryRelationships(ctx, &gallery)
|
|
}
|
|
|
|
return nil
|
|
}); err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
logf("... created %d galleries", i)
|
|
}
|
|
}
|
|
|
|
func generateZipFile(parentFolderID file.FolderID, path string) file.File {
|
|
return generateBaseFile(parentFolderID, path)
|
|
}
|
|
|
|
func makeZipFile(ctx context.Context, path string) (file.File, error) {
|
|
folderPath := fsutil.GetIntraDir(path, 2, 2)
|
|
parentFolder, err := getOrCreateFolder(ctx, folderPath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
f := generateZipFile(parentFolder.ID, path)
|
|
|
|
if err := repo.File.Create(ctx, f); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return f, nil
|
|
}
|
|
|
|
func generateGallery(i int) models.Gallery {
|
|
return models.Gallery{
|
|
Title: names[c.Naming.Galleries].generateName(rand.Intn(7) + 1),
|
|
Date: &models.Date{
|
|
Time: getDate(),
|
|
},
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
}
|
|
}
|
|
|
|
func makeChapters(n int) {
|
|
logf("creating %d chapters...", n)
|
|
for i := 0; i < n; {
|
|
// do in batches of 1000
|
|
batch := i + batchSize
|
|
if err := withTxn(func(ctx context.Context) error {
|
|
for ; i < batch && i < n; i++ {
|
|
chapter := generateChapter(i)
|
|
chapter.GalleryID = models.NullInt64(int64(getRandomGallery()))
|
|
|
|
created, err := repo.GalleryChapter.Create(ctx, chapter)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
logf("... created %d chapters", i)
|
|
|
|
return nil
|
|
}); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func generateChapter(i int) models.GalleryChapter {
|
|
return models.GalleryChapter{
|
|
Title: names[c.Naming.Galleries].generateName(rand.Intn(7) + 1),
|
|
ImageIndex: rand.Intn(200),
|
|
}
|
|
}
|
|
|
|
func makeMarkers(n int) {
|
|
logf("creating %d markers...", n)
|
|
for i := 0; i < n; {
|
|
// do in batches of 1000
|
|
batch := i + batchSize
|
|
if err := withTxn(func(ctx context.Context) error {
|
|
for ; i < batch && i < n; i++ {
|
|
marker := generateMarker(i)
|
|
marker.SceneID = models.NullInt64(int64(getRandomScene()))
|
|
marker.PrimaryTagID = getRandomTags(ctx, 1, 1)[0]
|
|
|
|
created, err := repo.SceneMarker.Create(ctx, marker)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
tags := getRandomTags(ctx, 0, 5)
|
|
// remove primary tag
|
|
tags = sliceutil.Exclude(tags, []int{marker.PrimaryTagID})
|
|
if err := repo.SceneMarker.UpdateTags(ctx, created.ID, tags); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
logf("... created %d markers", i)
|
|
|
|
return nil
|
|
}); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func generateMarker(i int) models.SceneMarker {
|
|
return models.SceneMarker{
|
|
Title: names[c.Naming.Scenes].generateName(rand.Intn(7) + 1),
|
|
}
|
|
}
|
|
|
|
func getRandomFilter(n int) *models.FindFilterType {
|
|
seed := math.Floor(rand.Float64() * math.Pow10(8))
|
|
sortBy := fmt.Sprintf("random_%.f", seed)
|
|
return &models.FindFilterType{
|
|
Sort: &sortBy,
|
|
PerPage: &n,
|
|
}
|
|
}
|
|
|
|
func getRandomStudioID(ctx context.Context) *int {
|
|
if rand.Intn(10) == 0 {
|
|
return nil
|
|
}
|
|
|
|
// s, _, err := r.Studio().Query(nil, getRandomFilter(1))
|
|
// if err != nil {
|
|
// panic(err)
|
|
// }
|
|
|
|
v := rand.Intn(c.Studios) + 1
|
|
return &v
|
|
}
|
|
|
|
func makeSceneRelationships(ctx context.Context, s *models.Scene) {
|
|
// add tags
|
|
s.TagIDs = models.NewRelatedIDs(getRandomTags(ctx, 0, 15))
|
|
|
|
// add performers
|
|
s.PerformerIDs = models.NewRelatedIDs(getRandomPerformers(ctx))
|
|
}
|
|
|
|
func makeImageRelationships(ctx context.Context, i *models.Image) {
|
|
// there are typically many more images. For performance reasons
|
|
// only a small proportion should have tags/performers
|
|
|
|
// add tags
|
|
if rand.Intn(100) == 0 {
|
|
i.TagIDs = models.NewRelatedIDs(getRandomTags(ctx, 1, 15))
|
|
}
|
|
|
|
// add performers
|
|
if rand.Intn(100) <= 1 {
|
|
i.PerformerIDs = models.NewRelatedIDs(getRandomPerformers(ctx))
|
|
}
|
|
}
|
|
|
|
func makeGalleryRelationships(ctx context.Context, g *models.Gallery) {
|
|
// add images
|
|
imageIDs := getRandomImages(ctx)
|
|
if len(imageIDs) > 0 {
|
|
if err := repo.Gallery.UpdateImages(ctx, g.ID, imageIDs); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func getRandomPerformers(ctx context.Context) []int {
|
|
n := rand.Intn(5)
|
|
|
|
var ret []int
|
|
// if n > 0 {
|
|
// p, _, err := r.Performer().Query(nil, getRandomFilter(n))
|
|
// if err != nil {
|
|
// panic(err)
|
|
// }
|
|
|
|
// for _, pp := range p {
|
|
// ret = sliceutil.AppendUnique(ret, pp.ID)
|
|
// }
|
|
// }
|
|
|
|
for i := 0; i < n; i++ {
|
|
ret = sliceutil.AppendUnique(ret, rand.Intn(c.Performers)+1)
|
|
}
|
|
|
|
return ret
|
|
}
|
|
|
|
func getRandomScene() int {
|
|
return rand.Intn(c.Scenes) + 1
|
|
}
|
|
|
|
func getRandomGallery() int {
|
|
return rand.Intn(c.Galleries) + 1
|
|
}
|
|
|
|
func getRandomTags(ctx context.Context, min, max int) []int {
|
|
var n int
|
|
if min == max {
|
|
n = min
|
|
} else {
|
|
n = rand.Intn(max-min) + min
|
|
}
|
|
|
|
var ret []int
|
|
// if n > 0 {
|
|
// t, _, err := r.Tag().Query(nil, getRandomFilter(n))
|
|
// if err != nil {
|
|
// panic(err)
|
|
// }
|
|
|
|
// for _, tt := range t {
|
|
// ret = sliceutil.AppendUnique(ret, tt.ID)
|
|
// }
|
|
// }
|
|
|
|
for i := 0; i < n; i++ {
|
|
ret = sliceutil.AppendUnique(ret, rand.Intn(c.Tags)+1)
|
|
}
|
|
|
|
return ret
|
|
}
|
|
|
|
func getRandomImages(ctx context.Context) []int {
|
|
n := rand.Intn(500)
|
|
|
|
var ret []int
|
|
// if n > 0 {
|
|
// t, _, err := r.Image().Query(nil, getRandomFilter(n))
|
|
// if err != nil {
|
|
// panic(err)
|
|
// }
|
|
|
|
// for _, tt := range t {
|
|
// ret = sliceutil.AppendUnique(ret, tt.ID)
|
|
// }
|
|
// }
|
|
|
|
for i := 0; i < n; i++ {
|
|
ret = sliceutil.AppendUnique(ret, rand.Intn(c.Images)+1)
|
|
}
|
|
|
|
return ret
|
|
}
|