Go database/sql 包查询结果行数精确判断与首行数据获取

package main

import (
    "database/sql"
    "fmt"
    _ "github.com/go-sql-driver/mysql" // 导入MySQL驱动，也可替换为Postgres等其他驱动
)

// RowStatus 定义了查询结果的行数状态
type RowStatus int

const (
    ZeroRows RowStatus = iota // 未找到任何行
    OneRow                    // 找到且仅找到一行
    MultipleRows              // 找到多行
)

// String 方法用于方便地打印 RowStatus
func (s RowStatus) String() string {
    switch s {
    case ZeroRows:
        return "ZeroRows"
    case OneRow:
        return "OneRow"
    case MultipleRows:
        return "MultipleRows"
    default:
        return "UnknownStatus"
    }
}

// QueryAndCountRows 执行SQL查询，并确定返回的行数，
// 同时将第一行数据扫描到 dest 参数中。
//
// db: 数据库连接对象。
// query: SQL查询字符串。
// args: 查询参数。
// dest: 可变参数，指针列表，用于接收第一行扫描的数据。
//
// 返回值:
// RowStatus: 指示查询结果的行数状态（ZeroRows, OneRow, MultipleRows）。
// error: 如果查询或扫描过程中发生错误。
func QueryAndCountRows(db *sql.DB, query string, args []interface{}, dest ...interface{}) (RowStatus, error) {
    rows, err := db.Query(query, args...)
    if err != nil {
        return ZeroRows, fmt.Errorf("执行查询失败: %w", err)
    }
    defer rows.Close() // 确保无论如何都关闭 rows 资源

    // 尝试获取第一行
    if !rows.Next() {
        // 如果没有下一行，检查是否有迭代错误
        if err := rows.Err(); err != nil {
            return ZeroRows, fmt.Errorf("遍历第一行时发生错误: %w", err)
        }
        // 没有错误且没有下一行，表示没有找到任何数据
        return ZeroRows, nil
    }

    // 成功获取到第一行，进行扫描
    if err := rows.Scan(dest...); err != nil {
        return ZeroRows, fmt.Errorf("扫描第一行数据失败: %w", err)
    }

    // 检查是否还有第二行，以判断是单行还是多行
    if rows.Next() {
        // 如果有第二行，则表示有多行数据
        return MultipleRows, nil
    }

    // 如果没有第二行，检查是否有迭代错误
    if err := rows.Err(); err != nil {
        return ZeroRows, fmt.Errorf("遍历第二行时发生错误: %w", err)
    }

    // 成功扫描第一行，且没有第二行，表示恰好只有一行数据
    return OneRow, nil
}

func main() {
    // 1. 初始化数据库连接 (请根据实际情况替换连接字符串)
    // 例如，使用 MySQL 驱动
    // db, err := sql.Open("mysql", "user:password@tcp(127.0.0.1:3306)/testdb?parseTime=true")
    // 这里使用一个模拟的数据库连接，实际应用中应正确初始化
    // 为了示例运行，我们假设 db 已经初始化并可用
    // 实际应用中需要处理 db 的初始化和错误
    db, err := sql.Open("mysql", "root:password@tcp(127.0.0.1:3306)/testdb") // 请替换为你的数据库连接字符串
    if err != nil {
        fmt.Printf("数据库连接失败: %v\n", err)
        return
    }
    defer db.Close()

    // 确保数据库连接有效
    err = db.Ping()
    if err != nil {
        fmt.Printf("无法连接到数据库: %v\n", err)
        return
    }
    fmt.Println("数据库连接成功。")

    // 示例：查询 ID 为 1 的用户
    var id int
    var name string
    var age int

    fmt.Println("\n--- 查询 ID = 1 的用户 ---")
    status, err := QueryAndCountRows(db, "SELECT id, name, age FROM test_users WHERE id = ?", []interface{}{1}, &id, &name, &age)
    if err != nil {
        fmt.Printf("查询出错: %v\n", err)
        return
    }

    switch status {
    case ZeroRows:
        fmt.Println("未找到 ID 为 1 的用户。")
    case OneRow:
        fmt.Printf("找到一个用户: ID=%d, Name=%s, Age=%d\n", id, name, age)
    case MultipleRows:
        // 根据业务逻辑，多行可能是一个错误
        fmt.Printf("错误: 找到多个 ID 为 1 的用户，期望最多一个。首行数据: ID=%d, Name=%s, Age=%d\n", id, name, age)
    }

    // 示例：查询 ID 不存在的用户 (例如 ID = 999)
    fmt.Println("\n--- 查询 ID = 999 的用户 ---")
    var idNotFound int
    var nameNotFound string
    var ageNotFound int
    statusNotFound, err := QueryAndCountRows(db, "SELECT id, name, age FROM test_users WHERE id = ?", []interface{}{999}, &idNotFound, &nameNotFound, &ageNotFound)
    if err != nil {
        fmt.Printf("查询出错: %v\n", err)
        return
    }
    fmt.Printf("查询结果状态: %s\n", statusNotFound)

    // 示例：查询年龄大于 25 的所有用户 (可能有多行)
    fmt.Println("\n--- 查询年龄 > 25 的用户 ---")
    var firstId int
    var firstName string
    var firstAge int
    statusMultiple, err := QueryAndCountRows(db, "SELECT id, name, age FROM test_users WHERE age > ?", []interface{}{25}, &firstId, &firstName, &firstAge)
    if err != nil {
        fmt.Printf("查询出错: %v\n", err)
        return
    }

    switch statusMultiple {
    case ZeroRows:
        fmt.Println("未找到年龄大于 25 的用户。")
    case OneRow:
        fmt.Printf("找到一个年龄大于 25 的用户: ID=%d, Name=%s, Age=%d\n", firstId, firstName, firstAge)
    case MultipleRows:
        fmt.Printf("找到多个年龄大于 25 的用户。首行数据: ID=%d, Name=%s, Age=%d\n", firstId, firstName, firstAge)
        // 如果需要处理所有行，则需要重新执行 Query() 并遍历
        fmt.Println("提示: 如果需要所有结果，请使用 db.Query() 进行完整迭代。")
    }
}