github.com/crowdsecurity/crowdsec@v1.6.1/pkg/database/database.go (about)

     1  package database
     2  
     3  import (
     4  	"context"
     5  	"database/sql"
     6  	"fmt"
     7  	"os"
     8  
     9  	entsql "entgo.io/ent/dialect/sql"
    10  	_ "github.com/go-sql-driver/mysql"
    11  	_ "github.com/jackc/pgx/v4/stdlib"
    12  	_ "github.com/mattn/go-sqlite3"
    13  	log "github.com/sirupsen/logrus"
    14  
    15  	"github.com/crowdsecurity/go-cs-lib/ptr"
    16  
    17  	"github.com/crowdsecurity/crowdsec/pkg/csconfig"
    18  	"github.com/crowdsecurity/crowdsec/pkg/database/ent"
    19  	"github.com/crowdsecurity/crowdsec/pkg/types"
    20  )
    21  
    22  type Client struct {
    23  	Ent              *ent.Client
    24  	CTX              context.Context
    25  	Log              *log.Logger
    26  	CanFlush         bool
    27  	Type             string
    28  	WalMode          *bool
    29  	decisionBulkSize int
    30  }
    31  
    32  func getEntDriver(dbtype string, dbdialect string, dsn string, config *csconfig.DatabaseCfg) (*entsql.Driver, error) {
    33  	db, err := sql.Open(dbtype, dsn)
    34  	if err != nil {
    35  		return nil, err
    36  	}
    37  	if config.MaxOpenConns == nil {
    38  		log.Warningf("MaxOpenConns is 0, defaulting to %d", csconfig.DEFAULT_MAX_OPEN_CONNS)
    39  		config.MaxOpenConns = ptr.Of(csconfig.DEFAULT_MAX_OPEN_CONNS)
    40  	}
    41  	db.SetMaxOpenConns(*config.MaxOpenConns)
    42  	drv := entsql.OpenDB(dbdialect, db)
    43  	return drv, nil
    44  }
    45  
    46  func NewClient(config *csconfig.DatabaseCfg) (*Client, error) {
    47  	var client *ent.Client
    48  	var err error
    49  	if config == nil {
    50  		return &Client{}, fmt.Errorf("DB config is empty")
    51  	}
    52  	/*The logger that will be used by db operations*/
    53  	clog := log.New()
    54  	if err := types.ConfigureLogger(clog); err != nil {
    55  		return nil, fmt.Errorf("while configuring db logger: %w", err)
    56  	}
    57  	if config.LogLevel != nil {
    58  		clog.SetLevel(*config.LogLevel)
    59  	}
    60  	entLogger := clog.WithField("context", "ent")
    61  
    62  	entOpt := ent.Log(entLogger.Debug)
    63  	typ, dia, err := config.ConnectionDialect()
    64  	if err != nil {
    65  		return &Client{}, err //unsupported database caught here
    66  	}
    67  	if config.Type == "sqlite" {
    68  		/*if it's the first startup, we want to touch and chmod file*/
    69  		if _, err := os.Stat(config.DbPath); os.IsNotExist(err) {
    70  			f, err := os.OpenFile(config.DbPath, os.O_CREATE|os.O_RDWR, 0600)
    71  			if err != nil {
    72  				return &Client{}, fmt.Errorf("failed to create SQLite database file %q: %w", config.DbPath, err)
    73  			}
    74  			if err := f.Close(); err != nil {
    75  				return &Client{}, fmt.Errorf("failed to create SQLite database file %q: %w", config.DbPath, err)
    76  			}
    77  		}
    78  		//Always try to set permissions to simplify a bit the code for windows (as the permissions set by OpenFile will be garbage)
    79  		if err := setFilePerm(config.DbPath, 0640); err != nil {
    80  			return &Client{}, fmt.Errorf("unable to set perms on %s: %v", config.DbPath, err)
    81  		}
    82  	}
    83  	drv, err := getEntDriver(typ, dia, config.ConnectionString(), config)
    84  	if err != nil {
    85  		return &Client{}, fmt.Errorf("failed opening connection to %s: %v", config.Type, err)
    86  	}
    87  	client = ent.NewClient(ent.Driver(drv), entOpt)
    88  	if config.LogLevel != nil && *config.LogLevel >= log.DebugLevel {
    89  		clog.Debugf("Enabling request debug")
    90  		client = client.Debug()
    91  	}
    92  	if err = client.Schema.Create(context.Background()); err != nil {
    93  		return nil, fmt.Errorf("failed creating schema resources: %v", err)
    94  	}
    95  
    96  	return &Client{
    97  		Ent: client,
    98  		CTX: context.Background(),
    99  		Log: clog,
   100  		CanFlush: true,
   101  		Type: config.Type,
   102  		WalMode: config.UseWal,
   103  		decisionBulkSize: config.DecisionBulkSize,
   104  	}, nil
   105  }