sql.go

package psql

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"reflect"
	"strings"
	"time"
	"unsafe"

	"github.com/gopsql/db"
	"github.com/gopsql/logger"
)

var (
	// ErrInvalidTarget is returned when Query receives a target that is not a
	// pointer to a struct, slice, or map.
	ErrInvalidTarget = errors.New("target must be pointer of a struct, slice or map")

	// ErrNoConnection is returned when attempting to execute a query without
	// a database connection set on the Model.
	ErrNoConnection = errors.New("no connection")

	// ErrNoSQL is returned when Execute is called with an empty SQL statement.
	ErrNoSQL = errors.New("no sql statements to execute")

	// ErrTypeAssertionFailed is returned when scanning a JSONB value fails
	// due to an unexpected source type.
	ErrTypeAssertionFailed = errors.New("type assertion failed")

	// ErrUnsupportedExplainTarget is returned when Explain is called with a
	// target type that is not supported (*string, io.Writer, logger.Logger,
	// func(string), or func(...interface{})).
	ErrUnsupportedExplainTarget = errors.New("unsupported explain target type")
)

type (
	// SQL represents a SQL statement with parameter values. It provides methods
	// for executing queries (Query, QueryRow) and statements (Execute). Create
	// instances using Model.NewSQL or the query builder methods (Find, Select,
	// Insert, Update, Delete).
	SQL struct {
		main interface {
			String() string
			StringValues() (string, []interface{})
		}
		model          *Model
		sql            string
		values         []interface{}
		explainTarget  interface{}
		explainOptions []string
	}

	// Tx is an alias for db.Tx, representing a database transaction.
	Tx = db.Tx

	jsonbRaw map[string]json.RawMessage

	fieldsFunc = func([]string, string) []string
)

// AddTableName is a helper function that prefixes field names with the table
// name. It can be passed to Find to disambiguate columns in JOIN queries.
// Fields that already contain a dot are left unchanged.
var AddTableName fieldsFunc = func(fields []string, tableName string) (out []string) {
	for _, field := range fields {
		if strings.Contains(field, ".") {
			out = append(out, field)
			continue
		}
		out = append(out, tableName+"."+field)
	}
	return
}

func (j *jsonbRaw) Scan(src interface{}) error { // necessary for github.com/lib/pq
	if src == nil {
		return nil
	}
	switch source := src.(type) {
	case string:
		return json.Unmarshal([]byte(source), j)
	case []byte:
		return json.Unmarshal(source, j)
	default:
		return ErrTypeAssertionFailed
	}
}

// NewSQL creates a new SQL statement with the given query string and parameter
// values. Use $1, $2, etc. or $? as placeholders for parameters.
func (m Model) NewSQL(sql string, values ...interface{}) *SQL {
	return &SQL{
		model:  &m,
		sql:    strings.TrimSpace(sql),
		values: values,
	}
}

// Tap applies a series of functions to the SQL object, allowing method
// chaining with custom transformations.
func (s *SQL) Tap(funcs ...func(*SQL) *SQL) *SQL {
	for i := range funcs {
		s = funcs[i](s)
	}
	return s
}

// Explain sets up EXPLAIN output collection. When Query, QueryRow, or Execute
// is called, an EXPLAIN statement will be executed first and the result will
// be written to the target. Target can be *string, io.Writer, logger.Logger,
// func(string), or func(...interface{}) (e.g. log.Println).
// Options can include ANALYZE, VERBOSE, BUFFERS, COSTS, TIMING, FORMAT JSON, etc.
func (s *SQL) Explain(target interface{}, options ...string) *SQL {
	s.explainTarget = target
	s.explainOptions = options
	return s
}

// ExplainAnalyze is a shorthand for Explain(target, "ANALYZE", ...).
// Target can be *string, io.Writer, logger.Logger, func(string), or func(...interface{}).
// Note: The ANALYZE option causes the statement to be actually executed,
// not just planned. Use with caution on INSERT, UPDATE, DELETE statements
// as they will modify your data.
func (s *SQL) ExplainAnalyze(target interface{}, options ...string) *SQL {
	return s.Explain(target, append([]string{"ANALYZE"}, options...)...)
}

func (s SQL) formattedSQL() string {
	sql := s.sql
	i := 1
	for range s.values {
		sql = strings.Replace(sql, "$?", fmt.Sprintf("$%d", i), 1)
		i += 1
	}
	return sql
}

func (s SQL) String() string {
	if s.main != nil {
		return s.main.String()
	}
	return s.formattedSQL()
}

func (s SQL) StringValues() (string, []interface{}) {
	if s.main != nil {
		return s.main.StringValues()
	}
	return s.formattedSQL(), s.values
}

func (s SQL) Values() []interface{} {
	return s.values
}

// MustQuery is like Query but panics if query operation fails.
func (s SQL) MustQuery(target interface{}) {
	if err := s.Query(target); err != nil {
		panic(err)
	}
}

// Query executes the SQL query and scans results into target. Target must be
// a pointer to a struct (single row), slice (multiple rows), or map (key-value
// pairs). See Find and Select for examples.
func (s SQL) Query(target interface{}) error {
	return s.QueryCtxTx(context.Background(), nil, target)
}

// MustQueryCtx is like QueryCtx but panics if query operation fails.
func (s SQL) MustQueryCtx(ctx context.Context, target interface{}) {
	if err := s.QueryCtx(ctx, target); err != nil {
		panic(err)
	}
}

// QueryCtx is like Query but accepts a context for cancellation and timeouts.
func (s SQL) QueryCtx(ctx context.Context, target interface{}) error {
	return s.QueryCtxTx(ctx, nil, target)
}

// MustQueryCtxTx is like QueryCtxTx but panics if query operation fails.
func (s SQL) MustQueryCtxTx(ctx context.Context, tx Tx, target interface{}) {
	if err := s.QueryCtxTx(ctx, tx, target); err != nil {
		panic(err)
	}
}

// QueryCtxTx is like Query but accepts a context and optional transaction.
// If tx is non-nil, the query executes within that transaction.
func (s SQL) QueryCtxTx(ctx context.Context, tx Tx, target interface{}) error {
	if s.model.connection == nil {
		return ErrNoConnection
	}

	sqlQuery, values := s.StringValues()
	if sqlQuery == "" {
		return nil
	}

	if err := s.runExplain(ctx, tx, sqlQuery, values); err != nil {
		return err
	}

	var rv reflect.Value
	var rt reflect.Type

	targetIsRV := false
	switch v := target.(type) {
	case *reflect.Value:
		rv = *v
		targetIsRV = true
	case reflect.Value:
		rv = v
		targetIsRV = true
	}

	if targetIsRV {
		rt = rv.Type()
		if rt.Kind() == reflect.Ptr {
			rv = reflect.Indirect(rv)
			rt = rv.Type()
		}
		if !rv.CanAddr() {
			return ErrInvalidTarget
		}
	} else {
		rv = reflect.Indirect(reflect.ValueOf(target))
		rt = reflect.TypeOf(target)
		if rt.Kind() != reflect.Ptr {
			return ErrInvalidTarget
		}
		rt = rt.Elem()
	}

	kind := rt.Kind()
	if kind == reflect.Slice {
		rt = rt.Elem()
	}

	var mi *modelInfo
	if s.model.structType != nil && rt == s.model.structType {
		// use model's existing info if type is the same
		mi = s.model.modelInfo
	} else {
		// different type of struct
		mi = &modelInfo{tableName: s.model.tableName}
		mi.setColumnNamer(s.model.columnNamer)
		mi.updateColumnNames(rt)
	}

	if kind == reflect.Struct { // if target is not a slice, use QueryRow instead
		start := time.Now()
		defer s.log(sqlQuery, values, start)
		if tx != nil {
			return mi.scan(rv, tx.QueryRowContext(ctx, sqlQuery, values...))
		}
		return mi.scan(rv, s.model.connection.QueryRowContext(ctx, sqlQuery, values...))
	} else if kind == reflect.Map {
		start := time.Now()
		defer s.log(sqlQuery, values, start)
		var rows db.Rows
		var err error
		if tx != nil {
			rows, err = tx.QueryContext(ctx, sqlQuery, values...)
		} else {
			rows, err = s.model.connection.QueryContext(ctx, sqlQuery, values...)
		}
		if err != nil {
			return err
		}

		defer rows.Close()
		columns, _ := rows.Columns()
		columnLen := len(columns)
		if rv.IsNil() {
			rv.Set(reflect.MakeMapWithSize(rt, 0))
		}
		mapKeyType, mapValueType := rt.Key(), rt.Elem()
		isSlice := mapValueType.Kind() == reflect.Slice
		valueTypes := mapValueTypes(rt)
		for rows.Next() {
			mapKeys, end, dests := newDestsForMapType(mapKeyType, mapValueType, columnLen)
			if err := rows.Scan(dests...); err != nil {
				return err
			}
			if isSlice {
				slice := rv.MapIndex(mapKeys[0])
				if !slice.IsValid() {
					slice = reflect.MakeSlice(valueTypes[0], 0, 0)
				}
				rv.SetMapIndex(mapKeys[0], reflect.Append(slice, end))
				continue
			}
			subMap := rv
			i := 0
			for ; i < len(mapKeys)-1; i++ { // map[type]map...
				if !subMap.MapIndex(mapKeys[i]).IsValid() {
					subMap.SetMapIndex(mapKeys[i], reflect.MakeMap(valueTypes[i]))
				}
				subMap = subMap.MapIndex(mapKeys[i])
			}
			subMap.SetMapIndex(mapKeys[i], end)
		}
		return rows.Err()
	} else if kind != reflect.Slice {
		return ErrInvalidTarget
	}

	start := time.Now()
	defer s.log(sqlQuery, values, start)
	var rows db.Rows
	var err error
	if tx != nil {
		rows, err = tx.QueryContext(ctx, sqlQuery, values...)
	} else {
		rows, err = s.model.connection.QueryContext(ctx, sqlQuery, values...)
	}
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		nv := reflect.New(rt).Elem()
		if err := mi.scan(nv, rows); err != nil {
			return err
		}
		rv.Set(reflect.Append(rv, nv))
	}
	return rows.Err()
}

// scan a scannable (Row or Rows) into every field of a struct
func (mi *modelInfo) scan(rv reflect.Value, scannable db.Scannable) error {
	if rv.Kind() != reflect.Struct || (len(mi.modelFields) == 0 && len(mi.jsonbColumns) == 0) {
		return scannable.Scan(rv.Addr().Interface())
	}
	f := rv.FieldByName(tableNameField)
	if f.Kind() == reflect.String {
		// hack
		reflect.NewAt(f.Type(), unsafe.Pointer(f.UnsafeAddr())).Elem().SetString(mi.tableName)
	}
	dests := []interface{}{}
	for _, field := range mi.modelFields {
		if field.Jsonb != "" {
			continue
		}
		pointer := field.getFieldValueAddrFromStruct(rv)
		dests = append(dests, pointer)
	}
	jsonbValues := []jsonbRaw{}
	for range mi.jsonbColumns {
		jsonb := jsonbRaw{}
		dests = append(dests, &jsonb)
		jsonbValues = append(jsonbValues, jsonb)
	}
	if err := scannable.Scan(dests...); err != nil {
		return err
	}
	for _, jsonb := range jsonbValues {
		for _, field := range mi.modelFields {
			if field.Jsonb == "" {
				continue
			}
			val, ok := jsonb[field.ColumnName]
			if !ok {
				continue
			}
			pointer := field.getFieldValueAddrFromStruct(rv)
			if err := json.Unmarshal(val, pointer); err != nil {
				if field.Strict {
					return fmt.Errorf("error unmarshaling field %s of %s: %v", field.ColumnName, field.Jsonb, err)
				}
				continue
			}
		}
	}
	return nil
}

// MustQueryRow is like QueryRow but panics if query row operation fails.
func (s SQL) MustQueryRow(dest ...interface{}) {
	if err := s.QueryRow(dest...); err != nil {
		panic(err)
	}
}

// QueryRow executes the query and scans the first row's columns into dest.
// Each dest must be a pointer to a variable that can hold the column value.
//
//	var name string
//	var id int
//	psql.NewModelTable("users", conn).Select("name", "id").MustQueryRow(&name, &id)
func (s SQL) QueryRow(dest ...interface{}) error {
	return s.QueryRowCtxTx(context.Background(), nil, dest...)
}

// MustQueryRowCtx is like QueryRowCtx but panics if query row operation fails.
func (s SQL) MustQueryRowCtx(ctx context.Context, dest ...interface{}) {
	if err := s.QueryRowCtx(ctx, dest...); err != nil {
		panic(err)
	}
}

// QueryRowCtx is like QueryRow but accepts a context for cancellation and
// timeouts.
func (s SQL) QueryRowCtx(ctx context.Context, dest ...interface{}) error {
	return s.QueryRowCtxTx(ctx, nil, dest...)
}

// MustQueryRowCtxTx is like QueryRowCtxTx but panics if query row operation
// fails.
func (s SQL) MustQueryRowCtxTx(ctx context.Context, tx Tx, dest ...interface{}) {
	if err := s.QueryRowCtxTx(ctx, tx, dest...); err != nil {
		panic(err)
	}
}

// QueryRowCtxTx is like QueryRow but accepts a context and optional
// transaction. If tx is non-nil, the query executes within that transaction.
func (s SQL) QueryRowCtxTx(ctx context.Context, tx Tx, dest ...interface{}) error {
	if s.model.connection == nil {
		return ErrNoConnection
	}
	sqlQuery, values := s.StringValues()
	if sqlQuery == "" {
		return nil
	}
	if err := s.runExplain(ctx, tx, sqlQuery, values); err != nil {
		return err
	}
	start := time.Now()
	defer s.log(sqlQuery, values, start)
	if tx != nil {
		return tx.QueryRowContext(ctx, sqlQuery, values...).Scan(dest...)
	}
	return s.model.connection.QueryRowContext(ctx, sqlQuery, values...).Scan(dest...)
}

// MustExecute is like Execute but panics if execute operation fails.
func (s SQL) MustExecute(dest ...interface{}) {
	if err := s.Execute(dest...); err != nil {
		panic(err)
	}
}

// Execute runs an INSERT, UPDATE, or DELETE statement. To get the number of
// rows affected, pass a pointer to an int or int64 as dest.
func (s SQL) Execute(dest ...interface{}) error {
	return s.ExecuteCtxTx(context.Background(), nil, dest...)
}

// MustExecuteCtx is like ExecuteCtx but panics if execute operation fails.
func (s SQL) MustExecuteCtx(ctx context.Context, dest ...interface{}) {
	if err := s.ExecuteCtx(ctx, dest...); err != nil {
		panic(err)
	}
}

// ExecuteCtx is like Execute but accepts a context for cancellation and
// timeouts.
func (s SQL) ExecuteCtx(ctx context.Context, dest ...interface{}) error {
	return s.ExecuteCtxTx(ctx, nil, dest...)
}

// MustExecuteCtxTx is like ExecuteCtxTx but panics if execute operation fails.
func (s SQL) MustExecuteCtxTx(ctx context.Context, tx Tx, dest ...interface{}) {
	if err := s.ExecuteCtxTx(ctx, tx, dest...); err != nil {
		panic(err)
	}
}

// ExecuteCtxTx is like Execute but accepts a context and optional transaction.
// If tx is non-nil, the statement executes within that transaction.
func (s SQL) ExecuteCtxTx(ctx context.Context, tx Tx, dest ...interface{}) error {
	if s.model.connection == nil {
		return ErrNoConnection
	}
	sqlQuery, values := s.StringValues()
	if sqlQuery == "" {
		return ErrNoSQL
	}
	if err := s.runExplain(ctx, tx, sqlQuery, values); err != nil {
		return err
	}
	start := time.Now()
	defer s.log(sqlQuery, values, start)
	if tx != nil {
		return returnRowsAffected(dest)(tx.ExecContext(ctx, sqlQuery, values...))
	}
	return returnRowsAffected(dest)(s.model.connection.ExecContext(ctx, sqlQuery, values...))
}

func (s SQL) log(sql string, args []interface{}, startTime time.Time) {
	s.model.log(sql, args, time.Since(startTime))
}

// runExplain executes EXPLAIN on the given SQL and writes result to explainTarget.
func (s SQL) runExplain(ctx context.Context, tx Tx, sqlQuery string, values []interface{}) error {
	if s.explainTarget == nil {
		return nil
	}
	explainSQL := "EXPLAIN"
	if len(s.explainOptions) > 0 {
		explainSQL += " (" + strings.Join(s.explainOptions, ", ") + ")"
	}
	explainSQL += " " + sqlQuery

	start := time.Now()
	defer s.log(explainSQL, values, start)

	var rows db.Rows
	var err error
	if tx != nil {
		rows, err = tx.QueryContext(ctx, explainSQL, values...)
	} else {
		rows, err = s.model.connection.QueryContext(ctx, explainSQL, values...)
	}
	if err != nil {
		return err
	}
	defer rows.Close()

	var lines []string
	for rows.Next() {
		var line string
		if err := rows.Scan(&line); err != nil {
			return err
		}
		lines = append(lines, line)
	}
	if err := rows.Err(); err != nil {
		return err
	}

	result := strings.Join(lines, "\n")
	switch t := s.explainTarget.(type) {
	case *string:
		*t = result
	case io.Writer:
		_, err = t.Write([]byte(result + "\n"))
		return err
	case logger.Logger:
		t.Debug(result)
	case func(string):
		t(result)
	case func(...interface{}):
		t(result)
	default:
		return ErrUnsupportedExplainTarget
	}
	return nil
}

func returnRowsAffected(dest []interface{}) func(db.Result, error) error {
	return func(result db.Result, err error) error {
		if err != nil {
			return err
		}
		if len(dest) == 0 {
			return nil
		}
		ra, err := result.RowsAffected()
		if err != nil {
			return err
		}
		switch x := dest[0].(type) {
		case *int:
			*x = int(ra)
		case *int64:
			*x = ra
		}
		return nil
	}
}

// Get all element types of a map recursively, for example:
// mapValueTypes(reflect.TypeOf(map[string]map[int]map[bool]int{})) returns:
// [ map[int]map[bool]int, map[bool]int, int ]
func mapValueTypes(mapType reflect.Type) (types []reflect.Type) {
	if mapType.Kind() != reflect.Map {
		return
	}
	mapValueType := mapType.Elem()
	types = append(types, mapValueType)
	types = append(types, mapValueTypes(mapValueType)...)
	return
}

// Make new destination pointers from map type for Scannable. The "end" is the
// last non-map type value. Map keys are paths to the "end" value.
func newDestsForMapType(mapKeyType, mapValueType reflect.Type, columnLen int) (mapKeys []reflect.Value, end reflect.Value, dests []interface{}) {
	isSlice := mapValueType.Kind() == reflect.Slice
	if isSlice {
		mapValueType = mapValueType.Elem()
	}
	newMapKey := reflect.New(mapKeyType).Elem()
	newMapVal := reflect.New(mapValueType).Elem()
	switch mapKeyType.Kind() {
	case reflect.Struct:
		for i := 0; i < columnLen && i < mapKeyType.NumField(); i++ {
			dests = append(dests, getAddrOfStructField(mapKeyType.Field(i), newMapKey.Field(i)))
		}
	case reflect.Array:
		for i := 0; i < columnLen && i < mapKeyType.Len(); i++ {
			dests = append(dests, newMapKey.Index(i).Addr().Interface())
		}
	default:
		dests = append(dests, newMapKey.Addr().Interface())
	}
	mapKeys = append(mapKeys, newMapKey)
	end = newMapVal
	size := columnLen - len(dests)
	switch mapValueType.Kind() {
	case reflect.Struct:
		if size == 1 {
			if dest, ok := newMapVal.Addr().Interface().(sql.Scanner); ok {
				dests = append(dests, dest)
				return
			}
		}
		for i := 0; i < size; i++ {
			dests = append(dests, getAddrOfStructField(mapValueType.Field(i), newMapVal.Field(i)))
		}
	case reflect.Map:
		if isSlice {
			// can't handle this kind of data structure at the moment
			panic("sorry, but map[type][]map... is not yet supported")
		}
		k, e, d := newDestsForMapType(mapValueType.Key(), mapValueType.Elem(), size)
		mapKeys = append(mapKeys, k...)
		end = e
		dests = append(dests, d...)
	case reflect.Slice:
		newMapVal.Set(reflect.MakeSlice(reflect.SliceOf(mapValueType.Elem()), size, size))
		fallthrough
	case reflect.Array:
		for i := 0; i < size; i++ {
			dests = append(dests, newMapVal.Index(i).Addr().Interface())
		}
	default:
		if size > 0 {
			dests = append(dests, newMapVal.Addr().Interface())
		}
	}
	return
}

func getAddrOfStructField(field reflect.StructField, value reflect.Value) interface{} {
	if field.PkgPath == "" {
		return value.Addr().Interface()
	}
	return reflect.NewAt(value.Type(), unsafe.Pointer(value.UnsafeAddr())).Interface()
}