diff --git a/helpers/dbHelpers.go b/helpers/dbHelpers.go index 5809ffa..9dd402a 100644 --- a/helpers/dbHelpers.go +++ b/helpers/dbHelpers.go @@ -5,6 +5,10 @@ import ( "gorm.io/gorm" ) +func AddMigrations(db *gorm.DB) { + db.AutoMigrate(&models.Post{}, &models.MediaAttachment{}, &models.Account{}) +} + func GetExistingPostIds(db *gorm.DB) []string { var existingPostIds []string db.Model(&models.Post{}).Pluck("id", &existingPostIds) diff --git a/main.go b/main.go index 1f60865..677ccba 100644 --- a/main.go +++ b/main.go @@ -2,17 +2,22 @@ package main import ( "CatsOfMastodonBotGo/helpers" - "CatsOfMastodonBotGo/models" "context" "log" + "os" "time" + "gorm.io/driver/sqlite" "gorm.io/gorm" + "gorm.io/gorm/logger" ) func main() { var tag = "catsofmastodon" - var instance = "https://haminoa.net" + var instance = os.Getenv("INSTANCE") + if instance == "" { + instance = "https://haminoa.net" + } ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() err, posts := helpers.GetPosts(ctx, tag, instance) @@ -21,13 +26,12 @@ func main() { } log.Println("Number of fetched posts: ", len(posts)) - db, err := gorm.Open(sqlite.Open("test.db"), &gorm.Config{}) + db, err := gorm.Open(sqlite.Open("test.db"), &gorm.Config{Logger: logger.Default.LogMode(logger.Warn)}) if err != nil { panic("failed to connect database") } - // Migrate the schema - db.AutoMigrate(&models.Post{}, &models.MediaAttachment{}, &models.Account{}) + helpers.AddMigrations(db) var existingPostIds = helpers.GetExistingPostIds(db) var existingAccountIds = helpers.GetExistingAccountIds(db) @@ -40,23 +44,14 @@ func main() { log.Println("Number of new accounts: ", len(newAccounts)) if newAccounts != nil { - log.Println("Number of inserted accounts: ", db.Create(&newAccounts).RowsAffected) + log.Println("Number of inserted accounts: ", helpers.InsertNewAccounts(db, newAccounts)) } else { log.Println("No new accounts inserted") } if newPosts != nil { - log.Println("Number of inserted posts: ", db.Create(&newPosts).RowsAffected) + log.Println("Number of inserted posts: ", helpers.InsertNewPosts(db, newPosts)) } else { log.Println("No new posts inserted") } } - -func arrayContains(arr []string, str string) bool { - for _, a := range arr { - if a == str { - return true - } - } - return false -}