golib/v2/features/sdb/sdb.go

package sdb

import (
	"context"
	"database/sql"
	"sync"
	"time"
)

const (
	driverName = "sqlite3"
)

type DB struct {
	FileName string
	db       *sql.DB
	mu       sync.Mutex
}

type TableInfo struct {
	Name        string
	ColumnsInfo []ColumnInfo
}

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

func Open(name string) (*DB, error) {
	db, err := sql.Open(driverName, name)
	if err != nil {
		return nil, err
	}
	sdb := &DB{
		FileName: name,
		db:       db,
	}
	return sdb, nil
}

func (s *DB) Close() error {
	return s.db.Close()
}

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

func (s *DB) RawDB() *sql.DB {
	return s.db
}

func (s *DB) Query(query string, args ...any) ([]M, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	ctx, cancel := s.createCtx()
	rows, err := Query(ctx, s.db, 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.db.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 {
	s.mu.Lock()
	defer s.mu.Unlock()
	ctx, cancel := s.createCtx()
	err := Exec(ctx, s.db, query, args...)
	cancel()
	return err
}

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

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

func (s *DB) Tables() ([]TableInfo, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	tblName, err := TableNames(s.db)
	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.db, name)
		if err != nil {
			return infos, err
		}
		infos[i] = TableInfo{
			Name:        name,
			ColumnsInfo: info,
		}
	}
	return infos, nil
}