github.com/choria-io/go-choria@v0.28.1-0.20240416190746-b3bf9c7d5a45/client/choria_provisionclient/requester.go (about)

     1  // generated code; DO NOT EDIT
     2  
     3  package choria_provisionclient
     4  
     5  import (
     6  	"context"
     7  	"fmt"
     8  	"time"
     9  
    10  	"github.com/choria-io/go-choria/protocol"
    11  	rpcclient "github.com/choria-io/go-choria/providers/agent/mcorpc/client"
    12  	"github.com/gosuri/uiprogress"
    13  )
    14  
    15  // requester is a generic request handler
    16  type requester struct {
    17  	client   *ChoriaProvisionClient
    18  	action   string
    19  	args     map[string]any
    20  	progress *uiprogress.Bar
    21  }
    22  
    23  // do performs the request
    24  func (r *requester) do(ctx context.Context, handler func(pr protocol.Reply, r *rpcclient.RPCReply)) (*rpcclient.Stats, error) {
    25  	targets := make([]string, len(r.client.targets))
    26  	var err error
    27  
    28  	r.client.Lock()
    29  	copy(targets, r.client.targets)
    30  	discoverer := r.client.ns
    31  	filters := r.client.filters
    32  	fw := r.client.fw
    33  
    34  	opts := []rpcclient.RequestOption{}
    35  	discoveryStart := time.Now()
    36  
    37  	if r.client.ddl.Metadata.Service {
    38  		opts = append(opts, rpcclient.ServiceRequest(), rpcclient.Workers(1))
    39  	} else if len(targets) == 0 {
    40  		if r.client.clientOpts.progress {
    41  			fmt.Print("Discovering nodes .... ")
    42  		} else {
    43  			r.client.infof("Starting discovery")
    44  		}
    45  
    46  		targets, err = discoverer.Discover(ctx, fw, filters)
    47  		if err != nil {
    48  			return nil, err
    49  		}
    50  
    51  		if len(targets) == 0 {
    52  			return nil, fmt.Errorf("no nodes were discovered")
    53  		}
    54  
    55  		if r.client.clientOpts.progress {
    56  			fmt.Printf("%d\n", len(targets))
    57  		} else {
    58  			r.client.infof("Discovered %d nodes", len(targets))
    59  		}
    60  	}
    61  	discoveryEnd := time.Now()
    62  
    63  	if r.client.workers > 0 {
    64  		opts = append(opts, rpcclient.Workers(r.client.workers))
    65  	}
    66  
    67  	if r.client.exprFilter != "" {
    68  		opts = append(opts, rpcclient.ReplyExprFilter(r.client.exprFilter))
    69  	}
    70  
    71  	if len(targets) > 0 {
    72  		opts = append(opts, rpcclient.Targets(targets))
    73  	}
    74  
    75  	opts = append(opts, r.client.clientRPCOpts...)
    76  
    77  	r.client.Unlock()
    78  
    79  	if r.client.clientOpts.progress {
    80  		fmt.Println()
    81  		r.configureProgress(len(targets))
    82  	}
    83  
    84  	agent, err := rpcclient.New(r.client.fw, r.client.ddl.Metadata.Name, rpcclient.DDL(r.client.ddl))
    85  	if err != nil {
    86  		return nil, fmt.Errorf("could not create client: %s", err)
    87  	}
    88  
    89  	if r.client.clientOpts.progress {
    90  		opts = append(opts, rpcclient.ReplyHandler(func(pr protocol.Reply, rpcr *rpcclient.RPCReply) {
    91  			r.progress.Incr()
    92  			handler(pr, rpcr)
    93  		}))
    94  	} else if !r.client.noReplies {
    95  		opts = append(opts, rpcclient.ReplyHandler(handler))
    96  	}
    97  
    98  	if !r.client.clientOpts.progress {
    99  		r.client.debugf("Invoking %s#%s action with %#v", r.client.ddl.Metadata.Name, r.action, r.args)
   100  	}
   101  
   102  	res, err := agent.Do(ctx, r.action, r.args, opts...)
   103  	if err != nil {
   104  		return nil, fmt.Errorf("could not perform request: %s", err)
   105  	}
   106  
   107  	if r.client.clientOpts.progress {
   108  		uiprogress.Stop()
   109  		fmt.Println()
   110  	}
   111  
   112  	if !discoveryStart.IsZero() && !discoveryEnd.IsZero() {
   113  		res.Stats().OverrideDiscoveryTime(discoveryStart, discoveryEnd)
   114  	}
   115  
   116  	return res.Stats(), nil
   117  }
   118  
   119  func (r *requester) configureProgress(count int) {
   120  	if !r.client.clientOpts.progress {
   121  		return
   122  	}
   123  
   124  	r.progress = uiprogress.AddBar(count).AppendCompleted().PrependElapsed()
   125  
   126  	width := r.client.fw.ProgressWidth()
   127  	if width == -1 {
   128  		r.client.clientOpts.progress = false
   129  		return
   130  	}
   131  
   132  	r.progress.Width = width
   133  
   134  	r.progress.PrependFunc(func(b *uiprogress.Bar) string {
   135  		if b.Current() < count {
   136  			return r.client.fw.Colorize("red", "%d / %d", b.Current(), count)
   137  		}
   138  
   139  		return r.client.fw.Colorize("green", "%d / %d", b.Current(), count)
   140  	})
   141  
   142  	uiprogress.Start()
   143  }