Generate DefaultCount function (#282)

* Generate DefaultCount function * dependent function changes * added demo example
infobloxopen · Dec 10, 2024 · 3f7c3c4 · 3f7c3c4
1 parent 1938937
commit 3f7c3c4
Show file tree

Hide file tree

Showing 2 changed files with 127 additions and 0 deletions.
diff --git a/example/feature_demo/demo_service.pb.gorm.go b/example/feature_demo/demo_service.pb.gorm.go
diff --git a/plugin/handlergen.go b/plugin/handlergen.go
@@ -28,6 +28,12 @@ func (p *OrmPlugin) generateDefaultHandlers(file *generator.FileDescriptor) {
 
 			p.generateApplyFieldMask(message)
 			p.generateListHandler(message)
+
+			typeName := p.TypeName(message)
+			ormable := p.getOrmable(typeName)
+			if p.listHasPagination(ormable) {
+				p.generateListCountHandler(message)
+			}
 		}
 	}
 }
@@ -632,6 +638,65 @@ func (p *OrmPlugin) generateListHandler(message *generator.Descriptor) {
 	p.generateAfterListHookDef(ormable)
 }
 
+func (p *OrmPlugin) generateListCountHandler(message *generator.Descriptor) {
+	typeName := p.TypeName(message)
+	ormable := p.getOrmable(typeName)
+
+	p.P(`// DefaultCount`, typeName, ` executes a gorm total record count call`)
+	listSign := fmt.Sprint(`func DefaultCount`, typeName, `(ctx context.Context, db *`, p.Import(gormImport), `.DB`)
+	var f string
+	if p.listHasFiltering(ormable) {
+		listSign += fmt.Sprint(`, f `, `*`, p.Import(queryImport), `.Filtering`)
+		f = "f"
+	} else {
+		f = "nil"
+	}
+	listSign += fmt.Sprint(`) (int64`, `, error) {`)
+	p.P(listSign)
+	p.P(`in := `, typeName, `{}`)
+	p.P(`ormObj, err := in.ToORM(ctx)`)
+	p.P(`if err != nil {`)
+	p.P(`return 0, err`)
+	p.P(`}`)
+	p.generateBeforeCountHookCall(ormable, "ApplyQuery")
+	p.P(`db, err = `, p.Import(tkgormImport), `.ApplyCollectionOperators(ctx, db, &`, ormable.Name, `{}, &`, typeName, `{}, `, f, `,nil,nil,nil`, `)`)
+	p.P(`if err != nil {`)
+	p.P(`return 0, err`)
+	p.P(`}`)
+	p.generateAfterCountHookCall(ormable, "ApplyQuery")
+	p.P(`db = db.Where(&ormObj)`)
+	p.P(`var total int64`)
+	p.P(`if err = db.Model(&ormObj).Count(&total).Error; err != nil {`)
+	p.P(`return 0, err`)
+	p.P(`}`)
+	p.P(`return total, nil`)
+	p.P(`}`)
+	p.generateBeforeCountHookDef(ormable, "ApplyQuery")
+	p.generateAfterCountHookDef(ormable, "ApplyQuery")
+}
+
+func (p *OrmPlugin) generateBeforeCountHookDef(orm *OrmableType, suffix string) {
+	p.P(`type `, orm.Name, `WithBeforeCount`, suffix, ` interface {`)
+	hookSign := fmt.Sprint(`BeforeCount`, suffix, `(context.Context, *`, p.Import(gormImport), `.DB`)
+	if p.listHasFiltering(orm) {
+		hookSign += fmt.Sprint(`, *`, p.Import(queryImport), `.Filtering`)
+	}
+	hookSign += fmt.Sprint(`) (*`, p.Import(gormImport), `.DB, error)`)
+	p.P(hookSign)
+	p.P(`}`)
+}
+
+func (p *OrmPlugin) generateAfterCountHookDef(orm *OrmableType, suffix string) {
+	p.P(`type `, orm.Name, `WithAfterCount`, suffix, ` interface {`)
+	hookSign := fmt.Sprint(`AfterCount`, suffix, `(context.Context, *`, p.Import(gormImport), `.DB`)
+	if p.listHasFiltering(orm) {
+		hookSign += fmt.Sprint(`, *`, p.Import(queryImport), `.Filtering`)
+	}
+	hookSign += fmt.Sprint(`) (*`, p.Import(gormImport), `.DB, error)`)
+	p.P(hookSign)
+	p.P(`}`)
+}
+
 func (p *OrmPlugin) generateBeforeListHookDef(orm *OrmableType, suffix string) {
 	p.P(`type `, orm.Name, `WithBeforeList`, suffix, ` interface {`)
 	hookSign := fmt.Sprint(`BeforeList`, suffix, `(context.Context, *`, p.Import(gormImport), `.DB`)
@@ -694,6 +759,32 @@ func (p *OrmPlugin) generateBeforeListHookCall(orm *OrmableType, suffix string)
 	p.P(`}`)
 }
 
+func (p *OrmPlugin) generateBeforeCountHookCall(orm *OrmableType, suffix string) {
+	p.P(`if hook, ok := interface{}(&ormObj).(`, orm.Name, `WithBeforeCount`, suffix, `); ok {`)
+	hookCall := fmt.Sprint(`if db, err = hook.BeforeCount`, suffix, `(ctx, db`)
+	if p.listHasFiltering(orm) {
+		hookCall += `,f`
+	}
+	hookCall += `); err != nil {`
+	p.P(hookCall)
+	p.P(`return 0, err`)
+	p.P(`}`)
+	p.P(`}`)
+}
+
+func (p *OrmPlugin) generateAfterCountHookCall(orm *OrmableType, suffix string) {
+	p.P(`if hook, ok := interface{}(&ormObj).(`, orm.Name, `WithAfterCount`, suffix, `); ok {`)
+	hookCall := fmt.Sprint(`if db, err = hook.AfterCount`, suffix, `(ctx, db`)
+	if p.listHasFiltering(orm) {
+		hookCall += `,f`
+	}
+	hookCall += `); err != nil {`
+	p.P(hookCall)
+	p.P(`return 0, err`)
+	p.P(`}`)
+	p.P(`}`)
+}
+
 func (p *OrmPlugin) generateAfterListHookCall(orm *OrmableType) {
 	p.P(`if hook, ok := interface{}(&ormObj).(`, orm.Name, `WithAfterListFind); ok {`)
 	hookCall := fmt.Sprint(`if err = hook.AfterListFind(ctx, db, &ormResponse`)