github.com/ddfisher/etcdctl@v0.1.2-0.20130925194301-eab7435d452d/sets.go (about)

     1  package main
     2  
     3  import (
     4  	"flag"
     5  	"fmt"
     6  	"crypto/md5"
     7  	"io"
     8  	"os"
     9  	"github.com/coreos/go-etcd/etcd"
    10  )
    11  
    12  const sUsage = `usage: etcdctl [etcd flags] <command>
    13  
    14  Commands:
    15  
    16    sadd <key> <value> [flags]
    17      --ttl to set add a value with a ttl to the set
    18    sdel <key> <value>
    19    smembers <key>
    20    sismember <key> <value>
    21  
    22  `
    23  
    24  var (
    25  	saddFlag = flag.NewFlagSet("sadd", flag.ExitOnError)
    26  	saddTtl = saddFlag.Int64("ttl", 0, "ttl of the key")
    27  )
    28  
    29  func hash(str string) string {
    30  	h := md5.New()
    31  	io.WriteString(h, str)
    32  	return fmt.Sprintf("%x", h.Sum(nil))
    33  }
    34  
    35  func getHeadKey(key string) string {
    36  	return fmt.Sprintf("%s/set-%s", key, hash(key))
    37  }
    38  
    39  func setExists(key string) bool {
    40  	headKey := getHeadKey(key)
    41  	_, err := client.Get(headKey)
    42  	return err == nil
    43  }
    44  
    45  func init() {
    46  	registerCommand("sadd", sUsage, 3, 4, sadd)
    47  	registerCommand("sdel", sUsage, 3, 3, sdel)
    48  	registerCommand("sismember", sUsage, 3, 3, sismember)
    49  	registerCommand("smembers", sUsage, 2, 2, smembers)
    50  }
    51  
    52  func sadd(args []string) error {
    53  
    54  	setKey := args[1]
    55  	value := args[2]
    56  	saddFlag.Parse(args[3:])
    57  
    58  	// Create the set unless it exists
    59  	if ! setExists(setKey) {
    60  		headKey := getHeadKey(setKey)
    61  		_, err := client.Set(headKey, "1", 0)
    62  		if err != nil {
    63  			return err
    64  		}
    65  	}
    66  
    67  	key := fmt.Sprintf("%s/%s", setKey, hash(value))
    68  	_, err := client.Set(key, value, uint64(*saddTtl))
    69  	if err != nil {
    70  		return err
    71  	}
    72  
    73  	fmt.Println(value)
    74  
    75  	return nil
    76  }
    77  
    78  func sdel(args []string) error {
    79  
    80  	setKey := args[1]
    81  
    82  	if ! setExists(setKey) {
    83  		return fmt.Errorf("%s is not a set", setKey)
    84  	}
    85  
    86  	value := args[2]
    87  	key := fmt.Sprintf("%s/%s", setKey, hash(value))
    88  	_, err := client.Delete(key)
    89  	if err != nil {
    90  		err := err.(etcd.EtcdError)
    91  		if err.ErrorCode == 100 {
    92  			return etcd.EtcdError{
    93  				ErrorCode: 100,
    94  				Message: "Not In Set",
    95  				Cause: setKey,
    96  			}
    97  		}
    98  		return err
    99  	}
   100  
   101  	return nil
   102  }
   103  
   104  func smembers(args []string) error {
   105  	setKey := args[1]
   106  
   107  	if ! setExists(setKey) {
   108  		return fmt.Errorf("%s is not a set", setKey)
   109  	}
   110  
   111  	resps, err := client.Get(setKey)
   112  	if err != nil {
   113  		return err
   114  	}
   115  
   116  	headKey := getHeadKey(setKey)
   117  	for _, resp := range resps {
   118  		if resp.Key != headKey {
   119  			fmt.Printf("%s\n", resp.Value)
   120  		}
   121  	}
   122  
   123  	return nil
   124  }
   125  
   126  func sismember(args []string) error {
   127  	setKey := args[1]
   128  	value := args[2]
   129  
   130  	if ! setExists(setKey) {
   131  		return fmt.Errorf("%s is not a set", setKey)
   132  	}
   133  
   134  	key := fmt.Sprintf("%s/%s", setKey, hash(value))
   135  	_, err := client.Get(key)
   136  	if err != nil {
   137  		fmt.Println("false")
   138  		os.Exit(1)
   139  	} else {
   140  		fmt.Println("true")
   141  		os.Exit(0)
   142  	}
   143  
   144  	return nil
   145  }