perf: add telegram.WithMiddlewareFunc

TBXark · Nov 6, 2024 · 32d80c5 · 32d80c5
1 parent c664ca0
commit 32d80c5
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/internal/server/bot/bot.go b/internal/server/bot/bot.go
@@ -37,7 +37,6 @@ func (b *Bot) Run(ctx context.Context) error {
 		b.BindCommand(CommandStart, route[botv1.BotHandlerBotServiceStart])
 		b.BindCommand(CommandCounter, route[botv1.BotHandlerBotServiceCounter], sfMid)
 		b.BindCallback(QueryCounter, route[botv1.BotHandlerBotServiceCounter], sfMid)
-
 		return nil
 	})
 }

diff --git a/pkg/telegram/handler.go b/pkg/telegram/handler.go
@@ -12,12 +12,15 @@ import (
 )
 
 type HandlerFunc = func(ctx context.Context, bot *bot.Bot, update *models.Update) error
+type MiddlewareFunc = func(next HandlerFunc) HandlerFunc
 type ErrorHandlerFunc = func(ctx context.Context, bot *bot.Bot, update *models.Update, err error)
 
 func WithMiddleware(h HandlerFunc, e ErrorHandlerFunc, middleware ...bot.Middleware) bot.HandlerFunc {
 	handler := func(ctx context.Context, bot *bot.Bot, update *models.Update) {
 		if err := h(ctx, bot, update); err != nil {
-			e(ctx, bot, update, err)
+			if e != nil {
+				e(ctx, bot, update, err)
+			}
 		}
 	}
 	for i := len(middleware) - 1; i >= 0; i-- {
@@ -26,6 +29,20 @@ func WithMiddleware(h HandlerFunc, e ErrorHandlerFunc, middleware ...bot.Middlew
 	return handler
 }
 
+func WithMiddlewareFunc(h HandlerFunc, e ErrorHandlerFunc, middleware ...MiddlewareFunc) bot.HandlerFunc {
+	handler := h
+	for i := len(middleware) - 1; i >= 0; i-- {
+		handler = middleware[i](handler)
+	}
+	return func(ctx context.Context, bot *bot.Bot, update *models.Update) {
+		if err := handler(ctx, bot, update); err != nil {
+			if e != nil {
+				e(ctx, bot, update, err)
+			}
+		}
+	}
+}
+
 func NewSingleFlightMiddleware() bot.Middleware {
 	sf := &singleflight.Group{}
 	return func(next bot.HandlerFunc) bot.HandlerFunc {