mydb/adapter/mongo/database.go

// Package mongo wraps the gopkg.in/mgo.v2 MongoDB driver. See
// https://github.com/upper/db/adapter/mongo for documentation, particularities and usage
// examples.
package mongo

import (
	"context"
	"database/sql"
	"strings"
	"sync"
	"time"

	"git.hexq.cn/tiglog/mydb"
	mgo "gopkg.in/mgo.v2"
)

// Adapter holds the name of the mongodb adapter.
const Adapter = `mongo`

var connTimeout = time.Second * 5

// Source represents a MongoDB database.
type Source struct {
	mydb.Settings

	ctx context.Context

	name          string
	connURL       mydb.ConnectionURL
	session       *mgo.Session
	database      *mgo.Database
	version       []int
	collections   map[string]*Collection
	collectionsMu sync.Mutex
}

type mongoAdapter struct {
}

func (mongoAdapter) Open(dsn mydb.ConnectionURL) (mydb.Session, error) {
	return Open(dsn)
}

func init() {
	mydb.RegisterAdapter(Adapter, mydb.Adapter(&mongoAdapter{}))
}

// Open stablishes a new connection to a SQL server.
func Open(settings mydb.ConnectionURL) (mydb.Session, error) {
	d := &Source{Settings: mydb.NewSettings(), ctx: context.Background()}
	if err := d.Open(settings); err != nil {
		return nil, err
	}
	return d, nil
}

func (s *Source) TxContext(context.Context, func(tx mydb.Session) error, *sql.TxOptions) error {
	return mydb.ErrNotSupportedByAdapter
}

func (s *Source) Tx(func(mydb.Session) error) error {
	return mydb.ErrNotSupportedByAdapter
}

func (s *Source) SQL() mydb.SQL {
	// Not supported
	panic("sql builder is not supported by mongodb")
}

func (s *Source) ConnectionURL() mydb.ConnectionURL {
	return s.connURL
}

// SetConnMaxLifetime is not supported.
func (s *Source) SetConnMaxLifetime(time.Duration) {
	s.Settings.SetConnMaxLifetime(time.Duration(0))
}

// SetMaxIdleConns is not supported.
func (s *Source) SetMaxIdleConns(int) {
	s.Settings.SetMaxIdleConns(0)
}

// SetMaxOpenConns is not supported.
func (s *Source) SetMaxOpenConns(int) {
	s.Settings.SetMaxOpenConns(0)
}

// Name returns the name of the database.
func (s *Source) Name() string {
	return s.name
}

// Open attempts to connect to the database.
func (s *Source) Open(connURL mydb.ConnectionURL) error {
	s.connURL = connURL
	return s.open()
}

// Clone returns a cloned mydb.Session session.
func (s *Source) Clone() (mydb.Session, error) {
	newSession := s.session.Copy()
	clone := &Source{
		Settings: mydb.NewSettings(),

		name:        s.name,
		connURL:     s.connURL,
		session:     newSession,
		database:    newSession.DB(s.database.Name),
		version:     s.version,
		collections: map[string]*Collection{},
	}
	return clone, nil
}

// Ping checks whether a connection to the database is still alive by pinging
// it, establishing a connection if necessary.
func (s *Source) Ping() error {
	return s.session.Ping()
}

func (s *Source) Reset() {
	s.collectionsMu.Lock()
	defer s.collectionsMu.Unlock()
	s.collections = make(map[string]*Collection)
}

// Driver returns the underlying *mgo.Session instance.
func (s *Source) Driver() interface{} {
	return s.session
}

func (s *Source) open() error {
	var err error

	if s.session, err = mgo.DialWithTimeout(s.connURL.String(), connTimeout); err != nil {
		return err
	}

	s.collections = map[string]*Collection{}
	s.database = s.session.DB("")

	return nil
}

// Close terminates the current database session.
func (s *Source) Close() error {
	if s.session != nil {
		s.session.Close()
	}
	return nil
}

// Collections returns a list of non-system tables from the database.
func (s *Source) Collections() (cols []mydb.Collection, err error) {
	var rawcols []string
	var col string

	if rawcols, err = s.database.CollectionNames(); err != nil {
		return nil, err
	}

	cols = make([]mydb.Collection, 0, len(rawcols))

	for _, col = range rawcols {
		if !strings.HasPrefix(col, "system.") {
			cols = append(cols, s.Collection(col))
		}
	}

	return cols, nil
}

func (s *Source) Delete(mydb.Record) error {
	return mydb.ErrNotImplemented
}

func (s *Source) Get(mydb.Record, interface{}) error {
	return mydb.ErrNotImplemented
}

func (s *Source) Save(mydb.Record) error {
	return mydb.ErrNotImplemented
}

func (s *Source) Context() context.Context {
	return s.ctx
}

func (s *Source) WithContext(ctx context.Context) mydb.Session {
	return &Source{
		ctx:      ctx,
		Settings: s.Settings,
		name:     s.name,
		connURL:  s.connURL,
		session:  s.session,
		database: s.database,
		version:  s.version,
	}
}

// Collection returns a collection by name.
func (s *Source) Collection(name string) mydb.Collection {
	s.collectionsMu.Lock()
	defer s.collectionsMu.Unlock()

	var col *Collection
	var ok bool

	if col, ok = s.collections[name]; !ok {
		col = &Collection{
			parent:     s,
			collection: s.database.C(name),
		}
		s.collections[name] = col
	}

	return col
}

func (s *Source) versionAtLeast(version ...int) bool {
	// only fetch this once - it makes a db call
	if len(s.version) == 0 {
		buildInfo, err := s.database.Session.BuildInfo()
		if err != nil {
			return false
		}
		s.version = buildInfo.VersionArray
	}

	// Check major version first
	if s.version[0] > version[0] {
		return true
	}

	for i := range version {
		if i == len(s.version) {
			return false
		}
		if s.version[i] < version[i] {
			return false
		}
	}
	return true
}