bitbucket.org/Aishee/synsec@v0.0.0-20210414005726-236fc01a153d/pkg/database/database.go (about)

     1  package database
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"time"
     8  
     9  	"bitbucket.org/Aishee/synsec/pkg/csconfig"
    10  	"bitbucket.org/Aishee/synsec/pkg/database/ent"
    11  	"bitbucket.org/Aishee/synsec/pkg/types"
    12  	"github.com/go-co-op/gocron"
    13  	_ "github.com/go-sql-driver/mysql"
    14  	_ "github.com/lib/pq"
    15  	_ "github.com/mattn/go-sqlite3"
    16  	"github.com/pkg/errors"
    17  	log "github.com/sirupsen/logrus"
    18  )
    19  
    20  type Client struct {
    21  	Ent *ent.Client
    22  	CTX context.Context
    23  	Log *log.Logger
    24  }
    25  
    26  func NewClient(config *csconfig.DatabaseCfg) (*Client, error) {
    27  	var client *ent.Client
    28  	var err error
    29  	if config == nil {
    30  		return &Client{}, fmt.Errorf("DB config is empty")
    31  	}
    32  	switch config.Type {
    33  	case "sqlite":
    34  
    35  		/*if it's the first startup, we want to touch and chmod file*/
    36  		if _, err := os.Stat(config.DbPath); os.IsNotExist(err) {
    37  			f, err := os.OpenFile(config.DbPath, os.O_CREATE|os.O_RDWR, 0600)
    38  			if err != nil {
    39  				return &Client{}, errors.Wrapf(err, "failed to create SQLite database file %q", config.DbPath)
    40  			}
    41  			if err := f.Close(); err != nil {
    42  				return &Client{}, errors.Wrapf(err, "failed to create SQLite database file %q", config.DbPath)
    43  			}
    44  		} else { /*ensure file perms*/
    45  			if err := os.Chmod(config.DbPath, 0660); err != nil {
    46  				return &Client{}, fmt.Errorf("unable to set perms on %s: %v", config.DbPath, err)
    47  			}
    48  		}
    49  		client, err = ent.Open("sqlite3", fmt.Sprintf("file:%s?_busy_timeout=100000&_fk=1", config.DbPath))
    50  		if err != nil {
    51  			return &Client{}, fmt.Errorf("failed opening connection to sqlite: %v", err)
    52  		}
    53  	case "mysql":
    54  		client, err = ent.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?parseTime=True", config.User, config.Password, config.Host, config.Port, config.DbName))
    55  		if err != nil {
    56  			return &Client{}, fmt.Errorf("failed opening connection to mysql: %v", err)
    57  		}
    58  	case "postgres", "postgresql":
    59  		client, err = ent.Open("postgres", fmt.Sprintf("host=%s port=%d user=%s dbname=%s password=%s", config.Host, config.Port, config.User, config.DbName, config.Password))
    60  		if err != nil {
    61  			return &Client{}, fmt.Errorf("failed opening connection to postgres: %v", err)
    62  		}
    63  	default:
    64  		return &Client{}, fmt.Errorf("unknown database type")
    65  	}
    66  
    67  	/*The logger that will be used by db operations*/
    68  	clog := log.New()
    69  	if err := types.ConfigureLogger(clog); err != nil {
    70  		return nil, errors.Wrap(err, "while configuring db logger")
    71  	}
    72  	if config.LogLevel != nil {
    73  		clog.SetLevel(*config.LogLevel)
    74  		if *config.LogLevel >= log.DebugLevel {
    75  			clog.Debugf("Enabling request debug")
    76  			client = client.Debug()
    77  		}
    78  	}
    79  	if err = client.Schema.Create(context.Background()); err != nil {
    80  		return nil, fmt.Errorf("failed creating schema resources: %v", err)
    81  	}
    82  	return &Client{Ent: client, CTX: context.Background(), Log: clog}, nil
    83  }
    84  
    85  func (c *Client) StartFlushScheduler(config *csconfig.FlushDBCfg) (*gocron.Scheduler, error) {
    86  	maxItems := 0
    87  	maxAge := ""
    88  	if config.MaxItems != nil && *config.MaxItems <= 0 {
    89  		return nil, fmt.Errorf("max_items can't be zero or negative number")
    90  	}
    91  	if config.MaxItems != nil {
    92  		maxItems = *config.MaxItems
    93  	}
    94  
    95  	if config.MaxAge != nil && *config.MaxAge != "" {
    96  		maxAge = *config.MaxAge
    97  	}
    98  	// Init & Start cronjob every minute
    99  	scheduler := gocron.NewScheduler(time.UTC)
   100  	scheduler.Every(1).Minute().Do(c.FlushAlerts, maxAge, maxItems)
   101  	scheduler.StartAsync()
   102  
   103  	return scheduler, nil
   104  }