139 lines
5.6 KiB
Go
139 lines
5.6 KiB
Go
package data
|
||
|
||
import (
|
||
"context"
|
||
"database/sql"
|
||
"time"
|
||
|
||
"rag/file-system/internal/conf"
|
||
|
||
"github.com/go-kratos/kratos/v2/log"
|
||
"github.com/google/wire"
|
||
"gorm.io/driver/postgres"
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
// Data holds the GORM database connection and provides transaction support.
|
||
type Data struct {
|
||
db *gorm.DB
|
||
log *log.Helper
|
||
}
|
||
|
||
// contextTxKey is the context key for storing the current transaction.
|
||
type contextTxKey struct{}
|
||
|
||
// NewData creates a new Data instance with GORM connected to PostgreSQL.
|
||
func NewData(c *conf.Data, logger log.Logger) (*Data, func(), error) {
|
||
helper := log.NewHelper(logger)
|
||
db, err := gorm.Open(postgres.Open(c.Database.Source), &gorm.Config{})
|
||
if err != nil {
|
||
return nil, nil, err
|
||
}
|
||
|
||
sqlDB, err := db.DB()
|
||
if err != nil {
|
||
return nil, nil, err
|
||
}
|
||
sqlDB.SetMaxOpenConns(25)
|
||
sqlDB.SetMaxIdleConns(5)
|
||
|
||
if err := db.AutoMigrate(&FolderPO{}, &FileMetaPO{}, &ShareLinkPO{}); err != nil {
|
||
return nil, nil, err
|
||
}
|
||
|
||
helper.Info("connected to PostgreSQL via GORM")
|
||
cleanup := func() {
|
||
sqlDB.Close()
|
||
}
|
||
return &Data{db: db, log: helper}, cleanup, nil
|
||
}
|
||
|
||
// SqlDB returns the underlying *sql.DB for use by Watermill.
|
||
func (d *Data) SqlDB() (*sql.DB, error) {
|
||
return d.db.DB()
|
||
}
|
||
|
||
// DB returns the *gorm.DB for the given context. If a transaction is active
|
||
// in the context, it returns the transaction DB; otherwise the global DB.
|
||
func (d *Data) DB(ctx context.Context) *gorm.DB {
|
||
tx, ok := ctx.Value(contextTxKey{}).(*gorm.DB)
|
||
if ok {
|
||
return tx
|
||
}
|
||
return d.db.WithContext(ctx)
|
||
}
|
||
|
||
// Transaction is the interface for executing operations within a database transaction.
|
||
type Transaction interface {
|
||
InTx(ctx context.Context, fn func(ctx context.Context) error) error
|
||
}
|
||
|
||
// InTx executes fn inside a database transaction.
|
||
func (d *Data) InTx(ctx context.Context, fn func(ctx context.Context) error) error {
|
||
return d.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||
ctx = context.WithValue(ctx, contextTxKey{}, tx)
|
||
return fn(ctx)
|
||
})
|
||
}
|
||
|
||
// ProviderSet is the Wire provider set for the data layer.
|
||
var ProviderSet = wire.NewSet(NewData, NewFileRepo, NewFolderRepo, NewFileMetaRepo, NewShareRepo)
|
||
|
||
// --- GORM Models (Persistence Objects) ---
|
||
|
||
// FolderPO maps to the "folders" table.
|
||
type FolderPO struct {
|
||
ID string `gorm:"primaryKey;type:uuid;default:gen_random_uuid();comment:主键ID"`
|
||
ParentID *string `gorm:"type:uuid;index:idx_folders_parent;comment:父文件夹ID"`
|
||
Name string `gorm:"type:varchar(255);not null;comment:文件夹名称"`
|
||
OwnerID string `gorm:"type:varchar(36);not null;index:idx_folders_owner;comment:所有者ID"`
|
||
CreatedBy string `gorm:"type:varchar(36);not null;default:'';comment:创建人ID"`
|
||
CreatedAt time.Time `gorm:"autoCreateTime;comment:创建时间"`
|
||
UpdatedBy string `gorm:"type:varchar(36);not null;default:'';comment:更新人ID"`
|
||
UpdatedAt time.Time `gorm:"autoUpdateTime;comment:更新时间"`
|
||
IsDeleted bool `gorm:"not null;default:false;comment:是否软删除"`
|
||
OperatorIP string `gorm:"type:varchar(500);comment:操作人IP地址"`
|
||
}
|
||
|
||
func (FolderPO) TableName() string { return "folders" }
|
||
|
||
// FileMetaPO maps to the "files" table.
|
||
type FileMetaPO struct {
|
||
ID string `gorm:"primaryKey;type:uuid;default:gen_random_uuid();comment:主键ID"`
|
||
FolderID string `gorm:"type:uuid;index:idx_files_folder;comment:所属文件夹ID"`
|
||
Name string `gorm:"type:varchar(255);not null;comment:文件名称"`
|
||
S3Key string `gorm:"type:varchar(512);not null;index:idx_files_s3_key;comment:S3对象键"`
|
||
S3Bucket string `gorm:"type:varchar(255);not null;comment:S3存储桶名称"`
|
||
Size int64 `gorm:"default:0;comment:文件大小（字节）"`
|
||
ContentType string `gorm:"type:varchar(255);default:'application/octet-stream';comment:文件MIME类型"`
|
||
OwnerID string `gorm:"type:varchar(36);not null;index:idx_files_owner;comment:所有者ID"`
|
||
CreatedBy string `gorm:"type:varchar(36);not null;default:'';comment:创建人ID"`
|
||
CreatedAt time.Time `gorm:"autoCreateTime;comment:创建时间"`
|
||
UpdatedBy string `gorm:"type:varchar(36);not null;default:'';comment:更新人ID"`
|
||
UpdatedAt time.Time `gorm:"autoUpdateTime;comment:更新时间"`
|
||
IsDeleted bool `gorm:"not null;default:false;comment:是否软删除"`
|
||
OperatorIP string `gorm:"type:varchar(500);comment:操作人IP地址"`
|
||
}
|
||
|
||
func (FileMetaPO) TableName() string { return "files" }
|
||
|
||
// ShareLinkPO maps to the "share_links" table.
|
||
type ShareLinkPO struct {
|
||
ID string `gorm:"primaryKey;type:uuid;default:gen_random_uuid();comment:主键ID"`
|
||
ResourceType string `gorm:"type:varchar(10);not null;comment:资源类型（folder/file）"`
|
||
ResourceID string `gorm:"type:uuid;not null;comment:资源ID"`
|
||
Token string `gorm:"type:varchar(32);not null;uniqueIndex:idx_share_token;comment:分享令牌"`
|
||
Password *string `gorm:"type:varchar(255);comment:访问密码"`
|
||
ExpiresAt *time.Time `gorm:"type:timestamptz;comment:过期时间"`
|
||
DownloadCount int `gorm:"default:0;comment:下载次数"`
|
||
MaxDownloads *int `gorm:"comment:最大下载次数"`
|
||
CreatedBy string `gorm:"type:varchar(36);not null;comment:创建人ID"`
|
||
CreatedAt time.Time `gorm:"autoCreateTime;comment:创建时间"`
|
||
UpdatedBy string `gorm:"type:varchar(36);not null;default:'';comment:更新人ID"`
|
||
UpdatedAt time.Time `gorm:"autoUpdateTime;comment:更新时间"`
|
||
IsDeleted bool `gorm:"not null;default:false;comment:是否软删除"`
|
||
OperatorIP string `gorm:"type:varchar(500);comment:操作人IP地址"`
|
||
}
|
||
|
||
func (ShareLinkPO) TableName() string { return "share_links" }
|