rnd-parser-ruptela/utility/logwriter.go

package utility

import (
	"database/sql"
	"io"
	"os"
	"time"

	"strings"

	"github.com/ipsusila/opt"
)

//SQLWriter write log to database i.e. SQLite
type SQLWriter struct {
	db          *sql.DB
	insertQuery string
	createQuery string
	appID       string
}

//LogWriter support multiple output
type LogWriter struct {
	writers []io.WriteCloser
	console bool
}

//NewLogWriter create multiple writer
func NewLogWriter(options *opt.Options) (*LogWriter, error) {
	logOpt := options.Get("log")

	//backward compatibility, get from server
	appID := options.Get("server").GetString("id", "")
	if appID == "" {
		//if failed, get from `log.id`
		appID = logOpt.GetString("id", "")
	}
	if appID == "" {
		//if failed, get from `id``
		appID = options.GetString("id", "*undefined*")
	}

	outTypes := logOpt.GetString("type", "")
	typeItems := strings.Split(outTypes, ",")
	writers := []io.WriteCloser{}

	console := false
	for _, item := range typeItems {
		item = strings.TrimSpace(item)
		switch item {
		case "console":
			console = true
		case "file":
			fileOpt := logOpt.Get("file")
			name := fileOpt.GetString("name", "")
			appendExisting := fileOpt.GetBool("append", true)
			flag := os.O_WRONLY | os.O_CREATE
			if appendExisting {
				flag |= os.O_APPEND
			} else {
				flag |= os.O_TRUNC
			}

			f, err := os.OpenFile(name, flag, 0666)
			if err != nil {
				return nil, err
			}
			writers = append(writers, f)
		case "sql":
			w, err := newSQLWriter(logOpt, appID)
			if err != nil {
				return nil, err
			}
			writers = append(writers, w)
		}
	}

	lw := &LogWriter{
		writers: writers,
		console: console,
	}

	return lw, nil
}

//NewSQLWriter open log writer to sql
func newSQLWriter(logOpt *opt.Options, appID string) (*SQLWriter, error) {
	sqlOpt := logOpt.Get("sql")
	dbDriver := sqlOpt.GetString("driver", "")
	dbInfo := sqlOpt.GetString("connection", "")
	createQuery := sqlOpt.GetString("createQuery", "")
	insertQuery := sqlOpt.GetString("insertQuery", "")
	dbMaxIdle := sqlOpt.GetInt("maxIdle", 0)
	dbMaxOpen := sqlOpt.GetInt("maxOpen", 5)
	dbMaxLifetime := time.Duration(sqlOpt.GetInt("maxLifetime", 60)) * time.Second

	//open databae connection
	db, err := sql.Open(dbDriver, dbInfo)
	if err != nil {
		return nil, err
	}
	err = db.Ping()
	if err != nil {
		db.Close()
		return nil, err
	}

	//connection pooling
	db.SetMaxIdleConns(dbMaxIdle)
	db.SetMaxOpenConns(dbMaxOpen)
	db.SetConnMaxLifetime(dbMaxLifetime)

	//ignore error
	//TODO, verify query (security, untrusted source)
	db.Exec(createQuery)

	w := &SQLWriter{
		db:          db,
		createQuery: createQuery,
		insertQuery: insertQuery,
		appID:       appID,
	}

	return w, nil
}

//Write to multiple log
func (lw *LogWriter) Write(data []byte) (n int, err error) {
	//now := time.Now().Format("[2006-01-02 15:04:05]")
	if lw.console {
		n, err = os.Stderr.Write(data)
	}
	for _, w := range lw.writers {
		n, err = w.Write(data)
	}

	return n, err
}

//Close all writer
func (lw *LogWriter) Close() error {
	for _, w := range lw.writers {
		w.Close()
	}
	return nil
}

//Write data with timestamp (ts, log)
func (w *SQLWriter) Write(data []byte) (int, error) {
	//TODO, verify query (security, untrusted source)
	result, err := w.db.Exec(w.insertQuery, time.Now(), w.appID, string(data))
	if err != nil {
		return 0, err
	}

	n := len(data)
	_, err = result.LastInsertId()
	if err != nil {
		return n, err
	}

	nrows, err := result.RowsAffected()
	if err != nil || nrows == 0 {
		return n, err
	}

	return n, nil
}

func (w *SQLWriter) Close() error {
	return w.db.Close()
}