<commit_before>package leveldb

import "github.com/syndtr/goleveldb/leveldb"

// key prefixes for each element type
var prefix = func() map[string][]byte {
	return map[string][]byte{
		"node":     []byte{'N'},
		"way":      []byte{'W'},
		"relation": []byte{'R'},
	}
}()

// Connection - Connection
type Connection struct {
	DB *leveldb.DB
}

// Open - open connection and set up
func (c *Connection) Open(path string) {
	db, err := leveldb.OpenFile(path, nil)
	if err != nil {
		panic(err)
	}
	c.DB = db
}

// Close - close connection and clean up
func (c *Connection) Close() {
	c.DB.Close()
}
<commit_msg>leveldb: disable snappy compression<commit_after>package leveldb

import (
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/opt"
)

// key prefixes for each element type
var prefix = func() map[string][]byte {
	return map[string][]byte{
		"node":     []byte{'N'},
		"way":      []byte{'W'},
		"relation": []byte{'R'},
	}
}()

// Connection - Connection
type Connection struct {
	DB *leveldb.DB
}

// Open - open connection and set up
func (c *Connection) Open(path string) {
	db, err := leveldb.OpenFile(path, &opt.Options{
		Compression: opt.NoCompression,
	})
	if err != nil {
		panic(err)
	}
	c.DB = db
}

// Close - close connection and clean up
func (c *Connection) Close() {
	c.DB.Close()
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package collate contains types for comparing and sorting Unicode strings
// according to a given collation order.  Package locale provides a high-level
// interface to collation. Users should typically use that package instead.
package collate

import (
	"bytes"
	"exp/norm"
)

// Level identifies the collation comparison level.
// The primary level corresponds to the basic sorting of text.
// The secondary level corresponds to accents and related linguistic elements.
// The tertiary level corresponds to casing and related concepts.
// The quaternary level is derived from the other levels by the
// various algorithms for handling variable elements.
type Level int

const (
	Primary Level = iota
	Secondary
	Tertiary
	Quaternary
	Identity
)

// AlternateHandling identifies the various ways in which variables are handled.
// A rune with a primary weight lower than the variable top is considered a
// variable.
// See http://www.unicode.org/reports/tr10/#Variable_Weighting for details.
type AlternateHandling int

const (
	// AltNonIgnorable turns off special handling of variables.
	AltNonIgnorable AlternateHandling = iota

	// AltBlanked sets variables and all subsequent primary ignorables to be
	// ignorable at all levels. This is identical to removing all variables
	// and subsequent primary ignorables from the input.
	AltBlanked

	// AltShifted sets variables to be ignorable for levels one through three and
	// adds a fourth level based on the values of the ignored levels.
	AltShifted

	// AltShiftTrimmed is a slight variant of AltShifted that is used to
	// emulate POSIX.
	AltShiftTrimmed
)

// Collator provides functionality for comparing strings for a given
// collation order.
type Collator struct {
	// Strength sets the maximum level to use in comparison.
	Strength Level

	// Alternate specifies an alternative handling of variables.
	Alternate AlternateHandling

	// Backwards specifies the order of sorting at the secondary level.
	// This option exists predominantly to support reverse sorting of accents in French.
	Backwards bool

	// TODO: implement:
	// With HiraganaQuaternary enabled, Hiragana codepoints will get lower values
	// than all the other non-variable code points. Strength must be greater or
	// equal to Quaternary for this to take effect.
	HiraganaQuaternary bool

	// If CaseLevel is true, a level consisting only of case characteristics will
	// be inserted in front of the tertiary level.  To ignore accents but take
	// cases into account, set Strength to Primary and CaseLevel to true.
	CaseLevel bool

	// If Numeric is true, any sequence of decimal digits (category is Nd) is sorted
	// at a primary level with its numeric value.  For example, "A-21" < "A-123".
	Numeric bool

	f norm.Form

	t *table
}

// Locales returns the list of locales for which collating differs from its parent locale.
func Locales() []string {
	return availableLocales
}

// New returns a new Collator initialized for the given locale.
func New(loc string) *Collator {
	// TODO: handle locale selection according to spec.
	t := &mainTable
	if loc != "" {
		if idx, ok := locales[loc]; ok {
			t = mainTable.indexedTable(idx)
		}
	}
	return &Collator{
		Strength: Quaternary,
		f:        norm.NFD,
		t:        t,
	}
}

// SetVariableTop sets all runes with primary strength less than the primary
// strength of r to be variable and thus affected by alternate handling.
func (c *Collator) SetVariableTop(r rune) {
	// TODO: implement
}

// Buffer holds reusable buffers that can be used during collation.
// Reusing a Buffer for the various calls that accept it may avoid
// unnecessary memory allocations.
type Buffer struct {
	// TODO: try various parameters and techniques, such as using
	// a chan of buffers for a pool.
	ba  [4096]byte
	wa  [512]weights
	key []byte
	ce  []weights
}

func (b *Buffer) init() {
	if b.ce == nil {
		b.ce = b.wa[:0]
		b.key = b.ba[:0]
	} else {
		b.ce = b.ce[:0]
	}
}

// ResetKeys clears the buffer used for generated keys. Calling ResetKeys
// invalidates keys previously obtained from Key or KeyFromString.
func (b *Buffer) ResetKeys() {
	b.ce = b.ce[:0]
	b.key = b.key[:0]
}

// Compare returns an integer comparing the two byte slices.
// The result will be 0 if a==b, -1 if a < b, and +1 if a > b.
// Compare calls ResetKeys, thereby invalidating keys
// previously generated using Key or KeyFromString using buf.
func (c *Collator) Compare(buf *Buffer, a, b []byte) int {
	// TODO: for now we simply compute keys and compare.  Once we
	// have good benchmarks, move to an implementation that works
	// incrementally for the majority of cases.
	// - Benchmark with long strings that only vary in modifiers.
	buf.ResetKeys()
	ka := c.Key(buf, a)
	kb := c.Key(buf, b)
	defer buf.ResetKeys()
	return bytes.Compare(ka, kb)
}

// CompareString returns an integer comparing the two strings.
// The result will be 0 if a==b, -1 if a < b, and +1 if a > b.
// CompareString calls ResetKeys, thereby invalidating keys
// previously generated using Key or KeyFromString using buf.
func (c *Collator) CompareString(buf *Buffer, a, b string) int {
	buf.ResetKeys()
	ka := c.KeyFromString(buf, a)
	kb := c.KeyFromString(buf, b)
	defer buf.ResetKeys()
	return bytes.Compare(ka, kb)
}

func (c *Collator) Prefix(buf *Buffer, s, prefix []byte) int {
	// iterate over s, track bytes consumed.
	return 0
}

// Key returns the collation key for str.
// Passing the buffer buf may avoid memory allocations.
// The returned slice will point to an allocation in Buffer and will remain
// valid until the next call to buf.ResetKeys().
func (c *Collator) Key(buf *Buffer, str []byte) []byte {
	// See http://www.unicode.org/reports/tr10/#Main_Algorithm for more details.
	buf.init()
	c.getColElems(buf, str)
	return c.key(buf, buf.ce)
}

// KeyFromString returns the collation key for str.
// Passing the buffer buf may avoid memory allocations.
// The returned slice will point to an allocation in Buffer and will retain
// valid until the next call to buf.ResetKeys().
func (c *Collator) KeyFromString(buf *Buffer, str string) []byte {
	// See http://www.unicode.org/reports/tr10/#Main_Algorithm for more details.
	buf.init()
	c.getColElemsString(buf, str)
	return c.key(buf, buf.ce)
}

func (c *Collator) key(buf *Buffer, w []weights) []byte {
	processWeights(c.Alternate, c.t.variableTop, w)
	kn := len(buf.key)
	c.keyFromElems(buf, w)
	return buf.key[kn:]
}

func (c *Collator) getColElems(buf *Buffer, str []byte) {
	i := c.iter()
	i.src.SetInput(c.f, str)
	for !i.done() {
		buf.ce = i.next(buf.ce)
	}
}

func (c *Collator) getColElemsString(buf *Buffer, str string) {
	i := c.iter()
	i.src.SetInputString(c.f, str)
	for !i.done() {
		buf.ce = i.next(buf.ce)
	}
}

type iter struct {
	src        norm.Iter
	ba         [1024]byte
	buf        []byte
	t          *table
	p          int
	minBufSize int
	_done, eof bool
}

func (c *Collator) iter() iter {
	i := iter{t: c.t, minBufSize: c.t.maxContractLen}
	i.buf = i.ba[:0]
	return i
}

func (i *iter) done() bool {
	return i._done
}

func (i *iter) next(ce []weights) []weights {
	if !i.eof && len(i.buf)-i.p < i.minBufSize {
		// replenish buffer
		n := copy(i.buf, i.buf[i.p:])
		n += i.src.Next(i.buf[n:cap(i.buf)])
		i.buf = i.buf[:n]
		i.p = 0
		i.eof = i.src.Done()
	}
	if i.p == len(i.buf) {
		i._done = true
		return ce
	}
	ce, sz := i.t.appendNext(ce, i.buf[i.p:])
	i.p += sz
	return ce
}

func appendPrimary(key []byte, p uint32) []byte {
	// Convert to variable length encoding; supports up to 23 bits.
	if p <= 0x7FFF {
		key = append(key, uint8(p>>8), uint8(p))
	} else {
		key = append(key, uint8(p>>16)|0x80, uint8(p>>8), uint8(p))
	}
	return key
}

// keyFromElems converts the weights ws to a compact sequence of bytes.
// The result will be appended to the byte buffer in buf.
func (c *Collator) keyFromElems(buf *Buffer, ws []weights) {
	for _, v := range ws {
		if w := v.primary; w > 0 {
			buf.key = appendPrimary(buf.key, w)
		}
	}
	if Secondary <= c.Strength {
		buf.key = append(buf.key, 0, 0)
		// TODO: we can use one 0 if we can guarantee that all non-zero weights are > 0xFF.
		if !c.Backwards {
			for _, v := range ws {
				if w := v.secondary; w > 0 {
					buf.key = append(buf.key, uint8(w>>8), uint8(w))
				}
			}
		} else {
			for i := len(ws) - 1; i >= 0; i-- {
				if w := ws[i].secondary; w > 0 {
					buf.key = append(buf.key, uint8(w>>8), uint8(w))
				}
			}
		}
	} else if c.CaseLevel {
		buf.key = append(buf.key, 0, 0)
	}
	if Tertiary <= c.Strength || c.CaseLevel {
		buf.key = append(buf.key, 0, 0)
		for _, v := range ws {
			if w := v.tertiary; w > 0 {
				buf.key = append(buf.key, w)
			}
		}
		// Derive the quaternary weights from the options and other levels.
		// Note that we represent maxQuaternary as 0xFF. The first byte of the
		// representation of a a primary weight is always smaller than 0xFF,
		// so using this single byte value will compare correctly.
		if Quaternary <= c.Strength {
			if c.Alternate == AltShiftTrimmed {
				lastNonFFFF := len(buf.key)
				buf.key = append(buf.key, 0)
				for _, v := range ws {
					if w := v.quaternary; w == maxQuaternary {
						buf.key = append(buf.key, 0xFF)
					} else if w > 0 {
						buf.key = appendPrimary(buf.key, w)
						lastNonFFFF = len(buf.key)
					}
				}
				buf.key = buf.key[:lastNonFFFF]
			} else {
				buf.key = append(buf.key, 0)
				for _, v := range ws {
					if w := v.quaternary; w == maxQuaternary {
						buf.key = append(buf.key, 0xFF)
					} else if w > 0 {
						buf.key = appendPrimary(buf.key, w)
					}
				}
			}
		}
	}
}

func processWeights(vw AlternateHandling, top uint32, wa []weights) {
	ignore := false
	switch vw {
	case AltShifted, AltShiftTrimmed:
		for i := range wa {
			if p := wa[i].primary; p <= top && p != 0 {
				wa[i] = weights{quaternary: p}
				ignore = true
			} else if p == 0 {
				if ignore {
					wa[i] = weights{}
				} else if wa[i].tertiary != 0 {
					wa[i].quaternary = maxQuaternary
				}
			} else {
				wa[i].quaternary = maxQuaternary
				ignore = false
			}
		}
	case AltBlanked:
		for i := range wa {
			if p := wa[i].primary; p <= top && (ignore || p != 0) {
				wa[i] = weights{}
				ignore = true
			} else {
				ignore = false
			}
		}
	}
}
<commit_msg>exp/locale/collate: clarification in comments on use of returned value.<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package collate contains types for comparing and sorting Unicode strings
// according to a given collation order.  Package locale provides a high-level
// interface to collation. Users should typically use that package instead.
package collate

import (
	"bytes"
	"exp/norm"
)

// Level identifies the collation comparison level.
// The primary level corresponds to the basic sorting of text.
// The secondary level corresponds to accents and related linguistic elements.
// The tertiary level corresponds to casing and related concepts.
// The quaternary level is derived from the other levels by the
// various algorithms for handling variable elements.
type Level int

const (
	Primary Level = iota
	Secondary
	Tertiary
	Quaternary
	Identity
)

// AlternateHandling identifies the various ways in which variables are handled.
// A rune with a primary weight lower than the variable top is considered a
// variable.
// See http://www.unicode.org/reports/tr10/#Variable_Weighting for details.
type AlternateHandling int

const (
	// AltNonIgnorable turns off special handling of variables.
	AltNonIgnorable AlternateHandling = iota

	// AltBlanked sets variables and all subsequent primary ignorables to be
	// ignorable at all levels. This is identical to removing all variables
	// and subsequent primary ignorables from the input.
	AltBlanked

	// AltShifted sets variables to be ignorable for levels one through three and
	// adds a fourth level based on the values of the ignored levels.
	AltShifted

	// AltShiftTrimmed is a slight variant of AltShifted that is used to
	// emulate POSIX.
	AltShiftTrimmed
)

// Collator provides functionality for comparing strings for a given
// collation order.
type Collator struct {
	// Strength sets the maximum level to use in comparison.
	Strength Level

	// Alternate specifies an alternative handling of variables.
	Alternate AlternateHandling

	// Backwards specifies the order of sorting at the secondary level.
	// This option exists predominantly to support reverse sorting of accents in French.
	Backwards bool

	// TODO: implement:
	// With HiraganaQuaternary enabled, Hiragana codepoints will get lower values
	// than all the other non-variable code points. Strength must be greater or
	// equal to Quaternary for this to take effect.
	HiraganaQuaternary bool

	// If CaseLevel is true, a level consisting only of case characteristics will
	// be inserted in front of the tertiary level.  To ignore accents but take
	// cases into account, set Strength to Primary and CaseLevel to true.
	CaseLevel bool

	// If Numeric is true, any sequence of decimal digits (category is Nd) is sorted
	// at a primary level with its numeric value.  For example, "A-21" < "A-123".
	Numeric bool

	f norm.Form

	t *table
}

// Locales returns the list of locales for which collating differs from its parent locale.
// The returned value should not be modified.
func Locales() []string {
	return availableLocales
}

// New returns a new Collator initialized for the given locale.
func New(loc string) *Collator {
	// TODO: handle locale selection according to spec.
	t := &mainTable
	if loc != "" {
		if idx, ok := locales[loc]; ok {
			t = mainTable.indexedTable(idx)
		}
	}
	return &Collator{
		Strength: Quaternary,
		f:        norm.NFD,
		t:        t,
	}
}

// SetVariableTop sets all runes with primary strength less than the primary
// strength of r to be variable and thus affected by alternate handling.
func (c *Collator) SetVariableTop(r rune) {
	// TODO: implement
}

// Buffer holds reusable buffers that can be used during collation.
// Reusing a Buffer for the various calls that accept it may avoid
// unnecessary memory allocations.
type Buffer struct {
	// TODO: try various parameters and techniques, such as using
	// a chan of buffers for a pool.
	ba  [4096]byte
	wa  [512]weights
	key []byte
	ce  []weights
}

func (b *Buffer) init() {
	if b.ce == nil {
		b.ce = b.wa[:0]
		b.key = b.ba[:0]
	} else {
		b.ce = b.ce[:0]
	}
}

// ResetKeys clears the buffer used for generated keys. Calling ResetKeys
// invalidates keys previously obtained from Key or KeyFromString.
func (b *Buffer) ResetKeys() {
	b.ce = b.ce[:0]
	b.key = b.key[:0]
}

// Compare returns an integer comparing the two byte slices.
// The result will be 0 if a==b, -1 if a < b, and +1 if a > b.
// Compare calls ResetKeys, thereby invalidating keys
// previously generated using Key or KeyFromString using buf.
func (c *Collator) Compare(buf *Buffer, a, b []byte) int {
	// TODO: for now we simply compute keys and compare.  Once we
	// have good benchmarks, move to an implementation that works
	// incrementally for the majority of cases.
	// - Benchmark with long strings that only vary in modifiers.
	buf.ResetKeys()
	ka := c.Key(buf, a)
	kb := c.Key(buf, b)
	defer buf.ResetKeys()
	return bytes.Compare(ka, kb)
}

// CompareString returns an integer comparing the two strings.
// The result will be 0 if a==b, -1 if a < b, and +1 if a > b.
// CompareString calls ResetKeys, thereby invalidating keys
// previously generated using Key or KeyFromString using buf.
func (c *Collator) CompareString(buf *Buffer, a, b string) int {
	buf.ResetKeys()
	ka := c.KeyFromString(buf, a)
	kb := c.KeyFromString(buf, b)
	defer buf.ResetKeys()
	return bytes.Compare(ka, kb)
}

func (c *Collator) Prefix(buf *Buffer, s, prefix []byte) int {
	// iterate over s, track bytes consumed.
	return 0
}

// Key returns the collation key for str.
// Passing the buffer buf may avoid memory allocations.
// The returned slice will point to an allocation in Buffer and will remain
// valid until the next call to buf.ResetKeys().
func (c *Collator) Key(buf *Buffer, str []byte) []byte {
	// See http://www.unicode.org/reports/tr10/#Main_Algorithm for more details.
	buf.init()
	c.getColElems(buf, str)
	return c.key(buf, buf.ce)
}

// KeyFromString returns the collation key for str.
// Passing the buffer buf may avoid memory allocations.
// The returned slice will point to an allocation in Buffer and will retain
// valid until the next call to buf.ResetKeys().
func (c *Collator) KeyFromString(buf *Buffer, str string) []byte {
	// See http://www.unicode.org/reports/tr10/#Main_Algorithm for more details.
	buf.init()
	c.getColElemsString(buf, str)
	return c.key(buf, buf.ce)
}

func (c *Collator) key(buf *Buffer, w []weights) []byte {
	processWeights(c.Alternate, c.t.variableTop, w)
	kn := len(buf.key)
	c.keyFromElems(buf, w)
	return buf.key[kn:]
}

func (c *Collator) getColElems(buf *Buffer, str []byte) {
	i := c.iter()
	i.src.SetInput(c.f, str)
	for !i.done() {
		buf.ce = i.next(buf.ce)
	}
}

func (c *Collator) getColElemsString(buf *Buffer, str string) {
	i := c.iter()
	i.src.SetInputString(c.f, str)
	for !i.done() {
		buf.ce = i.next(buf.ce)
	}
}

type iter struct {
	src        norm.Iter
	ba         [1024]byte
	buf        []byte
	t          *table
	p          int
	minBufSize int
	_done, eof bool
}

func (c *Collator) iter() iter {
	i := iter{t: c.t, minBufSize: c.t.maxContractLen}
	i.buf = i.ba[:0]
	return i
}

func (i *iter) done() bool {
	return i._done
}

func (i *iter) next(ce []weights) []weights {
	if !i.eof && len(i.buf)-i.p < i.minBufSize {
		// replenish buffer
		n := copy(i.buf, i.buf[i.p:])
		n += i.src.Next(i.buf[n:cap(i.buf)])
		i.buf = i.buf[:n]
		i.p = 0
		i.eof = i.src.Done()
	}
	if i.p == len(i.buf) {
		i._done = true
		return ce
	}
	ce, sz := i.t.appendNext(ce, i.buf[i.p:])
	i.p += sz
	return ce
}

func appendPrimary(key []byte, p uint32) []byte {
	// Convert to variable length encoding; supports up to 23 bits.
	if p <= 0x7FFF {
		key = append(key, uint8(p>>8), uint8(p))
	} else {
		key = append(key, uint8(p>>16)|0x80, uint8(p>>8), uint8(p))
	}
	return key
}

// keyFromElems converts the weights ws to a compact sequence of bytes.
// The result will be appended to the byte buffer in buf.
func (c *Collator) keyFromElems(buf *Buffer, ws []weights) {
	for _, v := range ws {
		if w := v.primary; w > 0 {
			buf.key = appendPrimary(buf.key, w)
		}
	}
	if Secondary <= c.Strength {
		buf.key = append(buf.key, 0, 0)
		// TODO: we can use one 0 if we can guarantee that all non-zero weights are > 0xFF.
		if !c.Backwards {
			for _, v := range ws {
				if w := v.secondary; w > 0 {
					buf.key = append(buf.key, uint8(w>>8), uint8(w))
				}
			}
		} else {
			for i := len(ws) - 1; i >= 0; i-- {
				if w := ws[i].secondary; w > 0 {
					buf.key = append(buf.key, uint8(w>>8), uint8(w))
				}
			}
		}
	} else if c.CaseLevel {
		buf.key = append(buf.key, 0, 0)
	}
	if Tertiary <= c.Strength || c.CaseLevel {
		buf.key = append(buf.key, 0, 0)
		for _, v := range ws {
			if w := v.tertiary; w > 0 {
				buf.key = append(buf.key, w)
			}
		}
		// Derive the quaternary weights from the options and other levels.
		// Note that we represent maxQuaternary as 0xFF. The first byte of the
		// representation of a a primary weight is always smaller than 0xFF,
		// so using this single byte value will compare correctly.
		if Quaternary <= c.Strength {
			if c.Alternate == AltShiftTrimmed {
				lastNonFFFF := len(buf.key)
				buf.key = append(buf.key, 0)
				for _, v := range ws {
					if w := v.quaternary; w == maxQuaternary {
						buf.key = append(buf.key, 0xFF)
					} else if w > 0 {
						buf.key = appendPrimary(buf.key, w)
						lastNonFFFF = len(buf.key)
					}
				}
				buf.key = buf.key[:lastNonFFFF]
			} else {
				buf.key = append(buf.key, 0)
				for _, v := range ws {
					if w := v.quaternary; w == maxQuaternary {
						buf.key = append(buf.key, 0xFF)
					} else if w > 0 {
						buf.key = appendPrimary(buf.key, w)
					}
				}
			}
		}
	}
}

func processWeights(vw AlternateHandling, top uint32, wa []weights) {
	ignore := false
	switch vw {
	case AltShifted, AltShiftTrimmed:
		for i := range wa {
			if p := wa[i].primary; p <= top && p != 0 {
				wa[i] = weights{quaternary: p}
				ignore = true
			} else if p == 0 {
				if ignore {
					wa[i] = weights{}
				} else if wa[i].tertiary != 0 {
					wa[i].quaternary = maxQuaternary
				}
			} else {
				wa[i].quaternary = maxQuaternary
				ignore = false
			}
		}
	case AltBlanked:
		for i := range wa {
			if p := wa[i].primary; p <= top && (ignore || p != 0) {
				wa[i] = weights{}
				ignore = true
			} else {
				ignore = false
			}
		}
	}
}
<|endoftext|>


<commit_before>package repository

import (
	"bufio"
	"database/sql"
	"time"

	"io"

	"github.com/lib/pq"
	_ "github.com/lib/pq"
	"github.com/syou6162/go-active-learning/lib/feature"
	"github.com/syou6162/go-active-learning/lib/model"
	"github.com/syou6162/go-active-learning/lib/util/file"
)

var exampleNotFoundError = model.NotFoundError("example")

func (r *repository) InsertOrUpdateExample(e *model.Example) error {
	now := time.Now()
	e.UpdatedAt = now
	_, err := r.db.NamedExec(`
INSERT INTO example
( url,  final_url,  title,  description,  og_description,  og_type,  og_image,  body,  score,  is_new,  status_code,  favicon,  label,  created_at,  updated_at)
VALUES
(:url, :final_url, :title, :description, :og_description, :og_type, :og_image, :body, :score, :is_new, :status_code, :favicon, :label, :created_at, :updated_at)
ON CONFLICT (url)
DO UPDATE SET
url = :url, final_url = :final_url, title = :title,
description = :description, og_description = :og_description, og_type = :og_type, og_image = :og_image,
body = :body, score = :score, is_new = :is_new, status_code = :status_code, favicon = :favicon,
label = :label, created_at = :created_at, updated_at = :updated_at
WHERE (:label != 0) AND (example.label != EXCLUDED.label)
;`, e)
	return err
}

func (r *repository) UpdateFeatureVector(e *model.Example) error {
	tmp, err := r.FindExampleByUlr(e.Url)
	if err != nil {
		return err
	}
	id := tmp.Id
	if _, err = r.db.Exec(`DELETE FROM feature WHERE example_id = $1;`, id); err != nil {
		return err
	}
	_, err = r.db.Exec(`INSERT INTO feature (example_id, feature) VALUES ($1, unnest(cast($2 AS TEXT[])));`, id, pq.Array(e.Fv))
	return err
}

func (r *repository) InsertExampleFromScanner(scanner *bufio.Scanner) (*model.Example, error) {
	line := scanner.Text()
	e, err := file.ParseLine(line)
	if err != nil {
		return nil, err
	}
	err = r.InsertOrUpdateExample(e)
	if err != nil {
		return nil, err
	}
	return e, nil
}

func (r *repository) InsertExamplesFromReader(reader io.Reader) error {
	scanner := bufio.NewScanner(reader)

	for scanner.Scan() {
		_, err := r.InsertExampleFromScanner(scanner)
		if err != nil {
			return err
		}
	}
	if err := scanner.Err(); err != nil {
		return err
	}
	return nil
}

func (r *repository) readExamples(query string, args ...interface{}) (model.Examples, error) {
	examples := model.Examples{}
	err := r.db.Select(&examples, query, args...)
	if err != nil {
		return nil, err
	}
	return examples, nil
}

func (r *repository) readExample(query string, args ...interface{}) (*model.Example, error) {
	e := model.Example{}

	err := r.db.Get(&e, query, args...)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, exampleNotFoundError
		}
		return nil, err
	}
	return &e, nil
}

func (r *repository) ReadExamples() (model.Examples, error) {
	query := `SELECT * FROM example;`
	return r.readExamples(query)
}

func (r *repository) ReadRecentExamples(from time.Time) (model.Examples, error) {
	query := `SELECT * FROM example WHERE created_at > $1 ORDER BY updated_at DESC;`
	return r.readExamples(query, from)
}

func (r *repository) ReadExamplesByLabel(label model.LabelType, limit int) (model.Examples, error) {
	query := `SELECT * FROM example WHERE label = $1 ORDER BY updated_at DESC LIMIT $2;`
	return r.readExamples(query, label, limit)
}

func (r *repository) ReadLabeledExamples(limit int) (model.Examples, error) {
	query := `SELECT * FROM example WHERE label != 0 ORDER BY updated_at DESC LIMIT $1;`
	return r.readExamples(query, limit)
}

func (r *repository) ReadPositiveExamples(limit int) (model.Examples, error) {
	return r.ReadExamplesByLabel(model.POSITIVE, limit)
}

func (r *repository) ReadNegativeExamples(limit int) (model.Examples, error) {
	return r.ReadExamplesByLabel(model.NEGATIVE, limit)
}

func (r *repository) ReadUnlabeledExamples(limit int) (model.Examples, error) {
	return r.ReadExamplesByLabel(model.UNLABELED, limit)
}

func (r *repository) FindExampleByUlr(url string) (*model.Example, error) {
	query := `SELECT * FROM example WHERE url = $1;`
	return r.readExample(query, url)
}

func (r *repository) SearchExamplesByUlrs(urls []string) (model.Examples, error) {
	// ref: https://godoc.org/github.com/lib/pq#Array
	query := `SELECT * FROM example WHERE url = ANY($1);`
	return r.readExamples(query, pq.Array(urls))
}

func (r *repository) FindFeatureVector(e *model.Example) (feature.FeatureVector, error) {
	fv := feature.FeatureVector{}
	tmp, err := r.FindExampleByUlr(e.Url)
	if err != nil {
		return fv, err
	}
	id := tmp.Id
	query := `SELECT feature FROM feature WHERE example_id = $1;`
	err = r.db.Select(&fv, query, id)
	if err != nil {
		return fv, err
	}
	return fv, nil
}

func (r *repository) SearchFeatureVector(examples model.Examples) ([]feature.FeatureVector, error) {
	type Pair struct {
		ExampleId int    `db:"example_id"`
		Feature   string `db:"feature"`
	}

	fvList := make([]feature.FeatureVector, 0)
	urls := make([]string, 0)
	for _, e := range examples {
		urls = append(urls, e.Url)
	}

	tmp, err := r.SearchExamplesByUlrs(urls)
	if err != nil {
		return fvList, err
	}
	ids := make([]int, 0)
	for _, e := range tmp {
		ids = append(ids, e.Id)
	}

	query := `SELECT example_id, feature FROM feature WHERE example_id = ANY($1);`
	pairs := make([]Pair, 0)
	err = r.db.Select(&pairs, query, pq.Array(ids))
	if err != nil {
		return fvList, err
	}
	m := make(map[int][]string)

	for _, pair := range pairs {
		m[pair.ExampleId] = append(m[pair.ExampleId], pair.Feature)
	}
	for _, id := range ids {
		fvList = append(fvList, m[id])
	}
	return fvList, nil
}

func (r *repository) DeleteAllExamples() error {
	_, err := r.db.Exec(`DELETE FROM example`)
	return err
}
<commit_msg>status codeが変わるときはupdateしてよい<commit_after>package repository

import (
	"bufio"
	"database/sql"
	"time"

	"io"

	"github.com/lib/pq"
	_ "github.com/lib/pq"
	"github.com/syou6162/go-active-learning/lib/feature"
	"github.com/syou6162/go-active-learning/lib/model"
	"github.com/syou6162/go-active-learning/lib/util/file"
)

var exampleNotFoundError = model.NotFoundError("example")

func (r *repository) InsertOrUpdateExample(e *model.Example) error {
	now := time.Now()
	e.UpdatedAt = now
	_, err := r.db.NamedExec(`
INSERT INTO example
( url,  final_url,  title,  description,  og_description,  og_type,  og_image,  body,  score,  is_new,  status_code,  favicon,  label,  created_at,  updated_at)
VALUES
(:url, :final_url, :title, :description, :og_description, :og_type, :og_image, :body, :score, :is_new, :status_code, :favicon, :label, :created_at, :updated_at)
ON CONFLICT (url)
DO UPDATE SET
url = :url, final_url = :final_url, title = :title,
description = :description, og_description = :og_description, og_type = :og_type, og_image = :og_image,
body = :body, score = :score, is_new = :is_new, status_code = :status_code, favicon = :favicon,
label = :label, created_at = :created_at, updated_at = :updated_at
WHERE
((EXCLUDED.label != 0) AND (example.label != EXCLUDED.label)) OR
((example.status_code != 200) AND (EXCLUDED.status_code = 200))
;`, e)
	return err
}

func (r *repository) UpdateFeatureVector(e *model.Example) error {
	tmp, err := r.FindExampleByUlr(e.Url)
	if err != nil {
		return err
	}
	id := tmp.Id
	if _, err = r.db.Exec(`DELETE FROM feature WHERE example_id = $1;`, id); err != nil {
		return err
	}
	_, err = r.db.Exec(`INSERT INTO feature (example_id, feature) VALUES ($1, unnest(cast($2 AS TEXT[])));`, id, pq.Array(e.Fv))
	return err
}

func (r *repository) InsertExampleFromScanner(scanner *bufio.Scanner) (*model.Example, error) {
	line := scanner.Text()
	e, err := file.ParseLine(line)
	if err != nil {
		return nil, err
	}
	err = r.InsertOrUpdateExample(e)
	if err != nil {
		return nil, err
	}
	return e, nil
}

func (r *repository) InsertExamplesFromReader(reader io.Reader) error {
	scanner := bufio.NewScanner(reader)

	for scanner.Scan() {
		_, err := r.InsertExampleFromScanner(scanner)
		if err != nil {
			return err
		}
	}
	if err := scanner.Err(); err != nil {
		return err
	}
	return nil
}

func (r *repository) readExamples(query string, args ...interface{}) (model.Examples, error) {
	examples := model.Examples{}
	err := r.db.Select(&examples, query, args...)
	if err != nil {
		return nil, err
	}
	return examples, nil
}

func (r *repository) readExample(query string, args ...interface{}) (*model.Example, error) {
	e := model.Example{}

	err := r.db.Get(&e, query, args...)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, exampleNotFoundError
		}
		return nil, err
	}
	return &e, nil
}

func (r *repository) ReadExamples() (model.Examples, error) {
	query := `SELECT * FROM example;`
	return r.readExamples(query)
}

func (r *repository) ReadRecentExamples(from time.Time) (model.Examples, error) {
	query := `SELECT * FROM example WHERE created_at > $1 ORDER BY updated_at DESC;`
	return r.readExamples(query, from)
}

func (r *repository) ReadExamplesByLabel(label model.LabelType, limit int) (model.Examples, error) {
	query := `SELECT * FROM example WHERE label = $1 ORDER BY updated_at DESC LIMIT $2;`
	return r.readExamples(query, label, limit)
}

func (r *repository) ReadLabeledExamples(limit int) (model.Examples, error) {
	query := `SELECT * FROM example WHERE label != 0 ORDER BY updated_at DESC LIMIT $1;`
	return r.readExamples(query, limit)
}

func (r *repository) ReadPositiveExamples(limit int) (model.Examples, error) {
	return r.ReadExamplesByLabel(model.POSITIVE, limit)
}

func (r *repository) ReadNegativeExamples(limit int) (model.Examples, error) {
	return r.ReadExamplesByLabel(model.NEGATIVE, limit)
}

func (r *repository) ReadUnlabeledExamples(limit int) (model.Examples, error) {
	return r.ReadExamplesByLabel(model.UNLABELED, limit)
}

func (r *repository) FindExampleByUlr(url string) (*model.Example, error) {
	query := `SELECT * FROM example WHERE url = $1;`
	return r.readExample(query, url)
}

func (r *repository) SearchExamplesByUlrs(urls []string) (model.Examples, error) {
	// ref: https://godoc.org/github.com/lib/pq#Array
	query := `SELECT * FROM example WHERE url = ANY($1);`
	return r.readExamples(query, pq.Array(urls))
}

func (r *repository) FindFeatureVector(e *model.Example) (feature.FeatureVector, error) {
	fv := feature.FeatureVector{}
	tmp, err := r.FindExampleByUlr(e.Url)
	if err != nil {
		return fv, err
	}
	id := tmp.Id
	query := `SELECT feature FROM feature WHERE example_id = $1;`
	err = r.db.Select(&fv, query, id)
	if err != nil {
		return fv, err
	}
	return fv, nil
}

func (r *repository) SearchFeatureVector(examples model.Examples) ([]feature.FeatureVector, error) {
	type Pair struct {
		ExampleId int    `db:"example_id"`
		Feature   string `db:"feature"`
	}

	fvList := make([]feature.FeatureVector, 0)
	urls := make([]string, 0)
	for _, e := range examples {
		urls = append(urls, e.Url)
	}

	tmp, err := r.SearchExamplesByUlrs(urls)
	if err != nil {
		return fvList, err
	}
	ids := make([]int, 0)
	for _, e := range tmp {
		ids = append(ids, e.Id)
	}

	query := `SELECT example_id, feature FROM feature WHERE example_id = ANY($1);`
	pairs := make([]Pair, 0)
	err = r.db.Select(&pairs, query, pq.Array(ids))
	if err != nil {
		return fvList, err
	}
	m := make(map[int][]string)

	for _, pair := range pairs {
		m[pair.ExampleId] = append(m[pair.ExampleId], pair.Feature)
	}
	for _, id := range ids {
		fvList = append(fvList, m[id])
	}
	return fvList, nil
}

func (r *repository) DeleteAllExamples() error {
	_, err := r.db.Exec(`DELETE FROM example`)
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"regexp"
	"runtime"
	"strings"

	"container/list"

	graphviz "code.google.com/p/gographviz"
	shellwords "github.com/mattn/go-shellwords"
)

type HandleR int

const (
	QUIT HandleR = iota
	SUCCESS
	FAILURE
)

func findEdge(network *graphviz.Graph, n1 string, n2 string) (uint, *graphviz.Edge) {
	n1e := network.Edges.SrcToDsts[n1]
	if n1e != nil && n1e[n2] != nil {
		return 0, n1e[n2]
	}

	n2e := network.Edges.SrcToDsts[n2]
	if n2e != nil && n2e[n1] != nil {
		return 1, n2e[n1]
	}

	return 0, nil
}

func qprocess(edge *graphviz.Edge) {
	if qs[edge].Len() == 0 {
		c, ok := ebrands[edge]

		delete(edge.Attrs, "arrowhead")

		if !ok || c.msg == "" {
			delete(edge.Attrs, "label")
		} else {
			edge.Attrs["label"] = c.msg
		}

		if !ok || c.color == "" {
			delete(edge.Attrs, "color")
		} else {
			edge.Attrs["color"] = c.color
		}
		return
	}

	edge.Attrs["labelfloat"] = "true"
	edge.Attrs["color"] = "blue"

	if qs[edge].Len() == 1 {
		edge.Attrs["label"] = q(fmt.Sprintf("%s", qs[edge].Front().Value))
		return
	}

	edge.Attrs["label"] = q(fmt.Sprintf("%s (+%d)", qs[edge].Front().Value, qs[edge].Len()-1))
}

type ebrand struct {
	color string
	msg   string
}

var e2d map[*graphviz.Edge]string
var ls map[*graphviz.Node]string
var qs map[*graphviz.Edge]*list.List
var ebrands map[*graphviz.Edge]ebrand
var display *os.Process = nil
var exit = -1

func handle(cmd string, args []string, network *graphviz.Graph) (HandleR, string) {
	switch cmd {
	case "emark":
		{
			if len(args) == 2 {
				args = append(args, "")
			}
			if len(args) == 3 {
				args = append(args, "")
			}
			if len(args) != 4 {
				return FAILURE, fmt.Sprintf("Usage: emark <from> <to> <color> <msg>")
			}
			_, edge := findEdge(network, args[0], args[1])
			if edge == nil {
				return FAILURE, fmt.Sprintf("[E] No edge from node %s to node %s", args[0], args[1])
			}

			eb := ebrand{color: args[2], msg: args[3]}
			if eb.color == "" && eb.msg == "" {
				delete(ebrands, edge)
			} else {
				ebrands[edge] = eb
			}
			qprocess(edge)
		}
	case "mark":
		{
			if len(args) != 2 {
				return FAILURE, fmt.Sprintf("Usage: mark <node> <color>")
			}
			node := network.Nodes.Lookup[args[0]]
			if node == nil {
				return FAILURE, fmt.Sprintf("[E] Unknown node: %s", args[0])
			}
			node.Attrs["color"] = q(args[1])
		}
	case "label":
		{
			if len(args) == 1 {
				args = append(args, "")
			}
			if len(args) != 2 {
				return FAILURE, fmt.Sprintf("Usage: label <node> <text>")
			}
			node := network.Nodes.Lookup[args[0]]
			if node == nil {
				return FAILURE, fmt.Sprintf("[E] Unknown node: %s", args[0])
			}
			if args[1] == "" {
				delete(node.Attrs, "shape")
				if ls[node] == "" {
					fmt.Println("deleting node label")
					delete(node.Attrs, "label")
				} else {
					fmt.Println("resetting node label to", ls[node])
					node.Attrs["label"] = q(ls[node])
				}
				break
			}

			ls[node] = node.Attrs["label"]
			node.Attrs["shape"] = "Mrecord"
			fmt.Println("setting node label to be", q(node.Name+" | "+args[1]))
			node.Attrs["label"] = q(node.Name + " | " + args[1])
		}
	case "unmark":
		{
			if len(args) != 1 {
				return FAILURE, fmt.Sprintf("Usage: unmark <node>")
			}
			node := network.Nodes.Lookup[args[0]]
			if node == nil {
				return FAILURE, fmt.Sprintf("[E] Unknown node: %s", args[0])
			}
			delete(node.Attrs, "color")
		}
	case "send":
		{
			if len(args) != 3 {
				return FAILURE, fmt.Sprintf("Usage: send <from> <to> <msg>")
			}
			src, edge := findEdge(network, args[0], args[1])
			if edge == nil {
				return FAILURE, fmt.Sprintf("[E] No edge from node %s to node %s", args[0], args[1])
			}

			qs[edge].PushBack(args[2])
			qprocess(edge)

			_, ok := e2d[edge]
			if !ok {
				e2d[edge] = edge.Attrs["dir"]
			}

			if src == 0 {
				edge.Attrs["dir"] = "forward"
			} else {
				edge.Attrs["dir"] = "back"
			}
			edge.Attrs["arrowhead"] = "normal"
		}
	case "recv":
		{
			if len(args) != 2 {
				return FAILURE, fmt.Sprintf("Usage: recv <receiver> <sender>")
			}
			_, edge := findEdge(network, args[1], args[0])
			if edge == nil {
				return FAILURE, fmt.Sprintf("[E] No edge to node %s from node %s", args[0], args[1])
			}

			qs[edge].Remove(qs[edge].Front())
			qprocess(edge)

			if e2d[edge] == "" {
				delete(edge.Attrs, "dir")
			} else {
				edge.Attrs["dir"] = e2d[edge]
			}
		}
	case "erase":
		{
			e2d = make(map[*graphviz.Edge]string)
			qs = make(map[*graphviz.Edge]*list.List)
			ebrands = make(map[*graphviz.Edge]ebrand)
			ls = make(map[*graphviz.Node]string)
			for _, e := range network.Edges.Edges {
				qs[e] = list.New()
				qprocess(e)
				delete(e.Attrs, "dir")
			}
			for _, n := range network.Nodes.Nodes {
				delete(n.Attrs, "color")
			}
		}
	case "quit":
		return QUIT, ""
	default:
		return FAILURE, fmt.Sprintf("[E] Unknown command: %s", cmd)
	}
	return SUCCESS, "! " + cmd + " " + strings.Join(args, " ")
}
func q(i string) string {
	o := strings.Replace(i, "\\", "\\\\", -1)
	o = strings.Replace(o, "\"", "\\\"", -1)
	return "\"" + o + "\""
}

func main() {
	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "No network graph description given\n")
		os.Exit(2)
		return
	}

	data, err := ioutil.ReadFile(os.Args[1])
	if err != nil {
		fmt.Fprintf(os.Stderr, "Could not read graph description: %s\n", err)
	}

	e2d = make(map[*graphviz.Edge]string)
	qs = make(map[*graphviz.Edge]*list.List)
	ebrands = make(map[*graphviz.Edge]ebrand)
	ls = make(map[*graphviz.Node]string)

	var spec []string
	var dot bytes.Buffer
	spec = nil
	indot := false
	for _, s := range strings.Split(string(data), "\n") {
		if indot {
			dot.WriteString("\n")
			dot.WriteString(s)

			indot = !strings.HasPrefix(s, "}")
			continue
		}

		m, _ := regexp.MatchString("(di)?graph\\s+{", s)
		if m {
			indot = true
			dot.WriteString(s)
		} else {
			spec = append(spec, s)
		}
	}

	tmpn, e := graphviz.Parse(dot.Bytes())
	if e != nil {
		panic(e)
	}
	network := graphviz.NewAnalysedGraph(tmpn).(*graphviz.Graph)

	for _, e := range network.Edges.Edges {
		qs[e] = list.New()
	}

	var start []*graphviz.Node
	start = nil
	for _, s := range spec {
		if strings.HasPrefix(s, "start ") {
			name := strings.TrimSpace(strings.TrimPrefix(s, "start"))
			start = append(start, network.Nodes.Lookup[name])
		}
	}

	for _, n := range start {
		if n == nil {
			continue
		}
		n.Attrs.Add("color", "red")
	}

	err = plot(network)
	if err != nil {
		panic(err)
	}

	in := bufio.NewReader(os.Stdin)
	for {
		cmd, err := in.ReadString('\n')

		if err != nil {
			if err == io.EOF {
				exit = 0
			} else {
				panic(err)
			}
		} else {
			iterate(network, strings.TrimSpace(cmd))
		}

		if exit >= 0 {
			display.Kill()
			os.Exit(exit)
		}
	}
}

func iterate(g *graphviz.Graph, command string) {
	args, _ := shellwords.Parse(command)

	var ok HandleR
	var out string
	switch len(args) {
	case 0:
		return
	case 1:
		ok, out = handle(args[0], args[0:0], g)
	default:
		ok, out = handle(args[0], args[1:], g)
	}

	fmt.Println(out)

	if ok == QUIT {
		exit = 0
		return
	}

	if ok == FAILURE {
		return
	}

	err := plot(g)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Plotting exited with status code %s\n", err)
		exit = 1
		return
	}

	// Redraw img
}

var started = false

func plot(g *graphviz.Graph) error {
	w, e := os.Create("next.svg")
	if e != nil {
		return e
	}

	var ebuf bytes.Buffer
	cmd := exec.Command("fdp", "-Tsvg")
	cmd.Stdin = bytes.NewReader([]byte(g.String()))
	cmd.Stdout = w
	cmd.Stderr = bufio.NewWriter(&ebuf)

	cmde := cmd.Run()
	w.Close()
	if cmde != nil {
		fmt.Fprintf(os.Stderr, "Failed to plot graph:\n---\n%s---\n%s", g.String(), ebuf.String())
		return cmde
	}

	e = os.Rename("next.svg", "current.svg")
	if e != nil {
		return e
	}

	if !started {
		started = true
		_, filename, _, _ := runtime.Caller(0)
		cmd := exec.Command(path.Join(path.Dir(filename), "show.py"), "current.svg")
		e := cmd.Start()
		if e != nil {
			fmt.Println(e)
		}
		display = cmd.Process
	}

	return nil
}
<commit_msg>Make edges less visible<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"regexp"
	"runtime"
	"strings"

	"container/list"

	graphviz "code.google.com/p/gographviz"
	shellwords "github.com/mattn/go-shellwords"
)

type HandleR int

const (
	QUIT HandleR = iota
	SUCCESS
	FAILURE
)

func findEdge(network *graphviz.Graph, n1 string, n2 string) (uint, *graphviz.Edge) {
	n1e := network.Edges.SrcToDsts[n1]
	if n1e != nil && n1e[n2] != nil {
		return 0, n1e[n2]
	}

	n2e := network.Edges.SrcToDsts[n2]
	if n2e != nil && n2e[n1] != nil {
		return 1, n2e[n1]
	}

	return 0, nil
}

func qprocess(edge *graphviz.Edge) {
	if qs[edge].Len() == 0 {
		c, ok := ebrands[edge]

		delete(edge.Attrs, "arrowhead")

		if !ok || c.msg == "" {
			delete(edge.Attrs, "label")
		} else {
			edge.Attrs["label"] = c.msg
		}

		if !ok || c.color == "" {
			edge.Attrs["color"] = "gainsboro"
		} else {
			edge.Attrs["color"] = c.color
		}
		return
	}

	edge.Attrs["labelfloat"] = "true"
	edge.Attrs["color"] = "blue"

	if qs[edge].Len() == 1 {
		edge.Attrs["label"] = q(fmt.Sprintf("%s", qs[edge].Front().Value))
		return
	}

	edge.Attrs["label"] = q(fmt.Sprintf("%s (+%d)", qs[edge].Front().Value, qs[edge].Len()-1))
}

type ebrand struct {
	color string
	msg   string
}

var e2d map[*graphviz.Edge]string
var ls map[*graphviz.Node]string
var qs map[*graphviz.Edge]*list.List
var ebrands map[*graphviz.Edge]ebrand
var display *os.Process = nil
var exit = -1

func handle(cmd string, args []string, network *graphviz.Graph) (HandleR, string) {
	switch cmd {
	case "emark":
		{
			if len(args) == 2 {
				args = append(args, "")
			}
			if len(args) == 3 {
				args = append(args, "")
			}
			if len(args) != 4 {
				return FAILURE, fmt.Sprintf("Usage: emark <from> <to> <color> <msg>")
			}
			_, edge := findEdge(network, args[0], args[1])
			if edge == nil {
				return FAILURE, fmt.Sprintf("[E] No edge from node %s to node %s", args[0], args[1])
			}

			eb := ebrand{color: args[2], msg: args[3]}
			if eb.color == "" && eb.msg == "" {
				delete(ebrands, edge)
			} else {
				ebrands[edge] = eb
			}
			qprocess(edge)
		}
	case "mark":
		{
			if len(args) != 2 {
				return FAILURE, fmt.Sprintf("Usage: mark <node> <color>")
			}
			node := network.Nodes.Lookup[args[0]]
			if node == nil {
				return FAILURE, fmt.Sprintf("[E] Unknown node: %s", args[0])
			}
			node.Attrs["color"] = q(args[1])
		}
	case "label":
		{
			if len(args) == 1 {
				args = append(args, "")
			}
			if len(args) != 2 {
				return FAILURE, fmt.Sprintf("Usage: label <node> <text>")
			}
			node := network.Nodes.Lookup[args[0]]
			if node == nil {
				return FAILURE, fmt.Sprintf("[E] Unknown node: %s", args[0])
			}
			if args[1] == "" {
				delete(node.Attrs, "shape")
				if ls[node] == "" {
					fmt.Println("deleting node label")
					delete(node.Attrs, "label")
				} else {
					fmt.Println("resetting node label to", ls[node])
					node.Attrs["label"] = q(ls[node])
				}
				break
			}

			ls[node] = node.Attrs["label"]
			node.Attrs["shape"] = "Mrecord"
			fmt.Println("setting node label to be", q(node.Name+" | "+args[1]))
			node.Attrs["label"] = q(node.Name + " | " + args[1])
		}
	case "unmark":
		{
			if len(args) != 1 {
				return FAILURE, fmt.Sprintf("Usage: unmark <node>")
			}
			node := network.Nodes.Lookup[args[0]]
			if node == nil {
				return FAILURE, fmt.Sprintf("[E] Unknown node: %s", args[0])
			}
			delete(node.Attrs, "color")
		}
	case "send":
		{
			if len(args) != 3 {
				return FAILURE, fmt.Sprintf("Usage: send <from> <to> <msg>")
			}
			src, edge := findEdge(network, args[0], args[1])
			if edge == nil {
				return FAILURE, fmt.Sprintf("[E] No edge from node %s to node %s", args[0], args[1])
			}

			qs[edge].PushBack(args[2])
			qprocess(edge)

			_, ok := e2d[edge]
			if !ok {
				e2d[edge] = edge.Attrs["dir"]
			}

			if src == 0 {
				edge.Attrs["dir"] = "forward"
			} else {
				edge.Attrs["dir"] = "back"
			}
			edge.Attrs["arrowhead"] = "normal"
		}
	case "recv":
		{
			if len(args) != 2 {
				return FAILURE, fmt.Sprintf("Usage: recv <receiver> <sender>")
			}
			_, edge := findEdge(network, args[1], args[0])
			if edge == nil {
				return FAILURE, fmt.Sprintf("[E] No edge to node %s from node %s", args[0], args[1])
			}

			qs[edge].Remove(qs[edge].Front())
			qprocess(edge)

			if e2d[edge] == "" {
				delete(edge.Attrs, "dir")
			} else {
				edge.Attrs["dir"] = e2d[edge]
			}
		}
	case "erase":
		{
			e2d = make(map[*graphviz.Edge]string)
			qs = make(map[*graphviz.Edge]*list.List)
			ebrands = make(map[*graphviz.Edge]ebrand)
			ls = make(map[*graphviz.Node]string)
			for _, e := range network.Edges.Edges {
				qs[e] = list.New()
				qprocess(e)
				delete(e.Attrs, "dir")
			}
			for _, n := range network.Nodes.Nodes {
				delete(n.Attrs, "color")
			}
		}
	case "quit":
		return QUIT, ""
	default:
		return FAILURE, fmt.Sprintf("[E] Unknown command: %s", cmd)
	}
	return SUCCESS, "! " + cmd + " " + strings.Join(args, " ")
}
func q(i string) string {
	o := strings.Replace(i, "\\", "\\\\", -1)
	o = strings.Replace(o, "\"", "\\\"", -1)
	return "\"" + o + "\""
}

func main() {
	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "No network graph description given\n")
		os.Exit(2)
		return
	}

	data, err := ioutil.ReadFile(os.Args[1])
	if err != nil {
		fmt.Fprintf(os.Stderr, "Could not read graph description: %s\n", err)
	}

	var spec []string
	var dot bytes.Buffer
	spec = nil
	indot := false
	for _, s := range strings.Split(string(data), "\n") {
		if indot {
			dot.WriteString("\n")
			dot.WriteString(s)

			indot = !strings.HasPrefix(s, "}")
			continue
		}

		m, _ := regexp.MatchString("(di)?graph\\s+{", s)
		if m {
			indot = true
			dot.WriteString(s)
		} else {
			spec = append(spec, s)
		}
	}

	tmpn, e := graphviz.Parse(dot.Bytes())
	if e != nil {
		panic(e)
	}
	network := graphviz.NewAnalysedGraph(tmpn).(*graphviz.Graph)
	handle("erase", nil, network)

	var start []*graphviz.Node
	start = nil
	for _, s := range spec {
		if strings.HasPrefix(s, "start ") {
			name := strings.TrimSpace(strings.TrimPrefix(s, "start"))
			start = append(start, network.Nodes.Lookup[name])
		}
	}

	for _, n := range start {
		if n == nil {
			continue
		}
		n.Attrs.Add("color", "red")
	}

	err = plot(network)
	if err != nil {
		panic(err)
	}

	in := bufio.NewReader(os.Stdin)
	for {
		cmd, err := in.ReadString('\n')

		if err != nil {
			if err == io.EOF {
				exit = 0
			} else {
				panic(err)
			}
		} else {
			iterate(network, strings.TrimSpace(cmd))
		}

		if exit >= 0 {
			display.Kill()
			os.Exit(exit)
		}
	}
}

func iterate(g *graphviz.Graph, command string) {
	args, _ := shellwords.Parse(command)

	var ok HandleR
	var out string
	switch len(args) {
	case 0:
		return
	case 1:
		ok, out = handle(args[0], args[0:0], g)
	default:
		ok, out = handle(args[0], args[1:], g)
	}

	fmt.Println(out)

	if ok == QUIT {
		exit = 0
		return
	}

	if ok == FAILURE {
		return
	}

	err := plot(g)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Plotting exited with status code %s\n", err)
		exit = 1
		return
	}

	// Redraw img
}

var started = false

func plot(g *graphviz.Graph) error {
	w, e := os.Create("next.svg")
	if e != nil {
		return e
	}

	var ebuf bytes.Buffer
	cmd := exec.Command("fdp", "-Tsvg")
	cmd.Stdin = bytes.NewReader([]byte(g.String()))
	cmd.Stdout = w
	cmd.Stderr = bufio.NewWriter(&ebuf)

	cmde := cmd.Run()
	w.Close()
	if cmde != nil {
		fmt.Fprintf(os.Stderr, "Failed to plot graph:\n---\n%s---\n%s", g.String(), ebuf.String())
		return cmde
	}

	e = os.Rename("next.svg", "current.svg")
	if e != nil {
		return e
	}

	if !started {
		started = true
		_, filename, _, _ := runtime.Caller(0)
		cmd := exec.Command(path.Join(path.Dir(filename), "show.py"), "current.svg")
		e := cmd.Start()
		if e != nil {
			fmt.Println(e)
		}
		display = cmd.Process
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/pflag"

	"github.com/kolibox/koli/pkg/clientset"
	"github.com/kolibox/koli/pkg/controller"
	"github.com/kolibox/koli/pkg/controller/informers"
	_ "github.com/kolibox/koli/pkg/spec/install"

	"k8s.io/client-go/1.5/kubernetes"
	"k8s.io/client-go/1.5/pkg/util/wait"
	"k8s.io/client-go/1.5/rest"
)

// Config defines configuration parameters for the Operator.
type Config struct {
	Host        string
	TLSInsecure bool
	TLSConfig   rest.TLSClientConfig
}

var cfg Config

func init() {
	pflag.CommandLine.AddGoFlagSet(flag.CommandLine)
	pflag.StringVar(&cfg.Host, "apiserver", "", "API Server addr, e.g. 'http://127.0.0.1:8080'. Omit parameter to run in on-cluster mode and utilize the service account token.")
	pflag.StringVar(&cfg.TLSConfig.CertFile, "cert-file", "", "Path to public TLS certificate file.")
	pflag.StringVar(&cfg.TLSConfig.KeyFile, "key-file", "", " Path to private TLS certificate file.")
	pflag.StringVar(&cfg.TLSConfig.CAFile, "ca-file", "", " Path to TLS CA file.")
	pflag.BoolVar(&cfg.TLSInsecure, "tls-insecure", false, " Don't verify API server's CA certificate.")
	pflag.Parse()
}

func startControllers(stop <-chan struct{}) error {
	cfg, err := clientset.NewClusterConfig(cfg.Host, cfg.TLSInsecure, &cfg.TLSConfig)
	if err != nil {
		return err
	}
	// if os.Getenv("SUPER_USER_TOKEN") == "" {
	// 	return fmt.Errorf("SUPER_USER_TOKEN env not defined")
	// }
	client, err := kubernetes.NewForConfig(cfg)
	if err != nil {
		return err
	}

	_, err = client.Discovery().ServerVersion()
	if err != nil {
		return fmt.Errorf("communicating with server failed: %s", err)
	}

	// Create required third party resources
	// controller.CreateAddonTPRs(cfg.Host, client)
	// controller.CreateServicePlan3PRs(cfg.Host, client)
	// controller.CreateServicePlanStatus3PRs(cfg.Host, client)

	sysClient, err := clientset.NewSysRESTClient(cfg)
	if err != nil {
		return err
	}

	sharedInformers := informers.NewSharedInformerFactory(client, 30*time.Second)

	// TODO: should we use the same client instance??
	// go controller.NewAddonController(
	// 	sharedInformers.Addons().Informer(sysClient),
	// 	sharedInformers.PetSets().Informer(),
	// 	client,
	// ).Run(1, wait.NeverStop)

	// go controller.NewNamespaceController(
	// 	sharedInformers.Namespaces().Informer(),
	// 	client,
	// ).Run(1, wait.NeverStop)

	go controller.NewServicePlanController(
		sharedInformers.ServicePlans().Informer(sysClient),
		client,
		sysClient,
	).Run(1, wait.NeverStop)

	sharedInformers.Start(stop)

	select {} // block forever
}

func main() {
	err := startControllers(make(chan struct{}))
	glog.Fatalf("error running controllers: %v", err)
}
<commit_msg>removing comments<commit_after>package main

import (
	"flag"
	"fmt"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/pflag"

	"github.com/kolibox/koli/pkg/clientset"
	"github.com/kolibox/koli/pkg/controller"
	"github.com/kolibox/koli/pkg/controller/informers"
	_ "github.com/kolibox/koli/pkg/spec/install"

	"k8s.io/client-go/1.5/kubernetes"
	"k8s.io/client-go/1.5/pkg/util/wait"
	"k8s.io/client-go/1.5/rest"
)

// Config defines configuration parameters for the Operator.
type Config struct {
	Host        string
	TLSInsecure bool
	TLSConfig   rest.TLSClientConfig
}

var cfg Config

func init() {
	pflag.CommandLine.AddGoFlagSet(flag.CommandLine)
	pflag.StringVar(&cfg.Host, "apiserver", "", "API Server addr, e.g. 'http://127.0.0.1:8080'. Omit parameter to run in on-cluster mode and utilize the service account token.")
	pflag.StringVar(&cfg.TLSConfig.CertFile, "cert-file", "", "Path to public TLS certificate file.")
	pflag.StringVar(&cfg.TLSConfig.KeyFile, "key-file", "", " Path to private TLS certificate file.")
	pflag.StringVar(&cfg.TLSConfig.CAFile, "ca-file", "", " Path to TLS CA file.")
	pflag.BoolVar(&cfg.TLSInsecure, "tls-insecure", false, " Don't verify API server's CA certificate.")
	pflag.Parse()
}

func startControllers(stop <-chan struct{}) error {
	cfg, err := clientset.NewClusterConfig(cfg.Host, cfg.TLSInsecure, &cfg.TLSConfig)
	if err != nil {
		return err
	}
	// if os.Getenv("SUPER_USER_TOKEN") == "" {
	// 	return fmt.Errorf("SUPER_USER_TOKEN env not defined")
	// }
	client, err := kubernetes.NewForConfig(cfg)
	if err != nil {
		return err
	}

	_, err = client.Discovery().ServerVersion()
	if err != nil {
		return fmt.Errorf("communicating with server failed: %s", err)
	}

	// Create required third party resources
	controller.CreateAddonTPRs(cfg.Host, client)
	controller.CreateServicePlan3PRs(cfg.Host, client)
	controller.CreateServicePlanStatus3PRs(cfg.Host, client)

	sysClient, err := clientset.NewSysRESTClient(cfg)
	if err != nil {
		return err
	}

	sharedInformers := informers.NewSharedInformerFactory(client, 30*time.Second)

	// TODO: should we use the same client instance??
	go controller.NewAddonController(
		sharedInformers.Addons().Informer(sysClient),
		sharedInformers.PetSets().Informer(),
		client,
	).Run(1, wait.NeverStop)

	go controller.NewNamespaceController(
		sharedInformers.Namespaces().Informer(),
		client,
	).Run(1, wait.NeverStop)

	go controller.NewServicePlanController(
		sharedInformers.ServicePlans().Informer(sysClient),
		client,
		sysClient,
	).Run(1, wait.NeverStop)

	sharedInformers.Start(stop)

	select {} // block forever
}

func main() {
	err := startControllers(make(chan struct{}))
	glog.Fatalf("error running controllers: %v", err)
}
<|endoftext|>


<commit_before>package restconf

import (
	"time"

	"context"

	"github.com/c2stack/c2g/c2"
	"github.com/c2stack/c2g/node"
	"golang.org/x/net/websocket"
)

// Determined using default websocket settings and Chrome 49 and stop watch when it
// timesout out w/o any pings from server.
const PingRate = 30 * time.Second

// websocket library will kill connection after this time. This is mostly unnec.
// for our usage because we actively ping so this just has to be larger than ping rate
const serverSocketTimeout = 2 * PingRate

type wsNotifyService struct {
	timeout int
	factory node.Subscriber
	cancel  context.CancelFunc
	conn    *wsconn
}

func (self *wsNotifyService) Handle(ws *websocket.Conn) {
	// ignore error, other-side is free to disappear at will
	defer ws.Close()
	defer self.cancel()

	var rate time.Duration
	if self.timeout == 0 {
		rate = PingRate
	} else {
		rate = time.Duration(self.timeout) * time.Millisecond
	}
	self.conn = &wsconn{
		pinger: time.NewTicker(rate),
		mgr:    node.NewSubscriptionManager(self.factory, ws, ws),
	}
	defer self.conn.close()
	ws.Request().Body.Close()
	go self.conn.keepAlive(ws)
	if err := self.conn.mgr.Run(); err != nil {
		c2.Info.Printf("unclean terminination of web socket: (%s). other side may have close browser. closing socket.", err)
	}
}

type wsconn struct {
	pinger *time.Ticker
	mgr    *node.SubscriptionManager
}

func (self *wsconn) keepAlive(ws *websocket.Conn) {
	for {
		ws.SetDeadline(time.Now().Add(serverSocketTimeout))
		if fw, err := ws.NewFrameWriter(websocket.PingFrame); err != nil {
			//self.Close()
			return
		} else if _, err = fw.Write([]byte{}); err != nil {
			//self.Close()
			return
		}
		if _, running := <-self.pinger.C; !running {
			return
		}
	}
}

func (self *wsconn) close() {
	self.pinger.Stop()
	self.mgr.Close()
}
<commit_msg>recover from panics caused by web browser improperly closing web socket<commit_after>package restconf

import (
	"time"

	"context"

	"github.com/c2stack/c2g/c2"
	"github.com/c2stack/c2g/node"
	"golang.org/x/net/websocket"
)

// Determined using default websocket settings and Chrome 49 and stop watch when it
// timesout out w/o any pings from server.
const PingRate = 30 * time.Second

// websocket library will kill connection after this time. This is mostly unnec.
// for our usage because we actively ping so this just has to be larger than ping rate
const serverSocketTimeout = 2 * PingRate

type wsNotifyService struct {
	timeout int
	factory node.Subscriber
	cancel  context.CancelFunc
	conn    *wsconn
}

func (self *wsNotifyService) Handle(ws *websocket.Conn) {
	eatPanics := false
	defer func() {
		// we want to ignore errors when other-side disappears w/o warning which is
		// outside our control.  We should still panic otherwise because it would
		// point out programtic errors which probably should be surfaced and fixed.
		if eatPanics {
			if r := recover(); r != nil {
				c2.Info.Printf("recovering from panic from unclear browser disconnect. %s", r)
			}
		}
	}()
	defer ws.Close()
	defer self.cancel()

	var rate time.Duration
	if self.timeout == 0 {
		rate = PingRate
	} else {
		rate = time.Duration(self.timeout) * time.Millisecond
	}
	self.conn = &wsconn{
		pinger: time.NewTicker(rate),
		mgr:    node.NewSubscriptionManager(self.factory, ws, ws),
	}
	defer self.conn.close()
	ws.Request().Body.Close()
	go self.conn.keepAlive(ws)
	if err := self.conn.mgr.Run(); err != nil {
		eatPanics = true
		c2.Info.Printf("unclean terminination of web socket: (%s). other side may have close browser. closing socket.", err)
	}
}

type wsconn struct {
	pinger *time.Ticker
	mgr    *node.SubscriptionManager
}

func (self *wsconn) keepAlive(ws *websocket.Conn) {
	for {
		ws.SetDeadline(time.Now().Add(serverSocketTimeout))
		if fw, err := ws.NewFrameWriter(websocket.PingFrame); err != nil {
			//self.Close()
			return
		} else if _, err = fw.Write([]byte{}); err != nil {
			//self.Close()
			return
		}
		if _, running := <-self.pinger.C; !running {
			return
		}
	}
}

func (self *wsconn) close() {
	self.pinger.Stop()
	self.mgr.Close()
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2016 Intel Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"

	"github.com/01org/ciao/ciao-controller/api"
	"github.com/01org/ciao/ciao-controller/types"
	"github.com/01org/ciao/openstack/compute"
	"github.com/01org/ciao/payloads"

	"gopkg.in/yaml.v2"
)

var workloadCommand = &command{
	SubCommands: map[string]subCommand{
		"list":   new(workloadListCommand),
		"create": new(workloadCreateCommand),
		"delete": new(workloadDeleteCommand),
	},
}

type workloadListCommand struct {
	Flag     flag.FlagSet
	template string
}

func (cmd *workloadListCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload list

List all workloads

`)
	cmd.Flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\n%s",
		generateUsageDecorated("f", compute.FlavorsDetails{}.Flavors))
	os.Exit(2)
}

func (cmd *workloadListCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.template, "f", "", "Template used to format output")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadListCommand) run(args []string) error {
	if *tenantID == "" {
		fatalf("Missing required -tenant-id parameter")
	}

	var flavors compute.FlavorsDetails
	if *tenantID == "" {
		*tenantID = "faketenant"
	}

	url := buildComputeURL("%s/flavors/detail", *tenantID)

	resp, err := sendHTTPRequest("GET", url, nil, nil)
	if err != nil {
		fatalf(err.Error())
	}

	err = unmarshalHTTPResponse(resp, &flavors)
	if err != nil {
		fatalf(err.Error())
	}

	if cmd.template != "" {
		return outputToTemplate("workload-list", cmd.template,
			&flavors.Flavors)
	}

	for i, flavor := range flavors.Flavors {
		fmt.Printf("Workload %d\n", i+1)
		fmt.Printf("\tName: %s\n\tUUID:%s\n\tImage UUID: %s\n\tCPUs: %d\n\tMemory: %d MB\n",
			flavor.Name, flavor.ID, flavor.Disk, flavor.Vcpus, flavor.RAM)
	}
	return nil
}

type workloadCreateCommand struct {
	Flag     flag.FlagSet
	yamlFile string
}

func (cmd *workloadCreateCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.yamlFile, "yaml", "", "filename for yaml which describes the workload")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadCreateCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload create [flags]

Create a new workload

The create flags are:

`)
	cmd.Flag.PrintDefaults()
	os.Exit(2)
}

func getCiaoWorkloadsResource() (string, error) {
	return getCiaoResource("workloads", api.WorkloadsV1)
}

type source struct {
	Type types.SourceType `yaml:"service"`
	ID   string           `yaml:"id"`
}

type disk struct {
	ID        *string `yaml:"volume_id"`
	Size      int     `yaml:"size"`
	Bootable  bool    `yaml:"bootable"`
	Source    *source `yaml:"source"`
	Ephemeral bool    `yaml:"ephemeral"`
}

type defaultResources struct {
	VCPUs int `yaml:"vcpus"`
	MemMB int `yaml:"mem_mb"`
}

// we currently only use the first disk due to lack of support
// in types.Workload for multiple storage resources.
type workloadOptions struct {
	Description     string           `yaml:"description"`
	VMType          string           `yaml:"vm_type"`
	FWType          string           `yaml:"fw_type"`
	ImageName       string           `yaml:"image_name"`
	ImageID         string           `yaml:"image_id"`
	Defaults        defaultResources `yaml:"defaults"`
	CloudConfigFile string           `yaml:"cloud_init"`
	Disks           []disk           `yaml:"disks"`
}

func optToReqStorage(opt workloadOptions) ([]types.StorageResource, error) {
	storage := make([]types.StorageResource, 0)
	bootableCount := 0
	for _, disk := range opt.Disks {
		res := types.StorageResource{
			Size:      disk.Size,
			Bootable:  disk.Bootable,
			Ephemeral: disk.Ephemeral,
		}

		if disk.Source != nil && disk.Source.Type == "" {
			disk.Source.Type = types.Empty
		}

		if disk.ID != nil {
			res.ID = *disk.ID
		} else if disk.Size == 0 {
			// source had better exist.
			if disk.Source == nil {
				return nil, errors.New("Invalid workload yaml: disk source may not be nil")
			}

			res.SourceType = disk.Source.Type
			res.SourceID = disk.Source.ID

			if res.SourceID != "" && res.SourceType == types.Empty {
				return nil, errors.New("Invalid workload yaml: when specifying a source ID a type must also be specified")
			}

			if res.SourceType != types.Empty && res.SourceID == "" {
				return nil, errors.New("Invalid workload yaml: when specifying a source type other than empty an id must also be specified")
			}
		} else {
			if disk.Bootable == true {
				// you may not request a bootable drive
				// from an empty source
				return nil, errors.New("Invalid workload yaml: empty disk source may not be bootable")
			}

			res.SourceType = types.Empty
		}

		if disk.Bootable {
			bootableCount++
		}

		storage = append(storage, res)
	}

	if payloads.Hypervisor(opt.VMType) == payloads.QEMU && bootableCount == 0 {
		return nil, errors.New("Invalid workload yaml: no bootable disks specified for a VM")
	}

	return storage, nil
}

func optToReq(opt workloadOptions, req *types.Workload) error {
	b, err := ioutil.ReadFile(opt.CloudConfigFile)
	if err != nil {
		return err
	}

	config := string(b)

	// this is where you'd validate that the options make
	// sense.
	req.Description = opt.Description
	req.VMType = payloads.Hypervisor(opt.VMType)
	req.FWType = opt.FWType
	req.ImageName = opt.ImageName
	req.ImageID = opt.ImageID
	req.Config = config
	req.Storage, err = optToReqStorage(opt)

	if err != nil {
		return err
	}

	// all default resources are required.
	defaults := opt.Defaults

	r := payloads.RequestedResource{
		Type:  payloads.VCPUs,
		Value: defaults.VCPUs,
	}
	req.Defaults = append(req.Defaults, r)

	r = payloads.RequestedResource{
		Type:  payloads.MemMB,
		Value: defaults.MemMB,
	}
	req.Defaults = append(req.Defaults, r)

	return nil
}

func (cmd *workloadCreateCommand) run(args []string) error {
	var opt workloadOptions
	var req types.Workload

	if cmd.yamlFile == "" {
		cmd.usage()
	}

	f, err := ioutil.ReadFile(cmd.yamlFile)
	if err != nil {
		fatalf("Unable to read workload config file: %s\n", err)
	}

	err = yaml.Unmarshal(f, &opt)
	if err != nil {
		fatalf("Config file invalid: %s\n", err)
	}

	err = optToReq(opt, &req)
	if err != nil {
		fatalf(err.Error())
	}

	b, err := json.Marshal(req)
	if err != nil {
		fatalf(err.Error())
	}

	body := bytes.NewReader(b)

	url, err := getCiaoWorkloadsResource()
	if err != nil {
		fatalf(err.Error())
	}

	ver := api.WorkloadsV1

	resp, err := sendCiaoRequest("POST", url, nil, body, &ver)
	if err != nil {
		fatalf(err.Error())
	}

	if resp.StatusCode != http.StatusCreated {
		fatalf("Workload creation failed: %s", resp.Status)
	}

	var workload types.WorkloadResponse

	err = unmarshalHTTPResponse(resp, &workload)
	if err != nil {
		fatalf(err.Error())
	}

	fmt.Printf("Created new workload: %s\n", workload.Workload.ID)

	return nil
}

type workloadDeleteCommand struct {
	Flag     flag.FlagSet
	workload string
}

func (cmd *workloadDeleteCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload delete [flags]

Deletes a given workload

The delete flags are:

`)
	cmd.Flag.PrintDefaults()
	os.Exit(2)
}

func (cmd *workloadDeleteCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.workload, "workload", "", "Workload UUID")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadDeleteCommand) run(args []string) error {
	if cmd.workload == "" {
		cmd.usage()
	}

	url, err := getCiaoWorkloadsResource()
	if err != nil {
		fatalf(err.Error())
	}

	ver := api.WorkloadsV1

	// you should do a get first and search for the workload,
	// then use the href - but not with the currently used
	// OpenStack API. Until we support GET with a ciao API,
	// just hard code the path.
	url = fmt.Sprintf("%s/%s", url, cmd.workload)

	resp, err := sendCiaoRequest("DELETE", url, nil, nil, &ver)
	if err != nil {
		fatalf(err.Error())
	}

	if resp.StatusCode != http.StatusNoContent {
		fatalf("Workload deletion failed: %s", resp.Status)
	}

	return nil
}
<commit_msg>ciao-cli: add workload show support<commit_after>//
// Copyright (c) 2016 Intel Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"

	"github.com/01org/ciao/ciao-controller/api"
	"github.com/01org/ciao/ciao-controller/types"
	"github.com/01org/ciao/openstack/compute"
	"github.com/01org/ciao/payloads"

	"gopkg.in/yaml.v2"
)

var workloadCommand = &command{
	SubCommands: map[string]subCommand{
		"list":   new(workloadListCommand),
		"create": new(workloadCreateCommand),
		"delete": new(workloadDeleteCommand),
		"show":   new(workloadShowCommand),
	},
}

type workloadListCommand struct {
	Flag     flag.FlagSet
	template string
}

func (cmd *workloadListCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload list

List all workloads

`)
	cmd.Flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\n%s",
		generateUsageDecorated("f", compute.FlavorsDetails{}.Flavors))
	os.Exit(2)
}

func (cmd *workloadListCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.template, "f", "", "Template used to format output")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadListCommand) run(args []string) error {
	if *tenantID == "" {
		fatalf("Missing required -tenant-id parameter")
	}

	var flavors compute.FlavorsDetails
	if *tenantID == "" {
		*tenantID = "faketenant"
	}

	url := buildComputeURL("%s/flavors/detail", *tenantID)

	resp, err := sendHTTPRequest("GET", url, nil, nil)
	if err != nil {
		fatalf(err.Error())
	}

	err = unmarshalHTTPResponse(resp, &flavors)
	if err != nil {
		fatalf(err.Error())
	}

	if cmd.template != "" {
		return outputToTemplate("workload-list", cmd.template,
			&flavors.Flavors)
	}

	for i, flavor := range flavors.Flavors {
		fmt.Printf("Workload %d\n", i+1)
		fmt.Printf("\tName: %s\n\tUUID:%s\n\tImage UUID: %s\n\tCPUs: %d\n\tMemory: %d MB\n",
			flavor.Name, flavor.ID, flavor.Disk, flavor.Vcpus, flavor.RAM)
	}
	return nil
}

type workloadCreateCommand struct {
	Flag     flag.FlagSet
	yamlFile string
}

func (cmd *workloadCreateCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.yamlFile, "yaml", "", "filename for yaml which describes the workload")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadCreateCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload create [flags]

Create a new workload

The create flags are:

`)
	cmd.Flag.PrintDefaults()
	os.Exit(2)
}

func getCiaoWorkloadsResource() (string, error) {
	return getCiaoResource("workloads", api.WorkloadsV1)
}

type source struct {
	Type types.SourceType `yaml:"service"`
	ID   string           `yaml:"id"`
}

type disk struct {
	ID        *string `yaml:"volume_id,omitempty"`
	Size      int     `yaml:"size"`
	Bootable  bool    `yaml:"bootable"`
	Source    *source `yaml:"source"`
	Ephemeral bool    `yaml:"ephemeral"`
}

type defaultResources struct {
	VCPUs int `yaml:"vcpus"`
	MemMB int `yaml:"mem_mb"`
}

// we currently only use the first disk due to lack of support
// in types.Workload for multiple storage resources.
type workloadOptions struct {
	Description     string           `yaml:"description"`
	VMType          string           `yaml:"vm_type"`
	FWType          string           `yaml:"fw_type,omitempty"`
	ImageName       string           `yaml:"image_name,omitempty"`
	ImageID         string           `yaml:"image_id,omitempty"`
	Defaults        defaultResources `yaml:"defaults"`
	CloudConfigFile string           `yaml:"cloud_init,omitempty"`
	Disks           []disk           `yaml:"disks,omitempty"`
}

func optToReqStorage(opt workloadOptions) ([]types.StorageResource, error) {
	storage := make([]types.StorageResource, 0)
	bootableCount := 0
	for _, disk := range opt.Disks {
		res := types.StorageResource{
			Size:      disk.Size,
			Bootable:  disk.Bootable,
			Ephemeral: disk.Ephemeral,
		}

		if disk.Source != nil && disk.Source.Type == "" {
			disk.Source.Type = types.Empty
		}

		if disk.ID != nil {
			res.ID = *disk.ID
		} else if disk.Size == 0 {
			// source had better exist.
			if disk.Source == nil {
				return nil, errors.New("Invalid workload yaml: disk source may not be nil")
			}

			res.SourceType = disk.Source.Type
			res.SourceID = disk.Source.ID

			if res.SourceID != "" && res.SourceType == types.Empty {
				return nil, errors.New("Invalid workload yaml: when specifying a source ID a type must also be specified")
			}

			if res.SourceType != types.Empty && res.SourceID == "" {
				return nil, errors.New("Invalid workload yaml: when specifying a source type other than empty an id must also be specified")
			}
		} else {
			if disk.Bootable == true {
				// you may not request a bootable drive
				// from an empty source
				return nil, errors.New("Invalid workload yaml: empty disk source may not be bootable")
			}

			res.SourceType = types.Empty
		}

		if disk.Bootable {
			bootableCount++
		}

		storage = append(storage, res)
	}

	if payloads.Hypervisor(opt.VMType) == payloads.QEMU && bootableCount == 0 {
		return nil, errors.New("Invalid workload yaml: no bootable disks specified for a VM")
	}

	return storage, nil
}

func optToReq(opt workloadOptions, req *types.Workload) error {
	b, err := ioutil.ReadFile(opt.CloudConfigFile)
	if err != nil {
		return err
	}

	config := string(b)

	// this is where you'd validate that the options make
	// sense.
	req.Description = opt.Description
	req.VMType = payloads.Hypervisor(opt.VMType)
	req.FWType = opt.FWType
	req.ImageName = opt.ImageName
	req.ImageID = opt.ImageID
	req.Config = config
	req.Storage, err = optToReqStorage(opt)

	if err != nil {
		return err
	}

	// all default resources are required.
	defaults := opt.Defaults

	r := payloads.RequestedResource{
		Type:  payloads.VCPUs,
		Value: defaults.VCPUs,
	}
	req.Defaults = append(req.Defaults, r)

	r = payloads.RequestedResource{
		Type:  payloads.MemMB,
		Value: defaults.MemMB,
	}
	req.Defaults = append(req.Defaults, r)

	return nil
}

func outputWorkload(w types.Workload) {
	var opt workloadOptions

	opt.Description = w.Description
	opt.VMType = string(w.VMType)
	opt.FWType = w.FWType
	opt.ImageName = w.ImageName
	opt.ImageID = w.ImageID
	for _, d := range w.Defaults {
		if d.Type == payloads.VCPUs {
			opt.Defaults.VCPUs = d.Value
		} else if d.Type == payloads.MemMB {
			opt.Defaults.MemMB = d.Value
		}
	}

	for _, s := range w.Storage {
		d := disk{
			Size:      s.Size,
			Bootable:  s.Bootable,
			Ephemeral: s.Ephemeral,
		}
		if s.ID != "" {
			d.ID = &s.ID
		}

		src := source{
			Type: s.SourceType,
			ID:   s.SourceID,
		}

		d.Source = &src

		opt.Disks = append(opt.Disks, d)
	}

	b, err := yaml.Marshal(opt)
	if err != nil {
		fatalf(err.Error())
	}

	fmt.Println(string(b))
	fmt.Println(w.Config)
}

func (cmd *workloadCreateCommand) run(args []string) error {
	var opt workloadOptions
	var req types.Workload

	if cmd.yamlFile == "" {
		cmd.usage()
	}

	f, err := ioutil.ReadFile(cmd.yamlFile)
	if err != nil {
		fatalf("Unable to read workload config file: %s\n", err)
	}

	err = yaml.Unmarshal(f, &opt)
	if err != nil {
		fatalf("Config file invalid: %s\n", err)
	}

	err = optToReq(opt, &req)
	if err != nil {
		fatalf(err.Error())
	}

	b, err := json.Marshal(req)
	if err != nil {
		fatalf(err.Error())
	}

	body := bytes.NewReader(b)

	url, err := getCiaoWorkloadsResource()
	if err != nil {
		fatalf(err.Error())
	}

	ver := api.WorkloadsV1

	resp, err := sendCiaoRequest("POST", url, nil, body, &ver)
	if err != nil {
		fatalf(err.Error())
	}

	if resp.StatusCode != http.StatusCreated {
		fatalf("Workload creation failed: %s", resp.Status)
	}

	var workload types.WorkloadResponse

	err = unmarshalHTTPResponse(resp, &workload)
	if err != nil {
		fatalf(err.Error())
	}

	fmt.Printf("Created new workload: %s\n", workload.Workload.ID)

	return nil
}

type workloadDeleteCommand struct {
	Flag     flag.FlagSet
	workload string
}

func (cmd *workloadDeleteCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload delete [flags]

Deletes a given workload

The delete flags are:

`)
	cmd.Flag.PrintDefaults()
	os.Exit(2)
}

func (cmd *workloadDeleteCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.workload, "workload", "", "Workload UUID")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadDeleteCommand) run(args []string) error {
	if cmd.workload == "" {
		cmd.usage()
	}

	url, err := getCiaoWorkloadsResource()
	if err != nil {
		fatalf(err.Error())
	}

	ver := api.WorkloadsV1

	// you should do a get first and search for the workload,
	// then use the href - but not with the currently used
	// OpenStack API. Until we support GET with a ciao API,
	// just hard code the path.
	url = fmt.Sprintf("%s/%s", url, cmd.workload)

	resp, err := sendCiaoRequest("DELETE", url, nil, nil, &ver)
	if err != nil {
		fatalf(err.Error())
	}

	if resp.StatusCode != http.StatusNoContent {
		fatalf("Workload deletion failed: %s", resp.Status)
	}

	return nil
}

type workloadShowCommand struct {
	Flag     flag.FlagSet
	template string
	workload string
}

func (cmd *workloadShowCommand) usage(...string) {
	fmt.Fprintf(os.Stderr, `usage: ciao-cli [options] workload show

Show workload details

`)
	cmd.Flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\n%s",
		generateUsageDecorated("f", types.Workload{}))
	os.Exit(2)
}

func (cmd *workloadShowCommand) parseArgs(args []string) []string {
	cmd.Flag.StringVar(&cmd.workload, "workload", "", "Workload UUID")
	cmd.Flag.StringVar(&cmd.template, "f", "", "Template used to format output")
	cmd.Flag.Usage = func() { cmd.usage() }
	cmd.Flag.Parse(args)
	return cmd.Flag.Args()
}

func (cmd *workloadShowCommand) run(args []string) error {
	var wl types.Workload

	if cmd.workload == "" {
		cmd.usage()
	}

	url, err := getCiaoWorkloadsResource()
	if err != nil {
		fatalf(err.Error())
	}

	ver := api.WorkloadsV1

	// you should do a get first and search for the workload,
	// then use the href - but not with the currently used
	// OpenStack API. Until we support GET with a ciao API,
	// just hard code the path.
	url = fmt.Sprintf("%s/%s", url, cmd.workload)

	resp, err := sendCiaoRequest("GET", url, nil, nil, &ver)
	if err != nil {
		fatalf(err.Error())
	}

	if resp.StatusCode != http.StatusOK {
		fatalf("Workload show failed: %s", resp.Status)
	}

	err = unmarshalHTTPResponse(resp, &wl)
	if err != nil {
		fatalf(err.Error())
	}

	if cmd.template != "" {
		return outputToTemplate("workload-show", cmd.template,
			&wl)
	}

	outputWorkload(wl)
	return nil
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"fmt"
	"io"
	"sync"

	"github.com/golang/snappy"

	"github.com/gogo/protobuf/types"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/admin"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/client/pkg/pbutil"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/log"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
)

type apiServer struct {
	log.Logger
	address        string
	pachClient     *client.APIClient
	pachClientOnce sync.Once
}

func (a *apiServer) Extract(request *admin.ExtractRequest, extractServer admin.API_ExtractServer) (retErr error) {
	pachClient := a.getPachClient()
	pachClient = pachClient.WithCtx(extractServer.Context())
	handleOp := extractServer.Send
	if request.URL != "" {
		url, err := obj.ParseURL(request.URL)
		if err != nil {
			return fmt.Errorf("error parsing url %v: %v", request.URL, err)
		}
		objClient, err := obj.NewClientFromURLAndSecret(extractServer.Context(), url)
		if err != nil {
			return err
		}
		objW, err := objClient.Writer(url.Object)
		if err != nil {
			return err
		}
		snappyW := snappy.NewBufferedWriter(objW)
		defer func() {
			if err := snappyW.Close(); err != nil && retErr == nil {
				retErr = err
			}
		}()
		w := pbutil.NewWriter(snappyW)
		handleOp = func(op *admin.Op) error { return w.Write(op) }
	}
	v, err := pachClient.VersionAPIClient.GetVersion(pachClient.Ctx(), &types.Empty{})
	if err != nil {
		return err
	}
	if err := handleOp(&admin.Op{
		Version: v,
	}); err != nil {
		return err
	}
	if !request.NoObjects {
		w := &extractObjectWriter{extractServer}
		if err := pachClient.ListObject(func(object *pfs.Object) error {
			if err := pachClient.GetObject(object.Hash, w); err != nil {
				return err
			}
			// empty PutObjectRequest to indicate EOF
			return handleOp(&admin.Op{Object: &pfs.PutObjectRequest{}})
		}); err != nil {
			return err
		}
		if err := pachClient.ListTag(func(resp *pfs.ListTagsResponse) error {
			return handleOp(&admin.Op{
				Tag: &pfs.TagObjectRequest{
					Object: resp.Object,
					Tags:   []*pfs.Tag{resp.Tag},
				},
			})
		}); err != nil {
			return err
		}
	}
	if !request.NoRepos {
		ris, err := pachClient.ListRepo(nil)
		if err != nil {
			return err
		}
		for _, ri := range ris {
			if len(ri.Provenance) > 0 {
				continue
			}
			if err := handleOp(&admin.Op{
				Repo: &pfs.CreateRepoRequest{
					Repo:        ri.Repo,
					Provenance:  ri.Provenance,
					Description: ri.Description,
				},
			}); err != nil {
				return err
			}
			cis, err := pachClient.ListCommit(ri.Repo.Name, "", "", 0)
			if err != nil {
				return err
			}
			for _, ci := range sortCommitInfos(cis) {
				// Even without a parent, ParentCommit is used to indicate which
				// repo to make the commit in.
				if ci.ParentCommit == nil {
					ci.ParentCommit = client.NewCommit(ri.Repo.Name, "")
				}
				if err := handleOp(&admin.Op{
					Commit: &pfs.BuildCommitRequest{
						Parent: ci.ParentCommit,
						Tree:   ci.Tree,
						ID:     ci.Commit.ID,
					},
				}); err != nil {
					return err
				}
			}
			bis, err := pachClient.ListBranch(ri.Repo.Name)
			if err != nil {
				return err
			}
			for _, bi := range bis {
				if err := handleOp(&admin.Op{
					Branch: &pfs.SetBranchRequest{
						Commit: bi.Head,
						Branch: bi.Name,
					},
				},
				); err != nil {
					return err
				}
			}
		}
	}
	if !request.NoPipelines {
		pis, err := pachClient.ListPipeline()
		if err != nil {
			return err
		}
		for _, pi := range pis {
			if err := handleOp(&admin.Op{
				Pipeline: &pps.CreatePipelineRequest{
					Pipeline:           pi.Pipeline,
					Transform:          pi.Transform,
					ParallelismSpec:    pi.ParallelismSpec,
					Egress:             pi.Egress,
					OutputBranch:       pi.OutputBranch,
					ScaleDownThreshold: pi.ScaleDownThreshold,
					ResourceRequests:   pi.ResourceRequests,
					ResourceLimits:     pi.ResourceLimits,
					Input:              pi.Input,
					Description:        pi.Description,
					Incremental:        pi.Incremental,
					CacheSize:          pi.CacheSize,
					EnableStats:        pi.EnableStats,
					Batch:              pi.Batch,
					MaxQueueSize:       pi.MaxQueueSize,
					Service:            pi.Service,
					ChunkSpec:          pi.ChunkSpec,
					DatumTimeout:       pi.DatumTimeout,
					JobTimeout:         pi.JobTimeout,
					Salt:               pi.Salt,
				},
			}); err != nil {
				return err
			}
		}
	}
	return nil
}

func sortCommitInfos(cis []*pfs.CommitInfo) []*pfs.CommitInfo {
	commitMap := make(map[string]*pfs.CommitInfo)
	for _, ci := range cis {
		commitMap[ci.Commit.ID] = ci
	}
	var result []*pfs.CommitInfo
	for _, ci := range cis {
		if commitMap[ci.Commit.ID] == nil {
			continue
		}
		var localResult []*pfs.CommitInfo
		for ci != nil {
			localResult = append(localResult, ci)
			delete(commitMap, ci.Commit.ID)
			if ci.ParentCommit != nil {
				ci = commitMap[ci.ParentCommit.ID]
			} else {
				ci = nil
			}
		}
		for i := range localResult {
			result = append(result, localResult[len(localResult)-i-1])
		}
	}
	return result
}

func (a *apiServer) Restore(restoreServer admin.API_RestoreServer) (retErr error) {
	ctx := restoreServer.Context()
	pachClient := a.getPachClient()
	defer func() {
		for {
			_, err := restoreServer.Recv()
			if err != nil {
				break
			}
		}
		if err := restoreServer.SendAndClose(&types.Empty{}); err != nil && retErr == nil {
			retErr = err
		}
	}()
	var r pbutil.Reader
	for {
		var op *admin.Op
		if r == nil {
			req, err := restoreServer.Recv()
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
			if req.URL != "" {
				url, err := obj.ParseURL(req.URL)
				if err != nil {
					return fmt.Errorf("error parsing url %v: %v", req.URL, err)
				}
				objClient, err := obj.NewClientFromURLAndSecret(restoreServer.Context(), url)
				if err != nil {
					return err
				}
				objR, err := objClient.Reader(url.Object, 0, 0)
				if err != nil {
					return err
				}
				snappyR := snappy.NewReader(objR)
				r = pbutil.NewReader(snappyR)
				continue
			} else {
				op = req.Op
			}
		} else {
			op = &admin.Op{}
			if err := r.Read(op); err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
		}
		switch {
		case op.Version != nil:
		case op.Object != nil:
			r := &extractObjectReader{adminAPIRestoreServer: restoreServer}
			r.buf.Write(op.Object.Value)
			if _, _, err := pachClient.PutObject(r); err != nil {
				return fmt.Errorf("error putting object: %v", err)
			}
		case op.Tag != nil:
			if _, err := pachClient.ObjectAPIClient.TagObject(ctx, op.Tag); err != nil {
				return fmt.Errorf("error tagging object: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Repo != nil:
			if _, err := pachClient.PfsAPIClient.CreateRepo(ctx, op.Repo); err != nil {
				return fmt.Errorf("error creating repo: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Commit != nil:
			if _, err := pachClient.PfsAPIClient.BuildCommit(ctx, op.Commit); err != nil {
				return fmt.Errorf("error creating commit: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Branch != nil:
			if _, err := pachClient.PfsAPIClient.SetBranch(ctx, op.Branch); err != nil {
				return fmt.Errorf("error creating branch: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Pipeline != nil:
			if _, err := pachClient.PpsAPIClient.CreatePipeline(ctx, op.Pipeline); err != nil {
				return fmt.Errorf("error creating pipeline: %v", grpcutil.ScrubGRPC(err))
			}
		}
	}
}

func (a *apiServer) getPachClient() *client.APIClient {
	a.pachClientOnce.Do(func() {
		var err error
		a.pachClient, err = client.NewFromAddress(a.address)
		if err != nil {
			panic(fmt.Sprintf("pps failed to initialize pach client: %v", err))
		}
	})
	return a.pachClient
}

type extractObjectWriter struct {
	admin.API_ExtractServer
}

func (w *extractObjectWriter) Write(p []byte) (int, error) {
	chunkSize := grpcutil.MaxMsgSize / 2
	var n int
	for i := 0; i*(chunkSize) < len(p); i++ {
		value := p[i*chunkSize:]
		if len(value) > chunkSize {
			value = value[:chunkSize]
		}
		if err := w.Send(&admin.Op{Object: &pfs.PutObjectRequest{Value: value}}); err != nil {
			return n, err
		}
		n += len(value)
	}
	return n, nil
}

type adminAPIRestoreServer admin.API_RestoreServer

type extractObjectReader struct {
	adminAPIRestoreServer
	buf bytes.Buffer
	eof bool
}

func (r *extractObjectReader) Read(p []byte) (int, error) {
	for len(p) > r.buf.Len() && !r.eof {
		request, err := r.Recv()
		if err != nil {
			return 0, grpcutil.ScrubGRPC(err)
		}
		op := request.Op
		if op.Object == nil {
			return 0, fmt.Errorf("expected an object, but got: %v", op)
		}
		r.buf.Write(op.Object.Value)
		if len(op.Object.Value) == 0 {
			r.eof = true
		}
	}
	return r.buf.Read(p)
}
<commit_msg>Migrate to new proto design.<commit_after>package server

import (
	"bytes"
	"fmt"
	"io"
	"sync"

	"github.com/golang/snappy"

	"github.com/gogo/protobuf/types"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/admin"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/client/pkg/pbutil"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/log"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
)

type apiServer struct {
	log.Logger
	address        string
	pachClient     *client.APIClient
	pachClientOnce sync.Once
}

func (a *apiServer) Extract(request *admin.ExtractRequest, extractServer admin.API_ExtractServer) (retErr error) {
	pachClient := a.getPachClient()
	pachClient = pachClient.WithCtx(extractServer.Context())
	handleOp := extractServer.Send
	if request.URL != "" {
		url, err := obj.ParseURL(request.URL)
		if err != nil {
			return fmt.Errorf("error parsing url %v: %v", request.URL, err)
		}
		objClient, err := obj.NewClientFromURLAndSecret(extractServer.Context(), url)
		if err != nil {
			return err
		}
		objW, err := objClient.Writer(url.Object)
		if err != nil {
			return err
		}
		snappyW := snappy.NewBufferedWriter(objW)
		defer func() {
			if err := snappyW.Close(); err != nil && retErr == nil {
				retErr = err
			}
		}()
		w := pbutil.NewWriter(snappyW)
		handleOp = func(op *admin.Op) error { return w.Write(op) }
	}
	if !request.NoObjects {
		w := &extractObjectWriter{extractServer}
		if err := pachClient.ListObject(func(object *pfs.Object) error {
			if err := pachClient.GetObject(object.Hash, w); err != nil {
				return err
			}
			// empty PutObjectRequest to indicate EOF
			return handleOp(&admin.Op{Op17: &admin.Op17{Object: &pfs.PutObjectRequest{}}})
		}); err != nil {
			return err
		}
		if err := pachClient.ListTag(func(resp *pfs.ListTagsResponse) error {
			return handleOp(&admin.Op{Op17: &admin.Op17{
				Tag: &pfs.TagObjectRequest{
					Object: resp.Object,
					Tags:   []*pfs.Tag{resp.Tag},
				},
			}})
		}); err != nil {
			return err
		}
	}
	if !request.NoRepos {
		ris, err := pachClient.ListRepo(nil)
		if err != nil {
			return err
		}
		for _, ri := range ris {
			if len(ri.Provenance) > 0 {
				continue
			}
			if err := handleOp(&admin.Op{Op17: &admin.Op17{
				Repo: &pfs.CreateRepoRequest{
					Repo:        ri.Repo,
					Provenance:  ri.Provenance,
					Description: ri.Description,
				}},
			}); err != nil {
				return err
			}
			cis, err := pachClient.ListCommit(ri.Repo.Name, "", "", 0)
			if err != nil {
				return err
			}
			for _, ci := range sortCommitInfos(cis) {
				// Even without a parent, ParentCommit is used to indicate which
				// repo to make the commit in.
				if ci.ParentCommit == nil {
					ci.ParentCommit = client.NewCommit(ri.Repo.Name, "")
				}
				if err := handleOp(&admin.Op{Op17: &admin.Op17{
					Commit: &pfs.BuildCommitRequest{
						Parent: ci.ParentCommit,
						Tree:   ci.Tree,
						ID:     ci.Commit.ID,
					},
				}}); err != nil {
					return err
				}
			}
			bis, err := pachClient.ListBranch(ri.Repo.Name)
			if err != nil {
				return err
			}
			for _, bi := range bis {
				if err := handleOp(&admin.Op{Op17: &admin.Op17{
					Branch: &pfs.SetBranchRequest{
						Commit: bi.Head,
						Branch: bi.Name,
					},
				}}); err != nil {
					return err
				}
			}
		}
	}
	if !request.NoPipelines {
		pis, err := pachClient.ListPipeline()
		if err != nil {
			return err
		}
		for _, pi := range pis {
			if err := handleOp(&admin.Op{Op17: &admin.Op17{
				Pipeline: &pps.CreatePipelineRequest{
					Pipeline:           pi.Pipeline,
					Transform:          pi.Transform,
					ParallelismSpec:    pi.ParallelismSpec,
					Egress:             pi.Egress,
					OutputBranch:       pi.OutputBranch,
					ScaleDownThreshold: pi.ScaleDownThreshold,
					ResourceRequests:   pi.ResourceRequests,
					ResourceLimits:     pi.ResourceLimits,
					Input:              pi.Input,
					Description:        pi.Description,
					Incremental:        pi.Incremental,
					CacheSize:          pi.CacheSize,
					EnableStats:        pi.EnableStats,
					Batch:              pi.Batch,
					MaxQueueSize:       pi.MaxQueueSize,
					Service:            pi.Service,
					ChunkSpec:          pi.ChunkSpec,
					DatumTimeout:       pi.DatumTimeout,
					JobTimeout:         pi.JobTimeout,
					Salt:               pi.Salt,
				}}}); err != nil {
				return err
			}
		}
	}
	return nil
}

func sortCommitInfos(cis []*pfs.CommitInfo) []*pfs.CommitInfo {
	commitMap := make(map[string]*pfs.CommitInfo)
	for _, ci := range cis {
		commitMap[ci.Commit.ID] = ci
	}
	var result []*pfs.CommitInfo
	for _, ci := range cis {
		if commitMap[ci.Commit.ID] == nil {
			continue
		}
		var localResult []*pfs.CommitInfo
		for ci != nil {
			localResult = append(localResult, ci)
			delete(commitMap, ci.Commit.ID)
			if ci.ParentCommit != nil {
				ci = commitMap[ci.ParentCommit.ID]
			} else {
				ci = nil
			}
		}
		for i := range localResult {
			result = append(result, localResult[len(localResult)-i-1])
		}
	}
	return result
}

func (a *apiServer) Restore(restoreServer admin.API_RestoreServer) (retErr error) {
	ctx := restoreServer.Context()
	pachClient := a.getPachClient()
	defer func() {
		for {
			_, err := restoreServer.Recv()
			if err != nil {
				break
			}
		}
		if err := restoreServer.SendAndClose(&types.Empty{}); err != nil && retErr == nil {
			retErr = err
		}
	}()
	var r pbutil.Reader
	for {
		var op *admin.Op
		if r == nil {
			req, err := restoreServer.Recv()
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
			if req.URL != "" {
				url, err := obj.ParseURL(req.URL)
				if err != nil {
					return fmt.Errorf("error parsing url %v: %v", req.URL, err)
				}
				objClient, err := obj.NewClientFromURLAndSecret(restoreServer.Context(), url)
				if err != nil {
					return err
				}
				objR, err := objClient.Reader(url.Object, 0, 0)
				if err != nil {
					return err
				}
				snappyR := snappy.NewReader(objR)
				r = pbutil.NewReader(snappyR)
				continue
			} else {
				op = req.Op
			}
		} else {
			op = &admin.Op{}
			if err := r.Read(op); err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
		}
		switch {
		case op.Op17 != nil && op.Op17.Object != nil:
			r := &extractObjectReader{adminAPIRestoreServer: restoreServer}
			r.buf.Write(op.Op17.Object.Value)
			if _, _, err := pachClient.PutObject(r); err != nil {
				return fmt.Errorf("error putting object: %v", err)
			}
		case op.Op17 != nil && op.Op17.Tag != nil:
			if _, err := pachClient.ObjectAPIClient.TagObject(ctx, op.Op17.Tag); err != nil {
				return fmt.Errorf("error tagging object: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op17 != nil && op.Op17.Repo != nil:
			if _, err := pachClient.PfsAPIClient.CreateRepo(ctx, op.Op17.Repo); err != nil {
				return fmt.Errorf("error creating repo: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op17 != nil && op.Op17.Commit != nil:
			if _, err := pachClient.PfsAPIClient.BuildCommit(ctx, op.Op17.Commit); err != nil {
				return fmt.Errorf("error creating commit: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op17 != nil && op.Op17.Branch != nil:
			if _, err := pachClient.PfsAPIClient.SetBranch(ctx, op.Op17.Branch); err != nil {
				return fmt.Errorf("error creating branch: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op17 != nil && op.Op17.Pipeline != nil:
			if _, err := pachClient.PpsAPIClient.CreatePipeline(ctx, op.Op17.Pipeline); err != nil {
				return fmt.Errorf("error creating pipeline: %v", grpcutil.ScrubGRPC(err))
			}
		}
	}
}

func (a *apiServer) getPachClient() *client.APIClient {
	a.pachClientOnce.Do(func() {
		var err error
		a.pachClient, err = client.NewFromAddress(a.address)
		if err != nil {
			panic(fmt.Sprintf("pps failed to initialize pach client: %v", err))
		}
	})
	return a.pachClient
}

type extractObjectWriter struct {
	admin.API_ExtractServer
}

func (w *extractObjectWriter) Write(p []byte) (int, error) {
	chunkSize := grpcutil.MaxMsgSize / 2
	var n int
	for i := 0; i*(chunkSize) < len(p); i++ {
		value := p[i*chunkSize:]
		if len(value) > chunkSize {
			value = value[:chunkSize]
		}
		if err := w.Send(&admin.Op{Op17: &admin.Op17{Object: &pfs.PutObjectRequest{Value: value}}}); err != nil {
			return n, err
		}
		n += len(value)
	}
	return n, nil
}

type adminAPIRestoreServer admin.API_RestoreServer

type extractObjectReader struct {
	adminAPIRestoreServer
	buf bytes.Buffer
	eof bool
}

func (r *extractObjectReader) Read(p []byte) (int, error) {
	for len(p) > r.buf.Len() && !r.eof {
		request, err := r.Recv()
		if err != nil {
			return 0, grpcutil.ScrubGRPC(err)
		}
		op := request.Op
		if op.Op17.Object == nil {
			return 0, fmt.Errorf("expected an object, but got: %v", op)
		}
		r.buf.Write(op.Op17.Object.Value)
		if len(op.Op17.Object.Value) == 0 {
			r.eof = true
		}
	}
	return r.buf.Read(p)
}
<|endoftext|>