package query import ( "context" "errors" "fmt" "github.com/bitmagnet-io/bitmagnet/internal/database/dao" "github.com/bitmagnet-io/bitmagnet/internal/maps" "github.com/bitmagnet-io/bitmagnet/internal/model" "gorm.io/gen" "gorm.io/gen/field" "gorm.io/gorm" "gorm.io/gorm/clause" "strings" "sync" ) type ResultItem struct { QueryStringRank float64 } type GenericResult[T interface{}] struct { TotalCount uint HasNextPage bool Items []T Aggregations Aggregations } type SubQueryFactory = func(context.Context, *dao.Query) SubQuery func GenericQuery[T interface{}]( ctx context.Context, daoQ *dao.Query, option Option, tableName string, factory SubQueryFactory, ) (r GenericResult[T], _ error) { dbCtx := dbContext{ q: daoQ, tableName: tableName, factory: factory, } builder, optionErr := option(newQueryContext(dbCtx)) ctx = builder.createContext(ctx) if optionErr != nil { return r, optionErr } newSubQuery := func() (SubQuery, error) { sq := factory(ctx, daoQ) if selectErr := builder.applySelect(sq); selectErr != nil { return sq, selectErr } if preErr := builder.applyPre(sq); preErr != nil { return sq, preErr } return sq, nil } var errs []error mtx := sync.Mutex{} addErr := func(err error) { mtx.Lock() defer mtx.Unlock() errs = append(errs, err) } wg := sync.WaitGroup{} wg.Add(3) go (func() { defer wg.Done() if builder.withTotalCount() { sq, sqErr := newSubQuery() if sqErr != nil { addErr(sqErr) return } if tc, countErr := sq.Count(); countErr != nil { addErr(countErr) } else { r.TotalCount = uint(tc) } } })() go (func() { defer wg.Done() sq, sqErr := newSubQuery() if sqErr != nil { addErr(sqErr) return } if postErr := builder.applyPost(sq); postErr != nil { addErr(postErr) return } if !builder.hasZeroLimit() || builder.needsNextPage() { var items []T if txErr := sq.UnderlyingDB().Find(&items).Error; txErr != nil { addErr(txErr) return } // copy items slice to avoid modifying cached results copiedItems := append([]T{}, items...) if builder.hasNextPage(len(copiedItems)) { r.HasNextPage = true copiedItems = copiedItems[:len(copiedItems)-1] } if len(copiedItems) > 0 { if cbErr := builder.applyCallbacks(ctx, copiedItems); cbErr != nil { addErr(cbErr) return } } r.Items = copiedItems } })() go (func() { defer wg.Done() if aggs, aggErr := builder.calculateAggregations(ctx); aggErr != nil { addErr(aggErr) } else { r.Aggregations = aggs } })() wg.Wait() return r, errors.Join(errs...) } type BaseSubQuery interface { Count() (int64, error) TableName() string UnderlyingDB() *gorm.DB } type SubQuery interface { BaseSubQuery Scopes(...GormScope) SubQuery } type GenericSubQuery[T BaseSubQuery] struct { SubQuery BaseSubQuery } func (sq GenericSubQuery[T]) Count() (int64, error) { return sq.SubQuery.Count() } func (sq GenericSubQuery[T]) TableName() string { return sq.SubQuery.TableName() } func (sq GenericSubQuery[T]) UnderlyingDB() *gorm.DB { return sq.SubQuery.UnderlyingDB() } type scoper[T BaseSubQuery] interface { Scopes(funcs ...GormScope) T } func (sq GenericSubQuery[T]) Scopes(fns ...func(gen.Dao) gen.Dao) SubQuery { sq.SubQuery = sq.SubQuery.(scoper[T]).Scopes(fns...) return sq } type Scope = func(SubQuery) error type GormScope = func(gen.Dao) gen.Dao type DbContext interface { Query() *dao.Query TableName() string NewSubQuery(context.Context) SubQuery } type dbContext struct { q *dao.Query tableName string factory SubQueryFactory } func (db dbContext) Query() *dao.Query { return db.q } func (db dbContext) TableName() string { return db.tableName } func (db dbContext) NewSubQuery(ctx context.Context) SubQuery { return db.factory(ctx, db.q) } type CallbackContext interface { DbContext Lock() Unlock() } type callbackContext struct { dbContext *sync.Mutex } type Callback func(ctx context.Context, cbCtx CallbackContext, results any) error type OptionBuilder interface { DbContext Table(string) OptionBuilder Join(...TableJoin) OptionBuilder RequireJoin(...string) OptionBuilder Scope(...Scope) OptionBuilder Select(...clause.Expr) OptionBuilder OrderBy(...clause.OrderByColumn) OptionBuilder Limit(uint) OptionBuilder Offset(uint) OptionBuilder Group(...clause.Column) OptionBuilder Facet(...Facet) OptionBuilder Preload(...field.RelationField) OptionBuilder Callback(...Callback) OptionBuilder Context(func(ctx context.Context) context.Context) OptionBuilder applySelect(SubQuery) error applyPre(SubQuery) error applyPost(SubQuery) error createFacetsFilterCriteria() (Criteria, error) calculateAggregations(context.Context) (Aggregations, error) WithTotalCount(bool) OptionBuilder WithHasNextPage(bool) OptionBuilder withTotalCount() bool applyCallbacks(context.Context, any) error hasZeroLimit() bool needsNextPage() bool hasNextPage(nItems int) bool withCurrentFacet(string) OptionBuilder createContext(context.Context) context.Context } type optionBuilder struct { dbContext joins map[string]TableJoin requiredJoins maps.InsertMap[string, struct{}] scopes []Scope selections []clause.Expr groupBy []clause.Column orderBy []clause.OrderByColumn limit model.NullUint nextPage bool offset uint facets []Facet currentFacet string preloads []field.RelationField totalCount bool callbacks []Callback contextFn func(context.Context) context.Context } type RawJoin struct { Query string Args []interface{} } func newQueryContext(dbCtx dbContext) OptionBuilder { return optionBuilder{ dbContext: dbCtx, joins: make(map[string]TableJoin), requiredJoins: maps.NewInsertMap[string, struct{}](), } } func (b optionBuilder) Table(name string) OptionBuilder { b.dbContext.tableName = name return b.Scope(func(sq SubQuery) error { sq.UnderlyingDB().Table(name) return nil }) } func (b optionBuilder) Join(joins ...TableJoin) OptionBuilder { bJoins := make(map[string]TableJoin, len(b.joins)) for _, j := range b.joins { bJoins[j.Table.TableName()] = j } bRequiredJoins := b.requiredJoins.Copy() for _, j := range joins { bJoins[j.Table.TableName()] = j if j.Required { bRequiredJoins.SetKey(j.Table.TableName()) } } b.joins = bJoins b.requiredJoins = bRequiredJoins return b } func (b optionBuilder) RequireJoin(names ...string) OptionBuilder { bRequiredJoins := b.requiredJoins.Copy() for _, name := range names { bRequiredJoins.SetKey(name) } b.requiredJoins = bRequiredJoins return b } func (b optionBuilder) Scope(scopes ...Scope) OptionBuilder { b.scopes = append(b.scopes, scopes...) return b } func (b optionBuilder) Select(selections ...clause.Expr) OptionBuilder { b.selections = append(b.selections, selections...) return b } func (b optionBuilder) Group(columns ...clause.Column) OptionBuilder { b.groupBy = append(b.groupBy, columns...) return b } func (b optionBuilder) OrderBy(columns ...clause.OrderByColumn) OptionBuilder { b.orderBy = append(b.orderBy, columns...) return b } func (b optionBuilder) Limit(limit uint) OptionBuilder { b.limit = model.NewNullUint(limit) return b } func (b optionBuilder) Offset(offset uint) OptionBuilder { b.offset = offset return b } func (b optionBuilder) Facet(facets ...Facet) OptionBuilder { b.facets = append(b.facets, facets...) return b } func (b optionBuilder) Preload(relations ...field.RelationField) OptionBuilder { b.preloads = append(b.preloads, relations...) return b } func (b optionBuilder) Callback(callbacks ...Callback) OptionBuilder { b.callbacks = append(b.callbacks, callbacks...) return b } func (b optionBuilder) Context(fn func(context.Context) context.Context) OptionBuilder { prevFn := b.contextFn b.contextFn = func(ctx context.Context) context.Context { if prevFn != nil { ctx = prevFn(ctx) } return fn(ctx) } return b } func (b optionBuilder) WithTotalCount(bl bool) OptionBuilder { b.totalCount = bl return b } func (b optionBuilder) WithHasNextPage(bl bool) OptionBuilder { b.nextPage = bl return b } func (b optionBuilder) withTotalCount() bool { return b.totalCount } func (b optionBuilder) hasZeroLimit() bool { return b.limit.Valid && b.limit.Uint == 0 } func (b optionBuilder) needsNextPage() bool { return b.limit.Valid && b.nextPage } func (b optionBuilder) hasNextPage(nItems int) bool { if !b.nextPage { return false } if !b.limit.Valid { return false } return nItems > int(b.limit.Uint) } func (b optionBuilder) withCurrentFacet(facet string) OptionBuilder { b.currentFacet = facet return b } func (b optionBuilder) createContext(ctx context.Context) context.Context { if b.contextFn != nil { return b.contextFn(ctx) } return ctx } func (b optionBuilder) applySelect(sq SubQuery) error { var selectQueryParts []string selectQueryArgs := make([]interface{}, 0) if len(b.selections) == 0 { selectQueryParts = append(selectQueryParts, "*") } else { for _, s := range b.selections { selectQueryParts = append(selectQueryParts, s.SQL) selectQueryArgs = append(selectQueryArgs, s.Vars...) } } sq.UnderlyingDB().Select(strings.Join(selectQueryParts, ", "), selectQueryArgs...) return nil } func (b optionBuilder) applyPre(sq SubQuery) error { for _, s := range b.scopes { if err := s(sq); err != nil { return err } } requiredJoins := b.requiredJoins.Copy() aggC, aggCErr := b.createFacetsFilterCriteria() if aggCErr != nil { return aggCErr } rawAggC, rawAggCErr := aggC.Raw(b) if rawAggCErr != nil { return rawAggCErr } requiredJoins.SetEntries(rawAggC.Joins.Entries()...) joins, joinsErr := extractRequiredJoins(b.dbContext.tableName, b.joins, requiredJoins) if joinsErr != nil { return joinsErr } applyJoins(sq, joins...) sq.UnderlyingDB().Where(rawAggC.Query, rawAggC.Args...) if len(b.groupBy) > 0 { sq.UnderlyingDB().Clauses(clause.GroupBy{ Columns: b.groupBy, }) } return nil } func extractRequiredJoins(tableName string, joins map[string]TableJoin, requiredJoins maps.InsertMap[string, struct{}]) ([]TableJoin, error) { resolvedJoins := maps.NewInsertMap[string, TableJoin]() var addJoin func(name string) error addJoin = func(name string) error { if name == tableName { return nil } j, ok := joins[name] if !ok { return fmt.Errorf("required join not found: %s", name) } for _, depName := range j.Dependencies.Keys() { if err := addJoin(depName); err != nil { return err } } resolvedJoins.Set(j.Table.TableName(), j) return nil } for _, joinName := range requiredJoins.Keys() { if err := addJoin(joinName); err != nil { return nil, err } } return resolvedJoins.Values(), nil } func applyJoins(sq SubQuery, joins ...TableJoin) { for _, j := range joins { join := j sq.Scopes(func(dao gen.Dao) gen.Dao { switch join.Type { case TableJoinTypeInner: return dao.Join(join.Table, join.On...) case TableJoinTypeLeft: return dao.LeftJoin(join.Table, join.On...) case TableJoinTypeRight: return dao.RightJoin(join.Table, join.On...) } panic("invalid join type") }) } } func (b optionBuilder) applyPost(sq SubQuery) error { if len(b.orderBy) > 0 { sq.UnderlyingDB().Statement.AddClause(clause.OrderBy{ Columns: b.orderBy, }) } if b.limit.Valid { limit := int(b.limit.Uint) if b.nextPage { limit++ } sq.UnderlyingDB().Limit(limit) } sq.UnderlyingDB().Offset(int(b.offset)) for _, p := range b.preloads { sq.UnderlyingDB().Preload(p.Name(), p) } return nil } func (b optionBuilder) applyCallbacks(ctx context.Context, results any) error { cbCtx := callbackContext{ dbContext: b.dbContext, Mutex: &sync.Mutex{}, } var errs []error wg := sync.WaitGroup{} wg.Add(len(b.callbacks)) for _, cb := range b.callbacks { go (func(cb Callback) { defer wg.Done() if err := cb(ctx, cbCtx, results); err != nil { cbCtx.Lock() defer cbCtx.Unlock() errs = append(errs, err) } })(cb) } wg.Wait() return errors.Join(errs...) }