github.com/kunlun-qilian/sqlx/v3@v3.0.0/database.go (about)

     1  package sqlx
     2  
     3  import (
     4  	"database/sql"
     5  	"database/sql/driver"
     6  	"fmt"
     7  	"os"
     8  
     9  	"github.com/kunlun-qilian/sqlx/v3/builder"
    10  )
    11  
    12  func NewFeatureDatabase(name string) *Database {
    13  	if projectFeature, exists := os.LookupEnv("PROJECT_FEATURE"); exists && projectFeature != "" {
    14  		name = name + "__" + projectFeature
    15  	}
    16  	return NewDatabase(name)
    17  }
    18  
    19  func NewDatabase(name string) *Database {
    20  	return &Database{
    21  		Name:   name,
    22  		Tables: builder.Tables{},
    23  	}
    24  }
    25  
    26  type Database struct {
    27  	Name   string
    28  	Schema string
    29  	Tables builder.Tables
    30  }
    31  
    32  func (database Database) WithSchema(schema string) *Database {
    33  	database.Schema = schema
    34  
    35  	tables := builder.Tables{}
    36  
    37  	database.Tables.Range(func(tab *builder.Table, idx int) {
    38  		tables.Add(tab.WithSchema(database.Schema))
    39  	})
    40  
    41  	database.Tables = tables
    42  
    43  	return &database
    44  }
    45  
    46  type DBNameBinder interface {
    47  	WithDBName(dbName string) driver.Connector
    48  }
    49  
    50  func (database *Database) OpenDB(connector driver.Connector) *DB {
    51  	if dbNameBinder, ok := connector.(DBNameBinder); ok {
    52  		connector = dbNameBinder.WithDBName(database.Name)
    53  	}
    54  	dialect, ok := connector.(builder.Dialect)
    55  	if !ok {
    56  		panic(fmt.Errorf("connector should implement builder.Dialect"))
    57  	}
    58  	return &DB{
    59  		Database:    database,
    60  		dialect:     dialect,
    61  		SqlExecutor: sql.OpenDB(connector),
    62  	}
    63  }
    64  
    65  func (database *Database) AddTable(table *builder.Table) {
    66  	database.Tables.Add(table)
    67  }
    68  
    69  func (database *Database) Register(model builder.Model) *builder.Table {
    70  	table := builder.TableFromModel(model)
    71  	table.Schema = database.Schema
    72  	database.AddTable(table)
    73  	return table
    74  }
    75  
    76  func (database *Database) Table(tableName string) *builder.Table {
    77  	return database.Tables.Table(tableName)
    78  }
    79  
    80  func (database *Database) T(model builder.Model) *builder.Table {
    81  	if td, ok := model.(builder.TableDefinition); ok {
    82  		return td.T()
    83  	}
    84  
    85  	if t, ok := model.(*builder.Table); ok {
    86  		return t
    87  	}
    88  
    89  	return database.Table(model.TableName())
    90  }