github.com/wfusion/gofusion@v1.1.14/db/db.go (about)

     1  package db
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/pkg/errors"
     9  	"gorm.io/gorm"
    10  
    11  	"github.com/wfusion/gofusion/common/infra/drivers/orm"
    12  	"github.com/wfusion/gofusion/common/utils"
    13  	"github.com/wfusion/gofusion/db/plugins"
    14  )
    15  
    16  const (
    17  	defaultSlowThreshold = 200 * time.Millisecond
    18  )
    19  
    20  var (
    21  	rwlock       = new(sync.RWMutex)
    22  	appInstances map[string]map[string]*Instance
    23  )
    24  
    25  type Instance struct {
    26  	name                 string
    27  	db                   *orm.DB
    28  	tableShardingPlugins map[string]plugins.TableSharding
    29  }
    30  
    31  func (d *Instance) GetProxy() *gorm.DB {
    32  	return d.db.GetProxy()
    33  }
    34  
    35  type DB struct {
    36  	*orm.DB
    37  	Name                 string
    38  	tableShardingPlugins map[string]plugins.TableSharding
    39  }
    40  
    41  type useOption struct {
    42  	appName string
    43  }
    44  
    45  func AppName(name string) utils.OptionFunc[useOption] {
    46  	return func(o *useOption) {
    47  		o.appName = name
    48  	}
    49  }
    50  
    51  func Use(ctx context.Context, name string, opts ...utils.OptionExtender) *DB {
    52  	opt := utils.ApplyOptions[useOption](opts...)
    53  
    54  	rwlock.RLock()
    55  	defer rwlock.RUnlock()
    56  	instances, ok := appInstances[opt.appName]
    57  	if !ok {
    58  		panic(errors.Errorf("db instance not found for app: %s", opt.appName))
    59  	}
    60  	instance, ok := instances[name]
    61  	if !ok {
    62  		panic(errors.Errorf("db instance not found for name: %s", name))
    63  	}
    64  
    65  	return &DB{DB: instance.db.WithContext(ctx), Name: name, tableShardingPlugins: instance.tableShardingPlugins}
    66  }