334 lines · 9.5 KB · Go
RawDownload
// Copyright 2019 The Gitea Authors. All rights reserved.// SPDX-License-Identifier: MIT package db import (	"context"	"database/sql" 	"xorm.io/builder"	"xorm.io/xorm") // DefaultContext is the default context to run xorm queries in// will be overwritten by Init with HammerContextvar DefaultContext context.Context // contextKey is a value for use with context.WithValue.type contextKey struct {	name string} // enginedContextKey is a context key. It is used with context.Value() to get the current Engined for the contextvar (	enginedContextKey         = &contextKey{"engined"}	_                 Engined = &Context{}) // Context represents a db contexttype Context struct {	context.Context	e           Engine	transaction bool} func newContext(ctx context.Context, e Engine, transaction bool) *Context {	return &Context{		Context:     ctx,		e:           e,		transaction: transaction,	}} // InTransaction if context is in a transactionfunc (ctx *Context) InTransaction() bool {	return ctx.transaction} // Engine returns db enginefunc (ctx *Context) Engine() Engine {	return ctx.e} // Value shadows Value for context.Context but allows us to get ourselves and an Engined objectfunc (ctx *Context) Value(key any) any {	if key == enginedContextKey {		return ctx	}	return ctx.Context.Value(key)} // WithContext returns this engine tied to this contextfunc (ctx *Context) WithContext(other context.Context) *Context {	return newContext(ctx, ctx.e.Context(other), ctx.transaction)} // Engined structs provide an Enginetype Engined interface {	Engine() Engine} // GetEngine will get a db Engine from this context or return an Engine restricted to this contextfunc GetEngine(ctx context.Context) Engine {	if e := getEngine(ctx); e != nil {		return e	}	return x.Context(ctx)} // getEngine will get a db Engine from this context or return nilfunc getEngine(ctx context.Context) Engine {	if engined, ok := ctx.(Engined); ok {		return engined.Engine()	}	enginedInterface := ctx.Value(enginedContextKey)	if enginedInterface != nil {		return enginedInterface.(Engined).Engine()	}	return nil} // Committer represents an interface to Commit or Close the Contexttype Committer interface {	Commit() error	Close() error} // halfCommitter is a wrapper of Committer.// It can be closed early, but can't be committed early, it is useful for reusing a transaction.type halfCommitter struct {	committer Committer	committed bool} func (c *halfCommitter) Commit() error {	c.committed = true	// should do nothing, and the parent committer will commit later	return nil} func (c *halfCommitter) Close() error {	if c.committed {		// it's "commit and close", should do nothing, and the parent committer will commit later		return nil	} 	// it's "rollback and close", let the parent committer rollback right now	return c.committer.Close()} // TxContext represents a transaction Context,// it will reuse the existing transaction in the parent context or create a new one.// Some tips to use:////	1 It's always recommended to use `WithTx` in new code instead of `TxContext`, since `WithTx` will handle the transaction automatically.//	2. To maintain the old code which uses `TxContext`://	  a. Always call `Close()` before returning regardless of whether `Commit()` has been called.//	  b. Always call `Commit()` before returning if there are no errors, even if the code did not change any data.//	  c. Remember the `Committer` will be a halfCommitter when a transaction is being reused.//	     So calling `Commit()` will do nothing, but calling `Close()` without calling `Commit()` will rollback the transaction.//	     And all operations submitted by the caller stack will be rollbacked as well, not only the operations in the current function.//	  d. It doesn't mean rollback is forbidden, but always do it only when there is an error, and you do want to rollback.func TxContext(parentCtx context.Context) (*Context, Committer, error) {	if sess, ok := inTransaction(parentCtx); ok {		return newContext(parentCtx, sess, true), &halfCommitter{committer: sess}, nil	} 	sess := x.NewSession()	if err := sess.Begin(); err != nil {		sess.Close()		return nil, nil, err	} 	return newContext(DefaultContext, sess, true), sess, nil} // WithTx represents executing database operations on a transaction, if the transaction exist,// this function will reuse it otherwise will create a new one and close it when finished.func WithTx(parentCtx context.Context, f func(ctx context.Context) error) error {	if sess, ok := inTransaction(parentCtx); ok {		err := f(newContext(parentCtx, sess, true))		if err != nil {			// rollback immediately, in case the caller ignores returned error and tries to commit the transaction.			_ = sess.Close()		}		return err	}	return txWithNoCheck(parentCtx, f)} func txWithNoCheck(parentCtx context.Context, f func(ctx context.Context) error) error {	sess := x.NewSession()	defer sess.Close()	if err := sess.Begin(); err != nil {		return err	} 	if err := f(newContext(parentCtx, sess, true)); err != nil {		return err	} 	return sess.Commit()} // Insert inserts records into databasefunc Insert(ctx context.Context, beans ...any) error {	_, err := GetEngine(ctx).Insert(beans...)	return err} // Exec executes a sql with argsfunc Exec(ctx context.Context, sqlAndArgs ...any) (sql.Result, error) {	return GetEngine(ctx).Exec(sqlAndArgs...)} func Get[T any](ctx context.Context, cond builder.Cond) (object *T, exist bool, err error) {	if !cond.IsValid() {		panic("cond is invalid in db.Get(ctx, cond). This should not be possible.")	} 	var bean T	has, err := GetEngine(ctx).Where(cond).NoAutoCondition().Get(&bean)	if err != nil {		return nil, false, err	} else if !has {		return nil, false, nil	}	return &bean, true, nil} func GetByID[T any](ctx context.Context, id int64) (object *T, exist bool, err error) {	var bean T	has, err := GetEngine(ctx).ID(id).NoAutoCondition().Get(&bean)	if err != nil {		return nil, false, err	} else if !has {		return nil, false, nil	}	return &bean, true, nil} func Exist[T any](ctx context.Context, cond builder.Cond) (bool, error) {	if !cond.IsValid() {		panic("cond is invalid in db.Exist(ctx, cond). This should not be possible.")	} 	var bean T	return GetEngine(ctx).Where(cond).NoAutoCondition().Exist(&bean)} func ExistByID[T any](ctx context.Context, id int64) (bool, error) {	var bean T	return GetEngine(ctx).ID(id).NoAutoCondition().Exist(&bean)} // DeleteByID deletes the given bean with the given IDfunc DeleteByID[T any](ctx context.Context, id int64) (int64, error) {	var bean T	return GetEngine(ctx).ID(id).NoAutoCondition().NoAutoTime().Delete(&bean)} func DeleteByIDs[T any](ctx context.Context, ids ...int64) error {	if len(ids) == 0 {		return nil	} 	var bean T	_, err := GetEngine(ctx).In("id", ids).NoAutoCondition().NoAutoTime().Delete(&bean)	return err} func Delete[T any](ctx context.Context, opts FindOptions) (int64, error) {	if opts == nil || !opts.ToConds().IsValid() {		panic("opts are empty or invalid in db.Delete(ctx, opts). This should not be possible.")	} 	var bean T	return GetEngine(ctx).Where(opts.ToConds()).NoAutoCondition().NoAutoTime().Delete(&bean)} // DeleteByBean deletes all records according non-empty fields of the bean as conditions.func DeleteByBean(ctx context.Context, bean any) (int64, error) {	return GetEngine(ctx).Delete(bean)} // FindIDs finds the IDs for the given table name satisfying the given condition// By passing a different value than "id" for "idCol", you can query for foreign IDs, i.e. the repo IDs which satisfy the conditionfunc FindIDs(ctx context.Context, tableName, idCol string, cond builder.Cond) ([]int64, error) {	ids := make([]int64, 0, 10)	if err := GetEngine(ctx).Table(tableName).		Cols(idCol).		Where(cond).		Find(&ids); err != nil {		return nil, err	}	return ids, nil} // DecrByIDs decreases the given column for entities of the "bean" type with one of the given ids by one// Timestamps of the entities won't be updatedfunc DecrByIDs(ctx context.Context, ids []int64, decrCol string, bean any) error {	if len(ids) == 0 {		return nil	}	_, err := GetEngine(ctx).Decr(decrCol).In("id", ids).NoAutoCondition().NoAutoTime().Update(bean)	return err} // DeleteBeans deletes all given beans, beans must contain delete conditions.func DeleteBeans(ctx context.Context, beans ...any) (err error) {	e := GetEngine(ctx)	for i := range beans {		if _, err = e.Delete(beans[i]); err != nil {			return err		}	}	return nil} // TruncateBeans deletes all given beans, beans may contain delete conditions.func TruncateBeans(ctx context.Context, beans ...any) (err error) {	e := GetEngine(ctx)	for i := range beans {		if _, err = e.Truncate(beans[i]); err != nil {			return err		}	}	return nil} // CountByBean counts the number of database records according non-empty fields of the bean as conditions.func CountByBean(ctx context.Context, bean any) (int64, error) {	return GetEngine(ctx).Count(bean)} // TableName returns the table name according a bean objectfunc TableName(bean any) string {	return x.TableName(bean)} // InTransaction returns true if the engine is in a transaction otherwise return falsefunc InTransaction(ctx context.Context) bool {	_, ok := inTransaction(ctx)	return ok} func inTransaction(ctx context.Context) (*xorm.Session, bool) {	e := getEngine(ctx)	if e == nil {		return nil, false	} 	switch t := e.(type) {	case *xorm.Engine:		return nil, false	case *xorm.Session:		if t.IsInTx() {			return t, true		}		return nil, false	default:		return nil, false	}}