exec.go

package sqlpro

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"maps"
	"reflect"
	"slices"
	"strings"

	"github.com/jackc/pgx/v5"
	"github.com/pkg/errors"
)

// checkData checks that the given data is either one of:
//
// *[]*strcut
// *[]struct
// []*struct
// []struct
// *struct
//
// For structs the function returns true, nil, for slices false, nil

func checkData(data any) (rv reflect.Value, structMode bool, err error) {
	erro := func() (reflect.Value, bool, error) {
		return rv, false, fmt.Errorf("Insert/Update needs a struct or slice of structs.")
	}

	rv = reflect.Indirect(reflect.ValueOf(data))
	switch rv.Type().Kind() {
	case reflect.Slice:
		switch rv.Type().Elem().Kind() {
		case reflect.Ptr:
			if rv.Type().Elem().Elem().Kind() != reflect.Struct {
				return erro()
			}
		case reflect.Interface, reflect.Struct:
		default:
			return rv, false, fmt.Errorf("Insert/Update needs a slice of structs. Have: %s", rv.Type().Elem().Kind())
		}
	case reflect.Struct:
		structMode = true
	default:
		return erro()
	}

	return rv, structMode, nil
}

func (db2 *db) Insert(table string, data any) error {
	return db2.InsertContext(context.Background(), table, data)
}

// Insert takes a table name and a struct and inserts
// the record in the DB.
// The given data needs to be:
//
// *[]*strcut
// *[]struct
// []*struct
// []struct
// struct
// *struct
//
// sqlpro will executes one INSERT statement per row.
// result.LastInsertId will be used to set the first primary
// key column.

func (db2 *db) InsertContext(ctx context.Context, table string, data any) error {
	var (
		rv         reflect.Value
		structMode bool
		err        error
	)

	rv, structMode, err = checkData(data)
	if err != nil {
		return err
	}

	if !structMode {
		for i := 0; i < rv.Len(); i++ {
			row := reflect.Indirect(rv.Index(i))
			insert_id, structInfo, err := db2.insertStruct(ctx, table, row.Interface())
			if err != nil {
				return err
			}
			pk := structInfo.onlyPrimaryKey()
			if pk != nil && pk.structField.Type.Kind() == reflect.Int64 {
				setPrimaryKey(row.FieldByName(pk.name), insert_id)
			}
		}
	} else {
		insert_id, structInfo, err := db2.insertStruct(ctx, table, rv.Interface())
		if err != nil {
			return err
		}
		pk := structInfo.onlyPrimaryKey()
		// log.Printf("PK: %d", insert_id)
		if pk != nil && rv.CanAddr() {
			switch pk.structField.Type.Kind() {
			case reflect.Int,
				reflect.Int8,
				reflect.Int16,
				reflect.Int32,
				reflect.Int64,
				reflect.Uint,
				reflect.Uint8,
				reflect.Uint16,
				reflect.Uint32,
				reflect.Uint64:
				setPrimaryKey(rv.FieldByName(pk.name), insert_id)
			}
		}
	}

	// data
	return nil
}

func setPrimaryKey(rv reflect.Value, id int64) {
	switch rv.Type().Kind() {
	case reflect.Int,
		reflect.Int8,
		reflect.Int16,
		reflect.Int32,
		reflect.Int64:
		rv.SetInt(id)
	case reflect.Uint,
		reflect.Uint8,
		reflect.Uint16,
		reflect.Uint32,
		reflect.Uint64:
		rv.SetUint(uint64(id))
	default:
		err := fmt.Errorf("Unknown type to set primary key: %s", rv.Type())
		panic(err)
	}
}

func (db2 *db) InsertBulk(table string, data any) error {
	return db2.InsertBulkContext(context.Background(), table, data)
}

// InsertBulk takes a table name and a slice of struct and inserts
// the record in the DB with one Exec.
// The given data needs to be:
//
// *[]*struct
// *[]struct
// []*struct
// []struct
//
// sqlpro will executes one INSERT statement per call.
func (db2 *db) InsertBulkContext(ctx context.Context, table string, data any) error {
	return db2.insertBulkContext(ctx, table, data, false, nil)
}

// InsertBulkOnConflictDoNothingContext works like InsertBulkContext but adds a
// "ON CONFLICT DO NOTHING" to the insert command.
func (db2 *db) InsertBulkOnConflictDoNothingContext(ctx context.Context, table string, data any, cols ...string) error {
	return db2.insertBulkContext(ctx, table, data, true, cols)
}

type copyFromData struct {
	keyMap      map[string]*fieldInfo
	columns     []string
	rows        []map[string]any
	values      []any
	nextCounter int
	db          *db
}

func newCopyFromData(db *db, keyMap map[string]*fieldInfo, rows []map[string]any) (cfd *copyFromData) {
	return &copyFromData{
		keyMap:      keyMap,
		columns:     slices.Collect(maps.Keys(keyMap)),
		rows:        rows,
		nextCounter: -1,
		db:          db,
	}
}

func (cfd copyFromData) Columns() []string {
	return cfd.columns
}

func (cfd copyFromData) Len() int64 {
	return int64(len(cfd.rows))
}

func (cfd *copyFromData) Next() bool {
	if len(cfd.rows)-1 > cfd.nextCounter {
		cfd.nextCounter++
		return true
	}
	return false
}

func (cfd copyFromData) Err() error {
	return nil
}

func (cfd copyFromData) Values() (values []any, err error) {
	defer func() {
		r := recover()
		if r == nil {
			return
		}
		err = fmt.Errorf("panic: %v", r)
	}()

	values = make([]any, len(cfd.columns))
	row := cfd.rows[cfd.nextCounter]
	for idx, col := range cfd.columns {
		values[idx] = cfd.db.valueForInsert(row[col], cfd.keyMap[col])
	}
	return values, nil
}

func (db2 *db) copyFrom(ctx context.Context, table string, data *copyFromData) error {
	pgxConn := db2.pgxConn()
	if pgxConn == nil {
		panic("copyFrom needs pgx connection")
	}

	rowsAffected, err := pgxConn.CopyFrom(ctx, pgx.Identifier{table}, data.Columns(), data)
	if err != nil {
		return err
	}
	if rowsAffected != data.Len() {
		err = ErrMismatchedRowsAffected
	}
	return nil
}

func (db2 *db) insertBulkContext(ctx context.Context, table string, data any, onConflictDoNothing bool, conflictCols []string) (err error) {
	var (
		rv         reflect.Value
		structMode bool
	)

	rv, structMode, err = checkData(data)
	if err != nil {
		return err
	}

	if structMode {
		return fmt.Errorf("InsertBulk: Need Slice to insert bulk.")
	}

	key_map := map[string]*fieldInfo{}
	rows := []map[string]any{}

	if rv.Len() == 0 {
		return nil
	}

	for i := 0; i < rv.Len(); i++ {
		row := reflect.Indirect(rv.Index(i)).Interface()
		values, structInfo, err := db2.valuesFromStruct(row)
		if err != nil {
			return fmt.Errorf("sqlpro.insertBulkContext: %w", err)
		}
		rows = append(rows, values)
		for key := range values {
			key_map[key] = structInfo[key]
		}
	}

	// Use faster COPY FROM for postgres if possible
	if !onConflictDoNothing && db2.pgxConn() != nil {
		return db2.copyFrom(ctx, table, newCopyFromData(db2, key_map, rows))
	}

	insert := strings.Builder{} // make([]string, 0)
	keys := make([]string, 0, len(key_map))

	insert.WriteString("INSERT INTO ")
	insert.WriteString(db2.Esc(table))
	insert.WriteString(" (")

	idx := 0
	for key := range key_map {
		if idx > 0 {
			insert.WriteRune(',')
		}
		insert.WriteString(db2.Esc(key))
		keys = append(keys, key)
		idx++
	}

	insert.WriteString(") VALUES \n")

	for idx, row := range rows {
		if idx > 0 {
			insert.WriteRune(',')
		}
		insert.WriteRune('(')
		for idx2, key := range keys {
			if idx2 > 0 {
				insert.WriteRune(',')
			}
			insert.WriteString(db2.escValueForInsert(row[key], key_map[key]))
		}
		insert.WriteRune(')')
		insert.WriteRune('\n')
	}

	if onConflictDoNothing {
		if len(conflictCols) > 0 {
			cCols := []string{}
			for _, cc := range conflictCols {
				cCols = append(cCols, db2.Esc(cc))
			}
			insert.WriteString(" ON CONFLICT (" + strings.Join(cCols, ",") + ") DO NOTHING")
		} else {
			insert.WriteString(" ON CONFLICT DO NOTHING")
		}
	}

	rowsAffected, _, err := db2.execContext(ctx, insert.String())
	if !onConflictDoNothing && err == nil && rowsAffected != int64(len(rows)) {
		err = ErrMismatchedRowsAffected
	}
	if err != nil {
		return db2.sqlError(err, insert.String(), []any{})
	}

	return nil
}

func (db2 *db) UpdateBulk(table string, data any) error {
	return db2.UpdateBulkContext(context.Background(), table, data)
}

// UpdateBulkContext updates all records of the passed slice. It using a single
// exec to send the data to the database. This is generally faster than calling Update
// with a slice (which sends individual update requests).
func (db2 *db) UpdateBulkContext(ctx context.Context, table string, data any) error {
	var (
		rv         reflect.Value
		structMode bool
		err        error
	)

	rv, structMode, err = checkData(data)
	if err != nil {
		return err
	}

	if structMode {
		return fmt.Errorf("UpdateBulk: Need Slice to update bulk.")
	}

	l := rv.Len()
	if l == 0 {
		return nil
	}

	update := strings.Builder{} // make([]string, 0)
	for i := 0; i < l; i++ {
		row := reflect.Indirect(rv.Index(i)).Interface()
		values, structInfo, err := db2.valuesFromStruct(row)
		if err != nil {
			return errors.Wrap(err, "sqlpro.UpdateBulk error.")
		}
		where := strings.Builder{}
		whereCount := 0
		update.WriteString("UPDATE ")
		update.WriteString(db2.Esc(table))
		update.WriteString(" SET ")
		idx2 := 0
		for key, value := range values {
			value2 := db2.nullValue(value, structInfo[key])
			if structInfo[key].primaryKey {
				// skip primary keys for update
				if value2 == nil {
					return fmt.Errorf("Unable to build UPDATE clause with <nil> primary key: %s", key)
				}
				if whereCount > 0 {
					where.WriteString(" AND ")
				}
				where.WriteString(db2.Esc(key))
				where.WriteRune('=')
				where.WriteString(db2.escValueForInsert(value2, structInfo[key]))
				whereCount++
			} else {
				if idx2 > 0 {
					update.WriteRune(',')
				}
				idx2++
				update.WriteString(db2.Esc(key))
				update.WriteRune('=')
				update.WriteString(db2.escValueForInsert(value2, structInfo[key]))
			}
		}
		update.WriteString(" WHERE ")
		update.Write([]byte(where.String()))
		update.WriteRune(';')
		update.WriteRune('\n')
	}

	rowsAffected, _, err := db2.execContext(ctx, update.String())
	if err == nil && rowsAffected != 1 {
		err = ErrMismatchedRowsAffected
	}
	if err != nil {
		return db2.sqlError(err, update.String(), []any{})
	}

	return nil
}

func (db2 *db) insertStruct(ctx context.Context, table string, row any) (int64, structInfo, error) {
	values, info, err := db2.valuesFromStruct(row)
	if err != nil {
		return 0, nil, err
	}

	sql, args, err := db2.insertClauseFromValues(table, values, info)
	if err != nil {
		return 0, nil, err
	}

	if db2.UseReturningForLastId {
		pk := info.onlyPrimaryKey()
		if pk != nil {
			// Fail if transaction present and not in write mode
			if db2.sqlTx != nil && !db2.txWriteMode {
				return 0, nil, fmt.Errorf("[%s] Trying to write into read-only transaction: %s", db2, sql)
			}

			sql = sql + " RETURNING " + db2.Esc(pk.dbName)
			var insert_id_any any
			if db2.Debug || db2.DebugExec {
				log.Printf("%s SQL: %s\nARGS:\n%s", db2, sql, argsToString(args...))
			}
			err := db2.QueryContext(ctx, &insert_id_any, sql, args...)
			if err != nil {
				return 0, nil, err
			}
			insert_id, _ := insert_id_any.(int64) // ignore conversion error, return 0 in that case
			// log.Printf("Returning ID: %T %v", insert_id_any, insert_id_any)
			return insert_id, info, nil
		}
	}

	// log.Printf("SQL: %s Debug: %v", sql, db.Debug)
	rowsAffected, insert_id, err := db2.execContext(ctx, sql, args...)
	if err == nil && rowsAffected != 1 {
		err = ErrMismatchedRowsAffected
	}
	if err != nil {
		return 0, nil, err
	}

	return insert_id, info, nil
}

func (db2 *db) insertClauseFromValues(table string, values map[string]any, info structInfo) (string, []any, error) {
	cols := make([]string, 0, len(values))
	vs := make([]string, 0, len(values))
	args := make([]any, 0, len(values))

	for col, value := range values {
		cols = append(cols, db2.Esc(col))
		vs = append(vs, "?")
		args = append(args, db2.nullValue(value, info[col]))
	}
	return fmt.Sprintf("INSERT INTO %s (%s) VALUES(%s)",
		db2.Esc(table),
		strings.Join(cols, ","),
		strings.Join(vs, ","),
	), args, nil
}

func (db2 *db) updateClauseFromRow(table string, row any) (string, []any, error) {

	var (
		valid     bool
		args      []any
		whereArgs []any
		pk_value  any
	)

	values, structInfo, err := db2.valuesFromStruct(row)
	if err != nil {
		return "", nil, err
	}

	update := strings.Builder{}
	where := strings.Builder{}

	update.WriteString("UPDATE ")
	update.WriteString(db2.Esc(table))
	update.WriteString(" SET ")

	where.WriteString(" WHERE ")

	for key, value := range values {
		if structInfo.primaryKey(key) {
			// skip primary keys for update
			pk_value = db2.nullValue(value, structInfo[key])
			if pk_value == nil {
				return "", args, fmt.Errorf("Unable to build UPDATE clause with <nil> key: %s", key)
			}
			if len(whereArgs) > 0 {
				where.WriteString(" AND ")
			}
			where.WriteString(db2.Esc(key))
			where.WriteString("=")
			where.WriteRune(db2.PlaceholderValue)

			whereArgs = append(whereArgs, pk_value)
			valid = true
		} else {
			if len(args) > 0 {
				update.WriteString(",")
			}
			update.WriteString(db2.Esc(key))
			update.WriteString("=")
			update.WriteRune(db2.PlaceholderValue)
			args = append(args, db2.nullValue(value, structInfo[key]))
		}
	}

	if !valid {
		return "", args, fmt.Errorf("Unable to build UPDATE clause, at least one key needed.")
	}

	args = append(args, whereArgs...)

	// Add where clause
	return update.String() + where.String(), args, nil
}

func (db2 *db) Update(table string, data any) error {
	return db2.UpdateContext(context.Background(), table, data)
}

// Update updates the given struct or slice of structs
// The WHERE clause is put together from the "pk" columns.
// If not all "pk" columns have non empty values, Update returns
// an error.
func (db2 *db) UpdateContext(ctx context.Context, table string, data any) error {
	var (
		rv         reflect.Value
		structMode bool
		err        error
		update     string
		args       []any
	)

	if db2 == nil {
		panic("Update on <nil> handle.")
	}

	rv, structMode, err = checkData(data)
	if err != nil {
		return err
	}

	if structMode {
		update, args, err = db2.updateClauseFromRow(table, rv.Interface())
		if err != nil {
			return err
		}
		rowsAffected, _, err := db2.execContext(ctx, update, args...)
		if err == nil && rowsAffected != 1 {
			err = ErrMismatchedRowsAffected
		}
		if err != nil {
			return err
		}
	} else {
		for i := 0; i < rv.Len(); i++ {
			row := reflect.Indirect(rv.Index(i))
			update, args, err = db2.updateClauseFromRow(table, row.Interface())
			if err != nil {
				return err
			}
			rowsAffected, _, err := db2.execContext(ctx, update, args...)
			if err == nil && rowsAffected != 1 {
				err = ErrMismatchedRowsAffected
			}
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Save saves the given data. It performs an INSERT if the only primary key is
// zero, and and UPDATE if it is not. It panics if it the record has no primary
// key or less than one
func (db2 *db) Save(table string, data any) error {

	rv, structMode, err := checkData(data)
	if err != nil {
		return err
	}

	if structMode {
		return db2.saveRow(table, data)
	} else {
		for i := 0; i < rv.Len(); i++ {
			err = db2.saveRow(table, rv.Index(i).Interface())
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func (db2 *db) saveRow(table string, data any) error {
	row := reflect.Indirect(reflect.ValueOf(data))

	values, info, err := db2.valuesFromStruct(row.Interface())
	if err != nil {
		return err
	}
	pk := info.onlyPrimaryKey()

	if pk == nil {
		return fmt.Errorf("Save needs a struct with exactly one 'pk' field.")
	}

	pk_value, ok := values[pk.dbName]

	if !ok || isZero(pk_value) {
		return db2.Insert(table, data)
	} else {
		return db2.Update(table, data)
	}
}

// valuesFromStruct returns the relevant values
// from struct, as map
func (db2 *db) valuesFromStruct(data any) (map[string]any, structInfo, error) {
	var (
		info   structInfo
		values map[string]any
		dataV  reflect.Value
		err    error
	)

	values = make(map[string]any, 0)
	dataV = reflect.ValueOf(data)

	info = getStructInfo(dataV.Type())

	for _, fieldInfo := range info {
		dataF := dataV.FieldByName(fieldInfo.name)

		actualData := dataF.Interface()
		isZero := isZero(actualData)

		if isZero && fieldInfo.omitEmpty {
			continue
		}

		if fieldInfo.readOnly {
			continue
		}

		if fieldInfo.isJson {
			if isZero {
				actualData = reflect.Zero(fieldInfo.structField.Type).Interface()
			}
			actualData, err = json.Marshal(actualData)
			if err != nil {
				if !fieldInfo.jsonIgnoreError {
					return nil, nil, errors.Wrap(err, "Unable to marshal as data as json.")
				}
			}
			// If the database accepts "null" we write NULL, if the db does not accept null
			// we write "null", if it is not specified we write NULL if the json renders to "null"
			if isZero && (fieldInfo.null || !fieldInfo.notNull && string(actualData.([]byte)) == "null") {
				actualData = nil
			}
		}

		values[fieldInfo.dbName] = actualData
		// log.Printf("Name: %s Value: %v %v", fieldInfo.name, dataF.Interface(), isZero)
	}
	return values, info, nil
}

// isZero returns true if given "x" equals Go's empty value.
func isZero(x any) bool {
	if x == nil {
		return true
	}
	return reflect.DeepEqual(x, reflect.Zero(reflect.TypeOf(x)).Interface())
}

// execContext wraps DB.Exec and returns the number of affected rows as reported
// by the driver as well as the ID inserted, if the driver supports it.
func (db2 *db) execContext(ctx context.Context, execSql string, args ...any) (rowsAffected, insertID int64, err error) {
	var (
		execSql0 string
		newArgs  []any
	)

	if db2.txExecQueryMtx != nil {
		db2.txExecQueryMtx.Lock()
		defer db2.txExecQueryMtx.Unlock()
	}

	if db2.Debug || db2.DebugExec {
		log.Printf("%s SQL: %s\nARGS:\n%s", db2, execSql, argsToString(args...))
	}

	// Fail if transaction present and not in write mode
	if db2.sqlTx != nil && !db2.txWriteMode {
		return 0, 0, fmt.Errorf("[%s] Trying to write into read-only transaction: %s", db2, execSql)
	}

	if len(args) > 0 {
		execSql0, newArgs, err = db2.replaceArgs(execSql, args...)
		if err != nil {
			return 0, 0, err
		}
	} else {
		execSql0 = execSql
		newArgs = args
	}

	// logrus.Infof("[%p] EXEC #%d %s %s", db.sqlDB, db.transID, aurora.Green(fmt.Sprintf("%p", db.db)), execSql0[0:10])

	var result sql.Result

	// tries := 0
	for {
		result, err = db2.db.ExecContext(ctx, execSql0, newArgs...)
		if err != nil {
			// pp.Println(err)
			// sqlErr, ok := err.(sqlite3.Error)
			// if ok {
			// 	if sqlErr.Code == 5 { // SQLITE_BUSY
			// 		tries++
			// 		time.Sleep(50 * time.Millisecond)
			// 		if tries < 3 {
			// 			continue
			// 		}
			// 	}
			// }
			return 0, 0, db2.debugError(db2.sqlError(err, execSql0, newArgs))
		}
		break
	}

	row_count, err := result.RowsAffected()
	if err != nil {
		// Ignore the error here, we might get
		// no RowsAffected available after the empty statement from pq driver
		// which is ok and not a real error (it happens with empty statements)
	}

	if !db2.SupportsLastInsertId {
		return row_count, 0, nil
	}

	last_insert_id, err := result.LastInsertId()
	if err != nil {
		return row_count, 0, db2.debugError(err)
	}
	return row_count, last_insert_id, nil
}