by smallnest

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

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

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

  • 定义和表相关的 struct (Entity)
1
2
3
4
type User struct {
	ID   int    `db:"id" json:"id,omitempty"`
	Name string `db:"name" json:"name,omitempty"`
}
  • 根据 dsn 创建 sql. DB
1
2
3
4
5
db, err := sql.Open("mysql","user:password@tcp(127.0.0.1:3306)/test")
if err != nil {
	log.Fatal(err)
}
defer db.Close()
  • 执行查询, 获得一组Row
1
2
3
4
5
6
7
8
9
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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

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 和最精简的代码，把查询这一块能抽象出一个通用的代码。

那么梳理一下我的需求：

  1. 提供一个 raw sql 语句
  2. 返回一个指定类型的 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 的方法以及只查某个字段的方法：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
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 , 同时支持泛型。

查询例子如下：

1
2
3
4
5
6
7
8
9
10
11
12
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:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
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:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
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即可。