path: root/query.go

// SPDX-FileCopyrightText: 2026 Stefan Majewsky <majewsky@gmx.net>
// SPDX-License-Identifier: Apache-2.0

package oblast

import (
	"context"
	"database/sql"
	"fmt"
	"reflect"
)

// PrepareThreshold is a tuning parameter for the strategy used by all methods of [Store] operating on batches of records provided by the caller
// (specifically, [Store.Insert], [Store.Update] and [Store.Delete]).
//
// For large amounts of records, it is obviously advantageous to build a prepared statement for the one query that will be used repeatedly on all of them.
// However, building a prepared statement is associated with some amount of bookkeeping on the level of the database/sql library.
// When operating on individual records or small amounts of records at a time (that is, in OLTP rather than OLAP workloads), this overhead becomes a measurable performance burden.
//
// This tuning parameter defines the minimum number of records that will justify maintaining a prepared statement.
// Our benchmarking with the mattn/go-sqlite3 driver (and last checked with Go 1.26.2 on x86_64) indicates that this becomes a worthwhile investment at 8 or more records, so this is our default.
// If your benchmarking indicates a different tradeoff depending on your choice of Go version or SQL driver, you may adjust this variable accordingly.
var PrepareThreshold int = 8

// preparedStatement behaves like sql.Stmt, but only uses *sql.Stmt when it is useful (see explanation above).
type preparedStatement struct {
	db    Handle
	query string
	stmt  *sql.Stmt // nil for input sizes below PrepareThreshold
}

// prepare behaves like [Handle.Prepare].
func prepare(ctx context.Context, db Handle, query, operation string, inputSize int) (preparedStatement, error) {
	if query == "" {
		return preparedStatement{}, fmt.Errorf("cannot execute %s() because query could not be autogenerated", operation)
	}

	if inputSize < PrepareThreshold {
		return preparedStatement{db, query, nil}, nil
	}
	stmt, err := db.PrepareContext(ctx, query)
	if err != nil {
		return preparedStatement{}, fmt.Errorf("during Prepare(): %w", err)
	}
	return preparedStatement{db, query, stmt}, nil
}

// Close behaves like [sql.Stmt.Close].
func (s preparedStatement) Close() error {
	if s.stmt == nil {
		return nil
	}
	return s.stmt.Close()
}

// ExecContext behaves like [sql.Stmt.ExecContext].
func (s preparedStatement) ExecContext(ctx context.Context, args ...any) (sql.Result, error) {
	if s.stmt == nil {
		return s.db.ExecContext(ctx, s.query, args...)
	}
	return s.stmt.ExecContext(ctx, args...)
}

// QueryRow behaves like [sql.Stmt.QueryRowContext].
func (s preparedStatement) QueryRowContext(ctx context.Context, args ...any) *sql.Row {
	if s.stmt == nil {
		return s.db.QueryRowContext(ctx, s.query, args...)
	}
	return s.stmt.QueryRowContext(ctx, args...)
}

// Insert executes an SQL INSERT statement for each of the provided records.
//
// Fields that are declared with the "auto" tag will not be written into the DB,
// and instead their value (as auto-generated by the DB on insert) will be placed in the record.
//
// Returns an error if [NewStore] was called without the [TableNameIs] option, which is required to generate a query for this method.
//
// Returns an error if any of the `records` has a non-zero value in any column marked as `db:",auto"`.
// Records that already exist in the database should be handled with [Store.Update] instead.
func (s Store[R]) Insert(ctx context.Context, db Handle, records ...*R) error {
	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
	//       Any expression that does not depend on type R should be factored out into a reusable function.

	stmt, err := prepare(ctx, db, s.plan.Insert.Query, "Insert", len(records))
	if err != nil {
		return err
	}
	return s.insertUsing(ctx, stmt, db, records)
}

func (s Store[R]) insertUsing(ctx context.Context, stmt preparedStatement, db Handle, records []*R) error {
	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
	//       Any expression that does not depend on type R should be factored out into a reusable function.

	var (
		argumentIndexes = s.plan.Insert.ArgumentIndexes
		argumentSlots   = make([]any, len(argumentIndexes))
		scanIndexes     = s.plan.Insert.ScanIndexes
		scanSlots       = make([]any, len(scanIndexes))
	)

	for idx, r := range records {
		v := reflect.ValueOf(r).Elem()
		err := insertRecord(ctx, v, idx, stmt, argumentIndexes, argumentSlots, scanIndexes, scanSlots)
		if err != nil {
			return newIOError(err, "Stmt.Close", stmt.Close())
		}
	}

	return newIOError(nil, "Stmt.Close", stmt.Close())
}

func insertRecord(ctx context.Context, v reflect.Value, recordIndex int, stmt preparedStatement, argumentIndexes [][]int, argumentSlots []any, scanIndexes [][]int, scanSlots []any) error {
	for idx, index := range argumentIndexes {
		argumentSlots[idx] = v.FieldByIndex(index).Interface()
	}
	for idx, index := range scanIndexes {
		f := v.FieldByIndex(index)
		if !f.IsZero() {
			return fmt.Errorf(`refusing to INSERT record with idx = %d that already has non-zero values in its "auto" columns`, recordIndex)
		}
		scanSlots[idx] = f.Addr().Interface()
	}
	var err error
	if len(scanSlots) == 0 {
		_, err = stmt.ExecContext(ctx, argumentSlots...)
	} else {
		// TODO: using QueryRow for inserting is extremely expensive because database/sql allocates a Rows instance under the hood; other libraries are doing better by limiting themselves to ExecContext() + LastInsertId()
		err = stmt.QueryRowContext(ctx, argumentSlots...).Scan(scanSlots...)
	}
	if err != nil {
		return fmt.Errorf("while inserting record with idx = %d: %w", recordIndex, err)
	}
	return nil
}

// Update executes an SQL UPDATE statement for each of the provided records, updating all non-primary-key columns with the values in the records.
// Returns [MissingRecordError] if any of the records does not exist in the database, that is, if for any of the records, the database contains no row with the same primary key values.
//
// Returns an error if [NewStore] was called without the [TableNameIs] or [PrimaryKeyIs] options, which are both required to generate a query for this method.
func (s Store[R]) Update(ctx context.Context, db Handle, records ...R) error {
	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
	//       Any expression that does not depend on type R should be factored out into a reusable function.

	stmt, err := prepare(ctx, db, s.plan.Update.Query, "Update", len(records))
	if err != nil {
		return err
	}

	var (
		argumentIndexes = s.plan.Update.ArgumentIndexes
		argumentSlots   = make([]any, len(argumentIndexes))
	)

	for idx := range records {
		v := reflect.ValueOf(&records[idx]).Elem()
		rowsAffected, err := updateRecord(ctx, v, idx, stmt, argumentIndexes, argumentSlots)
		if err == nil && rowsAffected == 0 {
			err = MissingRecordError[R]{records[idx], s.plan}
		}
		if err != nil {
			return newIOError(err, "Stmt.Close", stmt.Close())
		}
	}
	return newIOError(nil, "Stmt.Close", stmt.Close())
}

func updateRecord(ctx context.Context, v reflect.Value, recordIndex int, stmt preparedStatement, argumentIndexes [][]int, argumentSlots []any) (int64, error) {
	for idx, index := range argumentIndexes {
		argumentSlots[idx] = v.FieldByIndex(index).Interface()
	}
	result, err := stmt.ExecContext(ctx, argumentSlots...)
	if err != nil {
		return 0, fmt.Errorf("while updating record with idx = %d: %w", recordIndex, err)
	}
	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("during RowsAffected() for record with idx = %d: %w", recordIndex, err)
	}
	return rowsAffected, nil
}

// Delete executes an SQL DELETE statement for each of the provided records, using their primary keys to locate the respective table rows.
//
// Returns an error if [NewStore] was called without the [TableNameIs] or [PrimaryKeyIs] options, which are both required to generate a query for this method.
func (s Store[R]) Delete(ctx context.Context, db Handle, records ...R) error {
	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
	//       Any expression that does not depend on type R should be factored out into a reusable function.

	stmt, err := prepare(ctx, db, s.plan.Delete.Query, "Delete", len(records))
	if err != nil {
		return err
	}

	var (
		argumentIndexes = s.plan.Delete.ArgumentIndexes
		argumentSlots   = make([]any, len(argumentIndexes))
	)

	for idx := range records {
		v := reflect.ValueOf(&records[idx]).Elem()
		err := deleteRecord(ctx, v, idx, stmt, argumentIndexes, argumentSlots)
		if err != nil {
			return newIOError(err, "Stmt.Close", stmt.Close())
		}
	}

	return newIOError(nil, "Stmt.Close", stmt.Close())
}

func deleteRecord(ctx context.Context, v reflect.Value, recordIndex int, stmt preparedStatement, argumentIndexes [][]int, argumentSlots []any) error {
	for idx, index := range argumentIndexes {
		argumentSlots[idx] = v.FieldByIndex(index).Interface()
	}
	_, err := stmt.ExecContext(ctx, argumentSlots...)
	if err != nil {
		return fmt.Errorf("while deleting record with idx = %d: %w", recordIndex, err)
	}
	return nil
}

// Upsert executes either an SQL INSERT or UPDATE statement for each of the provided records,
// based on whether the record already exists in the DB or not.
//
//   - For record types that have fields declared with the "auto" tag, INSERT is chosen iff those fields hold zero values.
//     Returns an error if only some of the respective fields hold zero values while others don't.
//     Returns an error if [NewStore] was called without the [TableNameIs] or [PrimaryKeyIs] options, which are both required to generate the respective queries for this method.
//   - For record types that do not have fields declared with the "auto" tag, an INSERT ... ON CONFLICT statement is used.
//     Returns an error if [NewStore] was called without the [TableNameIs] option, which is required to generate a query for this method.
func (s Store[R]) Upsert(ctx context.Context, db Handle, records ...*R) error {
	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
	//       Any expression that does not depend on type R should be factored out into a reusable function.

	if len(s.plan.AutoColumnNames) == 0 {
		stmt, err := prepare(ctx, db, s.plan.Upsert.Query, "Upsert", len(records))
		if err != nil {
			return err
		}
		return s.insertUsing(ctx, stmt, db, records)
	}

	// TODO: respect PrepareThreshold (or not? may be too much bookkeeping overhead for not a whole lot of benefit)
	insertStmt, err := prepare(ctx, db, s.plan.Insert.Query, "Insert", 0)
	if err != nil {
		return err
	}
	updateStmt, err := prepare(ctx, db, s.plan.Update.Query, "Update", 0)
	if err != nil {
		return err
	}

	var (
		insertArgumentIndexes = s.plan.Insert.ArgumentIndexes
		insertArgumentSlots   = make([]any, len(insertArgumentIndexes))
		insertScanIndexes     = s.plan.Insert.ScanIndexes
		insertScanSlots       = make([]any, len(insertScanIndexes))
		updateArgumentIndexes = s.plan.Update.ArgumentIndexes
		updateArgumentSlots   = make([]any, len(updateArgumentIndexes))
	)

	for idx, r := range records {
		v := reflect.ValueOf(r).Elem()
		isInsert, err := upsertDecideStrategy(v, idx, insertScanIndexes)
		if err != nil {
			return err
		}

		if isInsert {
			err = insertRecord(ctx, v, idx, insertStmt, insertArgumentIndexes, insertArgumentSlots, insertScanIndexes, insertScanSlots)
		} else {
			var rowsAffected int64
			rowsAffected, err = updateRecord(ctx, v, idx, updateStmt, updateArgumentIndexes, updateArgumentSlots)
			if err == nil && rowsAffected == 0 {
				err = MissingRecordError[R]{*r, s.plan}
			}
		}
		if err != nil {
			err = newIOError(err, "InsertStmt.Close", insertStmt.Close())
			err = newIOError(err, "UpdateStmt.Close", updateStmt.Close())
			return err
		}
	}

	err = newIOError(err, "InsertStmt.Close", insertStmt.Close())
	err = newIOError(err, "UpdateStmt.Close", updateStmt.Close())
	return err
}

func upsertDecideStrategy(v reflect.Value, recordIndex int, scanIndexes [][]int) (isInsert bool, err error) {
	var isUpdate bool
	for _, index := range scanIndexes {
		if v.FieldByIndex(index).IsZero() {
			isInsert = true
		} else {
			isUpdate = true
		}
	}
	if isInsert && isUpdate {
		return false, fmt.Errorf(`cannot decide whether to INSERT or UPDATE record with idx = %d: some "auto" columns are zero, others are not`, recordIndex)
	}
	return isInsert, nil
}