golib/v2/features/sdb/db.go

package sdb

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"reflect"
	"strings"
)

func Query(ctx context.Context, db *sql.DB, query string, args ...any) ([]M, error) {
	rows, err := db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer func() {
		_ = rows.Close()
	}()
	columns, err := rows.ColumnTypes()
	if err != nil {
		return nil, err
	}
	rowList := make([]M, 0, 512)
	for rows.Next() {
		refs := make([]any, len(columns))
		for i, col := range columns {
			refs[i] = handleColumnType(col.DatabaseTypeName())
		}
		if err = rows.Scan(refs...); err != nil {
			return nil, err
		}
		row := make(M, len(columns))
		for i, k := range columns {
			row[k.Name()] = handleScanValue(refs[i])
		}
		rowList = append(rowList, row)
	}
	return rowList, nil
}

func Exec(ctx context.Context, db *sql.DB, query string, args ...interface{}) error {
	ret, err := db.ExecContext(ctx, query, args...)
	if err != nil {
		return err
	}
	if _, err = ret.RowsAffected(); err != nil {
		return err
	}
	return nil
}

func Execs(ctx context.Context, db *sql.DB, sql string, values ...[]any) error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	s, err := tx.Prepare(sql)
	if err != nil {
		return err
	}
	defer func() {
		_ = s.Close()
	}()
	for _, value := range values {
		_, err = s.ExecContext(ctx, value...)
		if err != nil {
			_ = tx.Rollback()
			return err
		}
	}
	return tx.Commit()
}

func TableNames(db *sql.DB) ([]string, error) {
	query := `SELECT Name FROM sqlite_master WHERE type = "table"`
	rows, err := db.Query(query)
	if err != nil {
		return nil, err
	}
	tables := make([]string, 0)
	for rows.Next() {
		var table sql.NullString
		if err = rows.Scan(&table); err != nil {
			return nil, err
		}
		if table.String != "" && table.String != "sqlite_sequence" {
			tables = append(tables, table.String)
		}
	}
	return tables, nil
}

func Columns(ctx context.Context, db *sql.DB, table string) ([]ColumnInfo, error) {
	query := fmt.Sprintf("pragma table_info('%s')", table)
	rows, err := db.QueryContext(ctx, query)
	if err != nil {
		return nil, err
	}
	cols := make([]ColumnInfo, 0)
	for rows.Next() {
		var tmp, name, types, notNull, dflt sql.NullString
		if err = rows.Scan(&tmp, &name, &types, &notNull, &dflt, &tmp); err != nil {
			return nil, err
		}
		var isNotNull bool
		if notNull.String == "1" {
			isNotNull = true
		} else {
			isNotNull = false
		}
		col := ColumnInfo{
			Name:         name.String,
			Type:         types.String,
			NotNull:      isNotNull,
			DefaultValue: dflt.String,
		}
		cols = append(cols, col)
	}
	return cols, nil
}

func DecodeRow(row M, v any) error {
	b, err := json.Marshal(row)
	if err != nil {
		return err
	}
	return json.Unmarshal(b, v)
}

func DecodeRows[T any](rows []M, dst []T) error {
	for i, row := range rows {
		var v T
		if err := DecodeRow(row, &v); err != nil {
			return err
		}
		dst[i] = v
	}
	return nil
}

// EncodeRow
// Deprecated, use Encode
func EncodeRow[T any](s T) (M, error) {
	return Encode(s)
}

// EncodeRows
// Deprecated, use Encodes
func EncodeRows[T any](s []T) ([]M, error) {
	return Encodes(s)
}

// Encode to M using v. The v Must be a json Kind
// in the after encoded, delete Tag has "none" Field.
// if v is a map Kind, Encode will be Deep copy params v in return value
func Encode(v any) (M, error) {
	var row M
	b, err := json.Marshal(v)
	if err != nil {
		return nil, err
	}
	if err = json.Unmarshal(b, &row); err != nil {
		return nil, err
	}
	if rt := reflect.TypeOf(v); rt.Kind() == reflect.Struct {
		handle := func(tags []string) (key string, skip bool) {
			if len(tags) < 2 {
				return "", false
			}
			for i, tag := range tags {
				if tag == "none" && i > 0 {
					return tags[0], true
				}
			}
			return
		}
		for i := 0; i < rt.NumField(); i++ {
			field := rt.Field(i)
			if !field.IsExported() {
				continue
			}
			value, ok := field.Tag.Lookup("json")
			if !ok {
				continue
			}
			tags := strings.Split(value, ",")
			if key, skip := handle(tags); skip {
				delete(row, key)
			}
		}
	}
	return row, nil
}

// Encodes encode to []M using v.
// Usually, the param v need be a list kind, but will be called Encode if v it's not it
func Encodes(v any) ([]M, error) {
	rt := reflect.TypeOf(v)
	if rt.Kind() != reflect.Slice && rt.Kind() != reflect.Array {
		row, err := Encode(v)
		if err != nil {
			return nil, err
		}
		return []M{row}, nil
	}
	rv := reflect.ValueOf(v)
	if rv.Type().Elem().Kind() != reflect.Struct {
		return nil, fmt.Errorf("unsupported element type: %s", rt.Kind().String())
	}
	rows := make([]M, rv.Len())
	for i := 0; i < rv.Len(); i++ {
		row, err := Encode(rv.Index(i).Interface())
		if err != nil {
			return nil, err
		}
		rows[i] = row
	}
	return rows, nil
}