From bb167ea360dda0925b7c62add74a827d9de3a4f9 Mon Sep 17 00:00:00 2001
From: Derrick Hammer <derrick@derrickhammer.com>
Date: Wed, 28 Feb 2024 21:20:05 -0500
Subject: [PATCH] refactor: use a dynamic model registration process so we can
 use build tags to exclude/include

---
 db/db.go                        | 20 +++-----------------
 db/models/api_key.go            |  4 ++++
 db/models/blocklist.go          |  7 ++++++-
 db/models/dnslink.go            |  4 ++++
 db/models/download.go           |  7 ++++++-
 db/models/email_verification.go |  4 ++++
 db/models/models.go             | 11 +++++++++++
 db/models/password_reset.go     |  4 ++++
 db/models/pin.go                |  4 ++++
 db/models/public_key.go         |  4 ++++
 db/models/s3_upload.go          |  4 ++++
 db/models/s5_challenge.go       |  5 +++++
 db/models/tus_lock.go           |  8 +++++++-
 db/models/tus_upload.go         |  5 +++++
 db/models/upload.go             |  4 ++++
 db/models/user.go               |  4 ++++
 16 files changed, 79 insertions(+), 20 deletions(-)
 create mode 100644 db/models/models.go

diff --git a/db/db.go b/db/db.go
index 4817a8d..982b4a6 100644
--- a/db/db.go
+++ b/db/db.go
@@ -4,13 +4,14 @@ import (
 	"context"
 	"fmt"
 
+	"git.lumeweb.com/LumeWeb/portal/db/models"
+
 	"github.com/redis/go-redis/v9"
 
 	"go.uber.org/zap"
 
 	"git.lumeweb.com/LumeWeb/portal/config"
 
-	"git.lumeweb.com/LumeWeb/portal/db/models"
 	"github.com/go-gorm/caches/v4"
 	"go.uber.org/fx"
 	"gorm.io/driver/mysql"
@@ -60,22 +61,7 @@ func NewDatabase(lc fx.Lifecycle, params DatabaseParams) *gorm.DB {
 
 	lc.Append(fx.Hook{
 		OnStart: func(ctx context.Context) error {
-			return db.AutoMigrate(
-				&models.APIKey{},
-				&models.Blocklist{},
-				&models.DNSLink{},
-				&models.Download{},
-				&models.EmailVerification{},
-				&models.PasswordReset{},
-				&models.Pin{},
-				&models.PublicKey{},
-				&models.Upload{},
-				&models.User{},
-				&models.S3Upload{},
-				&models.S5Challenge{},
-				&models.TusLock{},
-				&models.TusUpload{},
-			)
+			return db.AutoMigrate(models.GetModels()...)
 		},
 	})
 
diff --git a/db/models/api_key.go b/db/models/api_key.go
index 430b1e7..6c1a261 100644
--- a/db/models/api_key.go
+++ b/db/models/api_key.go
@@ -2,6 +2,10 @@ package models
 
 import "gorm.io/gorm"
 
+func init() {
+	registerModel(&APIKey{})
+}
+
 type APIKey struct {
 	gorm.Model
 	UserID uint
diff --git a/db/models/blocklist.go b/db/models/blocklist.go
index df9d013..0a1e64c 100644
--- a/db/models/blocklist.go
+++ b/db/models/blocklist.go
@@ -1,10 +1,15 @@
 package models
 
 import (
-	"gorm.io/gorm"
 	"time"
+
+	"gorm.io/gorm"
 )
 
+func init() {
+	registerModel(&Blocklist{})
+}
+
 type Blocklist struct {
 	gorm.Model
 	IP        string
diff --git a/db/models/dnslink.go b/db/models/dnslink.go
index 5d4424d..9da354b 100644
--- a/db/models/dnslink.go
+++ b/db/models/dnslink.go
@@ -2,6 +2,10 @@ package models
 
 import "gorm.io/gorm"
 
+func init() {
+	registerModel(&DNSLink{})
+}
+
 type DNSLink struct {
 	gorm.Model
 	UserID   uint `gorm:"uniqueIndex:idx_user_id_upload"`
diff --git a/db/models/download.go b/db/models/download.go
index 876c084..4f9456a 100644
--- a/db/models/download.go
+++ b/db/models/download.go
@@ -1,10 +1,15 @@
 package models
 
 import (
-	"gorm.io/gorm"
 	"time"
+
+	"gorm.io/gorm"
 )
 
+func init() {
+	registerModel(&Download{})
+}
+
 type Download struct {
 	gorm.Model
 	UserID       uint
diff --git a/db/models/email_verification.go b/db/models/email_verification.go
index ed386ec..17b7c80 100644
--- a/db/models/email_verification.go
+++ b/db/models/email_verification.go
@@ -6,6 +6,10 @@ import (
 	"gorm.io/gorm"
 )
 
+func init() {
+	registerModel(&EmailVerification{})
+}
+
 type EmailVerification struct {
 	gorm.Model
 
diff --git a/db/models/models.go b/db/models/models.go
new file mode 100644
index 0000000..f4d532a
--- /dev/null
+++ b/db/models/models.go
@@ -0,0 +1,11 @@
+package models
+
+var registered []interface{}
+
+func registerModel(model interface{}) {
+	registered = append(registered, model)
+}
+
+func GetModels() []interface{} {
+	return registered
+}
diff --git a/db/models/password_reset.go b/db/models/password_reset.go
index 539f216..5c044e1 100644
--- a/db/models/password_reset.go
+++ b/db/models/password_reset.go
@@ -6,6 +6,10 @@ import (
 	"gorm.io/gorm"
 )
 
+func init() {
+	registerModel(&PasswordReset{})
+}
+
 type PasswordReset struct {
 	gorm.Model
 
diff --git a/db/models/pin.go b/db/models/pin.go
index 1904cc6..b1b73a8 100644
--- a/db/models/pin.go
+++ b/db/models/pin.go
@@ -2,6 +2,10 @@ package models
 
 import "gorm.io/gorm"
 
+func init() {
+	registerModel(&Pin{})
+}
+
 type Pin struct {
 	gorm.Model
 	UploadID uint
diff --git a/db/models/public_key.go b/db/models/public_key.go
index f12a764..28b83aa 100644
--- a/db/models/public_key.go
+++ b/db/models/public_key.go
@@ -2,6 +2,10 @@ package models
 
 import "gorm.io/gorm"
 
+func init() {
+	registerModel(&PublicKey{})
+}
+
 type PublicKey struct {
 	gorm.Model
 	UserID uint
diff --git a/db/models/s3_upload.go b/db/models/s3_upload.go
index 348d1dc..2e5233f 100644
--- a/db/models/s3_upload.go
+++ b/db/models/s3_upload.go
@@ -2,6 +2,10 @@ package models
 
 import "gorm.io/gorm"
 
+func init() {
+	registerModel(&S3Upload{})
+}
+
 type S3Upload struct {
 	gorm.Model
 	UploadID string `gorm:"unique;not null"`
diff --git a/db/models/s5_challenge.go b/db/models/s5_challenge.go
index e317aaf..856b738 100644
--- a/db/models/s5_challenge.go
+++ b/db/models/s5_challenge.go
@@ -2,6 +2,11 @@ package models
 
 import "gorm.io/gorm"
 
+func init() {
+	registerModel(&S5Challenge{})
+
+}
+
 type S5Challenge struct {
 	gorm.Model
 	Challenge string
diff --git a/db/models/tus_lock.go b/db/models/tus_lock.go
index b6f1cef..0c90d9e 100644
--- a/db/models/tus_lock.go
+++ b/db/models/tus_lock.go
@@ -3,11 +3,17 @@ package models
 import (
 	"context"
 	"errors"
+	"time"
+
 	"gorm.io/gorm"
 	"gorm.io/gorm/clause"
-	"time"
 )
 
+func init() {
+	registerModel(&TusLock{})
+
+}
+
 var (
 	ErrTusLockBusy = errors.New("lock is currently held by another process")
 )
diff --git a/db/models/tus_upload.go b/db/models/tus_upload.go
index 56456ef..2114120 100644
--- a/db/models/tus_upload.go
+++ b/db/models/tus_upload.go
@@ -2,6 +2,11 @@ package models
 
 import "gorm.io/gorm"
 
+func init() {
+	registerModel(&TusUpload{})
+
+}
+
 type TusUpload struct {
 	gorm.Model
 	Hash       []byte `gorm:"type:binary(32);uniqueIndex:idx_hash_deleted"`
diff --git a/db/models/upload.go b/db/models/upload.go
index 5b78310..b272b89 100644
--- a/db/models/upload.go
+++ b/db/models/upload.go
@@ -2,6 +2,10 @@ package models
 
 import "gorm.io/gorm"
 
+func init() {
+	registerModel(&Upload{})
+}
+
 type Upload struct {
 	gorm.Model
 	UserID     uint
diff --git a/db/models/user.go b/db/models/user.go
index cf31d1e..529d708 100644
--- a/db/models/user.go
+++ b/db/models/user.go
@@ -8,6 +8,10 @@ import (
 	"gorm.io/gorm"
 )
 
+func init() {
+	registerModel(&User{})
+}
+
 type User struct {
 	gorm.Model
 	FirstName          string