GORM

model.go

package sqlite

import (
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/jinzhu/gorm"
	"github.com/maxwellforest/onefill/db/gorm/runner/pg"
	"github.com/maxwellforest/onefill/db/gorm/runner/sqlite"
)

type RunRecord struct {
	gorm.Model
	Log    string
	SiteId string // TODO Make SiteID - before deploy
	JobID  string
}

var (
	pg_cnf           *pg.Config
	migrated         bool
	connection_count = 2
	connections      = make(chan *gorm.DB, connection_count)
	initMutex        = &sync.Mutex{}
)

type ORM struct {
	DB *gorm.DB
}

func (o *ORM) Close() error {
	connections <- o.DB
	o.DB = nil
	return nil
}

func (o *ORM) GetRunRecordsForSite(site_id *string) ([]*RunRecord, error) {
	var rr []*RunRecord
	if o.DB == nil {
		return nil, fmt.Errorf("DB connection not available")
	}
	o.DB.Where("site_id = ?", *site_id).Find(&rr)
	return rr, nil
}

func (o *ORM) GetTotalRunRecords() (int, error) {
	var count int
	if o.DB == nil {
		if o.DB == nil {
			return 0, fmt.Errorf("DB connection not available")
		}
	}
	o.DB.Table("run_records").Count(&count)
	return count, nil
}

func (o *ORM) GetRunRecord(run_record_id int) (*RunRecord, error) {
	var rr RunRecord
	if o.DB == nil {
		if o.DB == nil {
			return nil, fmt.Errorf("DB connection not available")
		}
	}

	o.DB.First(&rr, run_record_id)

	return &rr, nil
}

func (o *ORM) GetRunRecords(page int, limit int) ([]*RunRecord, error) {
	var rr []*RunRecord
	if o.DB == nil {
		if o.DB == nil {
			return nil, fmt.Errorf("DB connection not available")
		}
	}
	offset := (page - 1) * limit
	o.DB.Limit(limit).Offset(offset).Find(&rr)
	return rr, nil
}

func (o *ORM) CreateRunRecordForSite(log *string, site_id *string, job_id *string) (*RunRecord, error) {
	if o.DB == nil {
		return nil, fmt.Errorf("DB connection not available")
	}

	rr := &RunRecord{Log: *log, SiteId: *site_id, JobID: *job_id}
	o.DB.Create(rr)
	return rr, nil
}

func InitPG(cnf pg.Config) error {
	db, err := pg.Open(&cnf)
	if err != nil {
		return err
	}
	defer db.Close()

	pg_cnf = &cnf
	return nil
}

func initDB() (*gorm.DB, error) {
	var (
		db  *gorm.DB
		err error
	)
	tries := 5
	attempts := 0

	initMutex.Lock()
	defer initMutex.Unlock()

	for db == nil {
		if pg_cnf != nil {
			db, err = pg.Open(pg_cnf)
			if err != nil {
				return nil, err
			}
		} else {
			db, err = sqlite.Open(&sqlite.PathToDB)
		}
		if err != nil && attempts > tries {
			return nil, fmt.Errorf("Exceeded retries. Error: %s", err.Error())
		}
		attempts += 1
		time.Sleep(time.Second)
	}

	if db != nil && !migrated {
		db.AutoMigrate(&RunRecord{})
		migrated = true
	}

	return db, nil
}

func CleanUpDB(pathToDB string) error {
	err := os.Remove(pathToDB)
	return err
}

func GetORMOrFail() (*ORM, error) {
	db := <-connections
	orm := &ORM{db}

	return orm, nil
}

func InitConnections() {
	for i := connection_count; i > 0; i-- {
		db, err := initDB()
		if err != nil {
			glog.Fatalf("Failed to init connections: %s", err.Error())
		}
		connections <- db
	}
}

model_test.go

package sqlite

import (
	"testing"

	"github.com/maxwellforest/onefill/db/gorm/runner/sqlite"
	"github.com/maxwellforest/onefill/test"
)

func TestCreateRunRecordsForSite(t *testing.T) {
	orm, _ := GetORMOrFail()
	defer orm.Close()

	testRR := RunRecord{
		Log:    "This is a test",
		SiteId: "99",
		JobID:  "Yolo",
	}

	rr, err := orm.CreateRunRecordForSite(&testRR.Log, &testRR.SiteId, &testRR.JobID)
	if err != nil {
		t.Errorf("Failed to saved record: %v", err)
	}
	if rr.ID == 0 {
		t.Errorf("Failed to save record")
	}
	if orm.DB.NewRecord(*rr) == true {
		t.Errorf("Failed to save record")
	}
}

func TestGetRunRecordsForSite(t *testing.T) {
	orm, _ := GetORMOrFail()
	defer orm.Close()

	site_id := "99"
	log1 := "log"
	job_id_1 := "100"
	_, _ = orm.CreateRunRecordForSite(&log1, &site_id, &job_id_1)
	job_id_2 := "100"
	log2 := "log"
	_, _ = orm.CreateRunRecordForSite(&log2, &site_id, &job_id_2)

	rr, err := orm.GetRunRecordsForSite(&site_id)
	if err != nil {
		t.Errorf("Didn't expect error, got %v\n", err)
	}
	if len(rr) != (2 + 1) {
		t.Errorf("Expected 2 records, got %d\n", len(rr))
	}
}

func TestCleanUpDB(t *testing.T) {
	_, err := initDB()
	if err != nil {
		t.Errorf("DB Failed to initialize: %s", err.Error())
	}

	exists, err := test.FileExists(sqlite.PathToDB)
	if !exists {
		t.Errorf("Expected file exist, it apparently doesn't")
	}

	CleanUpDB(sqlite.PathToDB)
	exists, err = test.FileExists(sqlite.PathToDB)
	if exists {
		t.Errorf("Expected file not to exist, it apparently does.")
	}
}