sql_select.go

package psql

import (
	"context"
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

type (
	// SelectSQL represents a SELECT query builder. Create instances using
	// Model.Find, Model.Select, or SQL.AsSelect.
	SelectSQL struct {
		*SQL
		sqlConditions
		sqlHavings
		fields  []string
		jfCount int // jsonb fields count
		from    string
		join    string
		with    string
		groupBy string
		orderBy string
		limit   string
		offset  string
	}

	sqlConditions struct {
		conditions []string
		args       []interface{}
	}

	sqlHavings struct {
		havings []string
	}
)

// AsSelect converts a raw SQL statement to a SelectSQL builder. Optional field
// names are used as the initial SELECT columns.
func (s SQL) AsSelect(fields ...string) *SelectSQL {
	f := &SelectSQL{
		SQL:    &s,
		fields: fields,
	}
	f.SQL.main = f
	return f
}

func (m Model) newSelect(fields ...string) *SelectSQL {
	return m.NewSQL("").AsSelect(fields...)
}

// Find creates a SELECT query for all columns defined in the Model's struct.
// Results are scanned into the same struct type or a slice of that type.
//
//	// Query multiple rows into a slice
//	var users []User
//	psql.NewModel(User{}, conn).Find().MustQuery(&users)
//
//	// Query single row into a struct
//	var user User
//	psql.NewModel(User{}, conn).Find().Where("id = $1", 1).MustQuery(&user)
//
// Options can include field transformation functions like AddTableName, which
// prefixes each column with the table name for use in JOIN queries.
func (m Model) Find(options ...interface{}) *SelectSQL {
	return m.newSelect().Find(options...)
}

// Select creates a SELECT query with the specified columns. Unlike Find, the
// result can be scanned into various types including slices, maps, and custom
// structs.
//
//	// Query into a slice of single values
//	var names []string
//	users.Select("name").OrderBy("id ASC").MustQuery(&names)
//
//	// Query into a slice of custom structs
//	var results []struct{ Name string; ID int }
//	users.Select("name", "id").MustQuery(&results)
//
//	// Query into a map for key-value lookup
//	var id2name map[int]string
//	users.Select("id", "name").MustQuery(&id2name)
//
//	// Query into a map with slice values for one-to-many grouping
//	var byCity map[string][]struct{ ID int; Name string }
//	users.Select("city", "id", "name").MustQuery(&byCity)
func (m Model) Select(fields ...string) *SelectSQL {
	return m.newSelect(fields...)
}

// From creates a SELECT query with additional FROM items (tables or subqueries).
func (m Model) From(items ...string) *SelectSQL {
	return m.newSelect().From(items...)
}

// Join creates a SELECT query with JOIN clauses.
func (m Model) Join(expressions ...string) *SelectSQL {
	return m.newSelect().Join(expressions...)
}

// With creates a SELECT query with a CTE (Common Table Expression).
func (m Model) With(expression string, args ...interface{}) *SelectSQL {
	return m.newSelect().With(expression, args...)
}

// WITH creates a SELECT query with a named CTE from another SelectSQL query.
func (m Model) WITH(name string, sql *SelectSQL) *SelectSQL {
	return m.newSelect().WITH(name, sql)
}

// Where creates a SELECT query with a WHERE condition. Use $1, $2 for
// positional parameters, or $? which is auto-replaced when a single argument
// is provided.
func (m Model) Where(condition string, args ...interface{}) *SelectSQL {
	return m.newSelect().Where(condition, args...)
}

// WHERE creates a SELECT query with conditions specified as field/operator/value
// tuples. Each tuple consists of three consecutive arguments.
//
//	// Single condition
//	users.WHERE("id", "=", 1)
//
//	// Multiple conditions (AND)
//	users.WHERE("status", "=", "active", "age", ">=", 18)
func (m Model) WHERE(args ...interface{}) *SelectSQL {
	return m.newSelect().WHERE(args...)
}

// Find populates the SELECT clause with all columns from the Model's struct.
// Options can include transformation functions like AddTableName or the string
// "--no-reset" to append rather than replace existing columns.
func (s *SelectSQL) Find(options ...interface{}) *SelectSQL {
	fields := []string{}
	for _, field := range s.model.modelFields {
		if field.Jsonb != "" {
			continue
		}
		fields = append(fields, field.ColumnName)
	}
	s.jfCount = 0
	for _, jsonbField := range s.model.jsonbColumns {
		fields = append(fields, jsonbField)
		s.jfCount += 1
	}
	var noReset bool
	for _, opts := range options {
		switch f := opts.(type) {
		case fieldsFunc:
			fields = f(fields, s.model.tableName)
		case string:
			if f == "--no-reset" {
				noReset = true
			}
		}
	}
	if noReset {
		return s.Select(fields...)
	}
	return s.ResetSelect(fields...)
}

// Update converts this SelectSQL to an UpdateSQL, preserving WHERE conditions.
func (s *SelectSQL) Update(lotsOfChanges ...interface{}) *UpdateSQL {
	n := s.model.Update(lotsOfChanges...)
	n.conditions = s.conditions
	n.args = s.args
	return n
}

// Delete converts this SelectSQL to a DeleteSQL, preserving WHERE conditions.
func (s *SelectSQL) Delete() *DeleteSQL {
	n := s.model.Delete()
	n.conditions = s.conditions
	n.args = s.args
	return n
}

// MustExists is like Exists but panics if existence check operation fails.
// Returns true if record exists, false if not exists.
func (s *SelectSQL) MustExists() bool {
	return s.MustExistsCtxTx(context.Background(), nil)
}

// MustExistsCtxTx is like ExistsCtxTx but panics if existence check operation fails.
// Returns true if record exists, false if not exists.
func (s *SelectSQL) MustExistsCtxTx(ctx context.Context, tx Tx) bool {
	exists, err := s.ExistsCtxTx(ctx, tx)
	if err != nil {
		panic(err)
	}
	return exists
}

// Exists executes a SELECT 1 query and returns true if at least one row
// matches the current conditions.
func (s *SelectSQL) Exists() (exists bool, err error) {
	return s.ExistsCtxTx(context.Background(), nil)
}

// ExistsCtxTx is like Exists but accepts a context and optional transaction.
func (s *SelectSQL) ExistsCtxTx(ctx context.Context, tx Tx) (exists bool, err error) {
	var ret int
	err = s.ResetSelect("1 AS one").QueryRowCtxTx(ctx, tx, &ret)
	if err == s.model.connection.ErrNoRows() {
		err = nil
		return
	}
	exists = ret == 1
	return
}

// MustCount is like Count but panics if count operation fails.
func (s *SelectSQL) MustCount(optional ...string) int {
	return s.MustCountCtxTx(context.Background(), nil, optional...)
}

// MustCountCtxTx is like CountCtxTx but panics if count operation fails.
func (s *SelectSQL) MustCountCtxTx(ctx context.Context, tx Tx, optional ...string) int {
	count, err := s.CountCtxTx(ctx, tx, optional...)
	if err != nil {
		panic(err)
	}
	return count
}

// Count executes a SELECT COUNT(*) query and returns the number of matching
// rows. Pass a custom expression for different counting, e.g.,
// Count("COUNT(DISTINCT author_id)").
func (s *SelectSQL) Count(optional ...string) (count int, err error) {
	return s.CountCtxTx(context.Background(), nil, optional...)
}

// CountCtxTx is like Count but accepts a context and optional transaction.
func (s *SelectSQL) CountCtxTx(ctx context.Context, tx Tx, optional ...string) (count int, err error) {
	var expr string
	if len(optional) > 0 && optional[0] != "" {
		expr = optional[0]
	} else {
		expr = "COUNT(*)"
	}
	err = s.ResetSelect(expr).QueryRowCtxTx(ctx, tx, &count)
	return
}

// ResetSelect replaces all SELECT columns with the given expressions.
func (s *SelectSQL) ResetSelect(expressions ...string) *SelectSQL {
	s.fields = expressions
	return s
}

// Select appends columns to the SELECT clause, inserting before JSONB columns.
func (s *SelectSQL) Select(expressions ...string) *SelectSQL {
	if s.jfCount > 0 {
		idx := len(s.fields) - s.jfCount
		s.fields = append(append(append([]string{}, s.fields[:idx]...), expressions...), s.fields[idx:]...)
	} else {
		s.fields = append(s.fields, expressions...)
	}
	return s
}

// ReplaceSelect replaces occurrences of old column name with new in the
// SELECT clause.
func (s *SelectSQL) ReplaceSelect(old, new string) *SelectSQL {
	for i := range s.fields {
		if s.fields[i] == old {
			s.fields[i] = new
		}
	}
	return s
}

// GroupBy adds a GROUP BY clause to the query.
func (s *SelectSQL) GroupBy(expressions ...string) *SelectSQL {
	s.groupBy = strings.Join(expressions, ", ")
	return s
}

// Having adds a HAVING clause to the query. Use $1, $2 for positional
// parameters, or $? which is auto-replaced when a single argument is provided.
func (s *SelectSQL) Having(condition string, args ...interface{}) *SelectSQL {
	s.args = append(s.args, args...)
	if len(args) == 1 {
		condition = strings.Replace(condition, "$?", fmt.Sprintf("$%d", len(s.args)), -1)
	}
	s.havings = append(s.havings, condition)
	return s
}

// OrderBy adds an ORDER BY clause to the query.
func (s *SelectSQL) OrderBy(expressions ...string) *SelectSQL {
	s.orderBy = strings.Join(expressions, ", ")
	return s
}

// Limit adds a LIMIT clause to the query. Pass nil to remove the limit.
func (s *SelectSQL) Limit(count interface{}) *SelectSQL {
	if count == nil {
		s.limit = ""
	} else {
		s.limit = fmt.Sprint(count)
	}
	return s
}

// Offset adds an OFFSET clause to the query. Pass nil to remove the offset.
func (s *SelectSQL) Offset(start interface{}) *SelectSQL {
	if start == nil {
		s.offset = ""
	} else {
		s.offset = fmt.Sprint(start)
	}
	return s
}

// Where adds a WHERE condition to the query. Multiple calls are combined with
// AND. Use $1, $2 for positional parameters, or $? for auto-replacement.
func (s *SelectSQL) Where(condition string, args ...interface{}) *SelectSQL {
	s.args = append(s.args, args...)
	if len(args) == 1 {
		condition = strings.Replace(condition, "$?", fmt.Sprintf("$%d", len(s.args)), -1)
	}
	s.conditions = append(s.conditions, condition)
	return s
}

// WHERE adds conditions from field/operator/value tuples. Each tuple consists
// of three consecutive arguments: field name, operator, and value. Multiple
// tuples are combined with AND.
func (s *SelectSQL) WHERE(args ...interface{}) *SelectSQL {
	for i := 0; i < len(args)/3; i++ {
		var column string
		if c, ok := args[i*3].(string); ok {
			column = c
		}
		var operator string
		if o, ok := args[i*3+1].(string); ok {
			operator = o
		}
		if column == "" || operator == "" {
			continue
		}
		s.args = append(s.args, args[i*3+2])
		s.conditions = append(s.conditions, fmt.Sprintf("%s %s $%d", s.model.ToColumnName(column), operator, len(s.args)))
	}
	return s
}

// ResetFrom replaces the FROM clause with the given items.
func (s *SelectSQL) ResetFrom(items ...string) *SelectSQL {
	s.from = strings.Join(items, ", ")
	return s
}

// From appends items to the FROM clause.
func (s *SelectSQL) From(items ...string) *SelectSQL {
	if s.from == "" {
		s.from = s.model.tableName
	}
	if s.from != "" {
		s.from += ", "
	}
	s.from += strings.Join(items, ", ")
	return s
}

// ResetJoin replaces all JOIN clauses with the given expressions.
func (s *SelectSQL) ResetJoin(expressions ...string) *SelectSQL {
	s.join = strings.Join(expressions, " ")
	return s
}

// Join appends JOIN clauses to the query.
func (s *SelectSQL) Join(expressions ...string) *SelectSQL {
	if s.join != "" && !strings.HasSuffix(s.join, " ") {
		s.join += " "
	}
	s.join += strings.Join(expressions, " ")
	return s
}

// With adds a CTE (Common Table Expression) to the query.
func (s *SelectSQL) With(expression string, args ...interface{}) *SelectSQL {
	i := 1
	for range args {
		expression = strings.Replace(expression, "$?", fmt.Sprintf("$%d", i), 1)
		i += 1
	}
	if offset := len(s.args); offset > 0 {
		re := regexp.MustCompile(`\$(\d+)`)
		expression = re.ReplaceAllStringFunc(expression, func(s string) string {
			num, err := strconv.Atoi(s[1:])
			if err != nil { // this should not happen
				panic(err)
			}
			return fmt.Sprintf("$%d", num+offset)
		})
	}
	if s.with != "" {
		s.with += ", "
	}
	s.with += expression
	s.args = append(s.args, args...)
	return s
}

// WITH adds a named CTE from another SelectSQL query. The name can include
// "AS MATERIALIZED" or "AS NOT MATERIALIZED" for PostgreSQL 12+.
func (s *SelectSQL) WITH(name string, sql *SelectSQL) *SelectSQL {
	sqlQuery := sql.String()
	if offset := len(s.args); offset > 0 {
		re := regexp.MustCompile(`\$(\d+)`)
		sqlQuery = re.ReplaceAllStringFunc(sqlQuery, func(s string) string {
			num, err := strconv.Atoi(s[1:])
			if err != nil { // this should not happen
				panic(err)
			}
			return fmt.Sprintf("$%d", num+offset)
		})
	}
	if s.with != "" {
		s.with += ", "
	}
	if strings.Contains(strings.ToLower(name), " as") {
		s.with += name + " (" + sqlQuery + ")"
	} else {
		s.with += name + " AS (" + sqlQuery + ")"
	}
	s.args = append(s.args, sql.args...)
	return s
}

// Tap applies transformation functions to this SelectSQL, enabling custom
// method chaining.
func (s *SelectSQL) Tap(funcs ...func(*SelectSQL) *SelectSQL) *SelectSQL {
	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 *SelectSQL) Explain(target interface{}, options ...string) *SelectSQL {
	s.SQL.Explain(target, 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 *SelectSQL) ExplainAnalyze(target interface{}, options ...string) *SelectSQL {
	s.SQL.ExplainAnalyze(target, options...)
	return s
}

func (s *SelectSQL) String() string {
	var sql string
	if s.with != "" {
		sql += "WITH " + s.with + " "
	}
	if s.sql != "" {
		sql += s.formattedSQL()
	} else {
		sql += "SELECT " + strings.Join(s.fields, ", ") + " FROM "
		if s.from != "" {
			sql += s.from
		} else {
			sql += s.model.tableName
		}
	}
	if s.join != "" {
		sql += " " + s.join
	}
	sql += s.where()
	if s.groupBy != "" {
		sql += " GROUP BY " + s.groupBy + s.having()
	}
	if s.orderBy != "" {
		sql += " ORDER BY " + s.orderBy
	}
	if s.limit != "" {
		sql += " LIMIT " + s.limit
	}
	if s.offset != "" {
		sql += " OFFSET " + s.offset
	}
	return sql
}

func (s *SelectSQL) StringValues() (string, []interface{}) {
	return s.model.convertValues(s.String(), s.args)
}

func (s sqlConditions) where() string {
	return conditionsToStr(s.conditions, " WHERE ")
}

func (s sqlHavings) having() string {
	return conditionsToStr(s.havings, " HAVING ")
}

func conditionsToStr(conds []string, prefix string) (out string) {
	moreThanOne := len(conds) > 1
	for i, conf := range conds {
		if i > 0 {
			out += " AND "
		}
		if moreThanOne {
			out += "(" + conf + ")"
		} else {
			out += conf
		}
	}
	if out != "" {
		out = prefix + out
	}
	return
}