<commit_before>package orm

import (
	"errors"
	"fmt"
	"sync"

	"gopkg.in/pg.v4/internal"
	"gopkg.in/pg.v4/types"
)

type Query struct {
	db    dber
	model tableModel
	err   error

	tableName  types.Q
	tableAlias string

	tables     []byte
	fields     []string
	columns    []byte
	set        []byte
	where      []byte
	join       []byte
	group      []byte
	order      []byte
	onConflict []byte
	returning  []byte
	limit      int
	offset     int
}

func NewQuery(db dber, v ...interface{}) *Query {
	q := Query{
		db: db,
	}
	switch l := len(v); {
	case l == 1:
		v0 := v[0]
		if v0 != nil {
			q.model, q.err = newTableModel(v0)
		}
	case l > 1:
		q.model, q.err = newTableModel(&v)
	}
	if q.model != nil {
		q.tableName = q.FormatQuery(q.tableName, string(q.model.Table().Name))
	}
	return &q
}

func (q *Query) copy() *Query {
	cp := *q
	return &cp
}

func (q *Query) setErr(err error) {
	if q.err == nil {
		q.err = err
	}
}

func (q *Query) Alias(alias string) *Query {
	q.tableAlias = alias
	return q
}

func (q *Query) Table(names ...string) *Query {
	for _, name := range names {
		q.tables = appendSep(q.tables, ", ")
		q.tables = types.AppendField(q.tables, name, 1)
	}
	return q
}

func (q *Query) Column(columns ...string) *Query {
loop:
	for _, column := range columns {
		if q.model != nil {
			if j := q.model.Join(column); j != nil {
				continue loop
			}
		}

		q.fields = append(q.fields, column)
		q.columns = appendSep(q.columns, ", ")
		q.columns = types.AppendField(q.columns, column, 1)
	}
	return q
}

func (q *Query) ColumnExpr(expr string, params ...interface{}) *Query {
	q.columns = appendSep(q.columns, ", ")
	q.columns = q.FormatQuery(q.columns, expr, params...)
	return q
}

func (q *Query) Set(set string, params ...interface{}) *Query {
	q.set = appendSep(q.set, ", ")
	q.set = q.FormatQuery(q.set, set, params...)
	return q
}

func (q *Query) Where(where string, params ...interface{}) *Query {
	q.where = appendSep(q.where, " AND ")
	q.where = append(q.where, '(')
	q.where = q.FormatQuery(q.where, where, params...)
	q.where = append(q.where, ')')
	return q
}

// WhereAnd joins passed conditions using AND operation.
func (q *Query) WhereAnd(conditions ...*SQL) *Query {
	q.where = appendSep(q.where, " AND ")
	q.where = append(q.where, '(')
	for i, cond := range conditions {
		q.where = cond.AppendFormat(q.where, q)
		if i != len(conditions)-1 {
			q.where = append(q.where, " AND "...)
		}
	}
	q.where = append(q.where, ')')
	return q
}

// WhereOr joins passed conditions using OR operation.
func (q *Query) WhereOr(conditions ...*SQL) *Query {
	q.where = appendSep(q.where, " AND ")
	q.where = append(q.where, '(')
	for i, cond := range conditions {
		q.where = cond.AppendFormat(q.where, q)
		if i != len(conditions)-1 {
			q.where = append(q.where, " OR "...)
		}
	}
	q.where = append(q.where, ')')
	return q
}

func (q *Query) Join(join string, params ...interface{}) *Query {
	q.join = appendSep(q.join, " ")
	q.join = q.FormatQuery(q.join, join, params...)
	return q
}

func (q *Query) Group(group string, params ...interface{}) *Query {
	q.group = appendSep(q.group, ", ")
	q.group = q.FormatQuery(q.group, group, params...)
	return q
}

func (q *Query) Order(order string, params ...interface{}) *Query {
	q.order = appendSep(q.order, ", ")
	q.order = q.FormatQuery(q.order, order, params...)
	return q
}

func (q *Query) Limit(n int) *Query {
	q.limit = n
	return q
}

func (q *Query) Offset(n int) *Query {
	q.offset = n
	return q
}

func (q *Query) OnConflict(s string, params ...interface{}) *Query {
	q.onConflict = q.FormatQuery(nil, s, params...)
	return q
}

func (q *Query) Returning(columns ...interface{}) *Query {
	for _, column := range columns {
		q.returning = appendSep(q.returning, ", ")

		switch column := column.(type) {
		case string:
			q.returning = types.AppendField(q.returning, column, 1)
		case types.ValueAppender:
			var err error
			q.returning, err = column.AppendValue(q.returning, 1)
			if err != nil {
				q.setErr(err)
			}
		default:
			q.setErr(fmt.Errorf("unsupported column type: %T", column))
		}
	}
	return q
}

// Count returns number of rows matching the query using count aggregate function.
func (q *Query) Count() (int, error) {
	if q.err != nil {
		return 0, q.err
	}

	q = q.copy()
	q.columns = types.Q("count(*)")
	q.order = nil
	q.limit = 0
	q.offset = 0

	sel := &selectQuery{
		Query: q,
	}
	var count int
	_, err := q.db.QueryOne(Scan(&count), sel, q.model)
	return count, err
}
<commit_msg>Add experimental Apply API.<commit_after>package orm

import (
	"errors"
	"fmt"
	"sync"

	"gopkg.in/pg.v4/internal"
	"gopkg.in/pg.v4/types"
)

type Query struct {
	db    dber
	model tableModel
	err   error

	tableName  types.Q
	tableAlias string

	tables     []byte
	fields     []string
	columns    []byte
	set        []byte
	where      []byte
	join       []byte
	group      []byte
	order      []byte
	onConflict []byte
	returning  []byte
	limit      int
	offset     int
}

func NewQuery(db dber, v ...interface{}) *Query {
	q := Query{
		db: db,
	}
	switch l := len(v); {
	case l == 1:
		v0 := v[0]
		if v0 != nil {
			q.model, q.err = newTableModel(v0)
		}
	case l > 1:
		q.model, q.err = newTableModel(&v)
	}
	if q.model != nil {
		q.tableName = q.FormatQuery(q.tableName, string(q.model.Table().Name))
	}
	return &q
}

func (q *Query) copy() *Query {
	cp := *q
	return &cp
}

func (q *Query) setErr(err error) {
	if q.err == nil {
		q.err = err
	}
}

func (q *Query) Alias(alias string) *Query {
	q.tableAlias = alias
	return q
}

func (q *Query) Table(names ...string) *Query {
	for _, name := range names {
		q.tables = appendSep(q.tables, ", ")
		q.tables = types.AppendField(q.tables, name, 1)
	}
	return q
}

func (q *Query) Column(columns ...string) *Query {
loop:
	for _, column := range columns {
		if q.model != nil {
			if j := q.model.Join(column); j != nil {
				continue loop
			}
		}

		q.fields = append(q.fields, column)
		q.columns = appendSep(q.columns, ", ")
		q.columns = types.AppendField(q.columns, column, 1)
	}
	return q
}

func (q *Query) ColumnExpr(expr string, params ...interface{}) *Query {
	q.columns = appendSep(q.columns, ", ")
	q.columns = q.FormatQuery(q.columns, expr, params...)
	return q
}

func (q *Query) Set(set string, params ...interface{}) *Query {
	q.set = appendSep(q.set, ", ")
	q.set = q.FormatQuery(q.set, set, params...)
	return q
}

func (q *Query) Where(where string, params ...interface{}) *Query {
	q.where = appendSep(q.where, " AND ")
	q.where = append(q.where, '(')
	q.where = q.FormatQuery(q.where, where, params...)
	q.where = append(q.where, ')')
	return q
}

// WhereOr joins passed conditions using OR operation.
func (q *Query) WhereOr(conditions ...*SQL) *Query {
	q.where = appendSep(q.where, " AND ")
	q.where = append(q.where, '(')
	for i, cond := range conditions {
		q.where = cond.AppendFormat(q.where, q)
		if i != len(conditions)-1 {
			q.where = append(q.where, " OR "...)
		}
	}
	q.where = append(q.where, ')')
	return q
}

func (q *Query) Join(join string, params ...interface{}) *Query {
	q.join = appendSep(q.join, " ")
	q.join = q.FormatQuery(q.join, join, params...)
	return q
}

func (q *Query) Group(group string, params ...interface{}) *Query {
	q.group = appendSep(q.group, ", ")
	q.group = q.FormatQuery(q.group, group, params...)
	return q
}

func (q *Query) Order(order string, params ...interface{}) *Query {
	q.order = appendSep(q.order, ", ")
	q.order = q.FormatQuery(q.order, order, params...)
	return q
}

func (q *Query) Limit(n int) *Query {
	q.limit = n
	return q
}

func (q *Query) Offset(n int) *Query {
	q.offset = n
	return q
}

func (q *Query) OnConflict(s string, params ...interface{}) *Query {
	q.onConflict = q.FormatQuery(nil, s, params...)
	return q
}

func (q *Query) Returning(columns ...interface{}) *Query {
	for _, column := range columns {
		q.returning = appendSep(q.returning, ", ")

		switch column := column.(type) {
		case string:
			q.returning = types.AppendField(q.returning, column, 1)
		case types.ValueAppender:
			var err error
			q.returning, err = column.AppendValue(q.returning, 1)
			if err != nil {
				q.setErr(err)
			}
		default:
			q.setErr(fmt.Errorf("unsupported column type: %T", column))
		}
	}
	return q
}

// Apply calls the fn passing the Query as an argument.
func (q *Query) Apply(fn func(*Query) *Query) *Query {
	return fn(q)
}

// Count returns number of rows matching the query using count aggregate function.
func (q *Query) Count() (int, error) {
	if q.err != nil {
		return 0, q.err
	}

	q = q.copy()
	q.columns = types.Q("count(*)")
	q.order = nil
	q.limit = 0
	q.offset = 0

	sel := &selectQuery{
		Query: q,
	}
	var count int
	_, err := q.db.QueryOne(Scan(&count), sel, q.model)
	return count, err
}

{'fraction_non_alphanumeric': 0.17226794695930497, 'fraction_numerical': 0.0037722908093278463, 'mean_word_length': 3.022528735632184, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1589623', 'n_tokens_mistral': 3474, 'n_tokens_neox': 3321, 'n_words': 904}