core/database/database.go

79 lines
1.6 KiB
Go

package database
import (
"context"
"crypto/tls"
"database/sql"
"fmt"
"os"
"system-trace/core/environment"
"time"
"github.com/gofiber/fiber/v2"
"github.com/uptrace/bun"
"github.com/uptrace/bun/dialect/pgdialect"
"github.com/uptrace/bun/driver/pgdriver"
"github.com/uptrace/bun/extra/bundebug"
)
var PG *bun.DB
func Connect() {
pgconn := pgdriver.NewConnector(
pgdriver.WithNetwork("tcp"),
pgdriver.WithAddr(fmt.Sprintf("%s:%s", os.Getenv("DB_HOST"), os.Getenv("DB_PORT"))),
pgdriver.WithTLSConfig(&tls.Config{InsecureSkipVerify: true}),
pgdriver.WithInsecure(true),
pgdriver.WithUser(os.Getenv("DB_USER")),
pgdriver.WithPassword(os.Getenv("DB_PASS")),
pgdriver.WithDatabase(os.Getenv("DB_NAME")),
pgdriver.WithApplicationName("st-controller"),
pgdriver.WithTimeout(5*time.Second),
pgdriver.WithDialTimeout(5*time.Second),
pgdriver.WithReadTimeout(5*time.Second),
pgdriver.WithWriteTimeout(5*time.Second),
)
sqldb := sql.OpenDB(pgconn)
db := bun.NewDB(sqldb, pgdialect.New())
if environment.IsDebug() {
db.AddQueryHook(bundebug.NewQueryHook(
bundebug.WithVerbose(true),
))
}
if !fiber.IsChild() {
err := createSchema(db)
if err != nil {
panic(err)
}
}
PG = db
}
func createSchema(db *bun.DB) error {
models := []interface{}{
(*Group)(nil),
(*GroupPermission)(nil),
(*User)(nil),
(*Session)(nil),
(*Server)(nil),
(*AuthToken)(nil),
}
ctx := context.Background()
for _, model := range models {
_, err := db.
NewCreateTable().
Model(model).
IfNotExists().
Exec(ctx)
if err != nil {
return err
}
}
return nil
}