arno/kube-bench
1
0
Fork 0
mirror of https://github.com/aquasecurity/kube-bench.git synced 2025-02-22 12:32:05 +00:00
Code Issues Releases Wiki Activity

feat(schema): add database schema config

Browse Source
This commit is contained in:
luisserra 2022-03-26 19:56:37 +00:00
parent ef40200f54
commit 3dddacac87
1 changed files with 23 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
cmd/database.go
View File

@ -9,6 +9,7 @@ import (
"github.com/spf13/viper" "github.com/spf13/viper"
"gorm.io/driver/postgres" "gorm.io/driver/postgres"
"gorm.io/gorm" "gorm.io/gorm"
"gorm.io/gorm/schema"
) )
type PsqlConnInfo struct { type PsqlConnInfo struct {
@ -17,8 +18,11 @@ type PsqlConnInfo struct {
DbName string DbName string
SslMode string SslMode string
Password string Password string
Schema string
} }
var db *gorm.DB
func getPsqlConnInfo() (PsqlConnInfo, error) { func getPsqlConnInfo() (PsqlConnInfo, error) {
var host string var host string
if value := viper.GetString("PGSQL_HOST"); value != "" { if value := viper.GetString("PGSQL_HOST"); value != "" {
@ -55,12 +59,21 @@ func getPsqlConnInfo() (PsqlConnInfo, error) {
return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_PASSWORD env var is required", envVarsPrefix) return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_PASSWORD env var is required", envVarsPrefix)
} }
var schema string
if value := viper.GetString("PGSQL_SCHEMA"); value != "" {
schema = value
} else {
schema = ""
fmt.Printf("No schema set.")
}
return PsqlConnInfo{ return PsqlConnInfo{
Host: host, Host: host,
User: user, User: user,
DbName: dbName, DbName: dbName,
SslMode: sslMode, SslMode: sslMode,
Password: password, Password: password,
Schema: schema,
}, nil }, nil
} }
@ -97,7 +110,16 @@ func savePgsql(jsonInfo string) {
exitWithError(err) exitWithError(err)
} }
db, err := gorm.Open(postgres.Open(PsqlConnInfo.toString()), &gorm.Config{}) if PsqlConnInfo.Schema == "" {
db, err = gorm.Open(postgres.Open(PsqlConnInfo.toString()), &gorm.Config{})
} else {
db, err = gorm.Open(postgres.Open(PsqlConnInfo.toString()), &gorm.Config{
NamingStrategy: schema.NamingStrategy{
TablePrefix: fmt.Sprintf("%s.", PsqlConnInfo.Schema), // schema name
SingularTable: false,
}})
}
if err != nil { if err != nil {
exitWithError(fmt.Errorf("received error connecting to database: %s", err)) exitWithError(fmt.Errorf("received error connecting to database: %s", err))
} }