From f68a5d54430b88f014943c2c7195d085299cacd3 Mon Sep 17 00:00:00 2001 From: iyear Date: Mon, 11 Sep 2023 10:43:31 +0800 Subject: [PATCH] feat(users): enhance output --- app/chat/users.go | 147 ++++++++++++++++++++++++++++------------------ cmd/chat.go | 2 +- 2 files changed, 92 insertions(+), 57 deletions(-) diff --git a/app/chat/users.go b/app/chat/users.go index 1cecb858be..72d15a82c7 100644 --- a/app/chat/users.go +++ b/app/chat/users.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "github.com/go-faster/errors" "os" "time" @@ -11,7 +12,6 @@ import ( "github.com/go-faster/jx" "github.com/gotd/contrib/middleware/ratelimit" "github.com/gotd/td/telegram/peers" - "github.com/gotd/td/telegram/query" "github.com/gotd/td/telegram/query/channels/participants" "github.com/gotd/td/tg" "github.com/jedib0t/go-pretty/v6/progress" @@ -32,107 +32,142 @@ type UsersOptions struct { type User struct { ID int64 `json:"id"` - Bot bool `json:"bot"` Username string `json:"username"` Phone string `json:"phone"` - FirstName string `json:"firstName"` - LastName string `json:"lastName"` + FirstName string `json:"first_name"` + LastName string `json:"last_name"` } -func Users(ctx context.Context, opts *UsersOptions) error { +func Users(ctx context.Context, opts UsersOptions) error { c, kvd, err := tgc.NoLogin(ctx, ratelimit.New(rate.Every(rateInterval), rateBucket)) if err != nil { return err } return tgc.RunWithAuth(ctx, c, func(ctx context.Context) (rerr error) { - var peer peers.Peer - manager := peers.Options{Storage: storage.NewPeers(kvd)}.Build(c.API()) - if opts.Chat == "" { // defaults to me(saved messages) + if opts.Chat == "" { return fmt.Errorf("missing domain id") } - peer, err = utils.Telegram.GetInputPeer(ctx, manager, opts.Chat) + peer, err := utils.Telegram.GetInputPeer(ctx, manager, opts.Chat) if err != nil { return fmt.Errorf("failed to get peer: %w", err) } + ch, ok := peer.(peers.Channel) + if !ok { + return fmt.Errorf("invalid type of chat. channels/groups are supported only") + } + color.Cyan("Occasional suspensions are due to Telegram rate limitations, please wait a moment.") fmt.Println() + f, err := os.Create(opts.Output) + if err != nil { + return err + } + defer multierr.AppendInvoke(&rerr, multierr.Close(f)) + + enc := jx.NewStreamingEncoder(f, 512) + defer multierr.AppendInvoke(&rerr, multierr.Close(enc)) + + enc.ObjStart() + defer enc.ObjEnd() + enc.Field("id", func(e *jx.Encoder) { e.Int64(peer.ID()) }) + pw := prog.New(progress.FormatNumber) pw.SetUpdateFrequency(200 * time.Millisecond) pw.Style().Visibility.TrackerOverall = false pw.Style().Visibility.ETA = false pw.Style().Visibility.Percentage = false - tracker := prog.AppendTracker(pw, progress.FormatNumber, fmt.Sprintf("%s-%d", peer.VisibleName(), peer.ID()), 0) - go pw.Render() - ch, ok := peer.(peers.Channel) - if !ok { - return fmt.Errorf("invalid type of chat. channels are supported only") + builder := func() *participants.GetParticipantsQueryBuilder { + return participants.NewQueryBuilder(c.API()). + GetParticipants(ch.InputChannel()). + BatchSize(100) } - usersList := []*tg.User{} - - iter := participants.NewIterator(query.NewQuery(c.API()).GetParticipants(ch.InputChannel()), 100) - for iter.Next(ctx) { - el := iter.Value() - us, ok := el.User() - if !ok { - continue - } - usersList = append(usersList, us) + fields := map[string]*participants.GetParticipantsQueryBuilder{ + "users": builder(), + "admins": builder().Admins(), + "kicked": builder().Kicked(""), + "banned": builder().Banned(""), + "bots": builder().Bots(), } - if err = iter.Err(); err != nil { - return err + for field, query := range fields { + iter := query.Iter() + if err = outputUsers(ctx, pw, peer, enc, field, iter, opts.Raw); err != nil { + return fmt.Errorf("failed to output %s: %w", field, err) + } } - f, err := os.Create(opts.Output) - if err != nil { - return err - } - defer multierr.AppendInvoke(&rerr, multierr.Close(f)) + prog.Wait(pw) + return nil + }) +} - enc := jx.NewStreamingEncoder(f, 512) - defer multierr.AppendInvoke(&rerr, multierr.Close(enc)) +func outputUsers(ctx context.Context, + pw progress.Writer, + peer peers.Peer, + enc *jx.Encoder, + field string, + iter *participants.Iterator, + raw bool) error { - enc.ObjStart() - defer enc.ObjEnd() - enc.Field("id", func(e *jx.Encoder) { e.Int64(peer.ID()) }) + total, err := iter.Total(ctx) + if err != nil { + return errors.Wrap(err, "get total count") + } - enc.FieldStart("users") - var output any = usersList - if !opts.Raw { - users := make([]User, len(usersList)) - for i := 0; i < len(usersList); i++ { - convertTelegramUser(&users[i], usersList[i]) - } + tracker := prog.AppendTracker(pw, + progress.FormatNumber, + fmt.Sprintf("%s-%d-%s", peer.VisibleName(), peer.ID(), field), + int64(total)) - output = users + enc.FieldStart(field) + enc.ArrStart() + defer enc.ArrEnd() + + for iter.Next(ctx) { + el := iter.Value() + u, ok := el.User() + if !ok { + continue + } + + var output any = u + if !raw { + output = convertTelegramUser(u) } buf, err := json.Marshal(output) if err != nil { - return fmt.Errorf("failed to marshal message: %w", err) + return errors.Wrap(err, "marshal user") } + enc.Raw(buf) - tracker.MarkAsDone() - prog.Wait(pw) - return nil - }) + tracker.Increment(1) + } + + if err = iter.Err(); err != nil { + return err + } + + tracker.MarkAsDone() + return nil } -func convertTelegramUser(dstUser *User, tgUser *tg.User) { - dstUser.ID = tgUser.ID - dstUser.Bot = tgUser.Bot - dstUser.FirstName = tgUser.FirstName - dstUser.LastName = tgUser.LastName - dstUser.Phone = tgUser.Phone - dstUser.Username = tgUser.Username +func convertTelegramUser(u *tg.User) User { + var dst User + dst.ID = u.ID + dst.Username = u.Username + dst.Phone = u.Phone + dst.FirstName = u.FirstName + dst.LastName = u.LastName + return dst } diff --git a/cmd/chat.go b/cmd/chat.go index 3a6492f29b..4d33293157 100644 --- a/cmd/chat.go +++ b/cmd/chat.go @@ -125,7 +125,7 @@ func NewChatUsers() *cobra.Command { Use: "users", Short: "export users from (protected) channels", RunE: func(cmd *cobra.Command, args []string) error { - return chat.Users(logger.Named(cmd.Context(), "users"), &opts) + return chat.Users(logger.Named(cmd.Context(), "users"), opts) }, }