VILLASweb-backend-go/common/database.go

package common

import (
	"flag"
	"fmt"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/postgres"
	"log"
)

var DB_HOST string
var DB_NAME string
var DB_DUMMY string
var DB_SSLMODE string

var DBpool *gorm.DB

// Initialize input command line flags
func init() {
	flag.StringVar(&DB_HOST, "dbhost", "/tmp", "Host of the PostgreSQL database (default is /tmp)")
	flag.StringVar(&DB_NAME, "dbname", "villasdb", "Name of the database to use (default is villasdb)")
	flag.StringVar(&DB_DUMMY, "dbdummy", "testvillasdb", "Name of the test database to use (default is testvillasdb)")
	flag.StringVar(&DB_SSLMODE, "dbsslmode", "disable", "SSL mode of DB (default is disable)") // TODO: change default for production
	flag.Parse()
	fmt.Println("DB_HOST has value ", DB_HOST)
	fmt.Println("DB_NAME has value ", DB_NAME)
	fmt.Println("DB_DUMMY has value ", DB_DUMMY)
	fmt.Println("DB_SSLMODE has value ", DB_SSLMODE)
}

// Initialize connection to the database
func InitDB() *gorm.DB {
	dbinfo := fmt.Sprintf("host=%s sslmode=%s dbname=%s",
		DB_HOST, DB_SSLMODE, DB_NAME)
	db, err := gorm.Open("postgres", dbinfo)
	checkErr(err)
	DBpool = db
	return db
}

// Connection pool to already opened DB
func GetDB() *gorm.DB {
	return DBpool
}

// Verify that the database connection is alive
func VerifyConnection(db *gorm.DB) error {
	return db.DB().Ping()
}

// Drop all the tables of the database
// TODO: Remove that function from the codebase and substitute the body
// to the Dummy*() where it is called
func DropTables(db *gorm.DB) {
	db.DropTableIfExists(&Simulator{})
	db.DropTableIfExists(&Sample{})
	db.DropTableIfExists(&Model{})
	db.DropTableIfExists(&File{})
	db.DropTableIfExists(&Project{})
	db.DropTableIfExists(&Simulation{})
	db.DropTableIfExists(&User{})
	db.DropTableIfExists(&Visualization{})
	db.DropTableIfExists(&Widget{})
}

// AutoMigrate the models
func MigrateModels(db *gorm.DB) {
	db.AutoMigrate(&Simulator{})
	db.AutoMigrate(&Sample{})
	db.AutoMigrate(&Model{})
	db.AutoMigrate(&File{})
	db.AutoMigrate(&Project{})
	db.AutoMigrate(&Simulation{})
	db.AutoMigrate(&User{})
	db.AutoMigrate(&Visualization{})
	db.AutoMigrate(&Widget{})
}

// Start a dummy database for testing
func DummyInitDB() *gorm.DB {

	dbinfo := fmt.Sprintf("host=%s sslmode=%s dbname=%s",
		DB_HOST, DB_SSLMODE, DB_DUMMY)
	test_db, err := gorm.Open("postgres", dbinfo)
	checkErr(err)
	DBpool = test_db
	// drop tables from previous tests
	DropTables(test_db)

	return test_db
}

// Migrates models and populates them with data
func DummyPopulateDB(test_db *gorm.DB) {

	MigrateModels(test_db)

	// Create two entries of each model

	simr_A := Simulator{UUID: "1", Host: "Host_A"}
	simr_B := Simulator{UUID: "2", Host: "Host_B"}
	checkErr(test_db.Create(&simr_A).Error)
	checkErr(test_db.Create(&simr_B).Error)

	outSmp_A := Sample{Name: "outSample_A"}
	outSmp_B := Sample{Name: "outSample_B"}
	inSmp_A := Sample{Name: "inSample_A"}
	inSmp_B := Sample{Name: "inSample_B"}
	checkErr(test_db.Create(&outSmp_A).Error)
	checkErr(test_db.Create(&outSmp_B).Error)
	checkErr(test_db.Create(&inSmp_A).Error)
	checkErr(test_db.Create(&inSmp_B).Error)

	mo_A := Model{Name: "Model_A"}
	mo_B := Model{Name: "Model_B"}
	checkErr(test_db.Create(&mo_A).Error)
	checkErr(test_db.Create(&mo_B).Error)

	file_A := File{Name: "File_A"}
	file_B := File{Name: "File_B"}
	checkErr(test_db.Create(&file_A).Error)
	checkErr(test_db.Create(&file_B).Error)

	proj_A := Project{Name: "Project_A"}
	proj_B := Project{Name: "Project_B"}
	checkErr(test_db.Create(&proj_A).Error)
	checkErr(test_db.Create(&proj_B).Error)

	simn_A := Simulation{Name: "Simulation_A"}
	simn_B := Simulation{Name: "Simulation_B"}
	checkErr(test_db.Create(&simn_A).Error)
	checkErr(test_db.Create(&simn_B).Error)

	usr_A := User{Username: "User_A"}
	usr_B := User{Username: "User_B"}
	checkErr(test_db.Create(&usr_A).Error)
	checkErr(test_db.Create(&usr_B).Error)

	vis_A := Visualization{Name: "Visualization_A"}
	vis_B := Visualization{Name: "Visualization_B"}
	checkErr(test_db.Create(&vis_A).Error)
	checkErr(test_db.Create(&vis_B).Error)

	widg_A := Widget{Name: "Widget_A"}
	widg_B := Widget{Name: "Widget_B"}
	checkErr(test_db.Create(&widg_A).Error)
	checkErr(test_db.Create(&widg_B).Error)

	// Associations betweend models
	// For `belongs to` use the model with id=1
	// For `has many` use the models with id=1 and id=2

	// Project HM Visualization, Visualization BT Project
	checkErr(test_db.Model(&vis_A).Association("Project").Append(&proj_A).Error)
	checkErr(test_db.Model(&vis_B).Association("Project").Append(&proj_A).Error)

	// User HM Project, Project BT User
	checkErr(test_db.Model(&proj_A).Association("User").Append(&usr_A).Error)
	checkErr(test_db.Model(&proj_B).Association("User").Append(&usr_A).Error)

	// Simulation HM Project, Project BT Simulation
	checkErr(test_db.Model(&proj_A).Association("Simulation").Append(&simn_A).Error)
	checkErr(test_db.Model(&proj_B).Association("Simulation").Append(&simn_A).Error)

	// Simulation HM SimModel, SimModel BT Simulation
	checkErr(test_db.Model(&mo_A).Association("BelongsToSimulation").Append(&simn_A).Error)
	checkErr(test_db.Model(&mo_B).Association("BelongsToSimulation").Append(&simn_A).Error)

	// User HM Simulation, Simulation BT User
	checkErr(test_db.Model(&simn_A).Association("User").Append(&usr_A).Error)
	checkErr(test_db.Model(&simn_B).Association("User").Append(&usr_A).Error)

	// Visualization HM Widget
	checkErr(test_db.Model(&vis_A).Association("Widgets").Append(&widg_A).Error)
	checkErr(test_db.Model(&vis_A).Association("Widgets").Append(&widg_B).Error)

	// SimModel HM Samples
	checkErr(test_db.Model(&mo_A).Association("InputMapping").Append(&inSmp_A).Error)
	checkErr(test_db.Model(&mo_A).Association("InputMapping").Append(&inSmp_B).Error)
	checkErr(test_db.Model(&mo_A).Association("OutputMapping").Append(&outSmp_A).Error)
	checkErr(test_db.Model(&mo_A).Association("OutputMapping").Append(&outSmp_B).Error)

	// Model HM Files
	checkErr(test_db.Model(&mo_A).Association("Files").Append(&file_A).Error)
	checkErr(test_db.Model(&mo_A).Association("Files").Append(&file_B).Error)

	// Visualization BT User
	checkErr(test_db.Model(&vis_A).Association("User").Append(&usr_A).Error)

	// Simulator BT SimModel
	checkErr(test_db.Model(&mo_A).Association("BelongsToSimulator").Append(&simr_A).Error)

	// Widget HM Files
	checkErr(test_db.Model(&widg_A).Association("Files").Append(&file_A).Error)
	checkErr(test_db.Model(&widg_A).Association("Files").Append(&file_B).Error)

}

// Erase tables and glose the testdb
func DummyCloseDB(test_db *gorm.DB) {
	test_db.Close()
}

// Quick error check
// NOTE: maybe this is not a good idea
func checkErr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}