feat: add chunk by query methods (#714)

Author: arashackdev

contracts/database/orm/orm.go

@@ -43,6 +43,12 @@ type Query interface {
 	Begin() (Query, error)
 	// BeginTransaction begins a new transaction
 	BeginTransaction() (Query, error)
+	// Chunk processes a given number of records in batches.
+	Chunk(count int, callback func([]any) error) error
+	// Chunk the results of a query by comparing numeric IDs.
+	ChunkByID(count int, callback func([]any) error) error
+	// ChunkByIDDesc processes a given number of records in batches, ordered by ID in descending order.
+	ChunkByIDDesc(count int, callback func([]any) error) error
 	// Commit commits the changes in a transaction.
 	Commit() error
 	// Count retrieve the "count" result of the query.
@@ -120,6 +126,8 @@ type Query interface {
 	OrderByDesc(column string) Query
 	// OrderByRaw specifies the order should be raw.
 	OrderByRaw(raw string) Query
+	// OrderedChunkByID processes a given number of records in batches, ordered by ID.
+	OrderedChunkByID(count int, callback func([]any) error, descending bool) error
 	// OrWhere add an "or where" clause to the query.
 	OrWhere(query any, args ...any) Query
 	// OrWhereBetween adds an "or where column between x and y" clause to the query.

database/gorm/query.go

@@ -189,6 +189,175 @@ func (r *Query) Cursor() chan contractsdb.Row {
 	return cursorChan
 }
 
+func (r *Query) Chunk(count int, callback func([]any) error) error {
+	if count <= 0 {
+		return errors.OrmQueryChunkZeroOrLess
+	}
+
+	var remaining *int
+	if r.conditions.limit != nil {
+		limit := *r.conditions.limit
+		remaining = &limit
+	}
+
+	query := r.addGlobalScopes().buildConditions()
+
+	destType := reflect.TypeOf(r.conditions.model)
+	sliceType := reflect.SliceOf(reflect.PointerTo(destType))
+	offset := 0
+
+	for {
+		if remaining != nil && *remaining == 0 {
+			break
+		}
+
+		chunkLimit := count
+		if remaining != nil && *remaining < count {
+			chunkLimit = *remaining
+		}
+
+		results := reflect.New(sliceType).Interface()
+
+		chunkQuery := query.Offset(offset).Limit(chunkLimit).(*Query)
+		err := chunkQuery.Find(results)
+		if err != nil {
+			return err
+		}
+
+		resultsValue := reflect.ValueOf(results).Elem()
+		length := resultsValue.Len()
+		if length == 0 {
+			return nil
+		}
+
+		if remaining != nil {
+			*remaining = max(*remaining-length, 0)
+		}
+
+		values := make([]any, length)
+		for i := 0; i < length; i++ {
+			values[i] = resultsValue.Index(i).Interface()
+		}
+
+		if err = callback(values); err != nil {
+			return err
+		}
+
+		if length < chunkLimit {
+			return nil
+		}
+
+		offset += chunkLimit
+	}
+
+	return nil
+}
+
+func (r *Query) ChunkByID(count int, callback func([]any) error) error {
+	return r.OrderedChunkByID(count, callback, false)
+}
+
+func (r *Query) ChunkByIDDesc(count int, callback func([]any) error) error {
+	return r.OrderedChunkByID(count, callback, true)
+}
+
+func (r *Query) OrderedChunkByID(count int, callback func([]any) error, descending bool) error {
+	if count <= 0 {
+		return errors.OrmQueryChunkZeroOrLess
+	}
+
+	column := "id"
+	initialOffset := 0
+	if r.conditions.offset != nil {
+		initialOffset = *r.conditions.offset
+	}
+	var remaining *int
+	if r.conditions.limit != nil {
+		limit := *r.conditions.limit
+		remaining = &limit
+	}
+
+	destType := reflect.TypeOf(r.conditions.model)
+	sliceType := reflect.SliceOf(reflect.PointerTo(destType))
+	var lastID any
+	page := 1
+
+	for {
+		if remaining != nil && *remaining == 0 {
+			break
+		}
+
+		chunkLimit := count
+		if remaining != nil && *remaining < count {
+			chunkLimit = *remaining
+		}
+
+		clone := r.addGlobalScopes()
+
+		if initialOffset > 0 {
+			if page > 1 {
+				clone = clone.Offset(0).(*Query)
+			} else {
+				clone = clone.Offset(initialOffset).(*Query)
+			}
+		}
+
+		if descending {
+			clone = clone.OrderByDesc(column).(*Query)
+		} else {
+			clone = clone.OrderBy(column).(*Query)
+		}
+
+		if lastID != nil {
+			if descending {
+				clone = clone.Where(column+" < ?", lastID).(*Query)
+			} else {
+				clone = clone.Where(column+" > ?", lastID).(*Query)
+			}
+		}
+		clone = clone.Limit(chunkLimit).(*Query)
+
+		query := clone.buildConditions()
+		results := reflect.New(sliceType).Interface()
+
+		err := query.Find(results)
+		if err != nil {
+			return err
+		}
+
+		resultsValue := reflect.ValueOf(results).Elem()
+		length := resultsValue.Len()
+
+		if length == 0 {
+			break
+		}
+
+		if remaining != nil {
+			*remaining = max(*remaining-length, 0)
+		}
+
+		values := make([]any, length)
+		for i := 0; i < length; i++ {
+			values[i] = resultsValue.Index(i).Interface()
+		}
+
+		lastRecord := values[length-1]
+		lastID = database.GetID(lastRecord)
+
+		if err = callback(values); err != nil {
+			return err
+		}
+
+		if length < chunkLimit {
+			break
+		}
+
+		page++
+	}
+
+	return nil
+}
+
 func (r *Query) DB() (*sql.DB, error) {
 	return r.instance.DB()
 }

errors/list.go

@@ -128,6 +128,7 @@ var (
 	OrmMissingWhereClause          = New("WHERE conditions required")
 	OrmNoDialectorsFound           = New("no dialectors found")
 	OrmQueryAssociationsConflict   = New("cannot set orm.Associations and other fields at the same time")
+	OrmQueryChunkZeroOrLess        = New("chunk count must be greater than 0")
 	OrmQueryConditionRequired      = New("query condition is required")
 	OrmQueryEmptyId                = New("id can't be empty")
 	OrmQueryEmptyRelation          = New("relation can't be empty")