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 WITH_AMQP bool

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.BoolVar(&WITH_AMQP, "amqp", false, "If AMQP client for simulators shall be enabled, set this option to true (default is false)")
	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)
	fmt.Println("WITH_AMQP has value ", WITH_AMQP)
}

// 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(&Signal{})
	db.DropTableIfExists(&SimulationModel{})
	db.DropTableIfExists(&File{})
	db.DropTableIfExists(&Scenario{})
	db.DropTableIfExists(&User{})
	db.DropTableIfExists(&Dashboard{})
	db.DropTableIfExists(&Widget{})
}

// AutoMigrate the models
func MigrateModels(db *gorm.DB) {
	db.AutoMigrate(&Simulator{})
	db.AutoMigrate(&Signal{})
	db.AutoMigrate(&SimulationModel{})
	db.AutoMigrate(&File{})
	db.AutoMigrate(&Scenario{})
	db.AutoMigrate(&User{})
	db.AutoMigrate(&Dashboard{})
	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 entries of each model (data defined in testdata.go)

	// Users
	checkErr(test_db.Create(&User0).Error) // Admin
	checkErr(test_db.Create(&UserA).Error) // Normal User
	checkErr(test_db.Create(&UserB).Error) // Normal User

	// Simulators
	checkErr(test_db.Create(&SimulatorA).Error)
	checkErr(test_db.Create(&SimulatorB).Error)

	// Scenarios
	checkErr(test_db.Create(&ScenarioA).Error)
	checkErr(test_db.Create(&ScenarioB).Error)

	// Signals
	checkErr(test_db.Create(&OutSignalA).Error)
	checkErr(test_db.Create(&OutSignalB).Error)
	checkErr(test_db.Create(&InSignalA).Error)
	checkErr(test_db.Create(&InSignalB).Error)

	// Simulation Models
	checkErr(test_db.Create(&SimulationModelA).Error)
	checkErr(test_db.Create(&SimulationModelB).Error)

	// Dashboards
	checkErr(test_db.Create(&DashboardA).Error)
	checkErr(test_db.Create(&DashboardB).Error)

	// Files
	checkErr(test_db.Create(&FileA).Error)
	checkErr(test_db.Create(&FileB).Error)
	checkErr(test_db.Create(&FileC).Error)
	checkErr(test_db.Create(&FileD).Error)

	// Widgets
	checkErr(test_db.Create(&WidgetA).Error)
	checkErr(test_db.Create(&WidgetB).Error)

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

	// User HM Scenarios, Scenario HM Users (Many-to-Many)
	checkErr(test_db.Model(&ScenarioA).Association("Users").Append(&UserA).Error)
	checkErr(test_db.Model(&ScenarioA).Association("Users").Append(&UserB).Error)
	checkErr(test_db.Model(&ScenarioB).Association("Users").Append(&UserA).Error)
	checkErr(test_db.Model(&ScenarioB).Association("Users").Append(&UserB).Error)

	// Scenario HM SimulationModels
	checkErr(test_db.Model(&ScenarioA).Association("SimulationModels").Append(&SimulationModelA).Error)
	checkErr(test_db.Model(&ScenarioA).Association("SimulationModels").Append(&SimulationModelB).Error)

	// Scenario HM Dashboards
	checkErr(test_db.Model(&ScenarioA).Association("Dashboards").Append(&DashboardA).Error)
	checkErr(test_db.Model(&ScenarioA).Association("Dashboards").Append(&DashboardB).Error)

	// Dashboard HM Widget
	checkErr(test_db.Model(&DashboardA).Association("Widgets").Append(&WidgetA).Error)
	checkErr(test_db.Model(&DashboardA).Association("Widgets").Append(&WidgetB).Error)

	// SimulationModel HM Signals
	checkErr(test_db.Model(&SimulationModelA).Association("InputMapping").Append(&InSignalA).Error)
	checkErr(test_db.Model(&SimulationModelA).Association("InputMapping").Append(&InSignalB).Error)
	checkErr(test_db.Model(&SimulationModelA).Association("OutputMapping").Append(&OutSignalA).Error)
	checkErr(test_db.Model(&SimulationModelA).Association("OutputMapping").Append(&OutSignalB).Error)

	// SimulationModel HM Files
	checkErr(test_db.Model(&SimulationModelA).Association("Files").Append(&FileC).Error)
	checkErr(test_db.Model(&SimulationModelA).Association("Files").Append(&FileD).Error)

	// Simulator HM SimulationModels
	checkErr(test_db.Model(&SimulatorA).Association("SimulationModels").Append(&SimulationModelA).Error)
	checkErr(test_db.Model(&SimulatorA).Association("SimulationModels").Append(&SimulationModelB).Error)

	// Widget HM Files
	checkErr(test_db.Model(&WidgetA).Association("Files").Append(&FileA).Error)
	checkErr(test_db.Model(&WidgetA).Association("Files").Append(&FileB).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)
	}
}