47 lines
1.2 KiB
Go
47 lines
1.2 KiB
Go
package find
|
|
|
|
import (
|
|
"errors"
|
|
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
const (
|
|
// Key under which result of `SELECT FOUND_ROWS()` should be stored in the
|
|
// driver context.
|
|
FOUND_ROWS_CTX_KEY = "maal:found_rows"
|
|
// Suggested name under which [find.FoundRowsCallback] can be registered.
|
|
FOUND_ROWS_CALLBACK = "maal:found_rows"
|
|
)
|
|
|
|
// Searches query clauses for presence of `SQL_CALC_FOUND_ROWS` and runs `SELECT
|
|
// FOUND_ROWS();` right after the query containing such clause. The result is
|
|
// put in the driver context under key [find.FOUND_ROWS_CTX_KEY]. For the
|
|
// callback to work correctly it must be registered and executed before the
|
|
// `gorm:preload` callback.
|
|
func FoundRowsCallback(d *gorm.DB) {
|
|
if _, ok := d.Statement.Clauses["SELECT"].AfterNameExpression.(sqlCalcFound); ok {
|
|
var count uint64
|
|
sqlDB, err := d.DB()
|
|
if err != nil {
|
|
_ = d.AddError(err)
|
|
return
|
|
}
|
|
res := sqlDB.QueryRowContext(d.Statement.Context, "SELECT FOUND_ROWS();")
|
|
if res == nil {
|
|
_ = d.AddError(errors.New(`fialed to issue SELECT FOUND_ROWS() query`))
|
|
return
|
|
}
|
|
if res.Err() != nil {
|
|
_ = d.AddError(res.Err())
|
|
return
|
|
}
|
|
err = res.Scan(&count)
|
|
if err != nil {
|
|
_ = d.AddError(err)
|
|
return
|
|
}
|
|
d.Set(FOUND_ROWS_CTX_KEY, count)
|
|
}
|
|
}
|