feat(RowIterator): add Count method

Useful when implementing Pagination with Total count.
einride · Aug 20, 2024 · 556d550 · 556d550
1 parent d14e895
commit 556d550
Show file tree

Hide file tree

Showing 10 changed files with 207 additions and 0 deletions.
diff --git a/internal/codegen/databasecodegen/rowiterator.go b/internal/codegen/databasecodegen/rowiterator.go
@@ -31,6 +31,7 @@ func (g RowIteratorCodeGenerator) GenerateCode(f *codegen.File) {
 	f.P("Next() (*", row.Type(), ", error)")
 	f.P("Do(f func(row *", row.Type(), ") error) error")
 	f.P("Stop()")
+	f.P("Count() int64")
 	f.P("}")
 	f.P()
 	f.P("type ", g.StreamingType(), " struct {")
@@ -59,6 +60,10 @@ func (g RowIteratorCodeGenerator) GenerateCode(f *codegen.File) {
 	f.P("})")
 	f.P("}")
 	f.P()
+	f.P("func (i *", g.StreamingType(), ") Count() int64 {")
+	f.P("return i.RowCount")
+	f.P("}")
+	f.P()
 	f.P("type ", g.BufferedType(), " struct {")
 	f.P("rows []*", row.Type())
 	f.P("err error")
@@ -76,6 +81,10 @@ func (g RowIteratorCodeGenerator) GenerateCode(f *codegen.File) {
 	f.P("return next, nil")
 	f.P("}")
 	f.P()
+	f.P("func (i *", g.BufferedType(), ") Count() int64 {")
+	f.P("return int64(len(i.rows))")
+	f.P("}")
+	f.P()
 	f.P("func (i *", g.BufferedType(), ") Do(f func(row *", row.Type(), ") error) error {")
 	f.P("for {")
 	f.P("row, err := i.Next()")

diff --git a/internal/codegen/databasecodegen/testdata/1.sql.database.go b/internal/codegen/databasecodegen/testdata/1.sql.database.go
diff --git a/internal/codegen/databasecodegen/testdata/2.sql.database.go b/internal/codegen/databasecodegen/testdata/2.sql.database.go
diff --git a/internal/codegen/databasecodegen/testdata/3.sql.database.go b/internal/codegen/databasecodegen/testdata/3.sql.database.go
diff --git a/internal/codegen/databasecodegen/testdata/4.sql.database.go b/internal/codegen/databasecodegen/testdata/4.sql.database.go
diff --git a/internal/codegen/databasecodegen/testdata/5.sql.database.go b/internal/codegen/databasecodegen/testdata/5.sql.database.go