github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/agents/db_sqlite.go (about)

     1  // Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  
     3  package agents
     4  
     5  import (
     6  	"errors"
     7  	"github.com/TeaOSLab/EdgeNode/internal/events"
     8  	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
     9  	"github.com/TeaOSLab/EdgeNode/internal/utils/dbs"
    10  	"github.com/iwind/TeaGo/Tea"
    11  	"github.com/iwind/TeaGo/types"
    12  	"log"
    13  	"os"
    14  	"path/filepath"
    15  	"strings"
    16  )
    17  
    18  const (
    19  	tableAgentIPs = "agentIPs"
    20  )
    21  
    22  type SQLiteDB struct {
    23  	db   *dbs.DB
    24  	path string
    25  
    26  	insertAgentIPStmt *dbs.Stmt
    27  	listAgentIPsStmt  *dbs.Stmt
    28  }
    29  
    30  func NewSQLiteDB(path string) *SQLiteDB {
    31  	var db = &SQLiteDB{path: path}
    32  
    33  	events.OnClose(func() {
    34  		_ = db.Close()
    35  	})
    36  
    37  	return db
    38  }
    39  
    40  func (this *SQLiteDB) Init() error {
    41  	// 检查目录是否存在
    42  	var dir = filepath.Dir(this.path)
    43  
    44  	_, err := os.Stat(dir)
    45  	if err != nil {
    46  		err = os.MkdirAll(dir, 0777)
    47  		if err != nil {
    48  			return err
    49  		}
    50  		remotelogs.Println("DB", "create database dir '"+dir+"'")
    51  	}
    52  
    53  	// TODO 思考 data.db 的数据安全性
    54  	db, err := dbs.OpenWriter("file:" + this.path + "?cache=shared&mode=rwc&_journal_mode=WAL&_locking_mode=EXCLUSIVE")
    55  	if err != nil {
    56  		return err
    57  	}
    58  	db.SetMaxOpenConns(1)
    59  
    60  	/**_, err = db.Exec("VACUUM")
    61  	if err != nil {
    62  		return err
    63  	}**/
    64  
    65  	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS "` + tableAgentIPs + `" (
    66    "id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
    67    "ip" varchar(64),
    68    "agentCode" varchar(128)
    69  );`)
    70  	if err != nil {
    71  		return err
    72  	}
    73  
    74  	// 预编译语句
    75  
    76  	// agent ip record statements
    77  	this.insertAgentIPStmt, err = db.Prepare(`INSERT INTO "` + tableAgentIPs + `" ("id", "ip", "agentCode") VALUES (?, ?, ?)`)
    78  	if err != nil {
    79  		return err
    80  	}
    81  
    82  	this.listAgentIPsStmt, err = db.Prepare(`SELECT "id", "ip", "agentCode" FROM "` + tableAgentIPs + `" ORDER BY "id" ASC LIMIT ? OFFSET ?`)
    83  	if err != nil {
    84  		return err
    85  	}
    86  
    87  	this.db = db
    88  
    89  	return nil
    90  }
    91  
    92  func (this *SQLiteDB) InsertAgentIP(ipId int64, ip string, agentCode string) error {
    93  	if this.db == nil {
    94  		return errors.New("db should not be nil")
    95  	}
    96  
    97  	_, err := this.insertAgentIPStmt.Exec(ipId, ip, agentCode)
    98  	if err != nil {
    99  		// 不提示ID重复错误
   100  		if strings.Contains(err.Error(), "UNIQUE constraint") {
   101  			return nil
   102  		}
   103  
   104  		return err
   105  	}
   106  
   107  	return nil
   108  }
   109  
   110  func (this *SQLiteDB) ListAgentIPs(offset int64, size int64) (agentIPs []*AgentIP, err error) {
   111  	if this.db == nil {
   112  		return nil, errors.New("db should not be nil")
   113  	}
   114  	rows, err := this.listAgentIPsStmt.Query(size, offset)
   115  	if err != nil {
   116  		return nil, err
   117  	}
   118  	defer func() {
   119  		_ = rows.Close()
   120  	}()
   121  	for rows.Next() {
   122  		var agentIP = &AgentIP{}
   123  		err = rows.Scan(&agentIP.Id, &agentIP.IP, &agentIP.AgentCode)
   124  		if err != nil {
   125  			return nil, err
   126  		}
   127  		agentIPs = append(agentIPs, agentIP)
   128  	}
   129  	return
   130  }
   131  
   132  func (this *SQLiteDB) Close() error {
   133  	if this.db == nil {
   134  		return nil
   135  	}
   136  
   137  	for _, stmt := range []*dbs.Stmt{
   138  		this.insertAgentIPStmt,
   139  		this.listAgentIPsStmt,
   140  	} {
   141  		if stmt != nil {
   142  			_ = stmt.Close()
   143  		}
   144  	}
   145  
   146  	return this.db.Close()
   147  }
   148  
   149  // 打印日志
   150  func (this *SQLiteDB) log(args ...any) {
   151  	if !Tea.IsTesting() {
   152  		return
   153  	}
   154  	if len(args) == 0 {
   155  		return
   156  	}
   157  
   158  	args[0] = "[" + types.String(args[0]) + "]"
   159  	log.Println(args...)
   160  }