<commit_before>package main

import (
  "fmt"
  "net"
  "bufio"
  "strings"
  "net/http"
  "time"
  "log/syslog"
  "log"
)

// BEGIN Configuration ===================================================================================================

const SleepTime = 30 * time.Minute

// Blocklist URLS need to return a list of IP addresses either in raw IP, I.e. 192.168.1.4 or
// CIDR format 192.168.0.0/16
var blocklists = []string{"https://www.spamhaus.org/drop/drop.txt",
        "https://www.spamhaus.org/drop/edrop.txt",
        "https://rules.emergingthreats.net/fwrules/emerging-Block-IPs.txt",
        "https://lists.blocklist.de/lists/all.txt"}

// Announcement & Withdrawl strings %s will be replaced by the CIDR formatted route to add/remove
const announce_template = "announce route %s next-hop 192.0.2.1 community [65332:666]\n"
const withdraw_template = "withdraw route %s next-hop 192.0.2.1 community [65332:666]\n"

// END Configuration ===================================================================================================

type IPSet struct {
  root BitNode
}

type BitNode struct {
  parent,zero,one *BitNode
  depth uint32
  full bool
  value uint32
}

// Process blocklists
func main () {

  var ipsetCurrent = createIPSet()

  sysLog, err := syslog.Dial("udp", "localhost:514",
		syslog.LOG_INFO|syslog.LOG_DAEMON, "blocklist")
	if err != nil {
		log.Fatal(err)
	}

  // Run forever, we sleep at the end of the loop
  for 1==1 {
    var announcements = 0
    var withdrawls = 0
    var ipsetNew = createIPSet()

    sysLog.Info("begin blocklist refresh")

    for _, url := range blocklists {
      data := downloadBlocklist(url)
      sysLog.Info(fmt.Sprintf("%d entries downloaded from %s\n",len(data), url))

      for _, ipnet := range data {
        ipsetNew.add(&ipnet)
      }
    }

    // If IPs in the new set aren't in the current set, announce them
    for _, newip := range ipsetNew.getAll() {
      current := ipsetCurrent.contains(&newip)
      if current == nil {
        fmt.Printf(announce_template, newip.String())
        announcements++
      } else {
        // However, if the thing that matched wasn't identical to the existing one
        // I.e. more specific IP, withdraw the existing one *and* annouce the new one
        if current.String() != newip.String() {
          fmt.Printf(withdraw_template, current.String())
          fmt.Printf(announce_template, newip.String())
          announcements++
          withdrawls++
        }
      }
    }

    // If IPs from the current set aren't in the new set
    for _, existing := range ipsetCurrent.getAll() {
      newip := ipsetNew.contains(&existing)

      if newip == nil {
        fmt.Printf(withdraw_template, existing.String())
        withdrawls++
      } else {
        // However, if the thing that matched wasn't identical to the existing one
        // I.e. more specific IP, withdraw the existing one *and* annouce the new one
        if newip != nil && newip.String() != existing.String() {
          fmt.Printf(withdraw_template, existing.String())
          fmt.Printf(announce_template, newip.String())
          announcements++
          withdrawls++
        }
      }
    }

    ipsetCurrent = ipsetNew
    sysLog.Info(fmt.Sprintf("completed with %d routes announced and %d routes withdrawn\n",announcements, withdrawls))
    time.Sleep(SleepTime)
  }
}

func downloadBlocklist(url string) []net.IPNet {
  var nets = make([]net.IPNet,0)
  resp, err := http.Get(url)

  if err != nil {
    fmt.Println(err)
    return nil
  }

  scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
    line := scanner.Text()
    items := strings.Split(line,";")
    ipnet := StringToIPNet(items[0])

    if ipnet != nil {
        nets = append(nets,*ipnet)
    }
	}
  return nets
}

func createIPSet() IPSet {
  return IPSet { root:  BitNode{parent: nil, zero :nil, one :nil, depth: 32, full: false, value: 0} }
}

func (ipset *IPSet) add(ipnet *net.IPNet) {
  if ipnet != nil {
      var bits,_ = ipnet.Mask.Size()
      var addr = IPtoInt(ipnet.IP)

      addIP(&ipset.root,addr,uint32(bits))
  }
}

func (ipset *IPSet) getAll() []net.IPNet {
  return collectIPs(&ipset.root)
}

func (ipset *IPSet) contains(ipnet *net.IPNet) *net.IPNet {
  if ipnet != nil {
      var bits,_ = ipnet.Mask.Size()
      var addr = IPtoInt(ipnet.IP)

      return containsIP(&ipset.root,addr,uint32(bits))
  } else {
    return nil
  }
}

func StringToIPNet(text string) *net.IPNet {
  var ip,ipnet,error = net.ParseCIDR(strings.TrimSpace(text))

  if error != nil {
    ip = net.ParseIP(text)
    if ip != nil {
      ipnet = &net.IPNet{IP: ip, Mask: net.CIDRMask(32,32)}
    }
  }

  return ipnet
}

func IPtoInt(addr net.IP) uint32 {
  return uint32(addr.To4()[0]) << 24 |
         uint32(addr.To4()[1]) << 16 |
         uint32(addr.To4()[2]) << 8 |
         uint32(addr.To4()[3])
}

func IntToIP(ip uint32) net.IP {
  result := make(net.IP, 4)
           result[3] = byte(ip)
           result[2] = byte(ip >>8)
           result[1] = byte(ip >>16)
           result[0] = byte(ip >> 24)
  return result
}

func CheckBit(num uint32, bit uint32) bool {
  return (num & (1 << bit)) != 0
}

func collectIPs(node *BitNode) []net.IPNet {
  var nets = make([]net.IPNet,0)

  if node.full {
    var ipnet = IPNetFromNode(node)
    return []net.IPNet{ipnet}
  } else {
    if node.zero != nil {
      nets = append(nets,collectIPs(node.zero)...)
    }

    if node.one != nil {
      nets = append(nets,collectIPs(node.one)...)
    }
  }
  return nets
}

func IPNetFromNode(node *BitNode) net.IPNet {
  var cur = node
  var accumulate uint32 = 0
  var mask = int(32 - node.depth)

  for cur.parent != nil {
    accumulate |= cur.value << cur.depth
    cur = cur.parent
  }

  return net.IPNet{IP: IntToIP(accumulate), Mask: net.CIDRMask(mask,32)}
}

func containsIP(node *BitNode, addr uint32, mask uint32) *net.IPNet {
  if node.full {
    ipnet := IPNetFromNode(node)

    return &ipnet
  }

  if 32-node.depth > mask {
    return nil
  }

  if CheckBit(addr,node.depth - 1) {
    if node.one == nil {
      return nil
    } else {
      return containsIP(node.one,addr,mask)
    }
  } else {
    if node.zero == nil {
      return nil
    } else {
      return containsIP(node.zero,addr,mask)
    }
  }
}

func addIP(node *BitNode, addr uint32, mask uint32) bool {

  if node.depth == 0 || 32 - node.depth == mask || node.full {
    node.full = true
    return node.full
  }

  var child *BitNode

  if CheckBit(addr,node.depth - 1) {
    if node.one == nil {
      child = &BitNode{parent: node, zero: nil, one: nil, depth: node.depth -1, full: false, value: 1}
      node.one = child
    } else {
      child = node.one
    }
  } else {
    if node.zero == nil {
      child = &BitNode{parent: node, zero: nil, one: nil, depth: node.depth -1, full: false, value: 0 }
      node.zero = child
    } else {
      child = node.zero
    }
  }

  addIP(child, addr, mask)

  node.full = (node.one != nil && node.one.full) && (node.zero != nil && node.zero.full)
  return node.full
}
<commit_msg>Syslog on HTTP failures Skip IPv6<commit_after>package main

import (
  "fmt"
  "net"
  "bufio"
  "strings"
  "net/http"
  "time"
  "log/syslog"
  "log"
)

// BEGIN Configuration ===================================================================================================

const SleepTime = 30 * time.Minute

// Blocklist URLS need to return a list of IP addresses either in raw IP, I.e. 192.168.1.4 or
// CIDR format 192.168.0.0/16
var blocklists = []string{"https://www.spamhaus.org/drop/drop.txt",
        "https://www.spamhaus.org/drop/edrop.txt",
        "https://rules.emergingthreats.net/fwrules/emerging-Block-IPs.txt",
        "https://lists.blocklist.de/lists/all.txt"}

// Announcement & Withdrawl strings %s will be replaced by the CIDR formatted route to add/remove
const announce_template = "announce route %s next-hop 192.0.2.1 community [65332:666]\n"
const withdraw_template = "withdraw route %s next-hop 192.0.2.1 community [65332:666]\n"

// END Configuration ===================================================================================================

type IPSet struct {
  root BitNode
}

type BitNode struct {
  parent,zero,one *BitNode
  depth uint32
  full bool
  value uint32
}

var sysLog *syslog.Writer

// Process blocklists
func main () {

  var ipsetCurrent = createIPSet()

  logger, err := syslog.Dial("udp", "localhost:514",
		syslog.LOG_INFO|syslog.LOG_DAEMON, "blocklist")
	if err != nil {
		log.Fatal(err)
	}

	sysLog = logger

  // Run forever, we sleep at the end of the loop
  for 1==1 {
    var announcements = 0
    var withdrawls = 0
    var ipsetNew = createIPSet()

    sysLog.Info("begin blocklist refresh")

    for _, url := range blocklists {
      data := downloadBlocklist(url)
      sysLog.Info(fmt.Sprintf("%d entries downloaded from %s\n",len(data), url))

      for _, ipnet := range data {
        ipsetNew.add(&ipnet)
      }
    }

    // If IPs in the new set aren't in the current set, announce them
    for _, newip := range ipsetNew.getAll() {
      current := ipsetCurrent.contains(&newip)
      if current == nil {
        fmt.Printf(announce_template, newip.String())
        announcements++
      } else {
        // However, if the thing that matched wasn't identical to the existing one
        // I.e. more specific IP, withdraw the existing one *and* annouce the new one
        if current.String() != newip.String() {
          fmt.Printf(withdraw_template, current.String())
          fmt.Printf(announce_template, newip.String())
          announcements++
          withdrawls++
        }
      }
    }

    // If IPs from the current set aren't in the new set
    for _, existing := range ipsetCurrent.getAll() {
      newip := ipsetNew.contains(&existing)

      if newip == nil {
        fmt.Printf(withdraw_template, existing.String())
        withdrawls++
      } else {
        // However, if the thing that matched wasn't identical to the existing one
        // I.e. more specific IP, withdraw the existing one *and* annouce the new one
        if newip != nil && newip.String() != existing.String() {
          fmt.Printf(withdraw_template, existing.String())
          fmt.Printf(announce_template, newip.String())
          announcements++
          withdrawls++
        }
      }
    }

    ipsetCurrent = ipsetNew
    sysLog.Info(fmt.Sprintf("completed with %d routes announced and %d routes withdrawn\n",announcements, withdrawls))
    time.Sleep(SleepTime)
  }
}

func downloadBlocklist(url string) []net.IPNet {
  var nets = make([]net.IPNet,0)
  resp, err := http.Get(url)

  if err != nil {
    sysLog.Err(fmt.Sprintf("Error downloading %s: %v",url,err))
    return nil
  }

  scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
    line := scanner.Text()
    items := strings.Split(line,";")
    ipnet := StringToIPNet(items[0])

    if ipnet != nil {
        nets = append(nets,*ipnet)
    }
	}
  return nets
}

func createIPSet() IPSet {
  return IPSet { root:  BitNode{parent: nil, zero :nil, one :nil, depth: 32, full: false, value: 0} }
}

func (ipset *IPSet) add(ipnet *net.IPNet) {

  if ipnet != nil {
      // We can't handle IPv6 now
      if ipnet.IP.To4() == nil {
        return
      }

      var bits,_ = ipnet.Mask.Size()
      var addr = IPtoInt(ipnet.IP)

      addIP(&ipset.root,addr,uint32(bits))
  }
}

func (ipset *IPSet) getAll() []net.IPNet {
  return collectIPs(&ipset.root)
}

func (ipset *IPSet) contains(ipnet *net.IPNet) *net.IPNet {
  if ipnet != nil {
      var bits,_ = ipnet.Mask.Size()
      var addr = IPtoInt(ipnet.IP)

      return containsIP(&ipset.root,addr,uint32(bits))
  } else {
    return nil
  }
}

func StringToIPNet(text string) *net.IPNet {
  var ip,ipnet,error = net.ParseCIDR(strings.TrimSpace(text))

  if error != nil {
    ip = net.ParseIP(text)
    if ip != nil {
      ipnet = &net.IPNet{IP: ip, Mask: net.CIDRMask(32,32)}
    }
  }

  return ipnet
}

func IPtoInt(addr net.IP) uint32 {
  return uint32(addr.To4()[0]) << 24 |
         uint32(addr.To4()[1]) << 16 |
         uint32(addr.To4()[2]) << 8 |
         uint32(addr.To4()[3])
}

func IntToIP(ip uint32) net.IP {
  result := make(net.IP, 4)
           result[3] = byte(ip)
           result[2] = byte(ip >>8)
           result[1] = byte(ip >>16)
           result[0] = byte(ip >> 24)
  return result
}

func CheckBit(num uint32, bit uint32) bool {
  return (num & (1 << bit)) != 0
}

func collectIPs(node *BitNode) []net.IPNet {
  var nets = make([]net.IPNet,0)

  if node.full {
    var ipnet = IPNetFromNode(node)
    return []net.IPNet{ipnet}
  } else {
    if node.zero != nil {
      nets = append(nets,collectIPs(node.zero)...)
    }

    if node.one != nil {
      nets = append(nets,collectIPs(node.one)...)
    }
  }
  return nets
}

func IPNetFromNode(node *BitNode) net.IPNet {
  var cur = node
  var accumulate uint32 = 0
  var mask = int(32 - node.depth)

  for cur.parent != nil {
    accumulate |= cur.value << cur.depth
    cur = cur.parent
  }

  return net.IPNet{IP: IntToIP(accumulate), Mask: net.CIDRMask(mask,32)}
}

func containsIP(node *BitNode, addr uint32, mask uint32) *net.IPNet {
  if node.full {
    ipnet := IPNetFromNode(node)

    return &ipnet
  }

  if 32-node.depth > mask {
    return nil
  }

  if CheckBit(addr,node.depth - 1) {
    if node.one == nil {
      return nil
    } else {
      return containsIP(node.one,addr,mask)
    }
  } else {
    if node.zero == nil {
      return nil
    } else {
      return containsIP(node.zero,addr,mask)
    }
  }
}

func addIP(node *BitNode, addr uint32, mask uint32) bool {

  if node.depth == 0 || 32 - node.depth == mask || node.full {
    node.full = true
    return node.full
  }

  var child *BitNode

  if CheckBit(addr,node.depth - 1) {
    if node.one == nil {
      child = &BitNode{parent: node, zero: nil, one: nil, depth: node.depth -1, full: false, value: 1}
      node.one = child
    } else {
      child = node.one
    }
  } else {
    if node.zero == nil {
      child = &BitNode{parent: node, zero: nil, one: nil, depth: node.depth -1, full: false, value: 0 }
      node.zero = child
    } else {
      child = node.zero
    }
  }

  addIP(child, addr, mask)

  node.full = (node.one != nil && node.one.full) && (node.zero != nil && node.zero.full)
  return node.full
}
<|endoftext|>


<commit_before>package bolt

import (
	"fmt"
	"io"
	"os"
	"time"

	"github.com/boltdb/bolt"
	"github.com/golang/protobuf/proto"
	"github.com/octavore/naga/service"
)

// Module bolt provides methods for writing and retrieving
// data from a bolt database.
type Module struct {
	DB *bolt.DB
}

// Init implements service.Init
func (m *Module) Init(c *service.Config) {
	c.Setup = func() (err error) {
		if c.Env().IsTest() {
			return
		}
		m.DB, err = bolt.Open("default.db", os.ModePerm, &bolt.Options{
			Timeout: 30 * time.Second,
		})
		if err != nil {
			return err
		}
		return m.init()
	}
	var testDB string
	c.SetupTest = func() {
		suffix := time.Now().Format("20060102150405.999")
		testDB = fmt.Sprintf("test-%s.db", suffix)
		err := os.RemoveAll(testDB)
		if err != nil {
			panic(err)
		}
		m.DB, err = bolt.Open(testDB, os.ModePerm, &bolt.Options{
			Timeout: 30 * time.Second,
		})
		if err != nil {
			panic(err)
		}
		m.init()
		return
	}
	c.Stop = func() {
		err := m.DB.Close()
		if err != nil {
			panic(err)
		}
		if testDB != "" {
			err := os.RemoveAll(testDB)
			if err != nil {
				panic(err)
			}
		}
	}
}

type ErrNoKey string

func (e ErrNoKey) Error() string {
	return fmt.Sprintf("key not found: %s", string(e))
}

func (m *Module) init() error {
	return m.DB.Update(func(tx *bolt.Tx) error {
		buckets := []string{PAGE_BUCKET, ROUTE_BUCKET, USER_BUCKET}
		for _, bucket := range buckets {
			_, err := tx.CreateBucketIfNotExists([]byte(bucket))
			if err != nil {
				return err
			}
		}
		return nil
	})
}

type AddressableProto interface {
	GetUuid() string
	proto.Message
}

func (m *Module) Get(bucket, key string, pb proto.Message) error {
	return m.DB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		data := b.Get([]byte(key))
		if data == nil {
			return ErrNoKey(key)
		}
		return proto.Unmarshal(data, pb)
	})
}

func (m *Module) Update(bucket string, pb AddressableProto) error {
	data, err := proto.Marshal(pb)
	if err != nil {
		return nil
	}
	key := []byte(pb.GetUuid())
	if len(key) == 0 {
		return fmt.Errorf("no uuid for proto")
	}
	return m.DB.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		return b.Put(key, data)
	})
}

func (m *Module) BackupToFile(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	return m.Backup(f)
}

func (m *Module) Backup(w io.Writer) error {
	return m.DB.View(func(tx *bolt.Tx) error {
		_, err := tx.WriteTo(w)
		return err
	})
}
<commit_msg>bolt: Missing err return.<commit_after>package bolt

import (
	"fmt"
	"io"
	"os"
	"time"

	"github.com/boltdb/bolt"
	"github.com/golang/protobuf/proto"
	"github.com/octavore/naga/service"
)

// Module bolt provides methods for writing and retrieving
// data from a bolt database.
type Module struct {
	DB *bolt.DB
}

// Init implements service.Init
func (m *Module) Init(c *service.Config) {
	c.Setup = func() (err error) {
		if c.Env().IsTest() {
			return
		}
		m.DB, err = bolt.Open("default.db", os.ModePerm, &bolt.Options{
			Timeout: 30 * time.Second,
		})
		if err != nil {
			return err
		}
		return m.init()
	}
	var testDB string
	c.SetupTest = func() {
		suffix := time.Now().Format("20060102150405.999")
		testDB = fmt.Sprintf("test-%s.db", suffix)
		err := os.RemoveAll(testDB)
		if err != nil {
			panic(err)
		}
		m.DB, err = bolt.Open(testDB, os.ModePerm, &bolt.Options{
			Timeout: 30 * time.Second,
		})
		if err != nil {
			panic(err)
		}
		m.init()
		return
	}
	c.Stop = func() {
		err := m.DB.Close()
		if err != nil {
			panic(err)
		}
		if testDB != "" {
			err := os.RemoveAll(testDB)
			if err != nil {
				panic(err)
			}
		}
	}
}

type ErrNoKey string

func (e ErrNoKey) Error() string {
	return fmt.Sprintf("key not found: %s", string(e))
}

func (m *Module) init() error {
	return m.DB.Update(func(tx *bolt.Tx) error {
		buckets := []string{PAGE_BUCKET, ROUTE_BUCKET, USER_BUCKET}
		for _, bucket := range buckets {
			_, err := tx.CreateBucketIfNotExists([]byte(bucket))
			if err != nil {
				return err
			}
		}
		return nil
	})
}

type AddressableProto interface {
	GetUuid() string
	proto.Message
}

func (m *Module) Get(bucket, key string, pb proto.Message) error {
	return m.DB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		data := b.Get([]byte(key))
		if data == nil {
			return ErrNoKey(key)
		}
		return proto.Unmarshal(data, pb)
	})
}

func (m *Module) Update(bucket string, pb AddressableProto) error {
	data, err := proto.Marshal(pb)
	if err != nil {
		return err
	}
	key := []byte(pb.GetUuid())
	if len(key) == 0 {
		return fmt.Errorf("no uuid for proto")
	}
	return m.DB.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		return b.Put(key, data)
	})
}

func (m *Module) BackupToFile(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	return m.Backup(f)
}

func (m *Module) Backup(w io.Writer) error {
	return m.DB.View(func(tx *bolt.Tx) error {
		_, err := tx.WriteTo(w)
		return err
	})
}
<|endoftext|>


<commit_before>package bot

import (
	"fmt"
	"strings"
	"time"

	"strconv"

	"github.com/nlopes/slack"
)

var (
	prefix        string = "party_"
	partyIndexKey string = prefix + "keys"
)

func rescheduleParty(bot *Meu) {
	items, err := bot.rc.Keys(prefix + "*")
	if err != nil {
		return
	}
	now := time.Now()
	for _, key := range items {
		parts := strings.Split(key, "_")
		sec, _ := strconv.Atoi(parts[1])
		saved_time := time.Unix(int64(sec), 0)
		if now.After(saved_time) {
			bot.rc.Erase(key)
		} else {
			scheduleParty(bot, &saved_time, parts[2])
			if _, err := bot.rc.SortedSetRank(partyIndexKey, key).Result(); err != nil {
				bot.rc.SortedSetAdd(partyIndexKey, sec, key)
			}
		}
	}

	bot.rc.SortedSetRemoveRange(partyIndexKey, 0, now.Unix())
}

func scheduleParty(bot *Meu, date *time.Time, keyword string) {
	bot.cr.AddFunc(fmt.Sprintf("0 %d %d %d %d *", date.Minute(), date.Hour(), date.Day(), date.Month()), alarmFuncGenerator(bot, keyword, partyKey(date, keyword)))
}

func partyKey(date *time.Time, keyword string) string {
	return prefix + strconv.FormatInt(date.Unix(), 10) + "_" + keyword
}

func alarmFuncGenerator(bot *Meu, keyword string, key string) func() {
	return func() {
		list, err := bot.rc.SetList(key)
		bot.rc.Erase(key)
		if err == nil {
			members := make([]string, len(list))
			for i, item := range list {
				members[i] = fmt.Sprintf("<@%s>", item)
			}
			bot.PostMessage("#random", fmt.Sprintf("'%s' 파티 10분 전이다 메우. %s", keyword, strings.Join(members, " ")), slack.PostMessageParameters{
				AsUser:    false,
				Username:  "파티 안내원 메우",
				IconEmoji: ":meu:",
			})
		}
	}
}

func register_party(bot *Meu, e *slack.MessageEvent, matched []string) {
	now := time.Now()
	month, err := strconv.Atoi(matched[1])
	var not_set struct {
		month bool
		day   bool
	}
	if err != nil {
		month = int(now.Month())
		not_set.month = true
	}
	day, err := strconv.Atoi(matched[2])
	if err != nil {
		day = now.Day()
		not_set.day = true
	}
	hour, err := strconv.Atoi(matched[3])
	min, err := strconv.Atoi(matched[4])
	if err != nil {
		min = 0
	}
	keyword := matched[5]

	date := time.Date(now.Year(), time.Month(month), day, hour, min, 0, 0, now.Location())
	if date.Before(now) {
		corrected := false
		// first try after 12 hour
		if not_set.day {
			if date.Hour() < 12 {
				date = date.Add(time.Hour * 12)
				if corrected = !date.Before(now); !corrected {
					// reset
					date = date.Add(time.Hour * -12)
				}
			}
			if !corrected {
				date = date.AddDate(0, 0, 1)
			}
		} else if not_set.month {
			date = date.AddDate(0, 1, 0)
		} else {
			date = date.AddDate(1, 0, 0)
		}
		if date.Before(now) {
			bot.replySimple(e, "과거에 대해서 파티를 모집할 수 없다 메우")
			return
		}
	}
	key := partyKey(&date, keyword)
	inserted := bot.rc.SetAdd(key, e.User)
	if inserted.Val() == 1 {
		bot.replySimple(e, fmt.Sprintf("파티 대기에 들어갔다 메우. - %s %s", date.String(), keyword))
		cardinal := bot.rc.SetCard(key)
		if cardinal.Val() == 1 {
			scheduleParty(bot, &date, keyword)
		}
	} else {
		bot.replySimple(e, fmt.Sprintf("이미 들어가있는 파티다 메우. - %s %s", date.String(), keyword))
	}
}
<commit_msg>Seperate register key into index function<commit_after>package bot

import (
	"fmt"
	"strings"
	"time"

	"strconv"

	"github.com/nlopes/slack"
)

var (
	prefix        string = "party_"
	partyIndexKey string = prefix + "keys"
)

func rescheduleParty(bot *Meu) {
	items, err := bot.rc.Keys(prefix + "*")
	if err != nil {
		return
	}
	now := time.Now()
	for _, key := range items {
		parts := strings.Split(key, "_")
		sec, _ := strconv.Atoi(parts[1])
		saved_time := time.Unix(int64(sec), 0)
		if now.After(saved_time) {
			bot.rc.Erase(key)
		} else {
			scheduleParty(bot, &saved_time, parts[2])
			registerToIndex(bot, &saved_time, key)
		}
	}

	bot.rc.SortedSetRemoveRange(partyIndexKey, 0, now.Unix())
}

func registerToIndex(bot *Meu, date *time.Time, key string) {
	if _, err := bot.rc.SortedSetRank(partyIndexKey, key).Result(); err != nil {
		bot.rc.SortedSetAdd(partyIndexKey, int(date.Unix()), key)
	}
}

func scheduleParty(bot *Meu, date *time.Time, keyword string) {
	bot.cr.AddFunc(fmt.Sprintf("0 %d %d %d %d *", date.Minute(), date.Hour(), date.Day(), date.Month()), alarmFuncGenerator(bot, keyword, partyKey(date, keyword)))
}

func partyKey(date *time.Time, keyword string) string {
	return prefix + strconv.FormatInt(date.Unix(), 10) + "_" + keyword
}

func alarmFuncGenerator(bot *Meu, keyword string, key string) func() {
	return func() {
		list, err := bot.rc.SetList(key)
		bot.rc.Erase(key)
		if err == nil {
			members := make([]string, len(list))
			for i, item := range list {
				members[i] = fmt.Sprintf("<@%s>", item)
			}
			bot.PostMessage("#random", fmt.Sprintf("'%s' 파티 10분 전이다 메우. %s", keyword, strings.Join(members, " ")), slack.PostMessageParameters{
				AsUser:    false,
				Username:  "파티 안내원 메우",
				IconEmoji: ":meu:",
			})
		}
	}
}

func register_party(bot *Meu, e *slack.MessageEvent, matched []string) {
	now := time.Now()
	month, err := strconv.Atoi(matched[1])
	var not_set struct {
		month bool
		day   bool
	}
	if err != nil {
		month = int(now.Month())
		not_set.month = true
	}
	day, err := strconv.Atoi(matched[2])
	if err != nil {
		day = now.Day()
		not_set.day = true
	}
	hour, err := strconv.Atoi(matched[3])
	min, err := strconv.Atoi(matched[4])
	if err != nil {
		min = 0
	}
	keyword := matched[5]

	date := time.Date(now.Year(), time.Month(month), day, hour, min, 0, 0, now.Location())
	if date.Before(now) {
		corrected := false
		// first try after 12 hour
		if not_set.day {
			if date.Hour() < 12 {
				date = date.Add(time.Hour * 12)
				if corrected = !date.Before(now); !corrected {
					// reset
					date = date.Add(time.Hour * -12)
				}
			}
			if !corrected {
				date = date.AddDate(0, 0, 1)
			}
		} else if not_set.month {
			date = date.AddDate(0, 1, 0)
		} else {
			date = date.AddDate(1, 0, 0)
		}
		if date.Before(now) {
			bot.replySimple(e, "과거에 대해서 파티를 모집할 수 없다 메우")
			return
		}
	}
	key := partyKey(&date, keyword)
	inserted := bot.rc.SetAdd(key, e.User)
	registerToIndex(bot, &date, key)
	if inserted.Val() == 1 {
		bot.replySimple(e, fmt.Sprintf("파티 대기에 들어갔다 메우. - %s %s", date.String(), keyword))
		cardinal := bot.rc.SetCard(key)
		if cardinal.Val() == 1 {
			scheduleParty(bot, &date, keyword)
		}
	} else {
		bot.replySimple(e, fmt.Sprintf("이미 들어가있는 파티다 메우. - %s %s", date.String(), keyword))
	}
}
<|endoftext|>


<commit_before>package worker

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/bitly/go-simplejson"
	"github.com/cenk/backoff"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/travis-ci/worker/backend"
	"github.com/travis-ci/worker/context"
	"github.com/travis-ci/worker/metrics"

	gocontext "context"
)

var (
	httpJobQueueNoJobsErr = fmt.Errorf("no jobs available")
)

// ProcessorEacherSizer is the minimal interface required by the HTTPJobQueue
// for interacting with the ProcessorPool
type ProcessorEacherSizer interface {
	Each(func(int, *Processor))
	Size() int
}

// HTTPJobQueue is a JobQueue that uses http
type HTTPJobQueue struct {
	processors   ProcessorEacherSizer
	jobBoardURL  *url.URL
	site         string
	providerName string
	queue        string
	workerID     string
	pollInterval time.Duration

	DefaultLanguage, DefaultDist, DefaultGroup, DefaultOS string
}

type httpFetchJobsRequest struct {
	Jobs []string `json:"jobs"`
}

type httpFetchJobsResponse struct {
	Jobs []string `json:"jobs"`
}

type jobBoardErrorResponse struct {
	Type          string `json:"@type"`
	Error         string `json:"error"`
	UpstreamError string `json:"upstream_error,omitempty"`
}

// NewHTTPJobQueue creates a new job-board job queue
func NewHTTPJobQueue(processors ProcessorEacherSizer, jobBoardURL *url.URL, site, providerName, queue, workerID string) (*HTTPJobQueue, error) {
	return &HTTPJobQueue{
		processors:   processors,
		jobBoardURL:  jobBoardURL,
		site:         site,
		providerName: providerName,
		queue:        queue,
		workerID:     workerID,
		// TODO: make pollInterval configurable
		pollInterval: time.Second,
	}, nil
}

// Jobs consumes new jobs from job-board
func (q *HTTPJobQueue) Jobs(ctx gocontext.Context) (outChan <-chan Job, err error) {
	buildJobChan := make(chan Job)
	outChan = buildJobChan
	logger := context.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"self": "http_job_queue",
		"inst": fmt.Sprintf("%p", q),
	})

	go func() {
		for {
			logger.Debug("polling for job tick")
			if !q.pollForJob(ctx, buildJobChan) {
				return
			}
			time.Sleep(q.pollInterval)
		}
	}()

	return outChan, nil
}

func (q *HTTPJobQueue) pollForJob(ctx gocontext.Context, buildJobChan chan Job) bool {
	logger := context.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"self": "http_job_queue",
		"inst": fmt.Sprintf("%p", q),
	})

	logger.Debug("fetching job id")
	jobID, err := q.fetchJobID(ctx, 1, []uint64{})
	if err != nil {
		logger.WithField("err", err).Debug("continuing after failing to get job id")
		return true
	}
	logger.WithField("job_id", jobID).Debug("fetching complete job")
	buildJob, err := q.fetchJob(ctx, jobID)
	if err != nil {
		logger.WithFields(logrus.Fields{
			"err": err,
			"id":  jobID,
		}).Warn("failed to get complete job")
		return true
	}

	logger.WithField("job_id", jobID).Debug("sending job to output channel")
	jobSendBegin := time.Now()
	select {
	case buildJobChan <- buildJob:
		metrics.TimeSince("travis.worker.job_queue.http.blocking_time", jobSendBegin)
		logger.WithFields(logrus.Fields{
			"source": "http",
			"dur":    time.Since(jobSendBegin),
		}).Info("sent job to output channel")
		return true
	case <-ctx.Done():
		logger.WithField("err", ctx.Err()).Warn("returning from jobs loop due to context done")
		return false
	}
}

func (q *HTTPJobQueue) fetchJobID(ctx gocontext.Context, desired uint64, running []uint64) (uint64, error) {
	logger := context.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"self": "http_job_queue",
		"inst": fmt.Sprintf("%p", q),
	})

	processorName, ok := context.ProcessorFromContext(ctx)
	if !ok {
		processorName = "unknown-processor"
	}
	fetchRequestPayload := &httpFetchJobsRequest{Jobs: []string{}}
	for _, jobID := range running {
		fetchRequestPayload.Jobs = append(fetchRequestPayload.Jobs, strconv.Itoa(int(jobID)))
	}

	jobIDsJSON, err := json.Marshal(fetchRequestPayload)
	if err != nil {
		return 0, errors.Wrap(err, "failed to marshal job board jobs request payload")
	}

	u := *q.jobBoardURL

	query := u.Query()
	query.Add("count", strconv.Itoa(int(desired)))
	query.Add("capacity", "1")
	query.Add("queue", q.queue)

	u.Path = "/jobs"
	u.RawQuery = query.Encode()

	client := &http.Client{}

	req, err := http.NewRequest("POST", u.String(), bytes.NewReader(jobIDsJSON))
	if err != nil {
		return 0, errors.Wrap(err, "failed to create job board jobs request")
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Travis-Site", q.site)
	req.Header.Add("From", fmt.Sprintf("%s+%s", processorName, q.workerID))
	req = req.WithContext(ctx)

	resp, err := client.Do(req)
	if err != nil {
		return 0, errors.Wrap(err, "failed to make job board jobs request")
	}

	defer resp.Body.Close()
	fetchResponsePayload := &httpFetchJobsResponse{}
	err = json.NewDecoder(resp.Body).Decode(&fetchResponsePayload)
	if err != nil {
		return 0, errors.Wrap(err, "failed to decode job board jobs response")
	}

	logger.WithField("jobs", fetchResponsePayload.Jobs).Debug("fetched raw jobs")
	var jobIDs []uint64
	for _, strID := range fetchResponsePayload.Jobs {
		alreadyRunning := false
		for _, prevStrID := range fetchRequestPayload.Jobs {
			if strID == prevStrID {
				alreadyRunning = true
			}
		}
		if alreadyRunning {
			logger.WithField("job_id", strID).Debug("skipping running job")
			continue
		}

		id, err := strconv.ParseUint(strID, 10, 64)
		if err != nil {
			return 0, errors.Wrap(err, "failed to parse job ID")
		}
		jobIDs = append(jobIDs, id)
	}

	if len(jobIDs) == 0 {
		return 0, httpJobQueueNoJobsErr
	}

	logger.WithField("job_id", jobIDs[0]).Debug("returning first filtered job ID")
	return jobIDs[0], nil
}

func (q *HTTPJobQueue) fetchJob(ctx gocontext.Context, id uint64) (Job, error) {
	logger := context.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"self": "http_job_queue",
		"inst": fmt.Sprintf("%p", q),
	})

	processorName, ok := context.ProcessorFromContext(ctx)
	if !ok {
		processorName = "unknown-processor"
	}

	buildJob := &httpJob{
		payload: &httpJobPayload{
			Data: &JobPayload{},
		},
		startAttributes: &backend.StartAttributes{},

		refreshClaim: q.generateJobRefreshClaimFunc(10*time.Second, id),

		jobBoardURL: q.jobBoardURL,
		site:        q.site,
		workerID:    fmt.Sprintf("%s+%s", processorName, q.workerID),
	}
	startAttrs := &httpJobPayloadStartAttrs{
		Data: &jobPayloadStartAttrs{
			Config: &backend.StartAttributes{},
		},
	}

	u := *q.jobBoardURL
	u.Path = fmt.Sprintf("/jobs/%d", id)

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return nil, errors.Wrap(err, "couldn't make job board job request")
	}

	// TODO: ensure infrastructure is not synonymous with providerName since
	// there's the possibility that a provider has multiple infrastructures, which
	// is expected to be the case with the future cloudbrain provider.
	req.Header.Add("Travis-Infrastructure", q.providerName)
	req.Header.Add("Travis-Site", q.site)
	req.Header.Add("From", fmt.Sprintf("%s+%s", processorName, q.workerID))
	req = req.WithContext(ctx)

	bo := backoff.NewExponentialBackOff()
	bo.MaxInterval = 10 * time.Second
	bo.MaxElapsedTime = 1 * time.Minute

	var resp *http.Response
	err = backoff.Retry(func() (err error) {
		resp, err = (&http.Client{}).Do(req)
		if resp != nil && resp.StatusCode != http.StatusOK {
			logger.WithFields(logrus.Fields{
				"expected_status": http.StatusOK,
				"actual_status":   resp.StatusCode,
			}).Debug("job fetch failed")

			if resp.Body != nil {
				resp.Body.Close()
			}

			return errors.Errorf("expected %d but got %d", http.StatusOK, resp.StatusCode)
		}
		return
	}, bo)

	if err != nil {
		return nil, errors.Wrap(err, "error making job board job request")
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, errors.Wrap(err, "error reading body from job board job request")
	}

	err = json.Unmarshal(body, buildJob.payload)
	if err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal job board payload")
	}

	err = json.Unmarshal(body, &startAttrs)
	if err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal start attributes from job board")
	}

	rawPayload, err := simplejson.NewJson(body)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse raw payload with simplejson")
	}
	buildJob.rawPayload = rawPayload.Get("data")

	buildJob.startAttributes = startAttrs.Data.Config
	buildJob.startAttributes.VMType = buildJob.payload.Data.VMType
	buildJob.startAttributes.SetDefaults(q.DefaultLanguage, q.DefaultDist, q.DefaultGroup, q.DefaultOS, VMTypeDefault)

	return buildJob, nil
}

func (q *HTTPJobQueue) generateJobRefreshClaimFunc(pollInterval time.Duration, id uint64) func(gocontext.Context) {
	return func(ctx gocontext.Context) {
		for {
			_, err := q.fetchJobID(ctx, 0, []uint64{id})
			if err != nil && err != httpJobQueueNoJobsErr {
				context.LoggerFromContext(ctx).WithFields(logrus.Fields{
					"err":    err,
					"job_id": id,
				}).Warn("failed to refresh claim")
			}
			select {
			case <-ctx.Done():
				return
			case <-time.After(pollInterval):
			}
		}
	}
}

// Name returns the name of this queue type, wow!
func (q *HTTPJobQueue) Name() string {
	return "http"
}

// Cleanup does not do anything!
func (q *HTTPJobQueue) Cleanup() error {
	return nil
}
<commit_msg>Add log debugging for http job fetch operation when refreshing claim<commit_after>package worker

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/bitly/go-simplejson"
	"github.com/cenk/backoff"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/travis-ci/worker/backend"
	"github.com/travis-ci/worker/context"
	"github.com/travis-ci/worker/metrics"

	gocontext "context"
)

var (
	httpJobQueueNoJobsErr = fmt.Errorf("no jobs available")
)

// ProcessorEacherSizer is the minimal interface required by the HTTPJobQueue
// for interacting with the ProcessorPool
type ProcessorEacherSizer interface {
	Each(func(int, *Processor))
	Size() int
}

// HTTPJobQueue is a JobQueue that uses http
type HTTPJobQueue struct {
	processors   ProcessorEacherSizer
	jobBoardURL  *url.URL
	site         string
	providerName string
	queue        string
	workerID     string
	pollInterval time.Duration

	DefaultLanguage, DefaultDist, DefaultGroup, DefaultOS string
}

type httpFetchJobsRequest struct {
	Jobs []string `json:"jobs"`
}

type httpFetchJobsResponse struct {
	Jobs []string `json:"jobs"`
}

type jobBoardErrorResponse struct {
	Type          string `json:"@type"`
	Error         string `json:"error"`
	UpstreamError string `json:"upstream_error,omitempty"`
}

// NewHTTPJobQueue creates a new job-board job queue
func NewHTTPJobQueue(processors ProcessorEacherSizer, jobBoardURL *url.URL, site, providerName, queue, workerID string) (*HTTPJobQueue, error) {
	return &HTTPJobQueue{
		processors:   processors,
		jobBoardURL:  jobBoardURL,
		site:         site,
		providerName: providerName,
		queue:        queue,
		workerID:     workerID,
		// TODO: make pollInterval configurable
		pollInterval: time.Second,
	}, nil
}

// Jobs consumes new jobs from job-board
func (q *HTTPJobQueue) Jobs(ctx gocontext.Context) (outChan <-chan Job, err error) {
	buildJobChan := make(chan Job)
	outChan = buildJobChan
	logger := context.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"self": "http_job_queue",
		"inst": fmt.Sprintf("%p", q),
	})

	go func() {
		for {
			logger.Debug("polling for job tick")
			if !q.pollForJob(ctx, buildJobChan) {
				return
			}
			time.Sleep(q.pollInterval)
		}
	}()

	return outChan, nil
}

func (q *HTTPJobQueue) pollForJob(ctx gocontext.Context, buildJobChan chan Job) bool {
	logger := context.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"self": "http_job_queue",
		"inst": fmt.Sprintf("%p", q),
	})

	logger.Debug("fetching job id")
	jobID, err := q.fetchJobID(ctx, 1, []uint64{})
	if err != nil {
		logger.WithField("err", err).Debug("continuing after failing to get job id")
		return true
	}
	logger.WithField("job_id", jobID).Debug("fetching complete job")
	buildJob, err := q.fetchJob(ctx, jobID)
	if err != nil {
		logger.WithFields(logrus.Fields{
			"err": err,
			"id":  jobID,
		}).Warn("failed to get complete job")
		return true
	}

	logger.WithField("job_id", jobID).Debug("sending job to output channel")
	jobSendBegin := time.Now()
	select {
	case buildJobChan <- buildJob:
		metrics.TimeSince("travis.worker.job_queue.http.blocking_time", jobSendBegin)
		logger.WithFields(logrus.Fields{
			"source": "http",
			"dur":    time.Since(jobSendBegin),
		}).Info("sent job to output channel")
		return true
	case <-ctx.Done():
		logger.WithField("err", ctx.Err()).Warn("returning from jobs loop due to context done")
		return false
	}
}

func (q *HTTPJobQueue) fetchJobID(ctx gocontext.Context, desired uint64, running []uint64) (uint64, error) {
	logger := context.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"self": "http_job_queue",
		"inst": fmt.Sprintf("%p", q),
	})

	processorName, ok := context.ProcessorFromContext(ctx)
	if !ok {
		processorName = "unknown-processor"
	}
	fetchRequestPayload := &httpFetchJobsRequest{Jobs: []string{}}
	for _, jobID := range running {
		fetchRequestPayload.Jobs = append(fetchRequestPayload.Jobs, strconv.Itoa(int(jobID)))
	}

	if len(fetchRequestPayload.Jobs) > 0 {
		logger.WithFields(logrus.Fields{
			"running": strings.Join(fetchRequestPayload.Jobs, ","),
		}).Debug("sending record of running job(s) to job-board")
	}

	jobIDsJSON, err := json.Marshal(fetchRequestPayload)
	if err != nil {
		return 0, errors.Wrap(err, "failed to marshal job-board jobs request payload")
	}

	u := *q.jobBoardURL

	query := u.Query()
	query.Add("count", strconv.Itoa(int(desired)))
	query.Add("capacity", "1")
	query.Add("queue", q.queue)

	u.Path = "/jobs"
	u.RawQuery = query.Encode()

	client := &http.Client{}

	req, err := http.NewRequest("POST", u.String(), bytes.NewReader(jobIDsJSON))
	if err != nil {
		return 0, errors.Wrap(err, "failed to create job-board jobs request")
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Travis-Site", q.site)
	req.Header.Add("From", fmt.Sprintf("%s+%s", processorName, q.workerID))
	req = req.WithContext(ctx)

	resp, err := client.Do(req)
	if err != nil {
		return 0, errors.Wrap(err, "failed to make job-board jobs request")
	}

	defer resp.Body.Close()
	fetchResponsePayload := &httpFetchJobsResponse{}
	err = json.NewDecoder(resp.Body).Decode(&fetchResponsePayload)
	if err != nil {
		return 0, errors.Wrap(err, "failed to decode job-board jobs response")
	}

	logger.WithField("jobs", fetchResponsePayload.Jobs).Debug("fetched raw jobs")
	var jobIDs []uint64
	for _, strID := range fetchResponsePayload.Jobs {
		alreadyRunning := false
		for _, prevStrID := range fetchRequestPayload.Jobs {
			if strID == prevStrID {
				alreadyRunning = true
			}
		}
		if alreadyRunning {
			logger.WithField("job_id", strID).Debug("skipping running job")
			continue
		}

		id, err := strconv.ParseUint(strID, 10, 64)
		if err != nil {
			return 0, errors.Wrap(err, "failed to parse job ID")
		}
		jobIDs = append(jobIDs, id)
	}

	if len(jobIDs) == 0 {
		return 0, httpJobQueueNoJobsErr
	}

	logger.WithField("job_id", jobIDs[0]).Debug("returning first filtered job ID")
	return jobIDs[0], nil
}

func (q *HTTPJobQueue) fetchJob(ctx gocontext.Context, id uint64) (Job, error) {
	logger := context.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"self": "http_job_queue",
		"inst": fmt.Sprintf("%p", q),
	})

	processorName, ok := context.ProcessorFromContext(ctx)
	if !ok {
		processorName = "unknown-processor"
	}

	buildJob := &httpJob{
		payload: &httpJobPayload{
			Data: &JobPayload{},
		},
		startAttributes: &backend.StartAttributes{},

		refreshClaim: q.generateJobRefreshClaimFunc(10*time.Second, id),

		jobBoardURL: q.jobBoardURL,
		site:        q.site,
		workerID:    fmt.Sprintf("%s+%s", processorName, q.workerID),
	}
	startAttrs := &httpJobPayloadStartAttrs{
		Data: &jobPayloadStartAttrs{
			Config: &backend.StartAttributes{},
		},
	}

	u := *q.jobBoardURL
	u.Path = fmt.Sprintf("/jobs/%d", id)

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return nil, errors.Wrap(err, "couldn't make job-board job request")
	}

	// TODO: ensure infrastructure is not synonymous with providerName since
	// there's the possibility that a provider has multiple infrastructures, which
	// is expected to be the case with the future cloudbrain provider.
	req.Header.Add("Travis-Infrastructure", q.providerName)
	req.Header.Add("Travis-Site", q.site)
	req.Header.Add("From", fmt.Sprintf("%s+%s", processorName, q.workerID))
	req = req.WithContext(ctx)

	bo := backoff.NewExponentialBackOff()
	bo.MaxInterval = 10 * time.Second
	bo.MaxElapsedTime = 1 * time.Minute

	var resp *http.Response
	err = backoff.Retry(func() (err error) {
		resp, err = (&http.Client{}).Do(req)
		if resp != nil && resp.StatusCode != http.StatusOK {
			logger.WithFields(logrus.Fields{
				"expected_status": http.StatusOK,
				"actual_status":   resp.StatusCode,
			}).Debug("job fetch failed")

			if resp.Body != nil {
				resp.Body.Close()
			}

			return errors.Errorf("expected %d but got %d", http.StatusOK, resp.StatusCode)
		}
		return
	}, bo)

	if err != nil {
		return nil, errors.Wrap(err, "error making job-board job request")
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, errors.Wrap(err, "error reading body from job-board job request")
	}

	err = json.Unmarshal(body, buildJob.payload)
	if err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal job-board payload")
	}

	err = json.Unmarshal(body, &startAttrs)
	if err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal start attributes from job-board")
	}

	rawPayload, err := simplejson.NewJson(body)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse raw payload with simplejson")
	}
	buildJob.rawPayload = rawPayload.Get("data")

	buildJob.startAttributes = startAttrs.Data.Config
	buildJob.startAttributes.VMType = buildJob.payload.Data.VMType
	buildJob.startAttributes.SetDefaults(q.DefaultLanguage, q.DefaultDist, q.DefaultGroup, q.DefaultOS, VMTypeDefault)

	return buildJob, nil
}

func (q *HTTPJobQueue) generateJobRefreshClaimFunc(pollInterval time.Duration, id uint64) func(gocontext.Context) {
	return func(ctx gocontext.Context) {
		for {
			_, err := q.fetchJobID(ctx, 0, []uint64{id})
			if err != nil && err != httpJobQueueNoJobsErr {
				context.LoggerFromContext(ctx).WithFields(logrus.Fields{
					"err":    err,
					"job_id": id,
				}).Warn("failed to refresh claim")
			}
			select {
			case <-ctx.Done():
				return
			case <-time.After(pollInterval):
			}
		}
	}
}

// Name returns the name of this queue type, wow!
func (q *HTTPJobQueue) Name() string {
	return "http"
}

// Cleanup does not do anything!
func (q *HTTPJobQueue) Cleanup() error {
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
)

func main() {
	http.HandleFunc("/", catHandler)
	http.ListenAndServe(":8080", nil)
}

func catHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("Got Req")
	w.WriteHeader(200)
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		panic(err)
	}

	fmt.Println(string(body))
}
<commit_msg>Allow $ADDR to specify address to bind to<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
)

func main() {
	addr := os.Getenv("ADDR")
	if addr == "" {
		addr = ":8080"
	}

	http.HandleFunc("/", catHandler)
	http.ListenAndServe(addr, nil)
}

func catHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("Got Req")
	w.WriteHeader(200)
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		panic(err)
	}

	fmt.Println(string(body))
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/bouk/httprouter"
	"github.com/influxdata/chronograf"
	"github.com/influxdata/chronograf/uuid"
)

// ValidTemplateRequest checks if the request sent to the server is the correct format.
func ValidTemplateRequest(template *chronograf.Template) error {
	switch template.Type {
	default:
		return fmt.Errorf("Unknown template type %s", template.Type)
	case "query", "constant", "csv", "fieldKeys", "tagKeys", "tagValues":
	}

	for _, v := range template.Values {
		switch v.Type {
		default:
			return fmt.Errorf("Unknown template variable type %s", v.Type)
		case "csv", "fieldKey", "tagKey", "tagValue":
		}
	}

	if template.Type == "query" && template.Query == nil {
		return fmt.Errorf("No query set for template of type 'query'")
	}
	return nil
}

type templateLinks struct {
	Self string `json:"self"` // Self link mapping to this resource
}

type templateResponse struct {
	chronograf.Template
	Links templateLinks `json:"links"`
}

func newTemplateResponses(dID chronograf.DashboardID, tmps []chronograf.Template) []templateResponse {
	res := make([]templateResponse, len(tmps))
	for i, t := range tmps {
		res[i] = newTemplateResponse(dID, t)
	}
	return res
}

func newTemplateResponse(dID chronograf.DashboardID, tmp chronograf.Template) templateResponse {
	base := "/chronograf/v1/dashboards"
	return templateResponse{
		Template: tmp,
		Links: templateLinks{
			Self: fmt.Sprintf("%s/%d/templates/%s", base, dID, tmp.ID),
		},
	}
}

// Templates returns all templates from a dashboard within the store
func (s *Service) Templates(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	e, err := s.DashboardsStore.Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	boards := newDashboardResponse(e)
	templates := boards.Templates
	encodeJSON(w, http.StatusOK, templates, s.Logger)
}

// NewTemplate adds a template to an existing dashboard
func (s *Service) NewTemplate(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.DashboardsStore.Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	var template chronograf.Template
	if err := json.NewDecoder(r.Body).Decode(&template); err != nil {
		invalidJSON(w, s.Logger)
		return
	}

	if err := ValidTemplateRequest(&template); err != nil {
		invalidData(w, err, s.Logger)
		return
	}

	ids := uuid.V4{}
	tid, err := ids.Generate()
	if err != nil {
		msg := fmt.Sprintf("Error creating template ID for dashboard %d: %v", id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}
	template.ID = chronograf.TemplateID(tid)

	dash.Templates = append(dash.Templates, template)
	if err := s.DashboardsStore.Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error adding template %s to dashboard %d: %v", tid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}

	res := newTemplateResponse(dash.ID, template)
	encodeJSON(w, http.StatusOK, res, s.Logger)
}

// TemplateID retrieves a specific template from a dashboard
func (s *Service) TemplateID(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	dash, err := s.DashboardsStore.Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	tid := httprouter.GetParamFromContext(ctx, "tid")
	for _, t := range dash.Templates {
		if t.ID == chronograf.TemplateID(tid) {
			res := newTemplateResponse(chronograf.DashboardID(id), t)
			encodeJSON(w, http.StatusOK, res, s.Logger)
		}
	}

	notFound(w, id, s.Logger)
}

// RemoveTemplate removes a specific template from an existing dashboard
func (s *Service) RemoveTemplate(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.DashboardsStore.Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	tid := httprouter.GetParamFromContext(ctx, "tid")
	pos := -1
	for i, t := range dash.Templates {
		if t.ID == chronograf.TemplateID(tid) {
			pos = i
			break
		}
	}
	if pos == -1 {
		notFound(w, id, s.Logger)
		return
	}

	dash.Templates = append(dash.Templates[:pos], dash.Templates[pos+1:]...)
	if err := s.DashboardsStore.Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error removing template %s from dashboard %d: %v", tid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

// ReplaceTemplate replaces a template entirely within an existing dashboard
func (s *Service) ReplaceTemplate(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.DashboardsStore.Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	tid := httprouter.GetParamFromContext(ctx, "tid")
	pos := -1
	for i, t := range dash.Templates {
		if t.ID == chronograf.TemplateID(tid) {
			pos = i
			break
		}
	}
	if pos == -1 {
		notFound(w, id, s.Logger)
		return
	}

	var template chronograf.Template
	if err := json.NewDecoder(r.Body).Decode(&template); err != nil {
		invalidJSON(w, s.Logger)
		return
	}

	if err := ValidTemplateRequest(&template); err != nil {
		invalidData(w, err, s.Logger)
		return
	}
	template.ID = chronograf.TemplateID(tid)

	dash.Templates[pos] = template
	if err := s.DashboardsStore.Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error updating template %s in dashboard %d: %v", tid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}

	res := newTemplateResponse(chronograf.DashboardID(id), template)
	encodeJSON(w, http.StatusOK, res, s.Logger)
}
<commit_msg>Update templates API response<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/bouk/httprouter"
	"github.com/influxdata/chronograf"
	"github.com/influxdata/chronograf/uuid"
)

// ValidTemplateRequest checks if the request sent to the server is the correct format.
func ValidTemplateRequest(template *chronograf.Template) error {
	switch template.Type {
	default:
		return fmt.Errorf("Unknown template type %s", template.Type)
	case "query", "constant", "csv", "fieldKeys", "tagKeys", "tagValues":
	}

	for _, v := range template.Values {
		switch v.Type {
		default:
			return fmt.Errorf("Unknown template variable type %s", v.Type)
		case "csv", "fieldKey", "tagKey", "tagValue":
		}
	}

	if template.Type == "query" && template.Query == nil {
		return fmt.Errorf("No query set for template of type 'query'")
	}
	return nil
}

type templateLinks struct {
	Self string `json:"self"` // Self link mapping to this resource
}

type templateResponse struct {
	chronograf.Template
	Links templateLinks `json:"links"`
}

func newTemplateResponses(dID chronograf.DashboardID, tmps []chronograf.Template) []templateResponse {
	res := make([]templateResponse, len(tmps))
	for i, t := range tmps {
		res[i] = newTemplateResponse(dID, t)
	}
	return res
}

type templatesResponses struct {
	Templates []templateResponse `json:"templates"`
}

func newTemplateResponse(dID chronograf.DashboardID, tmp chronograf.Template) templateResponse {
	base := "/chronograf/v1/dashboards"
	return templateResponse{
		Template: tmp,
		Links: templateLinks{
			Self: fmt.Sprintf("%s/%d/templates/%s", base, dID, tmp.ID),
		},
	}
}

// Templates returns all templates from a dashboard within the store
func (s *Service) Templates(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	d, err := s.DashboardsStore.Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	res := templatesResponses{
		Templates: newTemplateResponses(chronograf.DashboardID(id), d.Templates),
	}
	encodeJSON(w, http.StatusOK, res, s.Logger)
}

// NewTemplate adds a template to an existing dashboard
func (s *Service) NewTemplate(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.DashboardsStore.Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	var template chronograf.Template
	if err := json.NewDecoder(r.Body).Decode(&template); err != nil {
		invalidJSON(w, s.Logger)
		return
	}

	if err := ValidTemplateRequest(&template); err != nil {
		invalidData(w, err, s.Logger)
		return
	}

	ids := uuid.V4{}
	tid, err := ids.Generate()
	if err != nil {
		msg := fmt.Sprintf("Error creating template ID for dashboard %d: %v", id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}
	template.ID = chronograf.TemplateID(tid)

	dash.Templates = append(dash.Templates, template)
	if err := s.DashboardsStore.Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error adding template %s to dashboard %d: %v", tid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}

	res := newTemplateResponse(dash.ID, template)
	encodeJSON(w, http.StatusOK, res, s.Logger)
}

// TemplateID retrieves a specific template from a dashboard
func (s *Service) TemplateID(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	dash, err := s.DashboardsStore.Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	tid := httprouter.GetParamFromContext(ctx, "tid")
	for _, t := range dash.Templates {
		if t.ID == chronograf.TemplateID(tid) {
			res := newTemplateResponse(chronograf.DashboardID(id), t)
			encodeJSON(w, http.StatusOK, res, s.Logger)
		}
	}

	notFound(w, id, s.Logger)
}

// RemoveTemplate removes a specific template from an existing dashboard
func (s *Service) RemoveTemplate(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.DashboardsStore.Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	tid := httprouter.GetParamFromContext(ctx, "tid")
	pos := -1
	for i, t := range dash.Templates {
		if t.ID == chronograf.TemplateID(tid) {
			pos = i
			break
		}
	}
	if pos == -1 {
		notFound(w, id, s.Logger)
		return
	}

	dash.Templates = append(dash.Templates[:pos], dash.Templates[pos+1:]...)
	if err := s.DashboardsStore.Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error removing template %s from dashboard %d: %v", tid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

// ReplaceTemplate replaces a template entirely within an existing dashboard
func (s *Service) ReplaceTemplate(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.DashboardsStore.Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	tid := httprouter.GetParamFromContext(ctx, "tid")
	pos := -1
	for i, t := range dash.Templates {
		if t.ID == chronograf.TemplateID(tid) {
			pos = i
			break
		}
	}
	if pos == -1 {
		notFound(w, id, s.Logger)
		return
	}

	var template chronograf.Template
	if err := json.NewDecoder(r.Body).Decode(&template); err != nil {
		invalidJSON(w, s.Logger)
		return
	}

	if err := ValidTemplateRequest(&template); err != nil {
		invalidData(w, err, s.Logger)
		return
	}
	template.ID = chronograf.TemplateID(tid)

	dash.Templates[pos] = template
	if err := s.DashboardsStore.Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error updating template %s in dashboard %d: %v", tid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}

	res := newTemplateResponse(chronograf.DashboardID(id), template)
	encodeJSON(w, http.StatusOK, res, s.Logger)
}
<|endoftext|>


<commit_before>package service

import (
	"encoding/xml"
	"errors"
	"fmt"
	"github.com/coopernurse/gorp"
	"github.com/pascalj/disgo/models"
	"io"
	"time"
)

type disqus struct {
	Threads []thread `xml:"thread"`
	Posts   []post   `xml:"post"`
}

type thread struct {
	Id       string `xml:"id"`
	ThreadId string `xml:"http://disqus.com/disqus-internals id,attr"`
	Link     string `xml:"link"`
	Title    string `xml:"title"`
}

type post struct {
	Message   string    `xml:"message"`
	CreatedAt string    `xml:"createdAt"`
	IsDeleted string    `xml:"isDeleted"`
	IsSpam    string    `xml:"isSpam"`
	Author    author    `xml:"author"`
	IpAddress string    `xml:"ipAddress"`
	Thread    threadRef `xml:"thread"`
}

type author struct {
	Email string `xml:"email"`
	Name  string `xml:"name"`
}

type threadRef struct {
	Id string `xml:"http://disqus.com/disqus-internals id,attr"`
}

func Import(dbmap *gorp.DbMap, xmlReader io.Reader) error {
	parsed := &disqus{}
	decoder := xml.NewDecoder(xmlReader)
	if err := decoder.Decode(parsed); err != nil {
		return err
	}
	comments := make([]models.Comment, 0)
	for _, post := range parsed.Posts {
		thread, err := parsed.findThread(post.Thread.Id)
		if err != nil {
			fmt.Errorf("Could not find thread reference", post.Thread.Id)
		}
		if post.IsDeleted == "true" {
			continue
		}
		createdAt, err := time.Parse(time.RFC3339, post.CreatedAt)
		if err != nil {
			createdAt = time.Now()
		}
		comment := models.NewComment(post.Author.Email, post.Author.Name, "", post.Message, thread.Link, post.IpAddress, "")
		comment.Created = createdAt.Unix()
		comment.Approved = post.IsSpam == "false"
		comments = append(comments, comment)
	}

	total := len(parsed.Posts)
	fmt.Println("Read", total, "from the file.")

	for i, comment := range comments {
		if err := dbmap.Insert(&comment); err != nil {
			total--
			fmt.Errorf("Could not import comment", i)
		}
	}

	fmt.Println("Wrote", total, "comments to the database.")

	return nil
}

func (dis *disqus) findThread(id string) (thread, error) {
	for _, thread := range dis.Threads {
		if thread.Id == id {
			return thread, nil
		}
	}
	return thread{}, errors.New("Could not find thread")
}
<commit_msg>Remove HTML upon import<commit_after>package service

import (
	"encoding/xml"
	"errors"
	"fmt"
	"github.com/coopernurse/gorp"
	"github.com/kennygrant/sanitize"
	"github.com/pascalj/disgo/models"
	"io"
	"time"
)

type disqus struct {
	Threads []thread `xml:"thread"`
	Posts   []post   `xml:"post"`
}

type thread struct {
	Id       string `xml:"id"`
	ThreadId string `xml:"http://disqus.com/disqus-internals id,attr"`
	Link     string `xml:"link"`
	Title    string `xml:"title"`
}

type post struct {
	Message   string    `xml:"message"`
	CreatedAt string    `xml:"createdAt"`
	IsDeleted string    `xml:"isDeleted"`
	IsSpam    string    `xml:"isSpam"`
	Author    author    `xml:"author"`
	IpAddress string    `xml:"ipAddress"`
	Thread    threadRef `xml:"thread"`
}

type author struct {
	Email string `xml:"email"`
	Name  string `xml:"name"`
}

type threadRef struct {
	Id string `xml:"http://disqus.com/disqus-internals id,attr"`
}

func Import(dbmap *gorp.DbMap, xmlReader io.Reader) error {
	parsed := &disqus{}
	decoder := xml.NewDecoder(xmlReader)
	if err := decoder.Decode(parsed); err != nil {
		return err
	}
	comments := make([]models.Comment, 0)
	for _, post := range parsed.Posts {
		thread, err := parsed.findThread(post.Thread.Id)
		if err != nil {
			fmt.Println("Could not find thread reference", post.Thread.Id)
		}
		if post.IsDeleted == "true" {
			continue
		}
		createdAt, err := time.Parse(time.RFC3339, post.CreatedAt)
		if err != nil {
			createdAt = time.Now()
		}
		body := sanitize.HTML(post.Message)
		comment := models.NewComment(post.Author.Email, post.Author.Name, "", body, thread.Link, post.IpAddress, "")
		comment.Created = createdAt.Unix()
		comment.Approved = post.IsSpam == "false"
		comments = append(comments, comment)
	}

	total := len(parsed.Posts)
	fmt.Println("Read", total, "from the file.")

	for i, comment := range comments {
		if err := dbmap.Insert(&comment); err != nil {
			total--
			fmt.Println("Could not import comment", i)
		}
	}

	fmt.Println("Wrote", total, "comments to the database.")

	return nil
}

func (dis *disqus) findThread(id string) (thread, error) {
	for _, thread := range dis.Threads {
		if thread.ThreadId == id {
			return thread, nil
		}
	}
	return thread{}, errors.New("Could not find thread")
}
<|endoftext|>


<commit_before>package kafka

import (
	"context"
	"fmt"
	"net"

	"github.com/segmentio/kafka-go/protocol/createpartitions"
)

type CreatePartitionsRequest struct {
	// Address of the kafka broker to send the request to.
	Addr net.Addr

	Topic         string
	NewPartitions []CreatePartitionsRequestPartition
}

type CreatePartitionsRequestPartition struct {
	BrokerIDs []int32
}

type CreatePartitionsResponse struct {
	ErrorCode    int16
	ErrorMessage string
}

func (c *Client) CreatePartitions(
	ctx context.Context,
	req CreatePartitionsRequest,
) (*CreatePartitionsResponse, error) {
	assignments := []createpartitions.RequestAssignment{}
	for _, partition := range req.NewPartitions {
		assignments = append(assignments, createpartitions.RequestAssignment{
			BrokerIDs: partition.BrokerIDs,
		})
	}

	apiReq := &createpartitions.Request{
		Topics: []createpartitions.RequestTopic{
			{
				Name:        req.Topic,
				Count:       int32(len(req.NewPartitions)) + 3,
				Assignments: assignments,
			},
		},
		TimeoutMs: 5000,
	}

	protocolResp, err := c.roundTrip(
		ctx,
		req.Addr,
		apiReq,
	)
	if err != nil {
		return nil, err
	}
	apiResp := protocolResp.(*createpartitions.Response)

	fmt.Printf("Response: %+v", *apiResp)
	if len(apiResp.Results) == 0 {
		return nil, fmt.Errorf("Empty results")
	}

	return &CreatePartitionsResponse{
		ErrorCode:    apiResp.Results[0].ErrorCode,
		ErrorMessage: apiResp.Results[0].ErrorMessage,
	}, nil
}
<commit_msg>Keep iterating on new apis<commit_after>package kafka

import (
	"context"
	"fmt"
	"net"

	"github.com/segmentio/kafka-go/protocol/createpartitions"
)

type CreatePartitionsRequest struct {
	// Address of the kafka broker to send the request to.
	Addr net.Addr

	Topic         string
	NewPartitions []CreatePartitionsRequestPartition
	TotalCount    int32
}

type CreatePartitionsRequestPartition struct {
	BrokerIDs []int32
}

type CreatePartitionsResponse struct {
	ErrorCode    int16
	ErrorMessage string
}

func (c *Client) CreatePartitions(
	ctx context.Context,
	req CreatePartitionsRequest,
) (*CreatePartitionsResponse, error) {
	assignments := []createpartitions.RequestAssignment{}
	for _, partition := range req.NewPartitions {
		assignments = append(assignments, createpartitions.RequestAssignment{
			BrokerIDs: partition.BrokerIDs,
		})
	}

	apiReq := &createpartitions.Request{
		Topics: []createpartitions.RequestTopic{
			{
				Name:        req.Topic,
				Count:       req.TotalCount,
				Assignments: assignments,
			},
		},
		TimeoutMs: 5000,
	}

	protocolResp, err := c.roundTrip(
		ctx,
		req.Addr,
		apiReq,
	)
	if err != nil {
		return nil, err
	}
	apiResp := protocolResp.(*createpartitions.Response)

	fmt.Printf("Response: %+v", *apiResp)
	if len(apiResp.Results) == 0 {
		return nil, fmt.Errorf("Empty results")
	}

	return &CreatePartitionsResponse{
		ErrorCode:    apiResp.Results[0].ErrorCode,
		ErrorMessage: apiResp.Results[0].ErrorMessage,
	}, nil
}
<|endoftext|>


<commit_before>package bitswap

import (
	"time"

	process "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess"
	procctx "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess/context"
	ratelimit "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess/ratelimit"
	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"

	key "github.com/ipfs/go-ipfs/blocks/key"
	logging "github.com/ipfs/go-ipfs/vendor/QmXJkcEXB6C9h6Ytb6rrUTFU56Ro62zxgrbxTT3dgjQGA8/go-log"
)

var TaskWorkerCount = 8

func (bs *Bitswap) startWorkers(px process.Process, ctx context.Context) {
	// Start up a worker to handle block requests this node is making
	px.Go(func(px process.Process) {
		bs.providerConnector(ctx)
	})

	// Start up workers to handle requests from other nodes for the data on this node
	for i := 0; i < TaskWorkerCount; i++ {
		i := i
		px.Go(func(px process.Process) {
			bs.taskWorker(ctx, i)
		})
	}

	// Start up a worker to manage periodically resending our wantlist out to peers
	px.Go(func(px process.Process) {
		bs.rebroadcastWorker(ctx)
	})

	// Start up a worker to manage sending out provides messages
	px.Go(func(px process.Process) {
		bs.provideCollector(ctx)
	})

	// Spawn up multiple workers to handle incoming blocks
	// consider increasing number if providing blocks bottlenecks
	// file transfers
	px.Go(bs.provideWorker)
}

func (bs *Bitswap) taskWorker(ctx context.Context, id int) {
	idmap := logging.LoggableMap{"ID": id}
	defer log.Info("bitswap task worker shutting down...")
	for {
		log.Event(ctx, "Bitswap.TaskWorker.Loop", idmap)
		select {
		case nextEnvelope := <-bs.engine.Outbox():
			select {
			case envelope, ok := <-nextEnvelope:
				if !ok {
					continue
				}
				log.Event(ctx, "Bitswap.TaskWorker.Work", logging.LoggableMap{
					"ID":     id,
					"Target": envelope.Peer.Pretty(),
					"Block":  envelope.Block.Multihash.B58String(),
				})

				bs.wm.SendBlock(ctx, envelope)
			case <-ctx.Done():
				return
			}
		case <-ctx.Done():
			return
		}
	}
}

func (bs *Bitswap) provideWorker(px process.Process) {

	limiter := ratelimit.NewRateLimiter(px, provideWorkerMax)

	limitedGoProvide := func(k key.Key, wid int) {
		ev := logging.LoggableMap{"ID": wid}
		limiter.LimitedGo(func(px process.Process) {

			ctx := procctx.OnClosingContext(px) // derive ctx from px
			defer log.EventBegin(ctx, "Bitswap.ProvideWorker.Work", ev, &k).Done()

			ctx, cancel := context.WithTimeout(ctx, provideTimeout) // timeout ctx
			defer cancel()

			if err := bs.network.Provide(ctx, k); err != nil {
				log.Error(err)
			}
		})
	}

	// worker spawner, reads from bs.provideKeys until it closes, spawning a
	// _ratelimited_ number of workers to handle each key.
	limiter.Go(func(px process.Process) {
		for wid := 2; ; wid++ {
			ev := logging.LoggableMap{"ID": 1}
			log.Event(procctx.OnClosingContext(px), "Bitswap.ProvideWorker.Loop", ev)

			select {
			case <-px.Closing():
				return
			case k, ok := <-bs.provideKeys:
				if !ok {
					log.Debug("provideKeys channel closed")
					return
				}
				limitedGoProvide(k, wid)
			}
		}
	})
}

func (bs *Bitswap) provideCollector(ctx context.Context) {
	defer close(bs.provideKeys)
	var toProvide []key.Key
	var nextKey key.Key
	var keysOut chan key.Key

	for {
		select {
		case blk, ok := <-bs.newBlocks:
			if !ok {
				log.Debug("newBlocks channel closed")
				return
			}
			if keysOut == nil {
				nextKey = blk.Key()
				keysOut = bs.provideKeys
			} else {
				toProvide = append(toProvide, blk.Key())
			}
		case keysOut <- nextKey:
			if len(toProvide) > 0 {
				nextKey = toProvide[0]
				toProvide = toProvide[1:]
			} else {
				keysOut = nil
			}
		case <-ctx.Done():
			return
		}
	}
}

// connects to providers for the given keys
func (bs *Bitswap) providerConnector(parent context.Context) {
	defer log.Info("bitswap client worker shutting down...")

	for {
		log.Event(parent, "Bitswap.ProviderConnector.Loop")
		select {
		case req := <-bs.findKeys:
			keys := req.keys
			if len(keys) == 0 {
				log.Warning("Received batch request for zero blocks")
				continue
			}
			log.Event(parent, "Bitswap.ProviderConnector.Work", logging.LoggableMap{"Keys": keys})

			// NB: Optimization. Assumes that providers of key[0] are likely to
			// be able to provide for all keys. This currently holds true in most
			// every situation. Later, this assumption may not hold as true.
			child, cancel := context.WithTimeout(req.ctx, providerRequestTimeout)
			providers := bs.network.FindProvidersAsync(child, keys[0], maxProvidersPerRequest)
			for p := range providers {
				go bs.network.ConnectTo(req.ctx, p)
			}
			cancel()

		case <-parent.Done():
			return
		}
	}
}

func (bs *Bitswap) rebroadcastWorker(parent context.Context) {
	ctx, cancel := context.WithCancel(parent)
	defer cancel()

	broadcastSignal := time.NewTicker(rebroadcastDelay.Get())
	defer broadcastSignal.Stop()

	tick := time.NewTicker(10 * time.Second)
	defer tick.Stop()

	for {
		log.Event(ctx, "Bitswap.Rebroadcast.idle")
		select {
		case <-tick.C:
			n := bs.wm.wl.Len()
			if n > 0 {
				log.Debug(n, "keys in bitswap wantlist")
			}
		case <-broadcastSignal.C: // resend unfulfilled wantlist keys
			log.Event(ctx, "Bitswap.Rebroadcast.active")
			entries := bs.wm.wl.Entries()
			if len(entries) > 0 {
				bs.connectToProviders(ctx, entries)
			}
		case <-parent.Done():
			return
		}
	}
}
<commit_msg>fix panic in bitswap working limit spawning<commit_after>package bitswap

import (
	"time"

	process "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess"
	procctx "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess/context"
	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"

	key "github.com/ipfs/go-ipfs/blocks/key"
	logging "github.com/ipfs/go-ipfs/vendor/QmXJkcEXB6C9h6Ytb6rrUTFU56Ro62zxgrbxTT3dgjQGA8/go-log"
)

var TaskWorkerCount = 8

func (bs *Bitswap) startWorkers(px process.Process, ctx context.Context) {
	// Start up a worker to handle block requests this node is making
	px.Go(func(px process.Process) {
		bs.providerConnector(ctx)
	})

	// Start up workers to handle requests from other nodes for the data on this node
	for i := 0; i < TaskWorkerCount; i++ {
		i := i
		px.Go(func(px process.Process) {
			bs.taskWorker(ctx, i)
		})
	}

	// Start up a worker to manage periodically resending our wantlist out to peers
	px.Go(func(px process.Process) {
		bs.rebroadcastWorker(ctx)
	})

	// Start up a worker to manage sending out provides messages
	px.Go(func(px process.Process) {
		bs.provideCollector(ctx)
	})

	// Spawn up multiple workers to handle incoming blocks
	// consider increasing number if providing blocks bottlenecks
	// file transfers
	px.Go(bs.provideWorker)
}

func (bs *Bitswap) taskWorker(ctx context.Context, id int) {
	idmap := logging.LoggableMap{"ID": id}
	defer log.Info("bitswap task worker shutting down...")
	for {
		log.Event(ctx, "Bitswap.TaskWorker.Loop", idmap)
		select {
		case nextEnvelope := <-bs.engine.Outbox():
			select {
			case envelope, ok := <-nextEnvelope:
				if !ok {
					continue
				}
				log.Event(ctx, "Bitswap.TaskWorker.Work", logging.LoggableMap{
					"ID":     id,
					"Target": envelope.Peer.Pretty(),
					"Block":  envelope.Block.Multihash.B58String(),
				})

				bs.wm.SendBlock(ctx, envelope)
			case <-ctx.Done():
				return
			}
		case <-ctx.Done():
			return
		}
	}
}

func (bs *Bitswap) provideWorker(px process.Process) {

	limit := make(chan struct{}, provideWorkerMax)

	limitedGoProvide := func(k key.Key, wid int) {
		defer func() {
			// replace token when done
			<-limit
		}()
		ev := logging.LoggableMap{"ID": wid}

		ctx := procctx.OnClosingContext(px) // derive ctx from px
		defer log.EventBegin(ctx, "Bitswap.ProvideWorker.Work", ev, &k).Done()

		ctx, cancel := context.WithTimeout(ctx, provideTimeout) // timeout ctx
		defer cancel()

		if err := bs.network.Provide(ctx, k); err != nil {
			log.Error(err)
		}
	}

	// worker spawner, reads from bs.provideKeys until it closes, spawning a
	// _ratelimited_ number of workers to handle each key.
	for wid := 2; ; wid++ {
		ev := logging.LoggableMap{"ID": 1}
		log.Event(procctx.OnClosingContext(px), "Bitswap.ProvideWorker.Loop", ev)

		select {
		case <-px.Closing():
			return
		case k, ok := <-bs.provideKeys:
			if !ok {
				log.Debug("provideKeys channel closed")
				return
			}
			select {
			case <-px.Closing():
				return
			case limit <- struct{}{}:
				go limitedGoProvide(k, wid)
			}
		}
	}
}

func (bs *Bitswap) provideCollector(ctx context.Context) {
	defer close(bs.provideKeys)
	var toProvide []key.Key
	var nextKey key.Key
	var keysOut chan key.Key

	for {
		select {
		case blk, ok := <-bs.newBlocks:
			if !ok {
				log.Debug("newBlocks channel closed")
				return
			}
			if keysOut == nil {
				nextKey = blk.Key()
				keysOut = bs.provideKeys
			} else {
				toProvide = append(toProvide, blk.Key())
			}
		case keysOut <- nextKey:
			if len(toProvide) > 0 {
				nextKey = toProvide[0]
				toProvide = toProvide[1:]
			} else {
				keysOut = nil
			}
		case <-ctx.Done():
			return
		}
	}
}

// connects to providers for the given keys
func (bs *Bitswap) providerConnector(parent context.Context) {
	defer log.Info("bitswap client worker shutting down...")

	for {
		log.Event(parent, "Bitswap.ProviderConnector.Loop")
		select {
		case req := <-bs.findKeys:
			keys := req.keys
			if len(keys) == 0 {
				log.Warning("Received batch request for zero blocks")
				continue
			}
			log.Event(parent, "Bitswap.ProviderConnector.Work", logging.LoggableMap{"Keys": keys})

			// NB: Optimization. Assumes that providers of key[0] are likely to
			// be able to provide for all keys. This currently holds true in most
			// every situation. Later, this assumption may not hold as true.
			child, cancel := context.WithTimeout(req.ctx, providerRequestTimeout)
			providers := bs.network.FindProvidersAsync(child, keys[0], maxProvidersPerRequest)
			for p := range providers {
				go bs.network.ConnectTo(req.ctx, p)
			}
			cancel()

		case <-parent.Done():
			return
		}
	}
}

func (bs *Bitswap) rebroadcastWorker(parent context.Context) {
	ctx, cancel := context.WithCancel(parent)
	defer cancel()

	broadcastSignal := time.NewTicker(rebroadcastDelay.Get())
	defer broadcastSignal.Stop()

	tick := time.NewTicker(10 * time.Second)
	defer tick.Stop()

	for {
		log.Event(ctx, "Bitswap.Rebroadcast.idle")
		select {
		case <-tick.C:
			n := bs.wm.wl.Len()
			if n > 0 {
				log.Debug(n, "keys in bitswap wantlist")
			}
		case <-broadcastSignal.C: // resend unfulfilled wantlist keys
			log.Event(ctx, "Bitswap.Rebroadcast.active")
			entries := bs.wm.wl.Entries()
			if len(entries) > 0 {
				bs.connectToProviders(ctx, entries)
			}
		case <-parent.Done():
			return
		}
	}
}
<|endoftext|>


<commit_before>package expr

import (
	"testing"

	u "github.com/araddon/gou"
	"github.com/araddon/qlbridge/lex"
	"github.com/bmizerany/assert"
)

var (
	_ = u.EMPTY
)

func init() {
	lex.IDENTITY_CHARS = lex.IDENTITY_SQL_CHARS
}

func parseFilterQlTest(t *testing.T, ql string) {
	u.Debugf("before: %s", ql)
	req, err := ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	req2, err := ParseFilterQL(req.String())
	assert.Tf(t, err == nil, "must parse roundtrip %v", err)
	req.Raw = ""
	req2.Raw = ""
	u.Debugf("after:  %s", req2.String())
	assert.Equal(t, req, req2, "must roundtrip")
}

func TestFilterQlRoundTrip(t *testing.T) {

	parseFilterQlTest(t, `
		FILTER OR (
			AND (
				score NOT BETWEEN 5 and 10,
				email NOT IN ("abc")
			),
			NOT date > "now-3d"
		)`)
	parseFilterQlTest(t, `
		FILTER AND ( EXISTS user_id, NOT OR ( user_id like "a", user_id like "b", user_id like "c", user_id like "d", user_id like "e", user_id like "f" ) )
	`)
	parseFilterQlTest(t, `
		FILTER OR ( AND ( our_names like "2. has spaces", our_names like "1. has more spa'ces" ), INCLUDE 'f9f0dc74234af7e86ddeb660c50350e1' )
	`)
	parseFilterQlTest(t, `
		FILTER  AND ( NOT INCLUDE '791734b084019d99c82a475264464304',
			NOT INCLUDE 'd750a11e72b58778e302eb0893788680', NOT INCLUDE '61a624e5ca4153645ddc9e6ebaee8000' )
		`)
	parseFilterQlTest(t, `FILTER AND ( visitct >= "1", NOT INCLUDE 3d4240482815b9848caf2e6f )`)

	parseFilterQlTest(t, `
		FILTER AND (
			AND (
				score NOT BETWEEN 5 and 10,
				email NOT IN ("abc")
			),
			x > 7
		)`)
	parseFilterQlTest(t, `FILTER AND ( visitct >= "1", INCLUDE 3d4240482815b9848caf2e6f )`)

	parseFilterQlTest(t, `FILTER x > 7`)

	parseFilterQlTest(t, `FILTER AND ( NOT EXISTS email, email NOT IN ("abc") )`)

	parseFilterQlTest(t, `FILTER AND ( score NOT BETWEEN 5 and 10, email NOT IN ("abc") )`)
	parseFilterQlTest(t, `
		FILTER
			AND (
				NAME != NULL
				, tostring(fieldname) == "hello"
			)

			LIMIT 100
	`)
}

func TestFilterQLAstCheck(t *testing.T) {

	ql := `
		FILTER
			AND (
				NAME != NULL,
				tostring(fieldname) == "hello",
			)

		LIMIT 100
	`
	req, err := ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "expected 1 filters got:%d for %s", len(req.Filter.Filters), req.Filter.String())
	cf := req.Filter.Filters[0]
	assert.Tf(t, len(cf.Filter.Filters) == 2, "expected 2 filters got:%d for %s", len(cf.Filter.Filters), cf.String())
	f1 := cf.Filter.Filters[0]
	assert.Tf(t, f1.Expr != nil, "")
	assert.Tf(t, f1.Expr.String() == "NAME != NULL", "%v", f1.Expr)
	assert.Tf(t, req.Limit == 100, "wanted limit=100: %v", req.Limit)

	ql = `FILTER NOT AND ( name == "bob" ) ALIAS root`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filter expr: %#v", req.Filter.Filters)
	cf = req.Filter.Filters[0]
	assert.Tf(t, cf.Filter.Negate == true, "must negate")
	fex := cf.Filter.Filters[0]
	assert.Tf(t, fex.Expr.String() == `name == "bob"`, "Should have expr %v", fex)
	assert.Tf(t, req.String() == `FILTER NOT name == "bob" ALIAS root`, "roundtrip? %v", req.String())

	ql = `FILTER OR ( INCLUDE child_1, INCLUDE child_2 ) ALIAS root`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	cf = req.Filter.Filters[0]
	assert.Tf(t, len(cf.Filter.Filters) == 2, "has 2 filter expr: %#v", cf.Filter.Filters)
	assert.Tf(t, req.Filter.Op == lex.TokenOr, "must have or op %v", req.Filter.Op)
	f1 = cf.Filter.Filters[1]
	assert.Tf(t, f1.String() == `INCLUDE child_2`, "Should have include %q", f1.String())

	ql = `FILTER NOT ( name == "bob", OR ( NOT INCLUDE filter_xyz , NOT exists abc ) ) ALIAS root`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	cf = req.Filter.Filters[0]
	assert.Tf(t, len(cf.Filter.Filters) == 2, "has 2 filter expr: %#v", cf.Filter.Filters)
	assert.Tf(t, cf.Filter.Negate == true, "must negate")
	f1 = cf.Filter.Filters[1]
	assert.Tf(t, len(f1.Filter.Filters) == 2, "has 2 filter subfilter: %#v", f1.String())
	assert.Tf(t, f1.Filter.Op == lex.TokenOr, "is or %#v", f1.Filter.Op)
	f2 := f1.Filter.Filters[0]
	assert.T(t, f2.Negate == true)
	assert.Tf(t, f2.String() == `NOT INCLUDE filter_xyz`, "Should have include %v", f2)
	//assert.Tf(t, req.String() == ql, "roundtrip? %v", req.String())

	ql = `
    SELECT *
    FROM users
    WHERE
      domain(url) == "google.com"
      OR momentum > 20
    ALIAS my_filter_name
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, req.Alias == "my_filter_name", "has alias: %q", req.Alias)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filters: %#v", req.Filter)
	fs := req.Filter.Filters[0]
	assert.Tf(t, fs.Expr != nil, "")
	assert.Tf(t, fs.Expr.String() == `domain(url) == "google.com" OR momentum > 20`, "%v", fs.Expr)

	ql = `
    FILTER
      AND (
          -- Lets make sure the date is good
          daysago(datefield) < 100
          -- as well as domain
          , domain(url) == "google.com"
          , INCLUDE my_other_named_filter
          , OR (
              momentum > 20
             , propensity > 50
             , INCLUDE nested_filter
          )
          , NOT AND ( score > 20 , score < 50 )
       )
    ALIAS my_filter_name
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, req.Alias == "my_filter_name", "has alias: %q", req.Alias)
	u.Info(req.String())
	cf = req.Filter.Filters[0]
	assert.Equalf(t, len(cf.Filter.Filters), 5, "expected 5 filters: %#v", cf.Filter)
	f5 := cf.Filter.Filters[4]
	assert.Tf(t, f5.Negate || f5.Filter.Negate, "expr negated? %s", f5.String())
	assert.Tf(t, len(f5.Filter.Filters) == 2, "expr? %s", f5.String())
	assert.Equal(t, f5.String(), "NOT AND ( score > 20, score < 50 )")
	assert.Tf(t, len(req.Includes()) == 2, "has 2 includes: %v", req.Includes())
	//assert.Equalf(t, f5.Expr.NodeType(), UnaryNodeType, "%s != %s", f5.Expr.NodeType(), UnaryNodeType)

	ql = `
    FILTER
      AND (
          -- Lets make sure the date is good
          daysago(datefield) < 100
          -- as well as domain
          domain(url) == "google.com"
          INCLUDE my_other_named_filter
          OR (
              momentum > 20
             propensity > 50
          )
          NOT AND ( score > 20 , score < 50 )
       )
    ALIAS my_filter_name
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, req.Alias == "my_filter_name", "has alias: %q", req.Alias)
	u.Info(req.String())
	cf = req.Filter.Filters[0]
	assert.Equalf(t, len(cf.Filter.Filters), 5, "expected 5 filters: %#v", cf.Filter)
	f5 = cf.Filter.Filters[4]
	assert.Tf(t, f5.Negate || f5.Filter.Negate, "expr negated? %s", f5.String())
	assert.Tf(t, len(f5.Filter.Filters) == 2, "expr? %s", f5.String())
	assert.Equal(t, f5.String(), "NOT AND ( score > 20, score < 50 )")

	ql = `FILTER AND (
				INCLUDE child_1,
				INCLUDE child_2
			) ALIAS root`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	cf = req.Filter.Filters[0]
	for _, f := range cf.Filter.Filters {
		assert.Tf(t, f.Include != "", "has include filter %q", f.String())
	}
	assert.Tf(t, len(cf.Filter.Filters) == 2, "want 2 filter expr: %#v", cf.Filter.Filters)

	ql = `FILTER NOT INCLUDE child_1 ALIAS root`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filter expr: %#v", req.Filter.Filters)
	assert.Tf(t, req.Filter.Negate == true || req.Filter.Filters[0].Negate, "must negate %s", req.String())
	fInc := cf.Filter.Filters[0]
	assert.Tf(t, fInc.Include != "", "Should have include")

	ql = `
		FILTER *
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filter expr: %#v", req.Filter.Filters)
	fAll := req.Filter.Filters[0]
	assert.Tf(t, fAll.MatchAll, "Should have match all")

	ql = `
		FILTER match_all
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filter expr: %#v", req.Filter.Filters)
	fAll = req.Filter.Filters[0]
	assert.Tf(t, fAll.MatchAll, "Should have match all")

	// Make sure we support following features
	//  - naked single valid expressions that are compound
	//  - naked expression syntax (ie, not AND())
	ql = `
		FILTER
			NAME != NULL
			AND
			tostring(fieldname) == "hello"
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filter expr: %#v", req.Filter.Filters)

	ql = `
    FILTER
      AND (
          EXISTS datefield
       )
	FROM user
    ALIAS my_filter_name
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, req.Alias == "my_filter_name", "has alias: %q", req.Alias)
	assert.Tf(t, req.From == "user", "has FROM: %q", req.From)
	cf = req.Filter.Filters[0]
	assert.Tf(t, len(cf.Filter.Filters) == 1, "has 1 filters: %#v", cf.Filter)
	f1 = cf.Filter.Filters[0]
	assert.Tf(t, f1.Expr != nil, "")
	assert.Tf(t, f1.Expr.String() == "EXISTS datefield", "%#v", f1.Expr.String())
}

func TestFilterQLKeywords(t *testing.T) {
	ql := `
		FILTER
		  -- Test filter
			AND (
				created < "now-24h",
				deleted == false
			)
		FROM accounts
		ALIAS new_accounts
		LIMIT 100
	`
	req, err := ParseFilterQL(ql)
	assert.Equal(t, nil, err)
	assert.NotEqual(t, nil, req)
}
<commit_msg>update grammar for filterql<commit_after>package expr

import (
	"testing"

	u "github.com/araddon/gou"
	"github.com/araddon/qlbridge/lex"
	"github.com/bmizerany/assert"
)

var (
	_ = u.EMPTY
)

func init() {
	lex.IDENTITY_CHARS = lex.IDENTITY_SQL_CHARS
}

func parseFilterQlTest(t *testing.T, ql string) {
	u.Debugf("before: %s", ql)
	req, err := ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	req2, err := ParseFilterQL(req.String())
	assert.Tf(t, err == nil, "must parse roundtrip %v", err)
	req.Raw = ""
	req2.Raw = ""
	u.Debugf("after:  %s", req2.String())
	assert.Equal(t, req, req2, "must roundtrip")
}

func TestFilterQlRoundTrip(t *testing.T) {

	parseFilterQlTest(t, `
		FILTER OR (
			AND (
				score NOT BETWEEN 5 and 10,
				email NOT IN ("abc")
			),
			NOT date > "now-3d"
		)`)
	parseFilterQlTest(t, `
		FILTER AND ( EXISTS user_id, NOT OR ( user_id like "a", user_id like "b", user_id like "c", user_id like "d", user_id like "e", user_id like "f" ) )
	`)
	parseFilterQlTest(t, `
		FILTER OR ( AND ( our_names like "2. has spaces", our_names like "1. has more spa'ces" ), INCLUDE 'f9f0dc74234af7e86ddeb660c50350e1' )
	`)
	parseFilterQlTest(t, `
		FILTER  AND ( NOT INCLUDE '791734b084019d99c82a475264464304',
			NOT INCLUDE 'd750a11e72b58778e302eb0893788680', NOT INCLUDE '61a624e5ca4153645ddc9e6ebaee8000' )
		`)
	parseFilterQlTest(t, `FILTER AND ( visitct >= "1", NOT INCLUDE 3d4240482815b9848caf2e6f )`)

	parseFilterQlTest(t, `
		FILTER AND (
			AND (
				score NOT BETWEEN 5 and 10,
				email NOT IN ("abc")
			),
			x > 7
		)`)
	parseFilterQlTest(t, `FILTER AND ( visitct >= "1", INCLUDE 3d4240482815b9848caf2e6f )`)

	parseFilterQlTest(t, `FILTER x > 7`)

	parseFilterQlTest(t, `FILTER AND ( NOT EXISTS email, email NOT IN ("abc") )`)

	parseFilterQlTest(t, `FILTER AND ( score NOT BETWEEN 5 and 10, email NOT IN ("abc") )`)
	parseFilterQlTest(t, `
		FILTER
			AND (
				NAME != NULL
				, tostring(fieldname) == "hello"
			)

			LIMIT 100
	`)
}

func TestFilterQLAstCheck(t *testing.T) {

	ql := `
		FILTER
			AND (
				NAME != NULL,
				tostring(fieldname) == "hello",
			)

		LIMIT 100
	`
	req, err := ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "expected 1 filters got:%d for %s", len(req.Filter.Filters), req.Filter.String())
	cf := req.Filter.Filters[0]
	assert.Tf(t, len(cf.Filter.Filters) == 2, "expected 2 filters got:%d for %s", len(cf.Filter.Filters), cf.String())
	f1 := cf.Filter.Filters[0]
	assert.Tf(t, f1.Expr != nil, "")
	assert.Tf(t, f1.Expr.String() == "NAME != NULL", "%v", f1.Expr)
	assert.Tf(t, req.Limit == 100, "wanted limit=100: %v", req.Limit)

	ql = `FILTER NOT AND ( name == "bob" ) ALIAS root`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filter expr: %#v", req.Filter.Filters)
	cf = req.Filter.Filters[0]
	assert.Tf(t, cf.Filter.Negate == true, "must negate")
	fex := cf.Filter.Filters[0]
	assert.Tf(t, fex.Expr.String() == `name == "bob"`, "Should have expr %v", fex)
	assert.Tf(t, req.String() == `FILTER NOT name == "bob" ALIAS root`, "roundtrip? %v", req.String())

	ql = `FILTER OR ( INCLUDE child_1, INCLUDE child_2 ) ALIAS root`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	cf = req.Filter.Filters[0]
	assert.Tf(t, len(cf.Filter.Filters) == 2, "has 2 filter expr: %#v", cf.Filter.Filters)
	assert.Tf(t, req.Filter.Op == lex.TokenOr, "must have or op %v", req.Filter.Op)
	f1 = cf.Filter.Filters[1]
	assert.Tf(t, f1.String() == `INCLUDE child_2`, "Should have include %q", f1.String())

	ql = `FILTER NOT ( name == "bob", OR ( NOT INCLUDE filter_xyz , NOT exists abc ) ) ALIAS root`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	cf = req.Filter.Filters[0]
	assert.Tf(t, len(cf.Filter.Filters) == 2, "has 2 filter expr: %#v", cf.Filter.Filters)
	assert.Tf(t, cf.Filter.Negate == true, "must negate")
	f1 = cf.Filter.Filters[1]
	assert.Tf(t, len(f1.Filter.Filters) == 2, "has 2 filter subfilter: %#v", f1.String())
	assert.Tf(t, f1.Filter.Op == lex.TokenOr, "is or %#v", f1.Filter.Op)
	f2 := f1.Filter.Filters[0]
	assert.T(t, f2.Negate == true)
	assert.Tf(t, f2.String() == `NOT INCLUDE filter_xyz`, "Should have include %v", f2)
	//assert.Tf(t, req.String() == ql, "roundtrip? %v", req.String())

	ql = `
    SELECT *
    FROM users
    WHERE
      domain(url) == "google.com"
      OR momentum > 20
    ALIAS my_filter_name
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, req.Alias == "my_filter_name", "has alias: %q", req.Alias)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filters: %#v", req.Filter)
	fs := req.Filter.Filters[0]
	assert.Tf(t, fs.Expr != nil, "")
	assert.Tf(t, fs.Expr.String() == `domain(url) == "google.com" OR momentum > 20`, "%v", fs.Expr)

	ql = `
    FILTER
      AND (
          -- Lets make sure the date is good
          daysago(datefield) < 100
          -- as well as domain
          , domain(url) == "google.com"
          , INCLUDE my_other_named_filter
          , OR (
              momentum > 20
             , propensity > 50
             , INCLUDE nested_filter
          )
          , NOT AND ( score > 20 , score < 50 )
       )
    ALIAS my_filter_name
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, req.Alias == "my_filter_name", "has alias: %q", req.Alias)
	u.Info(req.String())
	cf = req.Filter.Filters[0]
	assert.Equalf(t, len(cf.Filter.Filters), 5, "expected 5 filters: %#v", cf.Filter)
	f5 := cf.Filter.Filters[4]
	assert.Tf(t, f5.Negate || f5.Filter.Negate, "expr negated? %s", f5.String())
	assert.Tf(t, len(f5.Filter.Filters) == 2, "expr? %s", f5.String())
	assert.Equal(t, f5.String(), "NOT AND ( score > 20, score < 50 )")
	assert.Tf(t, len(req.Includes()) == 2, "has 2 includes: %v", req.Includes())
	//assert.Equalf(t, f5.Expr.NodeType(), UnaryNodeType, "%s != %s", f5.Expr.NodeType(), UnaryNodeType)

	ql = `
    FILTER
      AND (
          -- Lets make sure the date is good
          daysago(datefield) < 100
          -- as well as domain
          domain(url) == "google.com"
          INCLUDE my_other_named_filter
          OR (
              momentum > 20
             propensity > 50
          )
          NOT AND ( score > 20 , score < 50 )
       )
    ALIAS my_filter_name
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, req.Alias == "my_filter_name", "has alias: %q", req.Alias)
	u.Info(req.String())
	cf = req.Filter.Filters[0]
	assert.Equalf(t, len(cf.Filter.Filters), 5, "expected 5 filters: %#v", cf.Filter)
	f5 = cf.Filter.Filters[4]
	assert.Tf(t, f5.Negate || f5.Filter.Negate, "expr negated? %s", f5.String())
	assert.Tf(t, len(f5.Filter.Filters) == 2, "expr? %s", f5.String())
	assert.Equal(t, f5.String(), "NOT AND ( score > 20, score < 50 )")

	ql = `FILTER AND (
				INCLUDE child_1,
				INCLUDE child_2
			) ALIAS root`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	cf = req.Filter.Filters[0]
	for _, f := range cf.Filter.Filters {
		assert.Tf(t, f.Include != "", "has include filter %q", f.String())
	}
	assert.Tf(t, len(cf.Filter.Filters) == 2, "want 2 filter expr: %#v", cf.Filter.Filters)

	ql = `FILTER NOT INCLUDE child_1 ALIAS root`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filter expr: %#v", req.Filter.Filters)
	assert.Tf(t, req.Filter.Negate == true || req.Filter.Filters[0].Negate, "must negate %s", req.String())
	fInc := cf.Filter.Filters[0]
	assert.Tf(t, fInc.Include != "", "Should have include")

	ql = `
		FILTER *
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filter expr: %#v", req.Filter.Filters)
	fAll := req.Filter.Filters[0]
	assert.Tf(t, fAll.MatchAll, "Should have match all")

	ql = `
		FILTER match_all
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filter expr: %#v", req.Filter.Filters)
	fAll = req.Filter.Filters[0]
	assert.Tf(t, fAll.MatchAll, "Should have match all")

	// Make sure we support following features
	//  - naked single valid expressions that are compound
	//  - naked expression syntax (ie, not AND())
	ql = `
		FILTER
			NAME != NULL
			AND
			tostring(fieldname) == "hello"
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, len(req.Filter.Filters) == 1, "has 1 filter expr: %#v", req.Filter.Filters)

	ql = `
    FILTER
      AND (
          EXISTS datefield
       )
	FROM user
    ALIAS my_filter_name
	`
	req, err = ParseFilterQL(ql)
	assert.Tf(t, err == nil && req != nil, "Must parse: %s  \n\t%v", ql, err)
	assert.Tf(t, req.Alias == "my_filter_name", "has alias: %q", req.Alias)
	assert.Tf(t, req.From == "user", "has FROM: %q", req.From)
	cf = req.Filter.Filters[0]
	assert.Tf(t, len(cf.Filter.Filters) == 1, "has 1 filters: %#v", cf.Filter)
	f1 = cf.Filter.Filters[0]
	assert.Tf(t, f1.Expr != nil, "")
	assert.Tf(t, f1.Expr.String() == "EXISTS datefield", "%#v", f1.Expr.String())
}

func TestFilterQLKeywords(t *testing.T) {
	ql := `
		FILTER
		  -- Test filter
			AND (
				created < "now-24h",
				deleted == false
			)
		FROM accounts
		LIMIT 100
		ALIAS new_accounts
	`
	req, err := ParseFilterQL(ql)
	assert.Equal(t, nil, err)
	assert.NotEqual(t, nil, req)
	assert.Equal(t, req.Limit, 100, "has limit 100")
}
<|endoftext|>