simpleorm/exec/exec.go

package exec

import (
	"database/sql"
	"errors"
	"reflect"

	"git.gdulai.com/gdulai/simpleorm"
	"git.gdulai.com/gdulai/simpleorm/schema"
	log "gitlab.com/gdulai/simpleloglvl"
)

// Created and executes the DDL created by [simpleorm.ORM]
func ExecuteDDL(conn *simpleorm.DBConnection, orm *simpleorm.ORM) {
	schema, err := orm.CreateSchema()
	if err != nil {
		log.LogError("Failed to run DDL: %s", schema)
		return
	}

	log.LogInfo("Executing DDL:\n%s", schema)

	if _, err := conn.Exec(schema); err != nil {
		log.LogFatalError("%", err)
	}
}

// Wraps and represents a DB transaction.
// Allows for all at once or separate execution of [exec.Exec] implementations.
type Transaction struct {
	conn       *simpleorm.DBConnection
	executions []*Exec[any]
	tx         *sql.Tx
	finished   bool
}

// Creates a new [exec.Transaction].
// Can be initialized with a set of [exec.Exec]s.
func NewTransaction(conn *simpleorm.DBConnection, executions ...*Exec[any]) *Transaction {
	return &Transaction{conn: conn, executions: executions, finished: false}
}

// The executions assigned to this transaction.
func (t *Transaction) Executions() []*Exec[any] {
	return t.executions
}

// Executes all the [exec.Exec] implementations assigned to this transactions.
// Begins, commits or rollbacks the transaction.
// This is a terminal operation, the [exec.Transaction] is considered finished after calling this.
func (t *Transaction) ExecuteAtOnce() error {
	if t.finished {
		return errors.New("Transaction already finished.")
	}

	if t.tx == nil {
		tx, err := t.conn.Begin()
		t.tx = tx
		if err != nil {
			t.finished = true
			return err
		}
	}

	for _, exec := range t.executions {
		err := (*exec).execute(nil, t.tx)
		if err != nil {
			rollbackErr := t.Rollback()
			if rollbackErr != nil {
				return rollbackErr
			}
			return err
		}

	}

	err := t.tx.Commit()
	if err != nil {
		return err
	}

	t.finished = true
	return nil
}

// Executes and assigns the passed [exec.Exec]s to the Transaction.
// Calls rollback in case of an error.
// This is NOT a terminal operation, transactions still has to be committed.
func (t *Transaction) Execute(execs ...Exec[any]) error {
	if t.finished {
		return errors.New("Transaction already finished.")
	}

	for _, exec := range execs {
		execPtr := &exec
		t.executions = append(t.executions, execPtr)
		if t.tx == nil {
			tx, err := t.conn.Begin()
			t.tx = tx
			if err != nil {
				t.finished = true
				return err
			}
		}

		err := (*execPtr).execute(nil, t.tx)
		if err != nil {
			rollbackErr := t.Rollback()
			if rollbackErr != nil {
				return rollbackErr
			}
			return err
		}

	}

	return nil
}

// Finishes the transaction, commits the changes.
// This is terminal operation.
func (t *Transaction) Finish() error {
	if t.finished {
		return errors.New("Transaction already finished.")
	}

	if t.tx == nil {
		return errors.New("Transaction is nil.")
	}

	t.tx.Commit()
	t.finished = true
	return nil
}

// Rollbacks and aborts the transaction.
// This is a terminal operation.
func (t *Transaction) Rollback() error {
	if t.finished {
		return errors.New("Transaction already finished.")
	}

	if t.tx == nil {
		return errors.New("Transaction is nil.")
	}

	rollbackErr := t.tx.Rollback()
	if rollbackErr != nil {
		return rollbackErr
	}
	t.finished = true
	return nil
}

type Exec[T any] interface {
	Execute(conn *simpleorm.DBConnection) error
	execute(conn *simpleorm.DBConnection, tx *sql.Tx) error
}

type Count[T any] struct {
	target    schema.Table
	whereStmt string
	args      []any
	result    int64
}

func CreateCount[T any](orm *simpleorm.ORM, whereStmt string, args ...any) (Count[T], error) {
	table, ok := orm.Cache().Get(reflect.TypeFor[T]().Name())
	if !ok {
		return Count[T]{}, errors.New("Failed to get table from schema cache")
	}

	return Count[T]{target: *table, whereStmt: whereStmt, args: args}, nil
}

func (c *Count[T]) Result() int64 {
	return c.result
}

func (c *Count[T]) Execute(conn *simpleorm.DBConnection) error {
	return c.execute(conn, nil)
}

func (c *Count[T]) execute(conn *simpleorm.DBConnection, tx *sql.Tx) error {
	// Reinit the results, new execution
	c.result = -1
	dml := c.target.GetCountDML()

	if c.whereStmt != "" {
		dml += " WHERE " + c.whereStmt
	}

	log.LogDebug("Preparing sql: %s, with args: %s", dml, c.args)

	var stmt *sql.Stmt
	var err error
	if tx != nil {
		stmt, err = tx.Prepare(dml)
	} else {
		stmt, err = conn.Prepare(dml)
	}

	if err != nil {
		return err
	}
	defer stmt.Close()

	log.LogDebug("Executing statement: %s", stmt)

	if len(c.args) == 0 {
		err = stmt.QueryRow().Scan(&c.result)
	} else {
		// Flattent args to make sure it can be parsed correctly
		var flatArgs []any
		for _, a := range c.args {
			if s, ok := a.([]any); ok {
				flatArgs = append(flatArgs, s...)
			} else {
				flatArgs = append(flatArgs, a)
			}
		}

		err = stmt.QueryRow(flatArgs...).Scan(&c.result)

	}
	if err != nil {
		return err
	}

	return nil
}

type Insert[T any] struct {
	target   schema.Table
	toInsert []T
	results  []T
}

func NewInsert[T any](orm *simpleorm.ORM, toInsert ...T) (Insert[T], error) {
	table, ok := orm.Cache().Get(reflect.TypeFor[T]().Name())
	if !ok {
		return Insert[T]{}, errors.New("Failed to get table from schema cache")
	}

	return Insert[T]{target: *table, toInsert: toInsert}, nil
}

func (ins *Insert[T]) Results() []T {
	return ins.results
}

func (ins *Insert[T]) Execute(conn *simpleorm.DBConnection) error {
	return ins.execute(conn, nil)
}

func (ins *Insert[T]) execute(conn *simpleorm.DBConnection, tx *sql.Tx) error {
	dml, err := ins.target.GetInsertDML(len(ins.toInsert))
	if err != nil {
		return err
	}

	var params []any
	for i := range len(ins.toInsert) {
		actualParams := prepareParams(ins.toInsert[i], ins.target)
		if len(params) == 0 {
			params = make([]any, len(ins.toInsert)*len(actualParams))
		}
		for j := range actualParams {
			params[(i*len(actualParams))+j] = actualParams[j]
		}
	}

	log.LogDebug("%s [%s]", dml, params)

	var stmt *sql.Stmt
	if tx != nil {
		stmt, err = tx.Prepare(dml)
	} else {
		stmt, err = conn.Prepare(dml)
	}
	defer stmt.Close()

	// Flattent args to make sure it can be parsed correctly
	var flatParams []any
	for _, param := range params {
		if s, ok := param.([]any); ok {
			flatParams = append(flatParams, s...)
		} else {
			flatParams = append(flatParams, param)
		}
	}

	rows, err := stmt.Query(flatParams...)

	if err != nil {
		return err
	}

	ins.results, err = readRows[T](ins.target, rows)
	if err != nil {
		return err
	}

	return nil
}

type Update[T any] struct {
	target   schema.Table
	toUpdate T
}

func NewUpdate[T any](orm *simpleorm.ORM, toUpdate T) (Update[T], error) {
	table, ok := orm.Cache().Get(reflect.TypeFor[T]().Name())
	if !ok {
		return Update[T]{}, errors.New("Failed to get table from schema cache")
	}

	return Update[T]{target: *table, toUpdate: toUpdate}, nil
}

func (u Update[T]) Execute(conn *simpleorm.DBConnection) error {
	return u.execute(conn, nil)
}

func (u Update[T]) execute(conn *simpleorm.DBConnection, tx *sql.Tx) error {
	dml, err := u.target.GetUpdateDML()
	if err != nil {
		return err
	}

	params := prepareParams(u.toUpdate, u.target)

	pkCols, err := getPk(u.toUpdate, u.target)
	if err != nil {
		return err
	}
	// Put the pks back at the end
	for _, pk := range pkCols {
		params = append(params, pk)
	}

	log.LogDebug("%s [%s]", dml, params)

	var stmt *sql.Stmt
	if tx != nil {
		stmt, err = tx.Prepare(dml)
	} else {
		stmt, err = conn.Prepare(dml)
	}
	defer stmt.Close()

	result, err := stmt.Exec(params...)
	if err != nil {
		return err
	} else {
		rowsAffected, _ := result.RowsAffected()
		log.LogInfo("Updated %s row", rowsAffected)
	}
	return nil
}

type Delete[T any] struct {
	target   schema.Table
	toDelete []T
}

func NewDelete[T any](orm *simpleorm.ORM, toDelete []T) (Delete[T], error) {
	table, ok := orm.Cache().Get(reflect.TypeFor[T]().Name())
	if !ok {
		return Delete[T]{}, errors.New("Failed to get table from schema cache")
	}

	return Delete[T]{target: *table, toDelete: toDelete}, nil
}

func (d *Delete[T]) Execute(conn *simpleorm.DBConnection) error {
	return d.execute(conn, nil)
}

func (d *Delete[T]) execute(conn *simpleorm.DBConnection, tx *sql.Tx) error {
	count := len(d.toDelete)
	dml, err := d.target.GetDeleteDML(count)
	if err != nil {
		return err
	}

	var params []any

	for _, del := range d.toDelete {
		pkCols, err := getPk(del, d.target)
		if err != nil {
			return err
		}

		for _, pk := range pkCols {
			params = append(params, pk)
		}
	}

	log.LogDebug("%s [%s]", dml, params)

	var stmt *sql.Stmt
	if tx != nil {
		stmt, err = tx.Prepare(dml)
	} else {
		stmt, err = conn.Prepare(dml)
	}
	defer stmt.Close()

	result, err := stmt.Exec(params...)
	if err != nil {
		return err
	} else {
		rowsAffected, _ := result.RowsAffected()
		log.LogInfo("Deleted %s row", rowsAffected)
	}

	return nil

}

func createSelectResultContainer(t schema.Table) []any {
	typ := t.Type()
	vals := make([]any, typ.NumField())
	for i := range vals {
		switch typ.Field(i).Type.Kind().String() {
		case "string":
			var fieldContainer string
			vals[i] = &fieldContainer
		case "int", "bool":
			var fieldContainer int
			vals[i] = &fieldContainer
		case "int64", "time.Time":
			var fieldContainer int64
			vals[i] = &fieldContainer
		}
	}
	return vals
}

func prepareParams(src any, t schema.Table) []any {
	typ := t.Type()
	var params []any
	for _, col := range t.Columns() {
		_, ok := col.Modifiers["pk"]
		if ok && t.IsPkAuto() {
			continue
		}
		field, ok := typ.FieldByName(col.FieldName)
		if !ok {
			continue
		}
		fieldValue := reflect.ValueOf(src).FieldByIndex(field.Index)

		params = append(params, col.Encode(fieldValue))
	}
	return params
}

func getPk(src any, t schema.Table) ([]any, error) {
	typ := t.Type()
	var values []any
	for _, constraint := range t.Constraints() {
		if constraint.Type == "pk" {
			for _, col := range constraint.Columns {
				field, ok := typ.FieldByName(col.FieldName)
				if !ok {
					continue

				}
				fieldValue := reflect.ValueOf(src).FieldByIndex(field.Index)
				values = append(values, fieldValue.Interface())
			}
		}
	}
	if len(values) == 0 {
		return nil, errors.New("Could not determine pk column!")
	}

	return values, nil
}

func readRows[T any](table schema.Table, rows *sql.Rows) ([]T, error) {
	rowContainer := createSelectResultContainer(table)

	var results []T

	for rows.Next() {
		err := rows.Scan(rowContainer...)
		if err != nil {
			return nil, err
		}

		targetType := table.Type()
		cols := table.Columns()
		parsedResult := reflect.New(targetType)
		for i, fieldVal := range rowContainer {
			col := cols[i]
			targetField := parsedResult.Elem().Field(i)

			rawValue := reflect.Indirect(reflect.ValueOf(fieldVal))
			decoded := reflect.ValueOf(col.Decode(targetField.Type().Name(), rawValue))

			targetField.Set(decoded)
		}
		parsedObj := reflect.Indirect(parsedResult).Interface().(T)
		results = append(results, parsedObj)
	}
	return results, nil
}