mirror of
https://github.com/aquasecurity/kube-bench.git
synced 2025-02-21 12:02:30 +00:00
feat(schema): add database schema config
This commit is contained in:
parent
ef40200f54
commit
3dddacac87
@ -9,6 +9,7 @@ import (
|
||||
"github.com/spf13/viper"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/schema"
|
||||
)
|
||||
|
||||
type PsqlConnInfo struct {
|
||||
@ -17,8 +18,11 @@ type PsqlConnInfo struct {
|
||||
DbName string
|
||||
SslMode string
|
||||
Password string
|
||||
Schema string
|
||||
}
|
||||
|
||||
var db *gorm.DB
|
||||
|
||||
func getPsqlConnInfo() (PsqlConnInfo, error) {
|
||||
var host string
|
||||
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)
|
||||
}
|
||||
|
||||
var schema string
|
||||
if value := viper.GetString("PGSQL_SCHEMA"); value != "" {
|
||||
schema = value
|
||||
} else {
|
||||
schema = ""
|
||||
fmt.Printf("No schema set.")
|
||||
}
|
||||
|
||||
return PsqlConnInfo{
|
||||
Host: host,
|
||||
User: user,
|
||||
DbName: dbName,
|
||||
SslMode: sslMode,
|
||||
Password: password,
|
||||
Schema: schema,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@ -97,7 +110,16 @@ func savePgsql(jsonInfo string) {
|
||||
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 {
|
||||
exitWithError(fmt.Errorf("received error connecting to database: %s", err))
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user