github.com/ngicks/gokugen@v0.0.5/impl/work_registry/cli.go (about)

     1  package workregistry
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os/exec"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/ngicks/gokugen"
    11  	"github.com/ngicks/gokugen/cron"
    12  	"github.com/ngicks/type-param-common/set"
    13  )
    14  
    15  var _ cron.WorkRegistry = &Cli{}
    16  
    17  // Cli builds work function from given string assuming it is a cli command.
    18  // Cli also holds whiteList and blackList.
    19  type Cli struct {
    20  	mu        sync.Mutex
    21  	whiteList set.Set[string]
    22  	backList  set.Set[string]
    23  }
    24  
    25  type CliOption func(c *Cli) *Cli
    26  
    27  func CliWhiteList(whilteList []string) CliOption {
    28  	return func(c *Cli) *Cli {
    29  		for _, w := range whilteList {
    30  			c.whiteList.Add(w)
    31  		}
    32  		return c
    33  	}
    34  }
    35  func CliBlackList(backList []string) func(c *Cli) *Cli {
    36  	return func(c *Cli) *Cli {
    37  		for _, w := range backList {
    38  			c.backList.Add(w)
    39  		}
    40  		return c
    41  	}
    42  }
    43  
    44  func NewCli(options ...CliOption) *Cli {
    45  	c := new(Cli)
    46  	for _, opt := range options {
    47  		c = opt(c)
    48  	}
    49  	return c
    50  }
    51  
    52  func (c *Cli) isKeyAllowed(key string) (ok bool) {
    53  	if c.backList.Len() == 0 && c.whiteList.Len() == 0 {
    54  		return true
    55  	}
    56  	if c.backList.Len() != 0 {
    57  		if c.backList.Has(key) {
    58  			return false
    59  		} else {
    60  			return false
    61  		}
    62  	}
    63  	if c.whiteList.Len() != 0 {
    64  		if c.whiteList.Has(key) {
    65  			return true
    66  		} else {
    67  			return false
    68  		}
    69  	}
    70  	// unreachable?
    71  	return false
    72  }
    73  
    74  func (c *Cli) Load(key string) (value gokugen.WorkFnWParam, ok bool) {
    75  	c.mu.Lock()
    76  	defer c.mu.Unlock()
    77  
    78  	ok = c.isKeyAllowed(key)
    79  	if !ok {
    80  		return
    81  	}
    82  	value = buildCliWorkFn(key)
    83  	return
    84  }
    85  
    86  func buildCliWorkFn(command string) cron.WorkFnWParam {
    87  	return func(taskCtx context.Context, scheduled time.Time, param any) (v any, err error) {
    88  
    89  		wg := sync.WaitGroup{}
    90  		wg.Add(1)
    91  
    92  		cmd := exec.CommandContext(taskCtx, command, serializeParam(param)...)
    93  		b, err := cmd.Output()
    94  		if b != nil {
    95  			v = string(b)
    96  		}
    97  
    98  		wg.Wait()
    99  		return
   100  	}
   101  }
   102  
   103  func serializeParam(input any) (params []string) {
   104  	if input == nil {
   105  		return []string{}
   106  	}
   107  	switch x := input.(type) {
   108  	case []string:
   109  		return x
   110  	case string:
   111  		return []string{x}
   112  	case []byte:
   113  		return []string{string(x)}
   114  	default:
   115  		str := fmt.Sprintf("%s", x)
   116  		if str != "" {
   117  			return []string{str}
   118  		} else {
   119  			return []string{}
   120  		}
   121  	}
   122  }