crowdsec/pkg/database/database.go

106 lines
3.2 KiB
Go
Raw Permalink Normal View History

package database
import (
"context"
"database/sql"
"fmt"
"os"
entsql "entgo.io/ent/dialect/sql"
_ "github.com/go-sql-driver/mysql"
_ "github.com/jackc/pgx/v4/stdlib"
_ "github.com/mattn/go-sqlite3"
log "github.com/sirupsen/logrus"
"github.com/crowdsecurity/go-cs-lib/ptr"
"github.com/crowdsecurity/crowdsec/pkg/csconfig"
"github.com/crowdsecurity/crowdsec/pkg/database/ent"
"github.com/crowdsecurity/crowdsec/pkg/types"
)
type Client struct {
Ent *ent.Client
CTX context.Context
Log *log.Logger
CanFlush bool
Type string
WalMode *bool
decisionBulkSize int
}
func getEntDriver(dbtype string, dbdialect string, dsn string, config *csconfig.DatabaseCfg) (*entsql.Driver, error) {
db, err := sql.Open(dbtype, dsn)
if err != nil {
return nil, err
}
if config.MaxOpenConns == nil {
log.Warningf("MaxOpenConns is 0, defaulting to %d", csconfig.DEFAULT_MAX_OPEN_CONNS)
config.MaxOpenConns = ptr.Of(csconfig.DEFAULT_MAX_OPEN_CONNS)
}
db.SetMaxOpenConns(*config.MaxOpenConns)
drv := entsql.OpenDB(dbdialect, db)
return drv, nil
}
func NewClient(config *csconfig.DatabaseCfg) (*Client, error) {
var client *ent.Client
var err error
if config == nil {
return &Client{}, fmt.Errorf("DB config is empty")
}
/*The logger that will be used by db operations*/
clog := log.New()
if err := types.ConfigureLogger(clog); err != nil {
return nil, fmt.Errorf("while configuring db logger: %w", err)
}
if config.LogLevel != nil {
clog.SetLevel(*config.LogLevel)
}
entLogger := clog.WithField("context", "ent")
entOpt := ent.Log(entLogger.Debug)
typ, dia, err := config.ConnectionDialect()
if err != nil {
return &Client{}, err //unsupported database caught here
}
if config.Type == "sqlite" {
/*if it's the first startup, we want to touch and chmod file*/
if _, err := os.Stat(config.DbPath); os.IsNotExist(err) {
f, err := os.OpenFile(config.DbPath, os.O_CREATE|os.O_RDWR, 0600)
if err != nil {
return &Client{}, fmt.Errorf("failed to create SQLite database file %q: %w", config.DbPath, err)
}
if err := f.Close(); err != nil {
return &Client{}, fmt.Errorf("failed to create SQLite database file %q: %w", config.DbPath, err)
}
}
//Always try to set permissions to simplify a bit the code for windows (as the permissions set by OpenFile will be garbage)
2022-11-21 09:49:56 +00:00
if err := setFilePerm(config.DbPath, 0640); err != nil {
return &Client{}, fmt.Errorf("unable to set perms on %s: %v", config.DbPath, err)
}
}
drv, err := getEntDriver(typ, dia, config.ConnectionString(), config)
if err != nil {
return &Client{}, fmt.Errorf("failed opening connection to %s: %v", config.Type, err)
}
client = ent.NewClient(ent.Driver(drv), entOpt)
if config.LogLevel != nil && *config.LogLevel >= log.DebugLevel {
clog.Debugf("Enabling request debug")
client = client.Debug()
}
if err = client.Schema.Create(context.Background()); err != nil {
return nil, fmt.Errorf("failed creating schema resources: %v", err)
}
return &Client{
Ent: client,
CTX: context.Background(),
Log: clog,
CanFlush: true,
Type: config.Type,
WalMode: config.UseWal,
decisionBulkSize: config.DecisionBulkSize,
}, nil
}