IT博客汇 | 重复了很多次，我终于不再忍了

重复了很多次，我终于不再忍了

smallnest发表于 2024-05-08 15:53:15

虽然我不做管理系统，但是在项目中和数据库打交道还是比较多的，经常会从数据库中 (比如 Mysql 、ClickHouse 等) 查询一些记录，偶尔也会写入一些数据，但是不多。

每次从数据库中查询一些数据，套路几乎是一样的，无非是：

定义和表相关的 struct (Entity)

type User struct {
	ID   int    `db:"id" json:"id,omitempty"`
	Name string `db:"name" json:"name,omitempty"`
}

根据 dsn 创建 sql. DB

db, err := sql.Open("mysql","user:password@tcp(127.0.0.1:3306)/test")
if err != nil {
	log.Fatal(err)
}
defer db.Close()

执行查询, 获得一组Row

var (
	id int
	name string
)
rows, err := db.Query("select id, name from users where id = ?", 1)
if err != nil {
	log.Fatal(err)
}
defer rows.Close()

遍历 rows, 读取数据,并填充struct


var users []User
for rows.Next() {
    var user User
	err := rows.Scan(&user.ID, &user.Name)
	if err != nil {
		log.Fatal(err)
	}
	
    users = append(users, user)
}
err = rows.Err()
if err != nil {
	log.Fatal(err)
}

每一次，都是这个套路，套路用多了，也烦了。

当然也可以用 orm 库，比如 gorm, 减少一些代码。但是我还是想用原始的 sql 和最精简的代码，把查询这一块能抽象出一个通用的代码。

那么梳理一下我的需求：

提供一个 raw sql 语句
返回一个指定类型的 struct、或者是一个指定类型的struct 的切片

但是等等，好歹也得提供数据库的连接信息吧，或者提供一个连接好的 sql. DB, 所以输入变成了两项：

那么是输入 dsn 还是创建好的 sql. DB 呢？我最终选择了 sql. DB, 原因有两点：

可以重用创建好的 sql. DB, 多个地方可以共享使用
用户可以在外部配置 sql. DB

输入输出确定了，那么就是实现了。相关的代码在 smallnest/exp 让我们看看它是怎么封装的。

查询多条记录

func Rows[T any](ctx context.Context, db *sql.DB, query string, args ...any) ([]T, error)

ctx 可以设置超时时间，或者不想设置的话用 context.Background 即可
db 外部已经创建好的数据库连接
query sql 查询语句
args sql 中的参数，可选

返回结果就是 []T, 如果查询失败，返回 error

本质上，这个函数也没啥，就是执行查询，遍历 rows, 利用反射将数据库记录转换成 T 类型的结构体，所以我把它称之为 helper 函数，减少我的重复工作量。

其实底层我没也没有从零去写，而是使用了 blockloop/scan , 封装的更方便使用，同时支持泛型。

一个例子如下, 演示了查询多个 person 的方法以及只查某个字段的方法：

type person struct {
	ID   int    `db:"id" json:"id,omitempty"`
	Name string `json:"name,omitempty"` // `db:"name" json:"name,omitempty"`
}
func TestRows(t *testing.T) {
	db := exampleDB(t)
	persons, err := Rows[person](context.Background(), db, "SELECT * FROM persons order by id")
	assert.NoError(t, err)
	require.Equal(t, 2, len(persons))
	assert.Equal(t, 1, persons[0].ID)
	assert.Equal(t, "brett", persons[0].Name)
	assert.Equal(t, 2, persons[1].ID)
	assert.Equal(t, "fred", persons[1].Name)
	names, err := Rows[string](context.Background(), db, "SELECT name FROM persons order by id")
	assert.NoError(t, err)
	assert.Equal(t, 2, len(names))
	assert.Equal(t, "brett", names[0])
	assert.Equal(t, "fred", names[1])
}

查询单条记录

如果查询单条记录，你可以使用 Row 函数：

1	func Row[T any](ctx context.Context, db *sql.DB, query string, args ...any) (T, error)

和查询多条记录类似，只不过它返回一个 struct 而已。

也是使用了 blockloop/scan , 同时支持泛型。

查询例子如下：

func TestRow(t *testing.T) {
	db := exampleDB(t)
	person, err := Row[person](context.Background(), db, "SELECT * FROM persons order by id limit 1")
	assert.NoError(t, err)
	assert.Equal(t, 1, person.ID)
	assert.Equal(t, "brett", person.Name)
	name, err := Row[string](context.Background(), db, "SELECT name FROM persons order by id limit 1")
	assert.NoError(t, err)
	assert.Equal(t, "brett", name)
}

总得来说，这两个函数基本满足了我的日常查询需求，也减少了很多重复的代码，同时也提高了代码的可读性，同时它们还实现了简单的ORM的功能，把数据库的记录转换成了结构体。

这里你可能注意到了，我最开始定义了所需的结构体，如果你想更懒一些，你可以不定义结构体，直接使用 map[string]any 来接收查询结果, 但是这样会失去类型检查，只推荐在特定的场景下使用。

下面的提供了两个函数，和上面的函数，但是返回的是 map[string]any。

返回 map 类型

为了方便，我给map[string]any起了一个别名:

1 2	// Record is a type alias for map[string]any. type Record = map[string]any

Record代表一条记录，key是字段名，value是字段的值。

查询多条记录的函数如下，这里我没有依赖第三方的库，而是直接实现，遍历 rows, 读取数据,并填充 map[string]any:

func RowsMap(ctx context.Context, db *sql.DB, query string, args ...any) ([]Record, error) {
	rows, err := db.QueryContext(ctx, query, args...)
	defer rows.Close()
	if err != nil {
		return nil, err
	}
	colNames, err := rows.Columns()
	if err != nil {
		return nil, err
	}
	cols := make([]any, len(colNames))
	colPtrs := make([]any, len(colNames))
	for i := 0; i < len(colNames); i++ {
		colPtrs[i] = &cols[i]
	}
	var ret []Record
	for rows.Next() {
		err = rows.Scan(colPtrs...)
		if err != nil {
			return nil, err
		}
		row := make(Record)
		for i, col := range cols {
			row[colNames[i]] = col
		}
		ret = append(ret, row)
	}
	return ret, nil
}

查询单条记录类似，只不过返回的是 Record:

1	func RowMap(ctx context.Context, db *sql.DB, query string, args ...any) (Record, error)

接下来。我们看一个例子。这个例子和上面的例子类似，只不过调用我们这里介绍的两个函数，返回的是 map[string]any:

func TestRowsMap(t *testing.T) {
	db := exampleDB(t)
	persons, err := RowsMap(context.Background(), db, "SELECT * FROM persons order by id")
	assert.NoError(t, err)
	require.Equal(t, 2, len(persons))
	assert.Equal(t, int64(1), persons[0]["id"])
	assert.Equal(t, "brett", persons[0]["name"])
	assert.Equal(t, int64(2), persons[1]["id"])
	assert.Equal(t, "fred", persons[1]["name"])
	names, err := RowsMap(context.Background(), db, "SELECT name FROM persons order by id")
	assert.NoError(t, err)
	assert.Equal(t, 2, len(names))
	assert.Equal(t, "brett", names[0]["name"])
	assert.Equal(t, "fred", names[1]["name"])
}
func TestRowMap(t *testing.T) {
	db := exampleDB(t)
	person, err := RowMap(context.Background(), db, "SELECT * FROM persons order by id limit 1")
	assert.NoError(t, err)
	assert.Equal(t, int64(1), person["id"])
	assert.Equal(t, "brett", person["name"])
	name, err := Row[string](context.Background(), db, "SELECT name FROM persons order by id limit 1")
	assert.NoError(t, err)
	assert.Equal(t, "brett", name)
}

使用这里介绍的函数，直接引用github.com/smallnest/exp/db即可。