<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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 metrics

import (
	"github.com/blang/semver"
	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	"k8s.io/klog"
	"sync"
)

/*
kubeCollector extends the prometheus.Collector interface to allow customization of the metric
registration process. Defer metric initialization until Create() is called, which then
delegates to the underlying metric's initializeMetric or initializeDeprecatedMetric
method call depending on whether the metric is deprecated or not.
*/
type kubeCollector interface {
	Collector
	lazyKubeMetric
	DeprecatedVersion() *semver.Version
	// Each collector metric should provide an initialization function
	// for both deprecated and non-deprecated variants of a metric. This
	// is necessary since metric instantiation will be deferred
	// until the metric is actually registered somewhere.
	initializeMetric()
	initializeDeprecatedMetric()
}

/*
lazyKubeMetric defines our metric registration interface. lazyKubeMetric objects are expected
to lazily instantiate metrics (i.e defer metric instantiation until when
the Create() function is explicitly called).
*/
type lazyKubeMetric interface {
	Create(*semver.Version) bool
	IsCreated() bool
	IsHidden() bool
	IsDeprecated() bool
}

/*
lazyMetric implements lazyKubeMetric. A lazy metric is lazy because it waits until metric
registration time before instantiation. Add it as an anonymous field to a struct that
implements kubeCollector to get deferred registration behavior. You must call lazyInit
with the kubeCollector itself as an argument.
*/
type lazyMetric struct {
	isDeprecated        bool
	isHidden            bool
	isCreated           bool
	markDeprecationOnce sync.Once
	createOnce          sync.Once
	self                kubeCollector
}

func (r *lazyMetric) IsCreated() bool {
	return r.isCreated
}

// lazyInit provides the lazyMetric with a reference to the kubeCollector it is supposed
// to allow lazy initialization for. It should be invoked in the factory function which creates new
// kubeCollector type objects.
func (r *lazyMetric) lazyInit(self kubeCollector) {
	r.self = self
}

// determineDeprecationStatus figures out whether the lazy metric should be deprecated or not.
// This method takes a Version argument which should be the version of the binary in which
// this code is currently being executed.
func (r *lazyMetric) determineDeprecationStatus(version semver.Version) {
	selfVersion := r.self.DeprecatedVersion()
	if selfVersion == nil {
		return
	}
	r.markDeprecationOnce.Do(func() {
		if selfVersion.LTE(version) {
			r.isDeprecated = true
		}
		if ShouldShowHidden() {
			klog.Warningf("Hidden metrics have been manually overridden, showing this very deprecated metric.")
			return
		}
		if selfVersion.LT(version) {
			klog.Warningf("This metric has been deprecated for more than one release, hiding.")
			r.isHidden = true
		}
	})
}

func (r *lazyMetric) IsHidden() bool {
	return r.isHidden
}

func (r *lazyMetric) IsDeprecated() bool {
	return r.isDeprecated
}

// Create forces the initialization of metric which has been deferred until
// the point at which this method is invoked. This method will determine whether
// the metric is deprecated or hidden, no-opting if the metric should be considered
// hidden. Furthermore, this function no-opts and returns true if metric is already
// created.
func (r *lazyMetric) Create(version *semver.Version) bool {
	if version != nil {
		r.determineDeprecationStatus(*version)
	}
	// let's not create if this metric is slated to be hidden
	if r.IsHidden() {
		return false
	}
	r.createOnce.Do(func() {
		r.isCreated = true
		if r.IsDeprecated() {
			r.self.initializeDeprecatedMetric()
		} else {
			r.self.initializeMetric()
		}
	})
	return r.IsCreated()
}

/*
This code is directly lifted from the prometheus codebase. It's a convenience struct which
allows you satisfy the Collector interface automatically if you already satisfy the Metric interface.

For reference: https://github.com/prometheus/client_golang/blob/v0.9.2/prometheus/collector.go#L98-L120
*/
type selfCollector struct {
	metric prometheus.Metric
}

func (c *selfCollector) initSelfCollection(m prometheus.Metric) {
	c.metric = m
}

func (c *selfCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- c.metric.Desc()
}

func (c *selfCollector) Collect(ch chan<- prometheus.Metric) {
	ch <- c.metric
}

// no-op vecs for convenience
var noopCounterVec = &prometheus.CounterVec{}
var noopHistogramVec = &prometheus.HistogramVec{}
var noopSummaryVec = &prometheus.SummaryVec{}
var noopGaugeVec = &prometheus.GaugeVec{}
var noopObserverVec = &noopObserverVector{}

// just use a convenience struct for all the no-ops
var noop = &noopMetric{}

type noopMetric struct{}

func (noopMetric) Inc()                             {}
func (noopMetric) Add(float64)                      {}
func (noopMetric) Dec()                             {}
func (noopMetric) Set(float64)                      {}
func (noopMetric) Sub(float64)                      {}
func (noopMetric) Observe(float64)                  {}
func (noopMetric) SetToCurrentTime()                {}
func (noopMetric) Desc() *prometheus.Desc           { return nil }
func (noopMetric) Write(*dto.Metric) error          { return nil }
func (noopMetric) Describe(chan<- *prometheus.Desc) {}
func (noopMetric) Collect(chan<- prometheus.Metric) {}

type noopObserverVector struct{}

func (noopObserverVector) GetMetricWith(prometheus.Labels) (prometheus.Observer, error) {
	return noop, nil
}
func (noopObserverVector) GetMetricWithLabelValues(...string) (prometheus.Observer, error) {
	return noop, nil
}
func (noopObserverVector) With(prometheus.Labels) prometheus.Observer    { return noop }
func (noopObserverVector) WithLabelValues(...string) prometheus.Observer { return noop }
func (noopObserverVector) CurryWith(prometheus.Labels) (prometheus.ObserverVec, error) {
	return noopObserverVec, nil
}
func (noopObserverVector) MustCurryWith(prometheus.Labels) prometheus.ObserverVec {
	return noopObserverVec
}
func (noopObserverVector) Describe(chan<- *prometheus.Desc) {}
func (noopObserverVector) Collect(chan<- prometheus.Metric) {}
<commit_msg>lock the mutation of the create boolean for safety<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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 metrics

import (
	"github.com/blang/semver"
	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	"k8s.io/klog"
	"sync"
)

/*
kubeCollector extends the prometheus.Collector interface to allow customization of the metric
registration process. Defer metric initialization until Create() is called, which then
delegates to the underlying metric's initializeMetric or initializeDeprecatedMetric
method call depending on whether the metric is deprecated or not.
*/
type kubeCollector interface {
	Collector
	lazyKubeMetric
	DeprecatedVersion() *semver.Version
	// Each collector metric should provide an initialization function
	// for both deprecated and non-deprecated variants of a metric. This
	// is necessary since metric instantiation will be deferred
	// until the metric is actually registered somewhere.
	initializeMetric()
	initializeDeprecatedMetric()
}

/*
lazyKubeMetric defines our metric registration interface. lazyKubeMetric objects are expected
to lazily instantiate metrics (i.e defer metric instantiation until when
the Create() function is explicitly called).
*/
type lazyKubeMetric interface {
	Create(*semver.Version) bool
	IsCreated() bool
	IsHidden() bool
	IsDeprecated() bool
}

/*
lazyMetric implements lazyKubeMetric. A lazy metric is lazy because it waits until metric
registration time before instantiation. Add it as an anonymous field to a struct that
implements kubeCollector to get deferred registration behavior. You must call lazyInit
with the kubeCollector itself as an argument.
*/
type lazyMetric struct {
	isDeprecated        bool
	isHidden            bool
	isCreated           bool
	createLock          sync.RWMutex
	markDeprecationOnce sync.Once
	createOnce          sync.Once
	self                kubeCollector
}

func (r *lazyMetric) IsCreated() bool {
	r.createLock.RLock()
	defer r.createLock.RUnlock()
	return r.isCreated
}

// lazyInit provides the lazyMetric with a reference to the kubeCollector it is supposed
// to allow lazy initialization for. It should be invoked in the factory function which creates new
// kubeCollector type objects.
func (r *lazyMetric) lazyInit(self kubeCollector) {
	r.self = self
}

// determineDeprecationStatus figures out whether the lazy metric should be deprecated or not.
// This method takes a Version argument which should be the version of the binary in which
// this code is currently being executed.
func (r *lazyMetric) determineDeprecationStatus(version semver.Version) {
	selfVersion := r.self.DeprecatedVersion()
	if selfVersion == nil {
		return
	}
	r.markDeprecationOnce.Do(func() {
		if selfVersion.LTE(version) {
			r.isDeprecated = true
		}
		if ShouldShowHidden() {
			klog.Warningf("Hidden metrics have been manually overridden, showing this very deprecated metric.")
			return
		}
		if selfVersion.LT(version) {
			klog.Warningf("This metric has been deprecated for more than one release, hiding.")
			r.isHidden = true
		}
	})
}

func (r *lazyMetric) IsHidden() bool {
	return r.isHidden
}

func (r *lazyMetric) IsDeprecated() bool {
	return r.isDeprecated
}

// Create forces the initialization of metric which has been deferred until
// the point at which this method is invoked. This method will determine whether
// the metric is deprecated or hidden, no-opting if the metric should be considered
// hidden. Furthermore, this function no-opts and returns true if metric is already
// created.
func (r *lazyMetric) Create(version *semver.Version) bool {
	if version != nil {
		r.determineDeprecationStatus(*version)
	}
	// let's not create if this metric is slated to be hidden
	if r.IsHidden() {
		return false
	}
	r.createOnce.Do(func() {
		r.createLock.Lock()
		defer r.createLock.Unlock()
		r.isCreated = true
		if r.IsDeprecated() {
			r.self.initializeDeprecatedMetric()
		} else {
			r.self.initializeMetric()
		}
	})
	return r.IsCreated()
}

/*
This code is directly lifted from the prometheus codebase. It's a convenience struct which
allows you satisfy the Collector interface automatically if you already satisfy the Metric interface.

For reference: https://github.com/prometheus/client_golang/blob/v0.9.2/prometheus/collector.go#L98-L120
*/
type selfCollector struct {
	metric prometheus.Metric
}

func (c *selfCollector) initSelfCollection(m prometheus.Metric) {
	c.metric = m
}

func (c *selfCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- c.metric.Desc()
}

func (c *selfCollector) Collect(ch chan<- prometheus.Metric) {
	ch <- c.metric
}

// no-op vecs for convenience
var noopCounterVec = &prometheus.CounterVec{}
var noopHistogramVec = &prometheus.HistogramVec{}
var noopSummaryVec = &prometheus.SummaryVec{}
var noopGaugeVec = &prometheus.GaugeVec{}
var noopObserverVec = &noopObserverVector{}

// just use a convenience struct for all the no-ops
var noop = &noopMetric{}

type noopMetric struct{}

func (noopMetric) Inc()                             {}
func (noopMetric) Add(float64)                      {}
func (noopMetric) Dec()                             {}
func (noopMetric) Set(float64)                      {}
func (noopMetric) Sub(float64)                      {}
func (noopMetric) Observe(float64)                  {}
func (noopMetric) SetToCurrentTime()                {}
func (noopMetric) Desc() *prometheus.Desc           { return nil }
func (noopMetric) Write(*dto.Metric) error          { return nil }
func (noopMetric) Describe(chan<- *prometheus.Desc) {}
func (noopMetric) Collect(chan<- prometheus.Metric) {}

type noopObserverVector struct{}

func (noopObserverVector) GetMetricWith(prometheus.Labels) (prometheus.Observer, error) {
	return noop, nil
}
func (noopObserverVector) GetMetricWithLabelValues(...string) (prometheus.Observer, error) {
	return noop, nil
}
func (noopObserverVector) With(prometheus.Labels) prometheus.Observer    { return noop }
func (noopObserverVector) WithLabelValues(...string) prometheus.Observer { return noop }
func (noopObserverVector) CurryWith(prometheus.Labels) (prometheus.ObserverVec, error) {
	return noopObserverVec, nil
}
func (noopObserverVector) MustCurryWith(prometheus.Labels) prometheus.ObserverVec {
	return noopObserverVec
}
func (noopObserverVector) Describe(chan<- *prometheus.Desc) {}
func (noopObserverVector) Collect(chan<- prometheus.Metric) {}
<|endoftext|>


<commit_before>package quiclatest

import (
	"encoding/binary"

	"github.com/ami-GS/gQUIC/latest/qtype"
)

type Packet interface {
	GetWire() ([]byte, error)
	GetHeader() PacketHeader
	SetHeader(ph PacketHeader)
	GetFrames() []Frame
	SetFrames(fs []Frame)
}

func ParsePacket(data []byte) (packet Packet, idx int, err error) {
	version := binary.BigEndian.Uint32(data[1:5])
	if version == 0 {
		// VersionNegotiationPacket use all UDP datagram (doesn't have frame)
		return ParseVersionNegotiationPacket(data)
	}
	header, idx, err := PacketHeaderParserMap[PacketHeaderType(data[0]&0x80)](data) // ParseHeader
	if err != nil {
		return nil, 0, err
	}

	fs, idxTmp, err := ParseFrames(data[idx:])
	if err != nil {
		return nil, 0, err
	}

	packet, err = newPacket(data[0], header, fs)
	if err != nil {
		return nil, 0, err
	}

	return packet, idx + idxTmp, err
}

type BasePacket struct {
	Header     PacketHeader
	Frames     []Frame
	PaddingNum int
}

func (bp *BasePacket) GetHeader() PacketHeader {
	return bp.Header
}
func (bp *BasePacket) SetHeader(h PacketHeader) {
	bp.Header = h
}

func (bp *BasePacket) SetFrames(fs []Frame) {
	bp.Frames = fs
}

func (bp *BasePacket) GetFrames() []Frame {
	return bp.Frames
}

func (bp *BasePacket) GetWire() (wire []byte, err error) {
	hWire := bp.Header.GetWire()
	if err != nil {
		return nil, err
	}
	fsWire, err := GetFrameWires(bp.Frames)
	if err != nil {
		return nil, err
	}
	if bp.PaddingNum != 0 {
		fsWire = append(fsWire, make([]byte, bp.PaddingNum)...)
	}
	// TODO: protect for short header?
	return append(hWire, fsWire...), nil
}

func newPacket(firstByte byte, ph PacketHeader, fs []Frame) (Packet, error) {
	// TODO: needs frame type validation for each packet type
	if firstByte&byte(LongHeaderType) == byte(LongHeaderType) {
		payloadLen := 0
		for _, frame := range fs {
			payloadLen += frame.GetWireSize()
		}

		switch LongHeaderPacketType(firstByte & 0x7f) {
		case InitialPacketType:
			// TODO: check whether it is stream frame or not
			return &InitialPacket{
				BasePacket: &BasePacket{
					Header:     ph,
					Frames:     fs,
					PaddingNum: InitialPacketMinimumPayloadSize - payloadLen,
				},
			}, nil
		case RetryPacketType:
			return &RetryPacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case HandshakePacketType:
			return &HandshakePacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case ZeroRTTProtectedPacketType:
			return &ProtectedPacket{
				BasePacket: &BasePacket{
					Header: ph,
					Frames: fs,
				},
				RTT: 0,
			}, nil
		default:
			// error
			return nil, nil
		}
	} else if firstByte&byte(ShortHeaderType) == byte(ShortHeaderType) {
		return &ProtectedPacket{
			BasePacket: &BasePacket{
				Header: ph,
				Frames: fs,
			},
			RTT: 1,
		}, nil
	}
	// error
	return nil, nil
}

// long header with type of 0x7F
type InitialPacket struct {
	*BasePacket
}

func NewInitialPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, payloadLen uint64) *InitialPacket {
	return &InitialPacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(InitialPacketType, version, destConnID, srcConnID, packetNumber, payloadLen),
		},
	}

}

// long header with type of 0x7E
type RetryPacket struct {
	*BasePacket
}

func NewRetryPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, payloadLen uint64) *RetryPacket {
	return &RetryPacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(RetryPacketType, version, destConnID, srcConnID, packetNumber, payloadLen),
		},
	}
}

// long header with type of 0x7D
type HandshakePacket struct {
	*BasePacket
}

func NewHandshakePacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, payloadLen uint64) *HandshakePacket {
	return &HandshakePacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(HandshakePacketType, version, destConnID, srcConnID, packetNumber, payloadLen),
		},
	}
}

// long header with 0-RTT (type:0x7C)
// short header with 1-RTT
type ProtectedPacket struct {
	*BasePacket
	RTT byte
}

func NewProtectedPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, payloadLen uint64, rtt byte) *ProtectedPacket {
	var header PacketHeader
	if rtt == 0 {
		header = NewShortHeader(destConnID, packetNumber)
	} else if rtt == 1 {
		header = NewLongHeader(ZeroRTTProtectedPacketType, version, destConnID, srcConnID, packetNumber, payloadLen)
	} else {
		// error
	}

	return &ProtectedPacket{
		BasePacket: &BasePacket{
			Header: header,
		},
		RTT: rtt,
	}
}

/*
    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
   +-+-+-+-+-+-+-+-+
   |1|  Unused (7) |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                          Version (32)                         |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |DCIL(4)|SCIL(4)|
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |               Destination Connection ID (0/32..144)         ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                 Source Connection ID (0/32..144)            ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                    Supported Version 1 (32)                 ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version 2 (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
                                  ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version N (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
*/

// Version negotiation doesn't use long header, but have similar form
type VersionNegotiationPacket struct {
	*BasePacketHeader
	Version           qtype.Version
	DCIL              byte
	SCIL              byte
	SupportedVersions []qtype.Version
}

func NewVersionNegotiationPacket(destConnID, srcConnID qtype.ConnectionID, supportedVersions []qtype.Version) *VersionNegotiationPacket {
	dcil := 0
	if destConnID != nil {
		dcil = len(destConnID) - 3
	}
	scil := 0
	if srcConnID != nil {
		scil = len(srcConnID) - 3
	}

	return &VersionNegotiationPacket{
		BasePacketHeader: &BasePacketHeader{
			DestConnID:   destConnID,
			SrcConnID:    srcConnID,
			PacketNumber: 0,
		},
		Version:           0,
		DCIL:              byte(dcil),
		SCIL:              byte(scil),
		SupportedVersions: supportedVersions,
	}
}

func ParseVersionNegotiationPacket(data []byte) (Packet, int, error) {
	idx := 0
	packet := NewVersionNegotiationPacket(nil, nil, nil)
	if data[0]&0x80 != 0x80 {
		//TODO: error
	}
	idx++
	packet.Version = qtype.Version(binary.BigEndian.Uint32(data[idx:]))
	if packet.Version != 0 {
		// must be zero
		// TODO: error
	}
	idx += 4
	packet.DCIL = data[idx] >> 4
	packet.SCIL = data[idx] & 0x0f
	idx++
	if packet.DCIL != 0 {
		dcil := int(packet.DCIL + 3)
		qtype.ReadConnectionID(data[idx:], dcil)
		idx += dcil
	}
	if packet.SCIL != 0 {
		scil := int(packet.SCIL + 3)
		qtype.ReadConnectionID(data[idx:], scil)
		idx += scil
	}
	numVersions := (len(data) - idx) / 4
	packet.SupportedVersions = make([]qtype.Version, numVersions)
	for i := 0; i < numVersions; i++ {
		packet.SupportedVersions[i] = qtype.Version(binary.BigEndian.Uint32(data[idx:]))
		idx += 4
	}
	return packet, idx, nil
}

func (p VersionNegotiationPacket) GetWire() (wire []byte, err error) {
	wire = make([]byte, int(6+p.DCIL+p.SCIL)+len(p.SupportedVersions)*4)
	wire[0] = 0x80
	binary.BigEndian.PutUint32(wire[1:], uint32(p.Version))
	wire[5] = (p.DCIL << 4) | p.SCIL
	idx := 5
	if p.DCIL != 0 {
		for i := 0; i < int(p.DCIL+3); i++ {
			wire[idx+i] = p.DestConnID[i]
		}
		idx += int(p.DCIL + 3)
	}
	if p.SCIL != 0 {
		for i := 0; i < int(p.SCIL+3); i++ {
			wire[idx+i] = p.SrcConnID[i]
		}
		idx += int(p.SCIL + 3)
	}

	for i, version := range p.SupportedVersions {
		binary.BigEndian.PutUint32(wire[idx+i*4:], uint32(version))
	}
	// TODO: VersionNegotiationPacket fills MTU
	return
}

func (p VersionNegotiationPacket) GetHeader() PacketHeader {
	return nil
}
func (p VersionNegotiationPacket) SetHeader(h PacketHeader) {
	// no op?
}
func (p VersionNegotiationPacket) SetFrames(fs []Frame) {
	// no op?
}
func (p VersionNegotiationPacket) GetFrames() []Frame {
	return nil
}
<commit_msg>change packet creation interface to more realistic<commit_after>package quiclatest

import (
	"encoding/binary"

	"github.com/ami-GS/gQUIC/latest/qtype"
)

type Packet interface {
	GetWire() ([]byte, error)
	GetHeader() PacketHeader
	SetHeader(ph PacketHeader)
	GetFrames() []Frame
	SetFrames(fs []Frame)
}

func ParsePacket(data []byte) (packet Packet, idx int, err error) {
	version := binary.BigEndian.Uint32(data[1:5])
	if version == 0 {
		// VersionNegotiationPacket use all UDP datagram (doesn't have frame)
		return ParseVersionNegotiationPacket(data)
	}
	header, idx, err := PacketHeaderParserMap[PacketHeaderType(data[0]&0x80)](data) // ParseHeader
	if err != nil {
		return nil, 0, err
	}

	fs, idxTmp, err := ParseFrames(data[idx:])
	if err != nil {
		return nil, 0, err
	}

	packet, err = newPacket(data[0], header, fs)
	if err != nil {
		return nil, 0, err
	}

	return packet, idx + idxTmp, err
}

type BasePacket struct {
	Header     PacketHeader
	Frames     []Frame
	PaddingNum int
}

func (bp *BasePacket) GetHeader() PacketHeader {
	return bp.Header
}
func (bp *BasePacket) SetHeader(h PacketHeader) {
	bp.Header = h
}

func (bp *BasePacket) SetFrames(fs []Frame) {
	bp.Frames = fs
}

func (bp *BasePacket) GetFrames() []Frame {
	return bp.Frames
}

func (bp *BasePacket) GetWire() (wire []byte, err error) {
	hWire := bp.Header.GetWire()
	if err != nil {
		return nil, err
	}
	fsWire, err := GetFrameWires(bp.Frames)
	if err != nil {
		return nil, err
	}
	if bp.PaddingNum != 0 {
		fsWire = append(fsWire, make([]byte, bp.PaddingNum)...)
	}
	// TODO: protect for short header?
	return append(hWire, fsWire...), nil
}

func newPacket(firstByte byte, ph PacketHeader, fs []Frame) (Packet, error) {
	// TODO: needs frame type validation for each packet type
	if firstByte&byte(LongHeaderType) == byte(LongHeaderType) {
		payloadLen := 0
		for _, frame := range fs {
			payloadLen += frame.GetWireSize()
		}

		switch LongHeaderPacketType(firstByte & 0x7f) {
		case InitialPacketType:
			// TODO: check whether it is stream frame or not
			return &InitialPacket{
				BasePacket: &BasePacket{
					Header:     ph,
					Frames:     fs,
					PaddingNum: InitialPacketMinimumPayloadSize - payloadLen,
				},
			}, nil
		case RetryPacketType:
			return &RetryPacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case HandshakePacketType:
			return &HandshakePacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case ZeroRTTProtectedPacketType:
			return &ProtectedPacket{
				BasePacket: &BasePacket{
					Header: ph,
					Frames: fs,
				},
				RTT: 0,
			}, nil
		default:
			// error
			return nil, nil
		}
	} else if firstByte&byte(ShortHeaderType) == byte(ShortHeaderType) {
		return &ProtectedPacket{
			BasePacket: &BasePacket{
				Header: ph,
				Frames: fs,
			},
			RTT: 1,
		}, nil
	}
	// error
	return nil, nil
}

// long header with type of 0x7F
type InitialPacket struct {
	*BasePacket
}

const InitialPacketMinimumPayloadSize = 1200

func NewInitialPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, sFrame *StreamFrame) *InitialPacket {
	sFrameLen := sFrame.GetWireSize()
	var frames []Frame
	var lh *LongHeader
	paddingNum := 0
	if InitialPacketMinimumPayloadSize <= sFrameLen {
		// TODO: need to check when sFrameLen is over MTUIPv4
		lh = NewLongHeader(InitialPacketType, version, destConnID, srcConnID, packetNumber, uint64(sFrameLen))
		frames = []Frame{sFrame}
	} else {
		lh = NewLongHeader(InitialPacketType, version, destConnID, srcConnID, packetNumber, InitialPacketMinimumPayloadSize)
		frames = []Frame{sFrame}
		paddingNum = InitialPacketMinimumPayloadSize - sFrameLen
	}
	p := &InitialPacket{
		BasePacket: &BasePacket{
			Header:     lh,
			Frames:     frames,
			PaddingNum: paddingNum,
		},
	}
	return p
}

// long header with type of 0x7E
type RetryPacket struct {
	*BasePacket
}

func NewRetryPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, frames []Frame) *RetryPacket {
	if len(frames) < 2 {
		return nil
	}
	if _, ok := frames[0].(*StreamFrame); !ok {
		return nil
	}
	if _, ok := frames[1].(*AckFrame); !ok {
		return nil
	}
	for i := 2; i < len(frames); i++ {
		if _, ok := frames[i].(*PaddingFrame); !ok {
			return nil
		}
	}

	payloadLen := 0
	for i := 0; i < len(frames); i++ {
		payloadLen += frames[i].GetWireSize()
	}

	return &RetryPacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(RetryPacketType, version, destConnID, srcConnID, packetNumber, uint64(payloadLen)),
			Frames: frames,
		},
	}
}

// long header with type of 0x7D
type HandshakePacket struct {
	*BasePacket
}

func NewHandshakePacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, frames []Frame) *HandshakePacket {
	minimumReq := false
	payloadLen := 0
	for i := 0; i < len(frames); i++ {
		switch frames[i].(type) {
		case *StreamFrame:
			minimumReq = true
			payloadLen += frames[i].GetWireSize()
		case *AckFrame, *PathChallengeFrame /*or*/, *PathResponseFrame, *ConnectionCloseFrame:
			payloadLen += frames[i].GetWireSize()
		default:
			return nil
		}
	}
	if !minimumReq {
		return nil
	}

	return &HandshakePacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(HandshakePacketType, version, destConnID, srcConnID, packetNumber, uint64(payloadLen)),
			Frames: frames,
		},
	}
}

// long header with 0-RTT (type:0x7C)
// short header with 1-RTT
type ProtectedPacket struct {
	*BasePacket
	RTT byte
}

func NewProtectedPacket(version qtype.Version, key bool, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, rtt byte, frames []Frame) *ProtectedPacket {
	var header PacketHeader
	if rtt == 0 {
		payloadLen := 0
		for _, frame := range frames {
			payloadLen += frame.GetWireSize()
		}
		header = NewLongHeader(ZeroRTTProtectedPacketType, version, destConnID, srcConnID, packetNumber, uint64(payloadLen))
	} else if rtt == 1 {
		header = NewShortHeader(key, destConnID, packetNumber)
	} else {
		// error
	}

	return &ProtectedPacket{
		BasePacket: &BasePacket{
			Header: header,
			Frames: frames,
		},
		RTT: rtt,
	}

}

/*
    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
   +-+-+-+-+-+-+-+-+
   |1|  Unused (7) |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                          Version (32)                         |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |DCIL(4)|SCIL(4)|
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |               Destination Connection ID (0/32..144)         ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                 Source Connection ID (0/32..144)            ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                    Supported Version 1 (32)                 ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version 2 (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
                                  ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version N (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
*/

// Version negotiation doesn't use long header, but have similar form
type VersionNegotiationPacket struct {
	*BasePacketHeader
	Version           qtype.Version
	DCIL              byte
	SCIL              byte
	SupportedVersions []qtype.Version
}

func NewVersionNegotiationPacket(destConnID, srcConnID qtype.ConnectionID, supportedVersions []qtype.Version) *VersionNegotiationPacket {
	dcil := 0
	if destConnID != nil {
		dcil = len(destConnID) - 3
	}
	scil := 0
	if srcConnID != nil {
		scil = len(srcConnID) - 3
	}

	return &VersionNegotiationPacket{
		BasePacketHeader: &BasePacketHeader{
			DestConnID:   destConnID,
			SrcConnID:    srcConnID,
			PacketNumber: 0,
		},
		Version:           0,
		DCIL:              byte(dcil),
		SCIL:              byte(scil),
		SupportedVersions: supportedVersions,
	}
}

func ParseVersionNegotiationPacket(data []byte) (Packet, int, error) {
	idx := 0
	packet := NewVersionNegotiationPacket(nil, nil, nil)
	if data[0]&0x80 != 0x80 {
		//TODO: error
	}
	idx++
	packet.Version = qtype.Version(binary.BigEndian.Uint32(data[idx:]))
	if packet.Version != 0 {
		// must be zero
		// TODO: error
	}
	idx += 4
	packet.DCIL = data[idx] >> 4
	packet.SCIL = data[idx] & 0x0f
	idx++
	if packet.DCIL != 0 {
		dcil := int(packet.DCIL + 3)
		qtype.ReadConnectionID(data[idx:], dcil)
		idx += dcil
	}
	if packet.SCIL != 0 {
		scil := int(packet.SCIL + 3)
		qtype.ReadConnectionID(data[idx:], scil)
		idx += scil
	}
	numVersions := (len(data) - idx) / 4
	packet.SupportedVersions = make([]qtype.Version, numVersions)
	for i := 0; i < numVersions; i++ {
		packet.SupportedVersions[i] = qtype.Version(binary.BigEndian.Uint32(data[idx:]))
		idx += 4
	}
	return packet, idx, nil
}

func (p VersionNegotiationPacket) GetWire() (wire []byte, err error) {
	wire = make([]byte, int(6+p.DCIL+p.SCIL)+len(p.SupportedVersions)*4)
	wire[0] = 0x80
	binary.BigEndian.PutUint32(wire[1:], uint32(p.Version))
	wire[5] = (p.DCIL << 4) | p.SCIL
	idx := 5
	if p.DCIL != 0 {
		for i := 0; i < int(p.DCIL+3); i++ {
			wire[idx+i] = p.DestConnID[i]
		}
		idx += int(p.DCIL + 3)
	}
	if p.SCIL != 0 {
		for i := 0; i < int(p.SCIL+3); i++ {
			wire[idx+i] = p.SrcConnID[i]
		}
		idx += int(p.SCIL + 3)
	}

	for i, version := range p.SupportedVersions {
		binary.BigEndian.PutUint32(wire[idx+i*4:], uint32(version))
	}
	// TODO: VersionNegotiationPacket fills MTU
	return
}

func (p VersionNegotiationPacket) GetHeader() PacketHeader {
	return nil
}
func (p VersionNegotiationPacket) SetHeader(h PacketHeader) {
	// no op?
}
func (p VersionNegotiationPacket) SetFrames(fs []Frame) {
	// no op?
}
func (p VersionNegotiationPacket) GetFrames() []Frame {
	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strings"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/redshift"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func init() {
	resource.AddTestSweepers("aws_redshift_snapshot_schedule", &resource.Sweeper{
		Name: "aws_redshift_snapshot_schedule",
		F:    testSweepRedshiftSnapshotSchedules,
	})
}

func testSweepRedshiftSnapshotSchedules(region string) error {
	client, err := sharedClientForRegion(region)
	if err != nil {
		return fmt.Errorf("error getting client: %s", err)
	}
	conn := client.(*AWSClient).redshiftconn

	req := &redshift.DescribeSnapshotSchedulesInput{}

	resp, err := conn.DescribeSnapshotSchedules(req)
	if err != nil {
		if testSweepSkipSweepError(err) {
			log.Printf("[WARN] Skipping Redshift Regional Snapshot Schedules sweep for %s: %s", region, err)
			return nil
		}
		return fmt.Errorf("Error describing Redshift Regional Snapshot Schedules: %s", err)
	}

	if len(resp.SnapshotSchedules) == 0 {
		log.Print("[DEBUG] No AWS Redshift Regional Snapshot Schedules to sweep")
		return nil
	}

	for _, snapshotSchedules := range resp.SnapshotSchedules {
		identifier := aws.StringValue(snapshotSchedules.ScheduleIdentifier)

		hasPrefix := false
		prefixes := []string{"tf-snapshot-schedule-"}

		for _, prefix := range prefixes {
			if strings.HasPrefix(identifier, prefix) {
				hasPrefix = true
				break
			}
		}

		if !hasPrefix {
			log.Printf("[INFO] Skipping Delete Redshift Snapshot Schedule: %s", identifier)
			continue
		}

		_, err := conn.DeleteSnapshotSchedule(&redshift.DeleteSnapshotScheduleInput{
			ScheduleIdentifier: snapshotSchedules.ScheduleIdentifier,
		})
		if isAWSErr(err, redshift.ErrCodeSnapshotScheduleNotFoundFault, "") {
			return nil
		}
		if err != nil {
			return fmt.Errorf("Error deleting Redshift Snapshot Schedule %s: %s", identifier, err)
		}
	}

	return nil
}

func TestAccAWSRedshiftSnapshotSchedule_basic(t *testing.T) {
	var v redshift.SnapshotSchedule

	rName := acctest.RandString(8)
	resourceName := "aws_redshift_snapshot_schedule.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfig(rName, "rate(12 hours)"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "definitions.#", "1"),
				),
			},
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfig(rName, "cron(30 12 *)"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "definitions.#", "1"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"force_destroy",
				},
			},
		},
	})
}

func TestAccAWSRedshiftSnapshotSchedule_withMultipleDefinition(t *testing.T) {
	var v redshift.SnapshotSchedule

	rName := acctest.RandString(8)
	resourceName := "aws_redshift_snapshot_schedule.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithMultipleDefinition(rName, "cron(30 12 *)", "cron(15 6 *)"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "definitions.#", "2"),
				),
			},
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithMultipleDefinition(rName, "cron(30 8 *)", "cron(15 10 *)"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "definitions.#", "2"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"force_destroy",
				},
			},
		},
	})

}

func TestAccAWSRedshiftSnapshotSchedule_withIdentifierPrefix(t *testing.T) {
	var v redshift.SnapshotSchedule

	resourceName := "aws_redshift_snapshot_schedule.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithIdentifierPrefix,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"identifier_prefix",
					"force_destroy",
				},
			},
		},
	})
}

func TestAccAWSRedshiftSnapshotSchedule_withDescription(t *testing.T) {
	var v redshift.SnapshotSchedule

	rName := acctest.RandString(8)
	resourceName := "aws_redshift_snapshot_schedule.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithDescription(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(
						resourceName, "description", "Test Schedule"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"force_destroy",
				},
			},
		},
	})
}

func TestAccAWSRedshiftSnapshotSchedule_withTags(t *testing.T) {
	var v redshift.SnapshotSchedule

	rName := acctest.RandString(8)
	resourceName := "aws_redshift_snapshot_schedule.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithTags(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "tags.%", "2"),
					resource.TestCheckResourceAttr(resourceName, "tags.foo", "bar"),
					resource.TestCheckResourceAttr(resourceName, "tags.fizz", "buzz"),
				),
			},
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithTagsUpdate(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "tags.%", "2"),
					resource.TestCheckResourceAttr(resourceName, "tags.foo", "bar2"),
					resource.TestCheckResourceAttr(resourceName, "tags.good", "bad"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"force_destroy",
				},
			},
		},
	})
}

func TestAccAWSRedshiftSnapshotSchedule_withForceDestroy(t *testing.T) {
	var snapshotSchedule redshift.SnapshotSchedule
	var cluster redshift.Cluster
	rInt := acctest.RandInt()
	rName := acctest.RandString(8)
	resourceName := "aws_redshift_snapshot_schedule.default"
	clusterResourceName := "aws_redshift_cluster.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithForceDestroy(rInt, rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &snapshotSchedule),
					testAccCheckAWSRedshiftClusterExists(clusterResourceName, &cluster),
					testAccCheckAWSRedshiftSnapshotScheduleCreateSnapshotScheduleAssociation(&cluster, &snapshotSchedule),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"force_destroy",
				},
			},
		},
	})
}

func testAccCheckAWSRedshiftSnapshotScheduleDestroy(s *terraform.State) error {
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_redshift_snapshot_schedule" {
			continue
		}

		conn := testAccProvider.Meta().(*AWSClient).redshiftconn
		resp, err := conn.DescribeSnapshotSchedules(&redshift.DescribeSnapshotSchedulesInput{
			ScheduleIdentifier: aws.String(rs.Primary.ID),
		})

		if err == nil {
			if len(resp.SnapshotSchedules) != 0 {
				for _, s := range resp.SnapshotSchedules {
					if *s.ScheduleIdentifier == rs.Primary.ID {
						return fmt.Errorf("Redshift Cluster Snapshot Schedule %s still exists", rs.Primary.ID)
					}
				}
			}
		}

		return err
	}

	return nil
}

func testAccCheckAWSRedshiftSnapshotScheduleExists(n string, v *redshift.SnapshotSchedule) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No Redshift Cluster Snapshot Schedule ID is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).redshiftconn
		resp, err := conn.DescribeSnapshotSchedules(&redshift.DescribeSnapshotSchedulesInput{
			ScheduleIdentifier: aws.String(rs.Primary.ID),
		})

		if err != nil {
			return err
		}

		for _, s := range resp.SnapshotSchedules {
			if *s.ScheduleIdentifier == rs.Primary.ID {
				*v = *s
				return nil
			}
		}

		return fmt.Errorf("Redshift Snapshot Schedule (%s) not found", rs.Primary.ID)
	}
}

func testAccCheckAWSRedshiftSnapshotScheduleCreateSnapshotScheduleAssociation(cluster *redshift.Cluster, snapshotSchedule *redshift.SnapshotSchedule) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		conn := testAccProvider.Meta().(*AWSClient).redshiftconn

		if _, err := conn.ModifyClusterSnapshotSchedule(&redshift.ModifyClusterSnapshotScheduleInput{
			ClusterIdentifier:    cluster.ClusterIdentifier,
			ScheduleIdentifier:   snapshotSchedule.ScheduleIdentifier,
			DisassociateSchedule: aws.Bool(false),
		}); err != nil {
			return fmt.Errorf("Error associate Redshift Cluster and Snapshot Schedule: %s", err)
		}

		if err := waitForRedshiftSnapshotScheduleAssociationActive(conn, 75*time.Minute, aws.StringValue(cluster.ClusterIdentifier), aws.StringValue(snapshotSchedule.ScheduleIdentifier)); err != nil {
			return err
		}

		return nil
	}
}

const testAccAWSRedshiftSnapshotScheduleConfigWithIdentifierPrefix = `
resource "aws_redshift_snapshot_schedule" "default" {
	identifier_prefix = "tf-snapshot-schedule-"
	definitions = [
		"rate(12 hours)",
	]
}
`

func testAccAWSRedshiftSnapshotScheduleConfig(rName, definition string) string {
	return fmt.Sprintf(`
resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[1]s"
	definitions = [
		"%[2]s",
	]
}
	`, rName, definition)
}

func testAccAWSRedshiftSnapshotScheduleConfigWithMultipleDefinition(rName, definition1, definition2 string) string {
	return fmt.Sprintf(`
resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[1]s"
	definitions = [
		"%[2]s",
		"%[3]s",
	]
}
	`, rName, definition1, definition2)
}

func testAccAWSRedshiftSnapshotScheduleConfigWithDescription(rName string) string {
	return fmt.Sprintf(`
resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[1]s"
	description = "Test Schedule"
	definitions = [
		"rate(12 hours)",
	]
}
	`, rName)
}

func testAccAWSRedshiftSnapshotScheduleConfigWithTags(rName string) string {
	return fmt.Sprintf(`
resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[1]s"
	definitions = [
		"rate(12 hours)",
	]

	tags {
		foo = "bar"
		fizz = "buzz"
	}
}
	`, rName)
}

func testAccAWSRedshiftSnapshotScheduleConfigWithTagsUpdate(rName string) string {
	return fmt.Sprintf(`
resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[1]s"
	definitions = [
		"rate(12 hours)",
	]

	tags {
		foo = "bar2"
		good = "bad"
	}
}
	`, rName)
}

func testAccAWSRedshiftSnapshotScheduleConfigWithForceDestroy(rInt int, rName string) string {
	return fmt.Sprintf(`
%s

resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[2]s"
	description = "Test Schedule"
	definitions = [
		"rate(12 hours)",
	]
	force_destroy = true
}
`, testAccAWSRedshiftClusterConfig_basic(rInt), rName)
}
<commit_msg>fix format for type Map<commit_after>package aws

import (
	"fmt"
	"log"
	"strings"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/redshift"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func init() {
	resource.AddTestSweepers("aws_redshift_snapshot_schedule", &resource.Sweeper{
		Name: "aws_redshift_snapshot_schedule",
		F:    testSweepRedshiftSnapshotSchedules,
	})
}

func testSweepRedshiftSnapshotSchedules(region string) error {
	client, err := sharedClientForRegion(region)
	if err != nil {
		return fmt.Errorf("error getting client: %s", err)
	}
	conn := client.(*AWSClient).redshiftconn

	req := &redshift.DescribeSnapshotSchedulesInput{}

	resp, err := conn.DescribeSnapshotSchedules(req)
	if err != nil {
		if testSweepSkipSweepError(err) {
			log.Printf("[WARN] Skipping Redshift Regional Snapshot Schedules sweep for %s: %s", region, err)
			return nil
		}
		return fmt.Errorf("Error describing Redshift Regional Snapshot Schedules: %s", err)
	}

	if len(resp.SnapshotSchedules) == 0 {
		log.Print("[DEBUG] No AWS Redshift Regional Snapshot Schedules to sweep")
		return nil
	}

	for _, snapshotSchedules := range resp.SnapshotSchedules {
		identifier := aws.StringValue(snapshotSchedules.ScheduleIdentifier)

		hasPrefix := false
		prefixes := []string{"tf-snapshot-schedule-"}

		for _, prefix := range prefixes {
			if strings.HasPrefix(identifier, prefix) {
				hasPrefix = true
				break
			}
		}

		if !hasPrefix {
			log.Printf("[INFO] Skipping Delete Redshift Snapshot Schedule: %s", identifier)
			continue
		}

		_, err := conn.DeleteSnapshotSchedule(&redshift.DeleteSnapshotScheduleInput{
			ScheduleIdentifier: snapshotSchedules.ScheduleIdentifier,
		})
		if isAWSErr(err, redshift.ErrCodeSnapshotScheduleNotFoundFault, "") {
			return nil
		}
		if err != nil {
			return fmt.Errorf("Error deleting Redshift Snapshot Schedule %s: %s", identifier, err)
		}
	}

	return nil
}

func TestAccAWSRedshiftSnapshotSchedule_basic(t *testing.T) {
	var v redshift.SnapshotSchedule

	rName := acctest.RandString(8)
	resourceName := "aws_redshift_snapshot_schedule.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfig(rName, "rate(12 hours)"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "definitions.#", "1"),
				),
			},
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfig(rName, "cron(30 12 *)"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "definitions.#", "1"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"force_destroy",
				},
			},
		},
	})
}

func TestAccAWSRedshiftSnapshotSchedule_withMultipleDefinition(t *testing.T) {
	var v redshift.SnapshotSchedule

	rName := acctest.RandString(8)
	resourceName := "aws_redshift_snapshot_schedule.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithMultipleDefinition(rName, "cron(30 12 *)", "cron(15 6 *)"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "definitions.#", "2"),
				),
			},
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithMultipleDefinition(rName, "cron(30 8 *)", "cron(15 10 *)"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "definitions.#", "2"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"force_destroy",
				},
			},
		},
	})

}

func TestAccAWSRedshiftSnapshotSchedule_withIdentifierPrefix(t *testing.T) {
	var v redshift.SnapshotSchedule

	resourceName := "aws_redshift_snapshot_schedule.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithIdentifierPrefix,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"identifier_prefix",
					"force_destroy",
				},
			},
		},
	})
}

func TestAccAWSRedshiftSnapshotSchedule_withDescription(t *testing.T) {
	var v redshift.SnapshotSchedule

	rName := acctest.RandString(8)
	resourceName := "aws_redshift_snapshot_schedule.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithDescription(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(
						resourceName, "description", "Test Schedule"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"force_destroy",
				},
			},
		},
	})
}

func TestAccAWSRedshiftSnapshotSchedule_withTags(t *testing.T) {
	var v redshift.SnapshotSchedule

	rName := acctest.RandString(8)
	resourceName := "aws_redshift_snapshot_schedule.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithTags(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "tags.%", "2"),
					resource.TestCheckResourceAttr(resourceName, "tags.foo", "bar"),
					resource.TestCheckResourceAttr(resourceName, "tags.fizz", "buzz"),
				),
			},
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithTagsUpdate(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &v),
					resource.TestCheckResourceAttr(resourceName, "tags.%", "2"),
					resource.TestCheckResourceAttr(resourceName, "tags.foo", "bar2"),
					resource.TestCheckResourceAttr(resourceName, "tags.good", "bad"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"force_destroy",
				},
			},
		},
	})
}

func TestAccAWSRedshiftSnapshotSchedule_withForceDestroy(t *testing.T) {
	var snapshotSchedule redshift.SnapshotSchedule
	var cluster redshift.Cluster
	rInt := acctest.RandInt()
	rName := acctest.RandString(8)
	resourceName := "aws_redshift_snapshot_schedule.default"
	clusterResourceName := "aws_redshift_cluster.default"

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSRedshiftSnapshotScheduleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSRedshiftSnapshotScheduleConfigWithForceDestroy(rInt, rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSRedshiftSnapshotScheduleExists(resourceName, &snapshotSchedule),
					testAccCheckAWSRedshiftClusterExists(clusterResourceName, &cluster),
					testAccCheckAWSRedshiftSnapshotScheduleCreateSnapshotScheduleAssociation(&cluster, &snapshotSchedule),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
				ImportStateVerifyIgnore: []string{
					"force_destroy",
				},
			},
		},
	})
}

func testAccCheckAWSRedshiftSnapshotScheduleDestroy(s *terraform.State) error {
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_redshift_snapshot_schedule" {
			continue
		}

		conn := testAccProvider.Meta().(*AWSClient).redshiftconn
		resp, err := conn.DescribeSnapshotSchedules(&redshift.DescribeSnapshotSchedulesInput{
			ScheduleIdentifier: aws.String(rs.Primary.ID),
		})

		if err == nil {
			if len(resp.SnapshotSchedules) != 0 {
				for _, s := range resp.SnapshotSchedules {
					if *s.ScheduleIdentifier == rs.Primary.ID {
						return fmt.Errorf("Redshift Cluster Snapshot Schedule %s still exists", rs.Primary.ID)
					}
				}
			}
		}

		return err
	}

	return nil
}

func testAccCheckAWSRedshiftSnapshotScheduleExists(n string, v *redshift.SnapshotSchedule) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No Redshift Cluster Snapshot Schedule ID is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).redshiftconn
		resp, err := conn.DescribeSnapshotSchedules(&redshift.DescribeSnapshotSchedulesInput{
			ScheduleIdentifier: aws.String(rs.Primary.ID),
		})

		if err != nil {
			return err
		}

		for _, s := range resp.SnapshotSchedules {
			if *s.ScheduleIdentifier == rs.Primary.ID {
				*v = *s
				return nil
			}
		}

		return fmt.Errorf("Redshift Snapshot Schedule (%s) not found", rs.Primary.ID)
	}
}

func testAccCheckAWSRedshiftSnapshotScheduleCreateSnapshotScheduleAssociation(cluster *redshift.Cluster, snapshotSchedule *redshift.SnapshotSchedule) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		conn := testAccProvider.Meta().(*AWSClient).redshiftconn

		if _, err := conn.ModifyClusterSnapshotSchedule(&redshift.ModifyClusterSnapshotScheduleInput{
			ClusterIdentifier:    cluster.ClusterIdentifier,
			ScheduleIdentifier:   snapshotSchedule.ScheduleIdentifier,
			DisassociateSchedule: aws.Bool(false),
		}); err != nil {
			return fmt.Errorf("Error associate Redshift Cluster and Snapshot Schedule: %s", err)
		}

		if err := waitForRedshiftSnapshotScheduleAssociationActive(conn, 75*time.Minute, aws.StringValue(cluster.ClusterIdentifier), aws.StringValue(snapshotSchedule.ScheduleIdentifier)); err != nil {
			return err
		}

		return nil
	}
}

const testAccAWSRedshiftSnapshotScheduleConfigWithIdentifierPrefix = `
resource "aws_redshift_snapshot_schedule" "default" {
	identifier_prefix = "tf-snapshot-schedule-"
	definitions = [
		"rate(12 hours)",
	]
}
`

func testAccAWSRedshiftSnapshotScheduleConfig(rName, definition string) string {
	return fmt.Sprintf(`
resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[1]s"
	definitions = [
		"%[2]s",
	]
}
	`, rName, definition)
}

func testAccAWSRedshiftSnapshotScheduleConfigWithMultipleDefinition(rName, definition1, definition2 string) string {
	return fmt.Sprintf(`
resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[1]s"
	definitions = [
		"%[2]s",
		"%[3]s",
	]
}
	`, rName, definition1, definition2)
}

func testAccAWSRedshiftSnapshotScheduleConfigWithDescription(rName string) string {
	return fmt.Sprintf(`
resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[1]s"
	description = "Test Schedule"
	definitions = [
		"rate(12 hours)",
	]
}
	`, rName)
}

func testAccAWSRedshiftSnapshotScheduleConfigWithTags(rName string) string {
	return fmt.Sprintf(`
resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[1]s"
	definitions = [
		"rate(12 hours)",
	]

	tags = {
		foo = "bar"
		fizz = "buzz"
	}
}
	`, rName)
}

func testAccAWSRedshiftSnapshotScheduleConfigWithTagsUpdate(rName string) string {
	return fmt.Sprintf(`
resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[1]s"
	definitions = [
		"rate(12 hours)",
	]

	tags = {
		foo = "bar2"
		good = "bad"
	}
}
	`, rName)
}

func testAccAWSRedshiftSnapshotScheduleConfigWithForceDestroy(rInt int, rName string) string {
	return fmt.Sprintf(`
%s

resource "aws_redshift_snapshot_schedule" "default" {
	identifier = "tf-snapshot-schedule-%[2]s"
	description = "Test Schedule"
	definitions = [
		"rate(12 hours)",
	]
	force_destroy = true
}
`, testAccAWSRedshiftClusterConfig_basic(rInt), rName)
}
<|endoftext|>


<commit_before>package storage

import (
	"database/sql"
	"encoding/base64"
	"errors"
	"fmt"
	log "github.com/golang/glog"
	"github.com/jrupac/goliath/models"
	_ "github.com/lib/pq"
	"strconv"
	"strings"
)

const (
	DIALECT               = "postgres"
	FOLDER_TABLE          = "Folder"
	FOLDER_CHILDREN_TABLE = "FolderChildren"
	FEED_TABLE            = "Feed"
	ARTICLE_TABLE         = "Article"
)

type Database struct {
	db *sql.DB
}

func Open(dbPath string) (*Database, error) {
	db := new(Database)

	d, err := sql.Open(DIALECT, dbPath)
	if err != nil {
		return nil, err
	}

	db.db = d
	return db, nil
}

func (d *Database) Close() error {
	return d.db.Close()
}

func (d *Database) InsertArticle(a models.Article) error {
	var articleId int64
	var count int
	err := d.db.QueryRow(
		`SELECT COUNT(*) FROM `+ARTICLE_TABLE+` WHERE hash = $1`, a.Hash()).Scan(&count)
	if err != nil {
		return err
	}
	if count > 0 {
		log.Infof("Duplicate article entry, skipping.")
		return nil
	}

	err = d.db.QueryRow(
		`INSERT INTO `+ARTICLE_TABLE+`
		(feed, folder, hash, title, summary, content, link, date, read)
		VALUES($1, $2, $3, $4, $5, $6, $7, $8, $9) RETURNING id`,
		a.FeedId, a.FolderId, a.Hash(), a.Title, a.Summary, a.Content, a.Link, a.Date, a.Read).Scan(&articleId)
	if err != nil {
		return err
	}
	a.Id = articleId
	return nil
}

func (d *Database) InsertFavicon(feedId int64, mime string, img []byte) error {
	var count int
	err := d.db.QueryRow(
		`SELECT COUNT(*) FROM `+FEED_TABLE+` WHERE id = $1`, feedId).Scan(&count)
	if err != nil {
		return err
	}
	if count == 0 {
		return errors.New(fmt.Sprintf("Original feed not in table: %d", feedId))
	}

	// Convert to a base64 encoded string before inserting
	h := base64.StdEncoding.EncodeToString(img)

	_, err = d.db.Query(
		`UPDATE `+FEED_TABLE+` SET favicon = $1, mime = $2 WHERE id = $3`,
		h, mime, feedId)
	return err
}

func (d *Database) MarkArticle(id int64, status string) error {
	// Perhaps add support for saving articles in the future.
	if status == "saved" || status == "unsaved" {
		return nil
	}

	_, err := d.db.Query(`UPDATE `+ARTICLE_TABLE+` SET read WHERE id = $1`, id)
	return err
}

func (d *Database) MarkFeed(id int64, status string) error {
	// Perhaps add support for saving articles in the future.
	if status == "saved" || status == "unsaved" {
		return nil
	}

	_, err := d.db.Query(`UPDATE `+ARTICLE_TABLE+` SET read WHERE feed = $1`, id)
	return err
}

func (d *Database) MarkFolder(id int64, status string) error {
	// Perhaps add support for saving articles in the future.
	if status == "saved" || status == "unsaved" {
		return nil
	}

	_, err := d.db.Query(`UPDATE `+ARTICLE_TABLE+` SET read WHERE folder = $1`, id)
	return err
}

func (d *Database) GetAllFolders() ([]models.Folder, error) {
	folders := []models.Folder{}
	rows, err := d.db.Query(`SELECT id, name FROM ` + FOLDER_TABLE)
	if err != nil {
		return folders, err
	}
	defer rows.Close()

	for rows.Next() {
		f := models.Folder{}
		if err = rows.Scan(&f.Id, &f.Name); err != nil {
			return folders, err
		}
		folders = append(folders, f)
	}

	return folders, err
}

func (d *Database) GetAllFeeds() ([]models.Feed, error) {
	feeds := []models.Feed{}
	rows, err := d.db.Query(`SELECT id, folder, title, description, url, text FROM ` + FEED_TABLE)
	if err != nil {
		return feeds, err
	}
	defer rows.Close()

	for rows.Next() {
		f := models.Feed{}
		if err = rows.Scan(&f.Id, &f.FolderId, &f.Title, &f.Description, &f.Url, &f.Text); err != nil {
			return feeds, err
		}
		feeds = append(feeds, f)
	}

	return feeds, err
}

func (d *Database) GetFeedsPerFolder() (map[int64]string, error) {
	agg := map[int64][]string{}
	resp := map[int64]string{}

	// CockroachDB doesn't have a concat-with-separator aggregation function
	rows, err := d.db.Query(`SELECT folder, id FROM ` + FEED_TABLE)
	if err != nil {
		return resp, err
	}
	defer rows.Close()

	var folderId, feed_id int64
	for rows.Next() {
		if err = rows.Scan(&folderId, &feed_id); err != nil {
			return resp, err
		}
		agg[folderId] = append(agg[folderId], strconv.FormatInt(feed_id, 10))
	}

	for k, v := range agg {
		resp[k] = strings.Join(v, ",")
	}
	return resp, err
}

func (d *Database) GetAllFavicons() (map[int64]string, error) {
	favicons := map[int64]string{}
	rows, err := d.db.Query(
		`SELECT id, mime, favicon FROM ` + FEED_TABLE + ` WHERE favicon IS NOT NULL`)
	if err != nil {
		return favicons, err
	}
	defer rows.Close()

	var id int64
	var mime string
	var favicon string
	for rows.Next() {
		if err = rows.Scan(&id, &mime, &favicon); err != nil {
			return favicons, err
		}
		favicons[id] = fmt.Sprintf("%s;base64,%s", mime, favicon)
	}
	return favicons, err
}

func (d *Database) GetAllUnreadArticles(limit int) ([]models.Article, error) {
	articles := []models.Article{}
	var limitStr string
	if limit == -1 {
		limitStr = "ALL"
	} else {
		limitStr = strconv.Itoa(limit)
	}
	rows, err := d.db.Query(
		`SELECT id, feed, folder, title, summary, content, link, date FROM `+ARTICLE_TABLE+`
		WHERE NOT read LIMIT $1`, limitStr)
	if err != nil {
		return articles, err
	}
	defer rows.Close()

	var dateStr string
	for rows.Next() {
		a := models.Article{}
		if err = rows.Scan(
			&a.Id, &a.FeedId, &a.FolderId, &a.Title, &a.Summary, &a.Content, &a.Link, &dateStr); err != nil {
			return articles, err
		}
		articles = append(articles, a)
	}
	return articles, err
}

func (d *Database) ImportOpml(opml *models.Opml) error {
	root := opml.Folders
	var rootId int64
	err := d.db.QueryRow(
		`INSERT INTO `+FOLDER_TABLE+`(name) VALUES($1) RETURNING id`, root.Name).Scan(&rootId)
	if err != nil {
		return err
	}
	root.Id = rootId
	return d.importChildren(root)
}

func (d *Database) importChildren(parent models.Folder) error {
	var err error
	var feedId int64
	var count int
	for _, f := range parent.Feed {
		err := d.db.QueryRow(
			`SELECT COUNT(*) FROM `+FEED_TABLE+` WHERE hash = $1`, f.Hash()).Scan(&count)
		if err != nil {
			return err
		}
		if count > 0 {
			log.Infof("Duplicate feed entry, skipping.")
			continue
		}

		err = d.db.QueryRow(
			`INSERT INTO `+FEED_TABLE+`(folder, hash, title, description, url, text)
			VALUES($1, $2, $3, $4, $5, $6) RETURNING id`,
			parent.Id, f.Hash(), f.Title, f.Description, f.Url, f.Text).Scan(&feedId)
		if err != nil {
			return err
		}
		f.Id = feedId
		f.FolderId = parent.Id
	}

	var childId int64
	for _, child := range parent.Folders {
		err = d.db.QueryRow(
			`INSERT INTO `+FOLDER_TABLE+`(name) VALUES($1) RETURNING id`, child.Name).Scan(&childId)
		if err != nil {
			return err
		}
		child.Id = childId
		_, err = d.db.Query(
			`INSERT INTO `+FOLDER_CHILDREN_TABLE+`(parent, child) VALUES($1, $2)`, parent.Id, child.Id)
		if err != nil {
			return err
		}
		if err = d.importChildren(child); err != nil {
			return err
		}
	}
	return err
}
<commit_msg>Fix bug where the pq driver doesn't support "LIMIT ALL" syntax.<commit_after>package storage

import (
	"database/sql"
	"encoding/base64"
	"errors"
	"fmt"
	log "github.com/golang/glog"
	"github.com/jrupac/goliath/models"
	_ "github.com/lib/pq"
	"strconv"
	"strings"
)

const (
	DIALECT               = "postgres"
	FOLDER_TABLE          = "Folder"
	FOLDER_CHILDREN_TABLE = "FolderChildren"
	FEED_TABLE            = "Feed"
	ARTICLE_TABLE         = "Article"
)

type Database struct {
	db *sql.DB
}

func Open(dbPath string) (*Database, error) {
	db := new(Database)

	d, err := sql.Open(DIALECT, dbPath)
	if err != nil {
		return nil, err
	}

	db.db = d
	return db, nil
}

func (d *Database) Close() error {
	return d.db.Close()
}

func (d *Database) InsertArticle(a models.Article) error {
	var articleId int64
	var count int
	err := d.db.QueryRow(
		`SELECT COUNT(*) FROM `+ARTICLE_TABLE+` WHERE hash = $1`, a.Hash()).Scan(&count)
	if err != nil {
		return err
	}
	if count > 0 {
		log.Infof("Duplicate article entry, skipping.")
		return nil
	}

	err = d.db.QueryRow(
		`INSERT INTO `+ARTICLE_TABLE+`
		(feed, folder, hash, title, summary, content, link, date, read)
		VALUES($1, $2, $3, $4, $5, $6, $7, $8, $9) RETURNING id`,
		a.FeedId, a.FolderId, a.Hash(), a.Title, a.Summary, a.Content, a.Link, a.Date, a.Read).Scan(&articleId)
	if err != nil {
		return err
	}
	a.Id = articleId
	return nil
}

func (d *Database) InsertFavicon(feedId int64, mime string, img []byte) error {
	var count int
	err := d.db.QueryRow(
		`SELECT COUNT(*) FROM `+FEED_TABLE+` WHERE id = $1`, feedId).Scan(&count)
	if err != nil {
		return err
	}
	if count == 0 {
		return errors.New(fmt.Sprintf("Original feed not in table: %d", feedId))
	}

	// Convert to a base64 encoded string before inserting
	h := base64.StdEncoding.EncodeToString(img)

	_, err = d.db.Query(
		`UPDATE `+FEED_TABLE+` SET favicon = $1, mime = $2 WHERE id = $3`,
		h, mime, feedId)
	return err
}

func (d *Database) MarkArticle(id int64, status string) error {
	// Perhaps add support for saving articles in the future.
	if status == "saved" || status == "unsaved" {
		return nil
	}

	_, err := d.db.Query(`UPDATE `+ARTICLE_TABLE+` SET read WHERE id = $1`, id)
	return err
}

func (d *Database) MarkFeed(id int64, status string) error {
	// Perhaps add support for saving articles in the future.
	if status == "saved" || status == "unsaved" {
		return nil
	}

	_, err := d.db.Query(`UPDATE `+ARTICLE_TABLE+` SET read WHERE feed = $1`, id)
	return err
}

func (d *Database) MarkFolder(id int64, status string) error {
	// Perhaps add support for saving articles in the future.
	if status == "saved" || status == "unsaved" {
		return nil
	}

	_, err := d.db.Query(`UPDATE `+ARTICLE_TABLE+` SET read WHERE folder = $1`, id)
	return err
}

func (d *Database) GetAllFolders() ([]models.Folder, error) {
	folders := []models.Folder{}
	rows, err := d.db.Query(`SELECT id, name FROM ` + FOLDER_TABLE)
	if err != nil {
		return folders, err
	}
	defer rows.Close()

	for rows.Next() {
		f := models.Folder{}
		if err = rows.Scan(&f.Id, &f.Name); err != nil {
			return folders, err
		}
		folders = append(folders, f)
	}

	return folders, err
}

func (d *Database) GetAllFeeds() ([]models.Feed, error) {
	feeds := []models.Feed{}
	rows, err := d.db.Query(`SELECT id, folder, title, description, url, text FROM ` + FEED_TABLE)
	if err != nil {
		return feeds, err
	}
	defer rows.Close()

	for rows.Next() {
		f := models.Feed{}
		if err = rows.Scan(&f.Id, &f.FolderId, &f.Title, &f.Description, &f.Url, &f.Text); err != nil {
			return feeds, err
		}
		feeds = append(feeds, f)
	}

	return feeds, err
}

func (d *Database) GetFeedsPerFolder() (map[int64]string, error) {
	agg := map[int64][]string{}
	resp := map[int64]string{}

	// CockroachDB doesn't have a concat-with-separator aggregation function
	rows, err := d.db.Query(`SELECT folder, id FROM ` + FEED_TABLE)
	if err != nil {
		return resp, err
	}
	defer rows.Close()

	var folderId, feed_id int64
	for rows.Next() {
		if err = rows.Scan(&folderId, &feed_id); err != nil {
			return resp, err
		}
		agg[folderId] = append(agg[folderId], strconv.FormatInt(feed_id, 10))
	}

	for k, v := range agg {
		resp[k] = strings.Join(v, ",")
	}
	return resp, err
}

func (d *Database) GetAllFavicons() (map[int64]string, error) {
	favicons := map[int64]string{}
	rows, err := d.db.Query(
		`SELECT id, mime, favicon FROM ` + FEED_TABLE + ` WHERE favicon IS NOT NULL`)
	if err != nil {
		return favicons, err
	}
	defer rows.Close()

	var id int64
	var mime string
	var favicon string
	for rows.Next() {
		if err = rows.Scan(&id, &mime, &favicon); err != nil {
			return favicons, err
		}
		favicons[id] = fmt.Sprintf("%s;base64,%s", mime, favicon)
	}
	return favicons, err
}

func (d *Database) GetAllUnreadArticles(limit int) ([]models.Article, error) {
	articles := []models.Article{}
	var rows *sql.Rows
	var err error

	if limit == -1 {
		rows, err = d.db.Query(
			`SELECT id, feed, folder, title, summary, content, link, date FROM ` + ARTICLE_TABLE + `
			 WHERE NOT read`)
	} else {
		rows, err = d.db.Query(
			`SELECT id, feed, folder, title, summary, content, link, date FROM `+ARTICLE_TABLE+`
			 WHERE NOT read LIMIT $1`, limit)
	}

	if err != nil {
		return articles, err
	}
	defer rows.Close()

	var dateStr string
	for rows.Next() {
		a := models.Article{}
		if err = rows.Scan(
			&a.Id, &a.FeedId, &a.FolderId, &a.Title, &a.Summary, &a.Content, &a.Link, &dateStr); err != nil {
			return articles, err
		}
		articles = append(articles, a)
	}
	return articles, err
}

func (d *Database) ImportOpml(opml *models.Opml) error {
	root := opml.Folders
	var rootId int64
	err := d.db.QueryRow(
		`INSERT INTO `+FOLDER_TABLE+`(name) VALUES($1) RETURNING id`, root.Name).Scan(&rootId)
	if err != nil {
		return err
	}
	root.Id = rootId
	return d.importChildren(root)
}

func (d *Database) importChildren(parent models.Folder) error {
	var err error
	var feedId int64
	var count int
	for _, f := range parent.Feed {
		err := d.db.QueryRow(
			`SELECT COUNT(*) FROM `+FEED_TABLE+` WHERE hash = $1`, f.Hash()).Scan(&count)
		if err != nil {
			return err
		}
		if count > 0 {
			log.Infof("Duplicate feed entry, skipping.")
			continue
		}

		err = d.db.QueryRow(
			`INSERT INTO `+FEED_TABLE+`(folder, hash, title, description, url, text)
			VALUES($1, $2, $3, $4, $5, $6) RETURNING id`,
			parent.Id, f.Hash(), f.Title, f.Description, f.Url, f.Text).Scan(&feedId)
		if err != nil {
			return err
		}
		f.Id = feedId
		f.FolderId = parent.Id
	}

	var childId int64
	for _, child := range parent.Folders {
		err = d.db.QueryRow(
			`INSERT INTO `+FOLDER_TABLE+`(name) VALUES($1) RETURNING id`, child.Name).Scan(&childId)
		if err != nil {
			return err
		}
		child.Id = childId
		_, err = d.db.Query(
			`INSERT INTO `+FOLDER_CHILDREN_TABLE+`(parent, child) VALUES($1, $2)`, parent.Id, child.Id)
		if err != nil {
			return err
		}
		if err = d.importChildren(child); err != nil {
			return err
		}
	}
	return err
}
<|endoftext|>


<commit_before>package common

import (
	"errors"
	"github.com/Unknwon/com"
	"github.com/belogik/goes"
	"github.com/mitchellh/mapstructure"
	"net/url"
)

type User struct {
	Id           string         `json:"id"`
	Email        string         `json:"email"`
	FirstName    string         `json:"firstName"`
	LastName     string         `json:"lastName"`
	Password     string         `json:"password"`
	ApiKey       string         `json:"apiKey"`
	SourceGroups []*SourceGroup `json:"sourceGroups"`
}

// Returns index name to use in Elastic
func (u *User) GetIndexName() string {
	return u.ApiKey
}

func (u *User) GetLogTypes() []string {
	userLogTypes, _ := GetTypes(u.GetIndexName())
	return userLogTypes
}

// Source group represent group of log types.
// Each log type can be in various groups at the same time.
type SourceGroup struct {
	Id          string   `json:"id"`
	Name        string   `json:"name"`
	Description string   `json:"description"`
	Types       []string `json:"types"`
}

func (u *User) AddSourceGroup(group *SourceGroup) *User {
	if group.Id == "" {
		key := com.RandomCreateBytes(5)
		group.Id = string(key)
	} else {
		u.DeleteSourceGroup(group.Id)
	}
	u.SourceGroups = append(u.SourceGroups, group)
	return u
}

func (u *User) DeleteSourceGroup(id string) {
	for i, val := range u.SourceGroups {
		if val.Id == id {
			copy(u.SourceGroups[i:], u.SourceGroups[i+1:])
			u.SourceGroups[len(u.SourceGroups)-1] = nil // or the zero value of T
			u.SourceGroups = u.SourceGroups[:len(u.SourceGroups)-1]
			return
		}
	}
}

func (u *User) GetSourceGroup(id string) (*SourceGroup, error) {
	for _, val := range u.SourceGroups {
		if val.Id == id {
			return val, nil
		}
	}
	return nil, errors.New("Source group not found")
}

func FindUserByEmail(email string) *User {
	return FindUserBy("email", email)
}

func FindUserByApiKey(apiKey string) *User {
	return FindUserBy("apiKey", apiKey)
}

func (this *User) Save() {
	doc := goes.Document{
		Index:  "users",
		Type:   "user",
		Id:     this.Id,
		Fields: this,
	}
	extraArgs := make(url.Values, 0)
	GetConnection().Index(doc, extraArgs)
}

func FindUserBy(key string, value string) *User {
	var query = map[string]interface{}{
		"query": map[string]interface{}{
			"bool": map[string]interface{}{
				"must": map[string]interface{}{
					"term": map[string]interface{}{
						key: map[string]interface{}{
							"value": value,
						},
					},
				},
			},
		},
	}

	searchResults, err := GetConnection().Search(query, []string{"users"}, []string{"user"}, url.Values{})

	if err != nil || searchResults.Hits.Total == 0 {
		return nil
	}

	user := &User{}
	mapstructure.Decode(searchResults.Hits.Hits[0].Source, user)
	user.Id = searchResults.Hits.Hits[0].Id

	return user
}
<commit_msg>Preserve order for updated groups<commit_after>package common

import (
	"errors"
	"github.com/Unknwon/com"
	"github.com/belogik/goes"
	"github.com/mitchellh/mapstructure"
	"net/url"
)

type User struct {
	Id           string         `json:"id"`
	Email        string         `json:"email"`
	FirstName    string         `json:"firstName"`
	LastName     string         `json:"lastName"`
	Password     string         `json:"password"`
	ApiKey       string         `json:"apiKey"`
	SourceGroups []*SourceGroup `json:"sourceGroups"`
}

// Returns index name to use in Elastic
func (u *User) GetIndexName() string {
	return u.ApiKey
}

func (u *User) GetLogTypes() []string {
	userLogTypes, _ := GetTypes(u.GetIndexName())
	return userLogTypes
}

////////////////////////
// Source groups
////////////////////////

// Source group represent group of log types.
// Each log type can be in various groups at the same time.
type SourceGroup struct {
	Id          string   `json:"id"`
	Name        string   `json:"name"`
	Description string   `json:"description"`
	Types       []string `json:"types"`
}

func (u *User) AddSourceGroup(group *SourceGroup) *User {
	if group.Id == "" {
		key := com.RandomCreateBytes(5)
		group.Id = string(key)
		u.SourceGroups = append(u.SourceGroups, group)
	} else {
		u.UpdateSourceGroup(group)
	}
	return u
}

func (u *User) UpdateSourceGroup(group *SourceGroup) {
	for key, g := range u.SourceGroups {
		if group.Id == g.Id {
			u.SourceGroups[key] = group
		}
	}
}

func (u *User) DeleteSourceGroup(id string) {
	for i, val := range u.SourceGroups {
		if val.Id == id {
			copy(u.SourceGroups[i:], u.SourceGroups[i+1:])
			u.SourceGroups[len(u.SourceGroups)-1] = nil // or the zero value of T
			u.SourceGroups = u.SourceGroups[:len(u.SourceGroups)-1]
			return
		}
	}
}

func (u *User) GetSourceGroup(id string) (*SourceGroup, error) {
	for _, val := range u.SourceGroups {
		if val.Id == id {
			return val, nil
		}
	}
	return nil, errors.New("Source group not found")
}

////////////////////////
// Finders
////////////////////////

func FindUserByEmail(email string) *User {
	return FindUserBy("email", email)
}

func FindUserByApiKey(apiKey string) *User {
	return FindUserBy("apiKey", apiKey)
}

func (this *User) Save() {
	doc := goes.Document{
		Index:  "users",
		Type:   "user",
		Id:     this.Id,
		Fields: this,
	}
	extraArgs := make(url.Values, 0)
	GetConnection().Index(doc, extraArgs)
}

func FindUserBy(key string, value string) *User {
	var query = map[string]interface{}{
		"query": map[string]interface{}{
			"bool": map[string]interface{}{
				"must": map[string]interface{}{
					"term": map[string]interface{}{
						key: map[string]interface{}{
							"value": value,
						},
					},
				},
			},
		},
	}

	searchResults, err := GetConnection().Search(query, []string{"users"}, []string{"user"}, url.Values{})

	if err != nil || searchResults.Hits.Total == 0 {
		return nil
	}

	user := &User{}
	mapstructure.Decode(searchResults.Hits.Hits[0].Source, user)
	user.Id = searchResults.Hits.Hits[0].Id

	return user
}
<|endoftext|>


<commit_before>package gocql

import (
	"flag"
	"fmt"
	"log"
	"strings"
	"sync"
	"testing"
	"time"
)

var (
	flagCluster      = flag.String("cluster", "127.0.0.1", "a comma-separated list of host:port tuples")
	flagProto        = flag.Int("proto", 2, "protcol version")
	flagCQL          = flag.String("cql", "3.0.0", "CQL version")
	flagRF           = flag.Int("rf", 1, "replication factor for test keyspace")
	clusterSize      = flag.Int("clusterSize", 1, "the expected size of the cluster")
	flagRetry        = flag.Int("retries", 5, "number of times to retry queries")
	flagAutoWait     = flag.Duration("autowait", 1000*time.Millisecond, "time to wait for autodiscovery to fill the hosts poll")
	flagRunSslTest   = flag.Bool("runssl", false, "Set to true to run ssl test")
	flagRunAuthTest  = flag.Bool("runauth", false, "Set to true to run authentication test")
	flagCompressTest = flag.String("compressor", "", "compressor to use")
	flagTimeout      = flag.Duration("gocql.timeout", 5*time.Second, "sets the connection `timeout` for all operations")
	clusterHosts     []string
)

func init() {
	flag.Parse()
	clusterHosts = strings.Split(*flagCluster, ",")
	log.SetFlags(log.Lshortfile | log.LstdFlags)
}

func addSslOptions(cluster *ClusterConfig) *ClusterConfig {
	if *flagRunSslTest {
		cluster.SslOpts = &SslOptions{
			CertPath:               "testdata/pki/gocql.crt",
			KeyPath:                "testdata/pki/gocql.key",
			CaPath:                 "testdata/pki/ca.crt",
			EnableHostVerification: false,
		}
	}
	return cluster
}

var initOnce sync.Once

func createTable(s *Session, table string) error {
	// lets just be really sure
	if err := s.control.awaitSchemaAgreement(); err != nil {
		return err
	}

	if err := s.control.query(table).Close(); err != nil {
		return err
	}

	if err := s.control.awaitSchemaAgreement(); err != nil {
		return err
	}

	return nil
}

func createCluster() *ClusterConfig {
	cluster := NewCluster(clusterHosts...)
	cluster.ProtoVersion = *flagProto
	cluster.CQLVersion = *flagCQL
	cluster.Timeout = *flagTimeout
	cluster.Consistency = Quorum
	cluster.MaxWaitSchemaAgreement = 2 * time.Minute // travis might be slow
	if *flagRetry > 0 {
		cluster.RetryPolicy = &SimpleRetryPolicy{NumRetries: *flagRetry}
	}

	switch *flagCompressTest {
	case "snappy":
		cluster.Compressor = &SnappyCompressor{}
	case "":
	default:
		panic("invalid compressor: " + *flagCompressTest)
	}

	cluster = addSslOptions(cluster)
	return cluster
}

func createKeyspace(tb testing.TB, cluster *ClusterConfig, keyspace string) {
	c := *cluster
	c.Keyspace = "system"
	c.Timeout = 20 * time.Second
	session, err := c.CreateSession()
	if err != nil {
		tb.Fatal("createSession:", err)
	}
	defer session.Close()
	defer log.Println("closing keyspace session")

	err = createTable(session, `DROP KEYSPACE IF EXISTS `+keyspace)
	if err != nil {
		tb.Fatal(err)
	}

	err = createTable(session, fmt.Sprintf(`CREATE KEYSPACE %s
	WITH replication = {
		'class' : 'SimpleStrategy',
		'replication_factor' : %d
	}`, keyspace, *flagRF))

	if err != nil {
		tb.Fatal(err)
	}
}

func createSessionFromCluster(cluster *ClusterConfig, tb testing.TB) *Session {
	// Drop and re-create the keyspace once. Different tests should use their own
	// individual tables, but can assume that the table does not exist before.
	initOnce.Do(func() {
		createKeyspace(tb, cluster, "gocql_test")
	})

	cluster.Keyspace = "gocql_test"
	session, err := cluster.CreateSession()
	if err != nil {
		tb.Fatal("createSession:", err)
	}

	if err := session.control.awaitSchemaAgreement(); err != nil {
		tb.Fatal(err)
	}

	return session
}

func createSession(tb testing.TB) *Session {
	cluster := createCluster()
	return createSessionFromCluster(cluster, tb)
}
<commit_msg>dont use the control connection to create tables and dont use a retry policy<commit_after>package gocql

import (
	"flag"
	"fmt"
	"log"
	"strings"
	"sync"
	"testing"
	"time"
)

var (
	flagCluster      = flag.String("cluster", "127.0.0.1", "a comma-separated list of host:port tuples")
	flagProto        = flag.Int("proto", 2, "protcol version")
	flagCQL          = flag.String("cql", "3.0.0", "CQL version")
	flagRF           = flag.Int("rf", 1, "replication factor for test keyspace")
	clusterSize      = flag.Int("clusterSize", 1, "the expected size of the cluster")
	flagRetry        = flag.Int("retries", 5, "number of times to retry queries")
	flagAutoWait     = flag.Duration("autowait", 1000*time.Millisecond, "time to wait for autodiscovery to fill the hosts poll")
	flagRunSslTest   = flag.Bool("runssl", false, "Set to true to run ssl test")
	flagRunAuthTest  = flag.Bool("runauth", false, "Set to true to run authentication test")
	flagCompressTest = flag.String("compressor", "", "compressor to use")
	flagTimeout      = flag.Duration("gocql.timeout", 5*time.Second, "sets the connection `timeout` for all operations")
	clusterHosts     []string
)

func init() {
	flag.Parse()
	clusterHosts = strings.Split(*flagCluster, ",")
	log.SetFlags(log.Lshortfile | log.LstdFlags)
}

func addSslOptions(cluster *ClusterConfig) *ClusterConfig {
	if *flagRunSslTest {
		cluster.SslOpts = &SslOptions{
			CertPath:               "testdata/pki/gocql.crt",
			KeyPath:                "testdata/pki/gocql.key",
			CaPath:                 "testdata/pki/ca.crt",
			EnableHostVerification: false,
		}
	}
	return cluster
}

var initOnce sync.Once

func createTable(s *Session, table string) error {
	// lets just be really sure
	if err := s.control.awaitSchemaAgreement(); err != nil {
		return err
	}

	if err := s.Query(table).RetryPolicy(nil).Exec(); err != nil {
		return err
	}

	if err := s.control.awaitSchemaAgreement(); err != nil {
		return err
	}

	return nil
}

func createCluster() *ClusterConfig {
	cluster := NewCluster(clusterHosts...)
	cluster.ProtoVersion = *flagProto
	cluster.CQLVersion = *flagCQL
	cluster.Timeout = *flagTimeout
	cluster.Consistency = Quorum
	cluster.MaxWaitSchemaAgreement = 2 * time.Minute // travis might be slow
	if *flagRetry > 0 {
		cluster.RetryPolicy = &SimpleRetryPolicy{NumRetries: *flagRetry}
	}

	switch *flagCompressTest {
	case "snappy":
		cluster.Compressor = &SnappyCompressor{}
	case "":
	default:
		panic("invalid compressor: " + *flagCompressTest)
	}

	cluster = addSslOptions(cluster)
	return cluster
}

func createKeyspace(tb testing.TB, cluster *ClusterConfig, keyspace string) {
	c := *cluster
	c.Keyspace = "system"
	c.Timeout = 20 * time.Second
	session, err := c.CreateSession()
	if err != nil {
		tb.Fatal("createSession:", err)
	}
	defer session.Close()
	defer log.Println("closing keyspace session")

	err = createTable(session, `DROP KEYSPACE IF EXISTS `+keyspace)
	if err != nil {
		tb.Fatal(err)
	}

	err = createTable(session, fmt.Sprintf(`CREATE KEYSPACE %s
	WITH replication = {
		'class' : 'SimpleStrategy',
		'replication_factor' : %d
	}`, keyspace, *flagRF))

	if err != nil {
		tb.Fatal(err)
	}
}

func createSessionFromCluster(cluster *ClusterConfig, tb testing.TB) *Session {
	// Drop and re-create the keyspace once. Different tests should use their own
	// individual tables, but can assume that the table does not exist before.
	initOnce.Do(func() {
		createKeyspace(tb, cluster, "gocql_test")
	})

	cluster.Keyspace = "gocql_test"
	session, err := cluster.CreateSession()
	if err != nil {
		tb.Fatal("createSession:", err)
	}

	if err := session.control.awaitSchemaAgreement(); err != nil {
		tb.Fatal(err)
	}

	return session
}

func createSession(tb testing.TB) *Session {
	cluster := createCluster()
	return createSessionFromCluster(cluster, tb)
}
<|endoftext|>


<commit_before>package atlas

import (
	"bytes"
	"fmt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"io/ioutil"
	"net/http"
	"net/url"
	"testing"
)

const testURL = "http://localhost:10000"

func TestGetPageNum(t *testing.T) {
	nurl := "https://foo.example.com/"

	n := getPageNum(nurl)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
	nurl = "https://foo.example.com/?page=42"
	n = getPageNum(nurl)
	if n != "42" {
		t.Errorf("n=%s should be 42", n)
	}
	nurl = "https://foo.example.com/?country=fr&page=43"
	n = getPageNum(nurl)
	if n != "43" {
		t.Errorf("n=%s should be 43", n)
	}
	nurl = "https://foo.example.com/?country=fr&page=666&bar=1"
	n = getPageNum(nurl)
	if n != "666" {
		t.Errorf("n=%s should be 666", n)
	}

	nurl = ""
	n = getPageNum(nurl)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
}

func TestClienthandleAPIResponsese(t *testing.T) {
	var (
		r http.Response
		b bytes.Buffer
	)

	client, err := NewClient()
	err = client.handleAPIResponsese(nil)
	assert.Error(t, err, "should be in error")

	r = http.Response{StatusCode: 0}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be no error")

	r = http.Response{StatusCode: 200}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be no error")

	var jsonErr = `error:{status: 501, code: 500, detail: "test"}`

	fmt.Fprintf(&b, "%v", jsonErr)
	r.StatusCode = 300
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be in error")

	r.StatusCode = 500
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.Error(t, err, "should be in error")
}

func TestAddQueryParameters(t *testing.T) {
	p := AddQueryParameters("", map[string]string{})
	assert.Equal(t, "", p)
}

func TestAddQueryParameters_1(t *testing.T) {
	p := AddQueryParameters("", map[string]string{"": ""})
	assert.Equal(t, "?=", p)
}

func TestAddQueryParameters_2(t *testing.T) {
	p := AddQueryParameters("foo", map[string]string{"bar": "baz"})
	assert.Equal(t, "foo?bar=baz", p)
}

func TestClient_AddAPIKey(t *testing.T) {
	c, err := NewClient(Config{APIKey: "foo"})
	require.NoError(t, err)
	assert.NotNil(t, c)
	assert.NotEmpty(t, c)

	opts := map[string]string{}

	newk := c.addAPIKey(opts)
	assert.NotEmpty(t, c.config.APIKey)
	assert.Equal(t, 1, len(newk))
	assert.EqualValues(t, map[string]string{"key": "foo"}, newk)
}

func TestClient_PrepareRequest(t *testing.T) {
	c, err := NewClient(Config{endpoint: testURL})
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_PrepareRequest_2(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_PrepareRequest_3(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("FETCH", testURL +"/foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_MergeGlobalOptions(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{"foo": "bar"}
	c.opts = map[string]string{"baz": "xyz"}
	res := map[string]string{"foo": "bar", "baz": "xyz"}

	c.mergeGlobalOptions(opts)
	assert.EqualValues(t, res, opts)
}
<commit_msg>Simplify test and fix json to match real APIError.<commit_after>package atlas

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const testURL = "http://localhost:10000"

func TestGetPageNum(t *testing.T) {
	nurl := "https://foo.example.com/"

	n := getPageNum(nurl)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
	nurl = "https://foo.example.com/?page=42"
	n = getPageNum(nurl)
	if n != "42" {
		t.Errorf("n=%s should be 42", n)
	}
	nurl = "https://foo.example.com/?country=fr&page=43"
	n = getPageNum(nurl)
	if n != "43" {
		t.Errorf("n=%s should be 43", n)
	}
	nurl = "https://foo.example.com/?country=fr&page=666&bar=1"
	n = getPageNum(nurl)
	if n != "666" {
		t.Errorf("n=%s should be 666", n)
	}

	nurl = ""
	n = getPageNum(nurl)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
}

func TestClienthandleAPIResponsese(t *testing.T) {
	var (
		r http.Response
		b bytes.Buffer
	)

	client, err := NewClient()
	err = client.handleAPIResponsese(nil)
	assert.Error(t, err)

	r = http.Response{StatusCode: 0}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err)

	r = http.Response{StatusCode: 200}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err)

	var jsonErr = `{"error":{"status": 501, "errors":[{"detail": "test"}],"code": 500, "detail":"issue"}}`

	fmt.Fprintf(&b, "%v", jsonErr)
	r.StatusCode = 300
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err)

	r.StatusCode = 500
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.Error(t, err)
}

func TestAddQueryParameters(t *testing.T) {
	p := AddQueryParameters("", map[string]string{})
	assert.Equal(t, "", p)
}

func TestAddQueryParameters_1(t *testing.T) {
	p := AddQueryParameters("", map[string]string{"": ""})
	assert.Equal(t, "?=", p)
}

func TestAddQueryParameters_2(t *testing.T) {
	p := AddQueryParameters("foo", map[string]string{"bar": "baz"})
	assert.Equal(t, "foo?bar=baz", p)
}

func TestClient_AddAPIKey(t *testing.T) {
	c, err := NewClient(Config{APIKey: "foo"})
	require.NoError(t, err)
	assert.NotNil(t, c)
	assert.NotEmpty(t, c)

	opts := map[string]string{}

	newk := c.addAPIKey(opts)
	assert.NotEmpty(t, c.config.APIKey)
	assert.Equal(t, 1, len(newk))
	assert.EqualValues(t, map[string]string{"key": "foo"}, newk)
}

func TestClient_PrepareRequest(t *testing.T) {
	c, err := NewClient(Config{endpoint: testURL})
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_PrepareRequest_2(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_PrepareRequest_3(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("FETCH", testURL+"/foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_MergeGlobalOptions(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{"foo": "bar"}
	c.opts = map[string]string{"baz": "xyz"}
	res := map[string]string{"foo": "bar", "baz": "xyz"}

	c.mergeGlobalOptions(opts)
	assert.EqualValues(t, res, opts)
}
<|endoftext|>