golib/features/sdb/sdb.go

package sdb

import (
	"context"
	"database/sql"
	"errors"
	"runtime"
	"time"
)

const (
	driverName = "sqlite3"
)

type DB struct {
	FileName string
	w        *sql.DB
	r        *sql.DB
}

type TableInfo struct {
	Name        string
	ColumnsInfo []ColumnInfo
}

func (t *TableInfo) Column(col string) *ColumnInfo {
	for i, c := range t.ColumnsInfo {
		if c.Name == col {
			return &t.ColumnsInfo[i]
		}
	}
	return nil
}

type ColumnInfo struct {
	Name         string
	Type         string
	NotNull      bool
	DefaultValue interface{}
}

func (c *ColumnInfo) TypePool() any {
	return handleColumnType(c.Type)
}

func Open(name string) (*DB, error) {
	w, err := sql.Open(driverName, name)
	if err != nil {
		return nil, err
	}
	w.SetMaxOpenConns(1) // 写线程设置为 1 个

	r, err := sql.Open(driverName, name)
	if err != nil {
		return nil, err
	}
	r.SetMaxOpenConns(runtime.NumCPU())

	sdb := &DB{
		FileName: name,
		w:        w,
		r:        r,
	}
	return sdb, nil
}

func (s *DB) Close() error {
	es := make([]error, 0, 2)
	if err := s.w.Close(); err != nil {
		es = append(es, err)
	}
	if err := s.r.Close(); err != nil {
		es = append(es, err)
	}
	return errors.Join(es...)
}

func (s *DB) createCtx() (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), 5*time.Second)
}

func (s *DB) Query(query string, args ...any) ([]M, error) {
	ctx, cancel := s.createCtx()
	rows, err := Query(ctx, s.r, query, args...)
	cancel()
	return rows, err
}

func (s *DB) QueryRow(query string, args ...any) (M, error) {
	rows, err := s.Query(query, args...)
	if err != nil {
		return nil, err
	}
	if len(rows) == 0 {
		return M{}, nil
	}
	return rows[0], nil
}

func (s *DB) Count(fieldNum int, query string, args ...any) ([]int64, error) {
	ctx, cancel := s.createCtx()
	defer cancel()
	row := s.r.QueryRowContext(ctx, query, args...)
	if err := row.Err(); err != nil {
		return nil, err
	}
	scan := func() (arg []any) {
		for i := 0; i < fieldNum; i++ {
			arg = append(arg, new(int64))
		}
		return
	}()
	if err := row.Scan(scan...); err != nil {
		return nil, err
	}
	count := make([]int64, fieldNum)
	for i, num := range scan {
		count[i] = *num.(*int64)
	}
	return count, nil
}

func (s *DB) Exec(query string, args ...any) error {
	ctx, cancel := s.createCtx()
	err := Exec(ctx, s.w, query, args...)
	cancel()
	return err
}

func (s *DB) Execs(query string, args ...[]any) error {
	ctx, cancel := s.createCtx()
	err := Execs(ctx, s.w, query, args...)
	cancel()
	return err
}

func (s *DB) Columns(table string) ([]ColumnInfo, error) {
	ctx, cancel := s.createCtx()
	cols, err := Columns(ctx, s.r, table)
	cancel()
	return cols, err
}

func (s *DB) Tables() ([]TableInfo, error) {
	tblName, err := TableNames(s.r)
	if err != nil {
		return nil, err
	}
	ctx, cancel := s.createCtx()
	defer cancel()
	infos := make([]TableInfo, len(tblName))
	for i, name := range tblName {
		info, err := Columns(ctx, s.r, name)
		if err != nil {
			return infos, err
		}
		infos[i] = TableInfo{
			Name:        name,
			ColumnsInfo: info,
		}
	}
	return infos, nil
}

func (s *DB) HasTable(tblName string) bool {
	tblList, _ := s.Tables()
	for _, tbl := range tblList {
		if tbl.Name == tblName {
			return true
		}
	}
	return false
}