diamond-orm/diamond.go

131 lines
2.7 KiB
Go

package orm
import (
"context"
"fmt"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"log/slog"
"time"
)
const LevelQuery = slog.Level(-6)
const defaultKey = "default"
type Engine struct {
modelMap *internalModelMap
conn *pgxpool.Pool
m2mSeen map[string]bool
dryRun bool
cfg *pgxpool.Config
ctx context.Context
logger *slog.Logger
levelVar *slog.LevelVar
connStr string
}
func (e *Engine) Models(v ...any) {
emm := makeModelMap(v...)
for k := range emm.Map {
if _, ok := e.modelMap.Map[k]; !ok {
e.modelMap.Mux.Lock()
e.modelMap.Map[k] = emm.Map[k]
e.modelMap.Mux.Unlock()
}
}
}
func (e *Engine) Model(val any) *Query {
qq := &Query{
engine: e,
ctx: context.Background(),
wheres: make(map[string][]any),
orders: make([]string, 0),
populationTree: make(map[string]any),
joins: make([]string, 0),
}
return qq.setModel(val)
}
func (e *Engine) QueryRaw(sql string, args ...any) (pgx.Rows, error) {
return e.conn.Query(e.ctx, sql, args...)
}
func (e *Engine) Migrate() error {
failedMigrations := make(map[string]*Model)
var err error
for mk, m := range e.modelMap.Map {
err = m.migrate(e)
if err != nil {
failedMigrations[mk] = m
}
}
for len(failedMigrations) > 0 {
e.m2mSeen = make(map[string]bool)
for mk, m := range failedMigrations {
err = m.migrate(e)
if err == nil {
delete(failedMigrations, mk)
}
}
}
return err
}
func (e *Engine) MigrateDropping() error {
for _, m := range e.modelMap.Map {
sql := fmt.Sprintf("DROP TABLE IF EXISTS %s CASCADE;", m.TableName)
if _, err := e.conn.Exec(e.ctx, sql); err != nil {
return err
}
for _, r := range m.Relationships {
if r.m2mIsh() || r.Type == ManyToMany {
jsql := fmt.Sprintf("DROP TABLE IF EXISTS %s CASCADE;", r.ComputeJoinTable())
if _, err := e.conn.Exec(e.ctx, jsql); err != nil {
return err
}
}
}
}
return e.Migrate()
}
func (e *Engine) Disconnect() {
e.conn.Close()
}
func Open(connString string) (*Engine, error) {
e := &Engine{
modelMap: &internalModelMap{
Map: make(map[string]*Model),
},
m2mSeen: make(map[string]bool),
dryRun: connString == "",
ctx: context.Background(),
}
if connString != "" {
engines.Mux.Lock()
if len(engines.Engines) == 0 || engines.Engines[defaultKey] == nil {
engines.Engines[defaultKey] = e
} else {
engines.Engines[connString] = e
}
e.connStr = ""
engines.Mux.Unlock()
var err error
e.cfg, err = pgxpool.ParseConfig(connString)
e.cfg.MinConns = 5
e.cfg.MaxConns = 10
e.cfg.MaxConnIdleTime = time.Minute * 2
if err != nil {
return nil, err
}
e.conn, err = pgxpool.NewWithConfig(e.ctx, e.cfg)
if err != nil {
return nil, err
}
}
return e, nil
}