dbmanager.go

// Use Case:
//
// Say you have a bunch of decoupled components which access the same set of databases.
// Since database/sql does connection pooling it makes sense for the components to share instances of *sql.DB
// This package lets you do that with minimal code changes.
package dbmanager

import (
	"database/sql"
	"time"
	"sync"
)

type Factory struct {
	connections map[string]*DB
	sync.Mutex
}

func NewFactory() *Factory {
	return &Factory{
		connections: make(map[string]*DB),
	}
}

func (f *Factory) open(driverName, dataSourceName string) (*DB, error) {
	db, err := sql.Open(driverName, dataSourceName)
	if err != nil {
		return nil, err
	}
	return &DB{
		key:     dataSourceName,
		factory: f,
		DB:      db,
	}, nil
}

// Open returns a DB connection.
func (f *Factory) Open(driverName, dataSourceName string) (*DB, error) {
	f.Lock()
	defer f.Unlock()
	db, ok := f.connections[dataSourceName]
	if !ok {
		var err error
		db, err = f.open(driverName, dataSourceName)
		if err != nil {
			return nil, err
		}
		f.connections[dataSourceName] = db
	}
	db.leases++
	return db, nil
}

// release is called by the DB's Close() method
func (f *Factory) release(key string) {
	f.Lock()
	defer f.Unlock()
	db, ok := f.connections[key]
	if !ok {
		return
	}
	db.leases--
	if db.leases == 0 {
		delete(f.connections, key)
		db.DB.Close()
	}
}

type DB struct {
	leases  int
	key     string
	factory *Factory
	*sql.DB
}

func (db *DB) Close() {
	db.factory.release(db.key)
}