EdgeNode/internal/utils/agents/db_sqlite.go

// Copyright 2021 GoEdge goedge.cdn@gmail.com. All rights reserved.

package agents

import (
	"errors"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/TeaOSLab/EdgeNode/internal/events"
	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
	"github.com/TeaOSLab/EdgeNode/internal/utils/dbs"
	"github.com/iwind/TeaGo/Tea"
	"github.com/iwind/TeaGo/types"
)

const (
	tableAgentIPs = "agentIPs"
)

type SQLiteDB struct {
	db   *dbs.DB
	path string

	insertAgentIPStmt *dbs.Stmt
	listAgentIPsStmt  *dbs.Stmt
}

func NewSQLiteDB(path string) *SQLiteDB {
	var db = &SQLiteDB{path: path}

	events.OnClose(func() {
		_ = db.Close()
	})

	return db
}

func (this *SQLiteDB) Init() error {
	// 检查目录是否存在
	var dir = filepath.Dir(this.path)

	_, err := os.Stat(dir)
	if err != nil {
		err = os.MkdirAll(dir, 0777)
		if err != nil {
			return err
		}
		remotelogs.Println("DB", "create database dir '"+dir+"'")
	}

	// TODO 思考 data.db 的数据安全性
	db, err := dbs.OpenWriter("file:" + this.path + "?cache=shared&mode=rwc&_journal_mode=WAL&_locking_mode=EXCLUSIVE")
	if err != nil {
		return err
	}
	db.SetMaxOpenConns(1)

	/**_, err = db.Exec("VACUUM")
	if err != nil {
		return err
	}**/

	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS "` + tableAgentIPs + `" (
  "id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
  "ip" varchar(64),
  "agentCode" varchar(128)
);`)
	if err != nil {
		return err
	}

	// 预编译语句

	// agent ip record statements
	this.insertAgentIPStmt, err = db.Prepare(`INSERT INTO "` + tableAgentIPs + `" ("id", "ip", "agentCode") VALUES (?, ?, ?)`)
	if err != nil {
		return err
	}

	this.listAgentIPsStmt, err = db.Prepare(`SELECT "id", "ip", "agentCode" FROM "` + tableAgentIPs + `" ORDER BY "id" ASC LIMIT ? OFFSET ?`)
	if err != nil {
		return err
	}

	this.db = db

	return nil
}

func (this *SQLiteDB) InsertAgentIP(ipId int64, ip string, agentCode string) error {
	if this.db == nil {
		return errors.New("db should not be nil")
	}

	_, err := this.insertAgentIPStmt.Exec(ipId, ip, agentCode)
	if err != nil {
		// 不提示ID重复错误
		if strings.Contains(err.Error(), "UNIQUE constraint") {
			return nil
		}

		return err
	}

	return nil
}

func (this *SQLiteDB) ListAgentIPs(offset int64, size int64) (agentIPs []*AgentIP, err error) {
	if this.db == nil {
		return nil, errors.New("db should not be nil")
	}
	rows, err := this.listAgentIPsStmt.Query(size, offset)
	if err != nil {
		return nil, err
	}
	defer func() {
		_ = rows.Close()
	}()
	for rows.Next() {
		var agentIP = &AgentIP{}
		err = rows.Scan(&agentIP.Id, &agentIP.IP, &agentIP.AgentCode)
		if err != nil {
			return nil, err
		}
		agentIPs = append(agentIPs, agentIP)
	}
	return
}

func (this *SQLiteDB) Close() error {
	if this.db == nil {
		return nil
	}

	for _, stmt := range []*dbs.Stmt{
		this.insertAgentIPStmt,
		this.listAgentIPsStmt,
	} {
		if stmt != nil {
			_ = stmt.Close()
		}
	}

	return this.db.Close()
}

// 打印日志
func (this *SQLiteDB) log(args ...any) {
	if !Tea.IsTesting() {
		return
	}
	if len(args) == 0 {
		return
	}

	args[0] = "[" + types.String(args[0]) + "]"
	log.Println(args...)
}