<commit_before>//
// Copyright (c) 2014 Juniper Networks, Inc. All rights reserved.
//

package config

import (
	"fmt"
	"net"
	"regexp"
	"strconv"
	"strings"

	"github.com/Juniper/contrail-go-api"
	"github.com/Juniper/contrail-go-api/types"
)

type NetworkInfo struct {
	Uuid         string
	Name         string
	AdminState   bool
	NetworkId    int
	Transit      bool
	Mode         string
	Subnets      []string
	Policies     []string
	RouteTargets []string
}

func buildNetworkInfo(net *types.VirtualNetwork, detail bool) (
	*NetworkInfo, error) {
	var subnets []string
	var policies []string

	refList, err := net.GetNetworkIpamRefs()
	if err != nil {
		return nil, err
	}

	for _, ref := range refList {
		attr := ref.Attr.(types.VnSubnetsType)
		for _, ipamSubnet := range attr.IpamSubnets {
			subnets = append(subnets, fmt.Sprintf("%s/%d",
				ipamSubnet.Subnet.IpPrefix,
				ipamSubnet.Subnet.IpPrefixLen))
		}
	}

	if detail {
		refList, err = net.GetNetworkPolicyRefs()
		for _, ref := range refList {
			policies = append(policies, strings.Join(ref.To, ":"))
		}
	}

	info := &NetworkInfo{
		net.GetUuid(),
		net.GetName(),
		net.GetIdPerms().Enable,
		net.GetVirtualNetworkProperties().NetworkId,
		net.GetVirtualNetworkProperties().AllowTransit,
		net.GetVirtualNetworkProperties().ForwardingMode,
		subnets,
		policies,
		net.GetRouteTargetList().RouteTarget,
	}
	return info, err
}

func NetworkShow(client contrail.ApiClient, uuid string, detail bool) (
	*NetworkInfo, error) {
	obj, err := client.FindByUuid("virtual-network", uuid)
	if err != nil {
		return nil, err
	}
	return buildNetworkInfo(obj.(*types.VirtualNetwork), detail)
}

func NetworkList(client contrail.ApiClient, project_id string, detail bool) (
	[]*NetworkInfo, error) {
	fields := []string{"network_ipams"}
	if detail {
		fields = append(fields, "network_policys")
	}
	networks, err := client.ListDetailByParent(
		"virtual-network", project_id, fields)
	if err != nil {
		return nil, err
	}

	var networkList []*NetworkInfo
	for _, reference := range networks {
		info, err := buildNetworkInfo(
			reference.(*types.VirtualNetwork), detail)
		if err != nil {
			return nil, err
		}
		networkList = append(networkList, info)
	}

	return networkList, nil
}

func makeSubnet(prefix string) (*types.IpamSubnetType, error) {
	expr := regexp.MustCompile(`(([0-9]{1,3}\.){3}[0-9]{1,3})/([0-9]{1,2})`)
	match := expr.FindStringSubmatch(prefix)
	if match == nil {
		return nil, fmt.Errorf("Invalid subnet prefix %s", prefix)
	}
	address := match[1]
	if net.ParseIP(address) == nil {
		return nil, fmt.Errorf("%s is not a valid IP address", address)
	}
	prefixlen, _ := strconv.Atoi(match[3])
	if prefixlen < 0 || prefixlen > 32 {
		return nil, fmt.Errorf("Invalid subnet prefix length %d", prefixlen)
	}

	subnet := &types.IpamSubnetType{
		Subnet: &types.SubnetType{address, prefixlen}}
	return subnet, nil
}

func networkAddSubnet(
	client contrail.ApiClient,
	project *types.Project, network *types.VirtualNetwork,
	subnet *types.IpamSubnetType) error {

	refList, err := project.GetNetworkIpams()
	if err != nil {
		return err
	}

	var ipam *types.NetworkIpam
	if len(refList) > 0 {
		obj, err := client.FindByUuid("network-ipam", refList[0].Uuid)
		if err != nil {
			return err
		}
		ipam = obj.(*types.NetworkIpam)
	} else {
		obj, err := client.FindByName("network-ipam",
			"default-domain:default-project:default-network-ipam")
		if err != nil {
			return err
		}
		ipam = obj.(*types.NetworkIpam)
	}

	refs, err := network.GetNetworkIpamRefs()
	if err != nil {
		return err
	}

	var subnets types.VnSubnetsType

	for _, ref := range refs {
		if ref.Uuid == ipam.GetUuid() {
			subnets = ref.Attr.(types.VnSubnetsType)
			network.DeleteNetworkIpam(ref.Uuid)
			break
		}
	}
	subnets.AddIpamSubnets(subnet)
	network.AddNetworkIpam(ipam, subnets)
	return nil
}

func CreateNetworkWithSubnet(
	client contrail.ApiClient, project_id, name, prefix string) (
	string, error) {

	obj, err := client.FindByUuid("project", project_id)
	if err != nil {
		return "", err
	}

	project := obj.(*types.Project)

	net := new(types.VirtualNetwork)
	net.SetParent(project)
	net.SetName(name)

	subnet, err := makeSubnet(prefix)
	if err != nil {
		return "", err
	}
	err = networkAddSubnet(client, project, net, subnet)
	if err != nil {
		return "", err
	}

	err = client.Create(net)
	if err != nil {
		return "", err
	}
	return net.GetUuid(), nil
}

// AddSubnet
// returns true if the network was modified, false if the subnet already exists
// in the network.
func AddSubnet(
	client contrail.ApiClient, network *types.VirtualNetwork, prefix string) (
	bool, error) {

	ipamRefs, err := network.GetNetworkIpamRefs()
	if err != nil {
		return false, err
	}
	subnet, err := makeSubnet(prefix)
	if err != nil {
		return false, err
	}
	for _, ref := range ipamRefs {
		attr := ref.Attr.(types.VnSubnetsType)
		for _, entry := range attr.IpamSubnets {
			if *subnet.Subnet == *entry.Subnet {
				return false, nil
			}
		}
	}

	fqn := network.GetFQName()
	project, err := types.ProjectByName(client, strings.Join(fqn[:len(fqn)-1], ":"))
	if err != nil {
		return false, err
	}
	err = networkAddSubnet(client, project, network, subnet)
	if err != nil {
		return false, err
	}
	err = client.Update(network)
	if err != nil {
		return false, err
	}
	return true, nil
}

func subnetTypeStringRepr(subnet *types.SubnetType) string {
	return fmt.Sprintf("%s/%d", subnet.IpPrefix, subnet.IpPrefixLen)
}

func RemoveSubnet(client contrail.ApiClient, network *types.VirtualNetwork, prefix string) error {
	ipamRefs, err := network.GetNetworkIpamRefs()
	if err != nil {
		return err
	}

	removeOp := func(ix int, uuid string, attr types.VnSubnetsType) error {
		ipam, err := types.NetworkIpamByUuid(client, uuid)
		if err != nil {
			return err
		}
		attr.IpamSubnets = append(attr.IpamSubnets[0:ix], attr.IpamSubnets[ix+1:]...)
		network.DeleteNetworkIpam(uuid)
		if len(attr.IpamSubnets) > 0 {
			network.AddNetworkIpam(ipam, attr)
		}
		return client.Update(network)
	}

	for _, ref := range ipamRefs {
		attr := ref.Attr.(types.VnSubnetsType)
		for ix, entry := range attr.IpamSubnets {
			entryPrefix := subnetTypeStringRepr(entry.Subnet)
			if entryPrefix == prefix {
				return removeOp(ix, ref.Uuid, attr)
			}
		}
	}

	return fmt.Errorf("Prefix %s not associated with network %s", prefix, network.GetName())
}

func CreateNetwork(client contrail.ApiClient, project_id, name string) (
	string, error) {

	obj, err := client.FindByUuid("project", project_id)
	if err != nil {
		return "", err
	}

	project := obj.(*types.Project)

	net := new(types.VirtualNetwork)
	net.SetParent(project)
	net.SetName(name)

	err = client.Create(net)
	if err != nil {
		return "", err
	}
	return net.GetUuid(), nil
}
<commit_msg>Add createNetwork with Ipam<commit_after>//
// Copyright (c) 2014 Juniper Networks, Inc. All rights reserved.
//

package config

import (
	"fmt"
	"net"
	"regexp"
	"strconv"
	"strings"

	"github.com/Juniper/contrail-go-api"
	"github.com/Juniper/contrail-go-api/types"
)

type NetworkInfo struct {
	Uuid         string
	Name         string
	AdminState   bool
	NetworkId    int
	Transit      bool
	Mode         string
	Subnets      []string
	Policies     []string
	RouteTargets []string
}

func buildNetworkInfo(net *types.VirtualNetwork, detail bool) (
	*NetworkInfo, error) {
	var subnets []string
	var policies []string

	refList, err := net.GetNetworkIpamRefs()
	if err != nil {
		return nil, err
	}

	for _, ref := range refList {
		attr := ref.Attr.(types.VnSubnetsType)
		for _, ipamSubnet := range attr.IpamSubnets {
			subnets = append(subnets, fmt.Sprintf("%s/%d",
				ipamSubnet.Subnet.IpPrefix,
				ipamSubnet.Subnet.IpPrefixLen))
		}
	}

	if detail {
		refList, err = net.GetNetworkPolicyRefs()
		for _, ref := range refList {
			policies = append(policies, strings.Join(ref.To, ":"))
		}
	}

	info := &NetworkInfo{
		net.GetUuid(),
		net.GetName(),
		net.GetIdPerms().Enable,
		net.GetVirtualNetworkProperties().NetworkId,
		net.GetVirtualNetworkProperties().AllowTransit,
		net.GetVirtualNetworkProperties().ForwardingMode,
		subnets,
		policies,
		net.GetRouteTargetList().RouteTarget,
	}
	return info, err
}

func NetworkShow(client contrail.ApiClient, uuid string, detail bool) (
	*NetworkInfo, error) {
	obj, err := client.FindByUuid("virtual-network", uuid)
	if err != nil {
		return nil, err
	}
	return buildNetworkInfo(obj.(*types.VirtualNetwork), detail)
}

func NetworkList(client contrail.ApiClient, project_id string, detail bool) (
	[]*NetworkInfo, error) {
	fields := []string{"network_ipams"}
	if detail {
		fields = append(fields, "network_policys")
	}
	networks, err := client.ListDetailByParent(
		"virtual-network", project_id, fields)
	if err != nil {
		return nil, err
	}

	var networkList []*NetworkInfo
	for _, reference := range networks {
		info, err := buildNetworkInfo(
			reference.(*types.VirtualNetwork), detail)
		if err != nil {
			return nil, err
		}
		networkList = append(networkList, info)
	}

	return networkList, nil
}

func makeSubnet(prefix string) (*types.IpamSubnetType, error) {
	expr := regexp.MustCompile(`(([0-9]{1,3}\.){3}[0-9]{1,3})/([0-9]{1,2})`)
	match := expr.FindStringSubmatch(prefix)
	if match == nil {
		return nil, fmt.Errorf("Invalid subnet prefix %s", prefix)
	}
	address := match[1]
	if net.ParseIP(address) == nil {
		return nil, fmt.Errorf("%s is not a valid IP address", address)
	}
	prefixlen, _ := strconv.Atoi(match[3])
	if prefixlen < 0 || prefixlen > 32 {
		return nil, fmt.Errorf("Invalid subnet prefix length %d", prefixlen)
	}

	subnet := &types.IpamSubnetType{
		Subnet: &types.SubnetType{address, prefixlen}}
	return subnet, nil
}

func networkAddSubnet(
	client contrail.ApiClient,
	project *types.Project, network *types.VirtualNetwork,
	subnet *types.IpamSubnetType,
	ipam *types.NetworkIpam) error {

	/*refList, err := project.GetNetworkIpams()
	if err != nil {
		return err
	}

	var ipam *types.NetworkIpam*/
	if ipam != nil {
		obj, err := client.FindByUuid("network-ipam", ipam.GetUuid())
		if err != nil {
			return err
		}
		ipam = obj.(*types.NetworkIpam)
	} else {
		obj, err := client.FindByName("network-ipam",
			"default-domain:default-project:default-network-ipam")
		if err != nil {
			return err
		}
		ipam = obj.(*types.NetworkIpam)
	}

	refs, err := network.GetNetworkIpamRefs()
	if err != nil {
		return err
	}

	var subnets types.VnSubnetsType

	for _, ref := range refs {
		if ref.Uuid == ipam.GetUuid() {
			subnets = ref.Attr.(types.VnSubnetsType)
			network.DeleteNetworkIpam(ref.Uuid)
			break
		}
	}
	subnets.AddIpamSubnets(subnet)
	network.AddNetworkIpam(ipam, subnets)
	return nil
}

func CreateNetworkWithSubnet(
client contrail.ApiClient, project_id, name, prefix string) (
string, error) {

	obj, err := client.FindByUuid("project", project_id)
	if err != nil {
		return "", err
	}

	project := obj.(*types.Project)

	net := new(types.VirtualNetwork)
	net.SetParent(project)
	net.SetName(name)

	subnet, err := makeSubnet(prefix)
	if err != nil {
		return "", err
	}
	err = networkAddSubnet(client, project, net, subnet, nil)
	if err != nil {
		return "", err
	}

	err = client.Create(net)
	if err != nil {
		return "", err
	}
	return net.GetUuid(), nil
}

func CreateNetworkWithIpam(
	client contrail.ApiClient, project_id, name, prefix string, ipam *types.NetworkIpam) (
	string, error) {

	obj, err := client.FindByUuid("project", project_id)
	if err != nil {
		return "", err
	}

	project := obj.(*types.Project)

	net := new(types.VirtualNetwork)
	net.SetParent(project)
	net.SetName(name)

	subnet, err := makeSubnet(prefix)
	if err != nil {
		return "", err
	}
	err = networkAddSubnet(client, project, net, subnet, ipam)
	if err != nil {
		return "", err
	}

	err = client.Create(net)
	if err != nil {
		return "", err
	}
	return net.GetUuid(), nil
}

// AddSubnet
// returns true if the network was modified, false if the subnet already exists
// in the network.
func AddSubnet(
	client contrail.ApiClient, network *types.VirtualNetwork, prefix string) (
	bool, error) {

	ipamRefs, err := network.GetNetworkIpamRefs()
	if err != nil {
		return false, err
	}
	subnet, err := makeSubnet(prefix)
	if err != nil {
		return false, err
	}
	for _, ref := range ipamRefs {
		attr := ref.Attr.(types.VnSubnetsType)
		for _, entry := range attr.IpamSubnets {
			if *subnet.Subnet == *entry.Subnet {
				return false, nil
			}
		}
	}

	fqn := network.GetFQName()
	project, err := types.ProjectByName(client, strings.Join(fqn[:len(fqn)-1], ":"))
	if err != nil {
		return false, err
	}
	err = networkAddSubnet(client, project, network, subnet, nil)
	if err != nil {
		return false, err
	}
	err = client.Update(network)
	if err != nil {
		return false, err
	}
	return true, nil
}

func subnetTypeStringRepr(subnet *types.SubnetType) string {
	return fmt.Sprintf("%s/%d", subnet.IpPrefix, subnet.IpPrefixLen)
}

func RemoveSubnet(client contrail.ApiClient, network *types.VirtualNetwork, prefix string) error {
	ipamRefs, err := network.GetNetworkIpamRefs()
	if err != nil {
		return err
	}

	removeOp := func(ix int, uuid string, attr types.VnSubnetsType) error {
		ipam, err := types.NetworkIpamByUuid(client, uuid)
		if err != nil {
			return err
		}
		attr.IpamSubnets = append(attr.IpamSubnets[0:ix], attr.IpamSubnets[ix+1:]...)
		network.DeleteNetworkIpam(uuid)
		if len(attr.IpamSubnets) > 0 {
			network.AddNetworkIpam(ipam, attr)
		}
		return client.Update(network)
	}

	for _, ref := range ipamRefs {
		attr := ref.Attr.(types.VnSubnetsType)
		for ix, entry := range attr.IpamSubnets {
			entryPrefix := subnetTypeStringRepr(entry.Subnet)
			if entryPrefix == prefix {
				return removeOp(ix, ref.Uuid, attr)
			}
		}
	}

	return fmt.Errorf("Prefix %s not associated with network %s", prefix, network.GetName())
}

func CreateNetwork(client contrail.ApiClient, project_id, name string) (
	string, error) {

	obj, err := client.FindByUuid("project", project_id)
	if err != nil {
		return "", err
	}

	project := obj.(*types.Project)

	net := new(types.VirtualNetwork)
	net.SetParent(project)
	net.SetName(name)

	err = client.Create(net)
	if err != nil {
		return "", err
	}
	return net.GetUuid(), nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 The Jaeger 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 integration

import (
	"bytes"
	"os"
	"strconv"
	"testing"
	"time"

	"github.com/Shopify/sarama"
	"github.com/gogo/protobuf/jsonpb"
	"github.com/stretchr/testify/require"
	"github.com/uber/jaeger-lib/metrics"
	"go.uber.org/zap"

	"github.com/jaegertracing/jaeger/model"
	"github.com/jaegertracing/jaeger/pkg/config"
	"github.com/jaegertracing/jaeger/pkg/testutils"
	"github.com/jaegertracing/jaeger/plugin/storage/kafka"
	"github.com/jaegertracing/jaeger/storage/spanstore"
)

const defaultLocalKafkaBroker = "127.0.0.1:9092"

type KafkaIntegrationTestSuite struct {
	StorageIntegration
}

func (s *KafkaIntegrationTestSuite) initialize() error {
	logger, _ := testutils.NewLogger()
	s.logger = logger
	// A new topic is generated per execution to avoid data overlap
	topic := "jaeger-kafka-integration-test-" + strconv.FormatInt(time.Now().UnixNano(), 10)

	f := kafka.NewFactory()
	v, command := config.Viperize(f.AddFlags)
	command.ParseFlags([]string{
		"--kafka.topic",
		topic,
		"--kafka.brokers",
		defaultLocalKafkaBroker,
		"--kafka.encoding",
		"json",
	})
	f.InitFromViper(v)
	if err := f.Initialize(metrics.NullFactory, s.logger); err != nil {
		return err
	}
	spanWriter, err := f.CreateSpanWriter()
	if err != nil {
		return err
	}
	spanReader, err := createSpanReader(topic)
	if err != nil {
		return err
	}
	s.SpanWriter = spanWriter
	s.SpanReader = spanReader
	s.Refresh = func() error { return nil }
	s.CleanUp = func() error { return nil }
	return nil
}

type spanReader struct {
	logger   *zap.Logger
	topic    string
	consumer sarama.PartitionConsumer
}

func createSpanReader(topic string) (spanstore.Reader, error) {
	logger, _ := testutils.NewLogger()
	c, err := sarama.NewConsumer([]string{defaultLocalKafkaBroker}, nil)
	if err != nil {
		return nil, err
	}
	pc, err := c.ConsumePartition(topic, 0, sarama.OffsetOldest)
	if err != nil {
		return nil, err
	}
	return &spanReader{
		consumer: pc,
		topic:    topic,
		logger:   logger,
	}, nil
}

func (r *spanReader) GetTrace(traceID model.TraceID) (*model.Trace, error) {
	result := &model.Trace{}
	var err error
	doneCh := make(chan struct{})
	go func() {
		for {
			select {
			case msg := <-r.consumer.Messages():
				newSpan := model.Span{}
				if err = jsonpb.Unmarshal(bytes.NewReader(msg.Value), &newSpan); err != nil {
					r.logger.Error("protobuf unmarshaling error", zap.Error(err))
				}
				if newSpan.TraceID == traceID {
					result.Spans = append(result.Spans, &newSpan)
				}
			case <-doneCh:
				return
			}
		}
	}()
	time.Sleep(100 * time.Millisecond)
	doneCh <- struct{}{}
	if err != nil {
		return nil, err
	}
	return result, nil
}

func (r *spanReader) GetServices() ([]string, error) {
	return nil, nil
}

func (r *spanReader) GetOperations(service string) ([]string, error) {
	return nil, nil
}

func (r *spanReader) FindTraces(query *spanstore.TraceQueryParameters) ([]*model.Trace, error) {
	return nil, nil
}

func TestKafkaStorage(t *testing.T) {
	if os.Getenv("STORAGE") != "kafka" {
		t.Skip("Integration test against kafka skipped; set STORAGE env var to kafka to run this")
	}
	s := &KafkaIntegrationTestSuite{}
	require.NoError(t, s.initialize())
	t.Run("GetTrace", s.testGetTrace)
}
<commit_msg>Use Ingester in kafka integration test (#968)<commit_after>// Copyright (c) 2018 The Jaeger 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 integration

import (
	"os"
	"strconv"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"github.com/uber/jaeger-lib/metrics"

	"github.com/jaegertracing/jaeger/cmd/ingester/app"
	"github.com/jaegertracing/jaeger/cmd/ingester/app/builder"
	"github.com/jaegertracing/jaeger/model"
	"github.com/jaegertracing/jaeger/pkg/config"
	"github.com/jaegertracing/jaeger/pkg/testutils"
	"github.com/jaegertracing/jaeger/plugin/storage/kafka"
	"github.com/jaegertracing/jaeger/plugin/storage/memory"
	"github.com/jaegertracing/jaeger/storage/spanstore"
)

const defaultLocalKafkaBroker = "127.0.0.1:9092"

type KafkaIntegrationTestSuite struct {
	StorageIntegration
}

func (s *KafkaIntegrationTestSuite) initialize() error {
	logger, _ := testutils.NewLogger()
	s.logger = logger
	// A new topic is generated per execution to avoid data overlap
	topic := "jaeger-kafka-integration-test-" + strconv.FormatInt(time.Now().UnixNano(), 10)

	f := kafka.NewFactory()
	v, command := config.Viperize(f.AddFlags, app.AddFlags)
	command.ParseFlags([]string{
		"--kafka.topic",
		topic,
		"--kafka.brokers",
		defaultLocalKafkaBroker,
		"--kafka.encoding",
		"json",
		"--ingester.brokers",
		defaultLocalKafkaBroker,
		"--ingester.topic",
		topic,
		"--ingester.group-id",
		"kafka-integration-test",
		"--ingester.parallelism",
		"1000",
		"--ingester.encoding",
		"json",
	})
	f.InitFromViper(v)
	if err := f.Initialize(metrics.NullFactory, s.logger); err != nil {
		return err
	}
	spanWriter, err := f.CreateSpanWriter()
	if err != nil {
		return err
	}

	options := app.Options{}
	options.InitFromViper(v)
	traceStore := memory.NewStore()
	spanConsumer, err := builder.CreateConsumer(logger, metrics.NullFactory, traceStore, options)
	if err != nil {
		return err
	}
	spanConsumer.Start()

	s.SpanWriter = spanWriter
	s.SpanReader = &ingester{traceStore}
	s.Refresh = func() error { return nil }
	s.CleanUp = func() error { return nil }
	return nil
}

// The ingester consumes spans from kafka and writes them to an in-memory traceStore
type ingester struct {
	traceStore *memory.Store
}

func (r *ingester) GetTrace(traceID model.TraceID) (*model.Trace, error) {
	return r.traceStore.GetTrace(traceID)
}

func (r *ingester) GetServices() ([]string, error) {
	return nil, nil
}

func (r *ingester) GetOperations(service string) ([]string, error) {
	return nil, nil
}

func (r *ingester) FindTraces(query *spanstore.TraceQueryParameters) ([]*model.Trace, error) {
	return nil, nil
}

func TestKafkaStorage(t *testing.T) {
	if os.Getenv("STORAGE") != "kafka" {
		t.Skip("Integration test against kafka skipped; set STORAGE env var to kafka to run this")
	}
	s := &KafkaIntegrationTestSuite{}
	require.NoError(t, s.initialize())
	t.Run("GetTrace", s.testGetTrace)
}
<|endoftext|>


<commit_before>package parser

import (
	"fmt"
	"strings"
)

type PrimaryKey struct {
	Name       string
	FieldNames []string
	Fields     []*Field
	Obj        *MetaObject
}

func NewPrimaryKey(obj *MetaObject) *PrimaryKey {
	return &PrimaryKey{Obj: obj}
}

func (pk *PrimaryKey) IsSingleField() bool {
	if len(pk.Fields) == 1 {
		return true
	}
	return false
}

func (pk *PrimaryKey) GetFuncParam() string {
	return Fields(pk.Fields).GetFuncParam()
}

func (pk *PrimaryKey) FirstField() *Field {
	if len(pk.Fields) > 0 {
		return pk.Fields[0]
	}
	return nil
}

func (pk *PrimaryKey) IsAutocrement() bool {
	if len(pk.Fields) == 1 {
		return pk.Fields[0].Flags.Contains("autoinc")
	}
	return false
}

func (pk *PrimaryKey) IsRange() bool {
	fs := make([]*Field, 0, len(pk.Fields))
	for _, f := range pk.Fields {
		if f.IsNorange() {
			continue
		}
		fs = append(fs, f)
	}
	c := len(fs)
	if c > 0 {
		return pk.Fields[c-1].IsNumber()
	}
	return false
}

func (pk *PrimaryKey) build() error {
	pk.Name = fmt.Sprintf("%sOf%sPK", strings.Join(pk.FieldNames, ""), pk.Obj.Name)
	for _, name := range pk.FieldNames {
		f := pk.Obj.FieldByName(name)
		if f == nil {
			return fmt.Errorf("%s field not exist", name)
		}
		f.Flags.Add("primary")
		pk.Fields = append(pk.Fields, f)
	}
	if len(pk.Fields) == 0 {
		return fmt.Errorf("primary key  not declare.")
	}
	return nil
}

func (pk *PrimaryKey) SQLColumn(driver string) string {
	switch strings.ToLower(driver) {
	case "mysql":
		columns := make([]string, 0, len(pk.Fields))
		for _, f := range pk.Fields {
			columns = append(columns, f.SQLName(driver))
		}
		return fmt.Sprintf("PRIMARY KEY(%s)", strings.Join(columns, ","))
	}
	return ""
}

func (pk *PrimaryKey) GetConstructor() string {
	return Fields(pk.Fields).GetConstructor()
}

func (pk *PrimaryKey) GetObjectParam() string {
	return Fields(pk.Fields).GetObjectParam()
}
<commit_msg>parser: fix misorder of the pk keys<commit_after>package parser

import (
	"fmt"
	"strings"
)

type PrimaryKey struct {
	Name       string
	FieldNames []string
	Fields     []*Field
	Obj        *MetaObject
}

func NewPrimaryKey(obj *MetaObject) *PrimaryKey {
	return &PrimaryKey{Obj: obj}
}

func (pk *PrimaryKey) IsSingleField() bool {
	if len(pk.Fields) == 1 {
		return true
	}
	return false
}

func (pk *PrimaryKey) GetFuncParam() string {
	return Fields(pk.Fields).GetFuncParam()
}

func (pk *PrimaryKey) FirstField() *Field {
	if len(pk.Fields) > 0 {
		return pk.Fields[0]
	}
	return nil
}

func (pk *PrimaryKey) IsAutocrement() bool {
	if len(pk.Fields) == 1 {
		return pk.Fields[0].Flags.Contains("autoinc")
	}
	return false
}

func (pk *PrimaryKey) IsRange() bool {
	fs := make([]*Field, 0, len(pk.Fields))
	for _, f := range pk.Fields {
		if f.IsNorange() {
			continue
		}
		fs = append(fs, f)
	}
	c := len(fs)
	if c > 0 {
		return fs[c-1].IsNumber()
	}
	return false
}

func (pk *PrimaryKey) build() error {
	pk.Name = fmt.Sprintf("%sOf%sPK", strings.Join(pk.FieldNames, ""), pk.Obj.Name)
	for _, name := range pk.FieldNames {
		f := pk.Obj.FieldByName(name)
		if f == nil {
			return fmt.Errorf("%s field not exist", name)
		}
		f.Flags.Add("primary")
		pk.Fields = append(pk.Fields, f)
	}
	if len(pk.Fields) == 0 {
		return fmt.Errorf("primary key  not declare.")
	}
	return nil
}

func (pk *PrimaryKey) SQLColumn(driver string) string {
	switch strings.ToLower(driver) {
	case "mysql":
		columns := make([]string, 0, len(pk.Fields))
		for _, f := range pk.Fields {
			columns = append(columns, f.SQLName(driver))
		}
		return fmt.Sprintf("PRIMARY KEY(%s)", strings.Join(columns, ","))
	}
	return ""
}

func (pk *PrimaryKey) GetConstructor() string {
	return Fields(pk.Fields).GetConstructor()
}

func (pk *PrimaryKey) GetObjectParam() string {
	return Fields(pk.Fields).GetObjectParam()
}
<|endoftext|>


<commit_before>package compile

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/hashicorp/otto/app"
	"github.com/hashicorp/otto/foundation"
	"github.com/hashicorp/otto/helper/bindata"
	"github.com/hashicorp/otto/helper/oneline"
)

// AppOptions are the options for compiling an application.
//
// These options may be modified during customization processing, and
// in fact that is an intended use case and common pattern. To do this,
// use the AppCustomizationFunc method. See some of the builtin types for
// examples.
type AppOptions struct {
	// Ctx is the app context of this compilation.
	Ctx *app.Context

	// Result is the base CompileResult that will be used to return the result.
	// You can set this if you want to override some settings.
	Result *app.CompileResult

	// FoundationConfig is the configuration for the foundation that
	// will be returned as the compilation result.
	FoundationConfig foundation.Config

	// Bindata is the data that is used for templating. This must be set.
	// Template data should also be set on this. This will be modified with
	// default template data if those keys are not set.
	Bindata *bindata.Data

	// Customizations is a list of helpers to process customizations
	// in the Appfile. See the Customization docs for more information.
	Customizations []*Customization

	// Callbacks are called just prior to compilation completing.
	Callbacks []CompileCallback
}

// CompileCallback is a callback that can be registered to be run after
// compilation. To access any data within this callback, it should be created
// as a closure around the AppOptions.
type CompileCallback func() error

// App is an opinionated compilation function to help implement
// app.App.Compile.
//
// AppOptions may be modified by this function during this call.
func App(opts *AppOptions) (*app.CompileResult, error) {
	// Write the test data in case we're running tests right now
	testLock.RLock()
	defer testLock.RUnlock()
	testAppOpts = opts

	ctx := opts.Ctx

	// Setup the basic templating data. We put this into the "data" local
	// var just so that it is easier to reference.
	//
	// The exact default data put into the context is documented above.
	data := opts.Bindata
	if data.Context == nil {
		data.Context = make(map[string]interface{})
		opts.Bindata = data
	}

	data.Context["app_type"] = ctx.Appfile.Application.Type
	data.Context["name"] = ctx.Appfile.Application.Name
	data.Context["dev_fragments"] = ctx.DevDepFragments
	data.Context["dev_ip_address"] = ctx.DevIPAddress

	if data.Context["path"] == nil {
		data.Context["path"] = make(map[string]string)
	}
	pathMap := data.Context["path"].(map[string]string)
	pathMap["cache"] = ctx.CacheDir
	pathMap["compiled"] = ctx.Dir
	pathMap["working"] = filepath.Dir(ctx.Appfile.Path)
	foundationDirsContext := map[string][]string{
		"dev":     make([]string, len(ctx.FoundationDirs)),
		"dev_dep": make([]string, len(ctx.FoundationDirs)),
		"build":   make([]string, len(ctx.FoundationDirs)),
		"deploy":  make([]string, len(ctx.FoundationDirs)),
	}
	for i, dir := range ctx.FoundationDirs {
		foundationDirsContext["dev"][i] = filepath.Join(dir, "app-dev")
		foundationDirsContext["dev_dep"][i] = filepath.Join(dir, "app-dev-dep")
		foundationDirsContext["build"][i] = filepath.Join(dir, "app-build")
		foundationDirsContext["deploy"][i] = filepath.Join(dir, "app-deploy")
	}
	data.Context["foundation_dirs"] = foundationDirsContext

	// Setup the shared data
	if data.SharedExtends == nil {
		data.SharedExtends = make(map[string]*bindata.Data)
	}
	data.SharedExtends["compile"] = &bindata.Data{
		Asset:    Asset,
		AssetDir: AssetDir,
	}

	// Process the customizations!
	err := processCustomizations(&processOpts{
		Customizations: opts.Customizations,
		Appfile:        ctx.Appfile,
		Bindata:        data,
	})
	if err != nil {
		return nil, err
	}

	// Create the directory list that we'll copy from, and copy those
	// directly into the compilation directory.
	bindirs := []string{
		"data/common",
		fmt.Sprintf("data/%s-%s", ctx.Tuple.Infra, ctx.Tuple.InfraFlavor),
	}
	for _, dir := range bindirs {
		// Copy all the common files that exist
		if err := data.CopyDir(ctx.Dir, dir); err != nil {
			// Ignore any directories that don't exist
			if strings.Contains(err.Error(), "not found") {
				continue
			}

			return nil, err
		}
	}

	if err := appFoundations(opts); err != nil {
		return nil, err
	}

	// Callbacks
	for _, cb := range opts.Callbacks {
		if err := cb(); err != nil {
			return nil, err
		}
	}

	// If the DevDep fragment exists, then use it
	fragmentPath := filepath.Join(ctx.Dir, "dev-dep", "Vagrantfile.fragment")
	if _, err := os.Stat(fragmentPath); err != nil {
		fragmentPath = ""
	}

	// Set some defaults here
	if opts.FoundationConfig.ServiceName == "" {
		opts.FoundationConfig.ServiceName = opts.Ctx.Application.Name
	}

	result := opts.Result
	if result == nil {
		result = new(app.CompileResult)
	}
	result.FoundationConfig = opts.FoundationConfig
	result.DevDepFragmentPath = fragmentPath
	return result, nil
}

// appFoundations compiles the app-specific foundation files.
func appFoundations(opts *AppOptions) error {
	// Setup the bindata for rendering
	dataCopy := Data
	data := &dataCopy
	data.Context = make(map[string]interface{})
	for k, v := range opts.Bindata.Context {
		data.Context[k] = v
	}

	// Go through each foundation and setup the layers
	log.Printf("[INFO] compile: looking for foundation layers for dev")
	for i, dir := range opts.Ctx.FoundationDirs {
		devDir := filepath.Join(dir, "app-dev")
		log.Printf("[DEBUG] compile: checking foundation dir: %s", devDir)

		_, err := os.Stat(filepath.Join(devDir, "layer.sh"))
		if err != nil {
			// If the file doesn't exist then this foundation just
			// doesn't have a layer. Not a big deal.
			if os.IsNotExist(err) {
				log.Printf("[DEBUG] compile: dir %s has no layers", devDir)
				continue
			}

			// The error is something else, return it...
			return err
		}

		log.Printf("[DEBUG] compile: dir %s has a layer!")

		// We have a layer! Read the ID.
		id, err := oneline.Read(filepath.Join(devDir, "layer.id"))
		if err != nil {
			return err
		}

		// Setup the data for this render
		data.Context["foundation_id"] = id
		data.Context["foundation_dir"] = devDir

		// Create the directory where this will be stored
		renderDir := filepath.Join(
			opts.Ctx.Dir, "foundation-layers", fmt.Sprintf("%d-%s", i, id))
		if err := os.MkdirAll(renderDir, 0755); err != nil {
			return err
		}

		// Render our standard template for a foundation layer
		err = data.RenderAsset(
			filepath.Join(renderDir, "Vagrantfile"),
			"data/internal/foundation-layer.Vagrantfile.tpl")
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>helper/compile: fix vet issue<commit_after>package compile

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/hashicorp/otto/app"
	"github.com/hashicorp/otto/foundation"
	"github.com/hashicorp/otto/helper/bindata"
	"github.com/hashicorp/otto/helper/oneline"
)

// AppOptions are the options for compiling an application.
//
// These options may be modified during customization processing, and
// in fact that is an intended use case and common pattern. To do this,
// use the AppCustomizationFunc method. See some of the builtin types for
// examples.
type AppOptions struct {
	// Ctx is the app context of this compilation.
	Ctx *app.Context

	// Result is the base CompileResult that will be used to return the result.
	// You can set this if you want to override some settings.
	Result *app.CompileResult

	// FoundationConfig is the configuration for the foundation that
	// will be returned as the compilation result.
	FoundationConfig foundation.Config

	// Bindata is the data that is used for templating. This must be set.
	// Template data should also be set on this. This will be modified with
	// default template data if those keys are not set.
	Bindata *bindata.Data

	// Customizations is a list of helpers to process customizations
	// in the Appfile. See the Customization docs for more information.
	Customizations []*Customization

	// Callbacks are called just prior to compilation completing.
	Callbacks []CompileCallback
}

// CompileCallback is a callback that can be registered to be run after
// compilation. To access any data within this callback, it should be created
// as a closure around the AppOptions.
type CompileCallback func() error

// App is an opinionated compilation function to help implement
// app.App.Compile.
//
// AppOptions may be modified by this function during this call.
func App(opts *AppOptions) (*app.CompileResult, error) {
	// Write the test data in case we're running tests right now
	testLock.RLock()
	defer testLock.RUnlock()
	testAppOpts = opts

	ctx := opts.Ctx

	// Setup the basic templating data. We put this into the "data" local
	// var just so that it is easier to reference.
	//
	// The exact default data put into the context is documented above.
	data := opts.Bindata
	if data.Context == nil {
		data.Context = make(map[string]interface{})
		opts.Bindata = data
	}

	data.Context["app_type"] = ctx.Appfile.Application.Type
	data.Context["name"] = ctx.Appfile.Application.Name
	data.Context["dev_fragments"] = ctx.DevDepFragments
	data.Context["dev_ip_address"] = ctx.DevIPAddress

	if data.Context["path"] == nil {
		data.Context["path"] = make(map[string]string)
	}
	pathMap := data.Context["path"].(map[string]string)
	pathMap["cache"] = ctx.CacheDir
	pathMap["compiled"] = ctx.Dir
	pathMap["working"] = filepath.Dir(ctx.Appfile.Path)
	foundationDirsContext := map[string][]string{
		"dev":     make([]string, len(ctx.FoundationDirs)),
		"dev_dep": make([]string, len(ctx.FoundationDirs)),
		"build":   make([]string, len(ctx.FoundationDirs)),
		"deploy":  make([]string, len(ctx.FoundationDirs)),
	}
	for i, dir := range ctx.FoundationDirs {
		foundationDirsContext["dev"][i] = filepath.Join(dir, "app-dev")
		foundationDirsContext["dev_dep"][i] = filepath.Join(dir, "app-dev-dep")
		foundationDirsContext["build"][i] = filepath.Join(dir, "app-build")
		foundationDirsContext["deploy"][i] = filepath.Join(dir, "app-deploy")
	}
	data.Context["foundation_dirs"] = foundationDirsContext

	// Setup the shared data
	if data.SharedExtends == nil {
		data.SharedExtends = make(map[string]*bindata.Data)
	}
	data.SharedExtends["compile"] = &bindata.Data{
		Asset:    Asset,
		AssetDir: AssetDir,
	}

	// Process the customizations!
	err := processCustomizations(&processOpts{
		Customizations: opts.Customizations,
		Appfile:        ctx.Appfile,
		Bindata:        data,
	})
	if err != nil {
		return nil, err
	}

	// Create the directory list that we'll copy from, and copy those
	// directly into the compilation directory.
	bindirs := []string{
		"data/common",
		fmt.Sprintf("data/%s-%s", ctx.Tuple.Infra, ctx.Tuple.InfraFlavor),
	}
	for _, dir := range bindirs {
		// Copy all the common files that exist
		if err := data.CopyDir(ctx.Dir, dir); err != nil {
			// Ignore any directories that don't exist
			if strings.Contains(err.Error(), "not found") {
				continue
			}

			return nil, err
		}
	}

	if err := appFoundations(opts); err != nil {
		return nil, err
	}

	// Callbacks
	for _, cb := range opts.Callbacks {
		if err := cb(); err != nil {
			return nil, err
		}
	}

	// If the DevDep fragment exists, then use it
	fragmentPath := filepath.Join(ctx.Dir, "dev-dep", "Vagrantfile.fragment")
	if _, err := os.Stat(fragmentPath); err != nil {
		fragmentPath = ""
	}

	// Set some defaults here
	if opts.FoundationConfig.ServiceName == "" {
		opts.FoundationConfig.ServiceName = opts.Ctx.Application.Name
	}

	result := opts.Result
	if result == nil {
		result = new(app.CompileResult)
	}
	result.FoundationConfig = opts.FoundationConfig
	result.DevDepFragmentPath = fragmentPath
	return result, nil
}

// appFoundations compiles the app-specific foundation files.
func appFoundations(opts *AppOptions) error {
	// Setup the bindata for rendering
	dataCopy := Data
	data := &dataCopy
	data.Context = make(map[string]interface{})
	for k, v := range opts.Bindata.Context {
		data.Context[k] = v
	}

	// Go through each foundation and setup the layers
	log.Printf("[INFO] compile: looking for foundation layers for dev")
	for i, dir := range opts.Ctx.FoundationDirs {
		devDir := filepath.Join(dir, "app-dev")
		log.Printf("[DEBUG] compile: checking foundation dir: %s", devDir)

		_, err := os.Stat(filepath.Join(devDir, "layer.sh"))
		if err != nil {
			// If the file doesn't exist then this foundation just
			// doesn't have a layer. Not a big deal.
			if os.IsNotExist(err) {
				log.Printf("[DEBUG] compile: dir %s has no layers", devDir)
				continue
			}

			// The error is something else, return it...
			return err
		}

		log.Printf("[DEBUG] compile: dir %s has a layer!", devDir)

		// We have a layer! Read the ID.
		id, err := oneline.Read(filepath.Join(devDir, "layer.id"))
		if err != nil {
			return err
		}

		// Setup the data for this render
		data.Context["foundation_id"] = id
		data.Context["foundation_dir"] = devDir

		// Create the directory where this will be stored
		renderDir := filepath.Join(
			opts.Ctx.Dir, "foundation-layers", fmt.Sprintf("%d-%s", i, id))
		if err := os.MkdirAll(renderDir, 0755); err != nil {
			return err
		}

		// Render our standard template for a foundation layer
		err = data.RenderAsset(
			filepath.Join(renderDir, "Vagrantfile"),
			"data/internal/foundation-layer.Vagrantfile.tpl")
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/lib/pq"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json?extensions=true&page_size=500"

type Open311Request struct {
	Lat, Long                                                                                float64
	Ward, Police_district                                                                                int
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address, Channel, Media_url string
	Requested_datetime, Updated_datetime                                                                    string // FIXME: should these be proper time objects?
	Extended_attributes                                                                                map[string]interface{}
	Notes                                                                                []map[string]interface{}
}

type Worker struct {
	Db        *sql.DB
	LastRunAt time.Time
}

var worker Worker
var sr_number string
var backfill bool
var backfill_date string

func init() {
	// open database
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	worker.Db = db

	// fetch SR num from command line, if present
	flag.StringVar(&sr_number, "sr-number", "", "SR number to fetch")
	flag.BoolVar(&backfill, "backfill", false, "run in reverse and backfill data")
	flag.StringVar(&backfill_date, "backfill-from", "", "date to start backfilling data from. Use RFC3339 format. Default will be the tiem of the least recently updated SR in the database.")
}

func main() {
	defer worker.Db.Close()
	flag.Parse()

	if sr_number != "" {
		sr := fetchSingleRequest(sr_number)
		sr.Save()
		return
	}

	start_backfill_from := backfill_date
	for {
		switch {
		case backfill:
			requests := backFillRequests(start_backfill_from)
			for _, request := range requests {
				request.Save()
			}

			start_backfill_from = requests[len(requests)-1].Updated_datetime // FIXME: is it safe to assume the items are sorted?

		case time.Since(worker.LastRunAt) > (30 * time.Second):
			// load requests from open311
			for _, request := range fetchRequests() {
				request.Save()
			}
			worker.LastRunAt = time.Now()
		default:
			log.Print("sleeping for 10 seconds")
			time.Sleep(10 * time.Second)
		}
	}
}

func (req Open311Request) String() string {
	// pretty print SR information
	return fmt.Sprintf("%s: %s at %s %f,%f, last update %s", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long, req.Updated_datetime)
}

func (req Open311Request) Save() (persisted bool) {
	// create or update a SR

	// open311 says we should always ignore a SR that does not have a SR# assigned
	if req.Service_request_id == "" {
		log.Printf("cowardly refusing to create a new SR record because of empty SR#. Request type is %s", req.Service_name)
		return false
	}

	persisted = false

	// find existing record if exists
	var existing_id int
	err := worker.Db.QueryRow("SELECT id FROM service_requests WHERE service_request_id = $1", req.Service_request_id).Scan(&existing_id)
	switch {
	case err == sql.ErrNoRows:
		// log.Printf("did not find existing record %s", req.Service_request_id)
	case err != nil:
		// log.Print("error searching for existing SR", err)
	default:
		persisted = true
		// log.Printf("found existing sr %s", req.Service_request_id)
	}

	var stmt *sql.Stmt

	if !persisted {
		// create new record
		stmt, err = worker.Db.Prepare("INSERT INTO service_requests(service_request_id," +
			"status, service_name, service_code, agency_responsible, " +
			"address, requested_datetime, updated_datetime, lat, long," +
			"ward, police_district, media_url, channel, duplicate, parent_service_request_id, closed_datetime) " +
			"VALUES ($1::varchar, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17); ")

		// "WHERE NOT EXISTS (SELECT 1 FROM service_requests WHERE service_request_id = $1);")

		if err != nil {
			log.Fatal("error preparing database insert statement", err)
		}

	} else {
		// update existing record
		stmt, err = worker.Db.Prepare("UPDATE service_requests SET " +
			"status = $2, service_name = $3, service_code = $4, agency_responsible = $5, " +
			"address = $6, requested_datetime = $7, updated_datetime = $8, lat = $9, long = $10," +
			"ward = $11, police_district = $12, media_url = $13, channel = $14, duplicate = $15, " +
			"parent_service_request_id = $16, updated_at = NOW(), closed_datetime = $17 WHERE service_request_id = $1;")

		if err != nil {
			log.Fatal("error preparing database update statement", err)
		}
	}

	tx, err := worker.Db.Begin()

	if err != nil {
		log.Fatal("error beginning transaction", err)
	}

	t := req.ExtractClosedDatetime()
	closed_time := pq.NullTime{Time: t, Valid: !t.IsZero()}

	_, err = tx.Stmt(stmt).Exec(req.Service_request_id,
		req.Status,
		req.Service_name,
		req.Service_code,
		req.Agency_responsible,
		req.Address,
		req.Requested_datetime,
		req.Updated_datetime,
		req.Lat,
		req.Long,
		req.Extended_attributes["ward"],
		req.Extended_attributes["police_district"],
		req.Media_url,
		req.Extended_attributes["channel"],
		req.Extended_attributes["duplicate"],
		req.Extended_attributes["parent_service_request_id"],
		closed_time,
	)

	if err != nil {
		log.Fatalf("could not update %s because %s", req.Service_request_id, err)
	} else {
		var verb string
		switch {
		case !persisted && closed_time.Time.IsZero():
			verb = "CREATED"
		case !persisted && !closed_time.Time.IsZero():
			verb = "CREATED/CLOSED"
		case persisted && closed_time.Time.IsZero():
			verb = "UPDATED"
		case persisted && !closed_time.Time.IsZero():
			verb = "UPDATED/CLOSED"
		}

		log.Printf("[%s] %s", verb, req)
		persisted = true
	}

	err = tx.Commit()
	if err != nil {
		log.Fatal("error closing transaction", err)
	}

	return persisted

	// calculate closed time if necessary

}

func (req Open311Request) ExtractClosedDatetime() time.Time {
	// given an extended_attributes JSON blob, pluck out the closed time, if present
	// req.PrintNotes()

	var closed_at time.Time
	for _, note := range req.Notes {
		if note["type"] == "closed" {
			parsed_date, err := time.Parse("2006-01-02T15:04:05-07:00", note["datetime"].(string))
			if err != nil {
				log.Print("error parsing date", err)
			}
			log.Printf("SR %s closed at: %s", req, parsed_date)
			closed_at = parsed_date
		}
	}

	return closed_at
}

func (req Open311Request) PrintNotes() {
	fmt.Printf("Notes for SR %s:\n", req.Service_request_id)

	for _, note := range req.Notes {
		fmt.Printf("%+v\n", note)
	}
}

func fetchSingleRequest(sr_number string) (request Open311Request) {
	// given an SR, fetch the record
	log.Printf("fetching single SR %s", sr_number)
	open311_api_endpoint := fmt.Sprintf("http://311api.cityofchicago.org/open311/v2/requests/%s.json?extensions=true", sr_number)

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	var requests []Open311Request

	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests[0]
}

func fetchRequests() (requests []Open311Request) {
	// find the most recent SR that we know about in the database
	rows, err := worker.Db.Query("SELECT MAX(updated_datetime) FROM service_requests;")
	if err != nil {
		log.Fatal("error finding most recent service request", err)
	}

	last_updated_at := time.Now()
	for rows.Next() {
		if err := rows.Scan(&last_updated_at); err != nil {
			log.Print("error finding most recent SR", err)
		}

		log.Printf("most recent SR timestamp %s", last_updated_at)
	}

	// janky hack to transform the last updated timestamp into
	// a format that plays nicely with the Open311 API
	// FIXME: there HAS to be a better way to handle this.
	formatted_date_string := last_updated_at.Format(time.RFC3339)
	formatted_date_string_with_tz := formatted_date_string[0:len(formatted_date_string)-1] + "-0500" // trunc the trailing 'Z' and tack on timezone

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_after=" + formatted_date_string_with_tz

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests
}

func backFillRequests(start_from string) (requests []Open311Request) {
	var fetch_from time.Time

	if start_from == "" {
		err := worker.Db.QueryRow("SELECT updated_datetime FROM service_requests ORDER BY updated_datetime ASC LIMIT 1").Scan(&fetch_from)
		if err != nil {
			log.Println("error fetching oldest SR:", err)
		}
		log.Printf("no start_from value provided, so falling back to oldest (by last update) SR in the database: %s", start_from)
	} else {
		t, err := time.Parse(time.RFC3339, start_from)
		if err != nil {
			log.Fatal("[backfill] error parsing date to start from", err)
		}
		fetch_from = t
	}

	formatted_date_string_with_tz := fetch_from.Format(time.RFC3339)

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_before=" + formatted_date_string_with_tz

	log.Printf("[backfill] fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("[backfill] error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("[backfill] error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("[backfill] error parsing JSON:", err)
	}

	log.Printf("[backfill] received %d requests from Open311", len(requests))

	return requests
}
<commit_msg>gofmt<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/lib/pq"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json?extensions=true&page_size=500"

type Open311Request struct {
	Lat, Long                                                                                float64
	Ward, Police_district                                                                                int
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address, Channel, Media_url string
	Requested_datetime, Updated_datetime                                                                    string // FIXME: should these be proper time objects?
	Extended_attributes                                                                                map[string]interface{}
	Notes                                                                                []map[string]interface{}
}

type Worker struct {
	Db        *sql.DB
	LastRunAt time.Time
}

var worker Worker
var sr_number string
var backfill bool
var backfill_date string

func init() {
	// open database
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	worker.Db = db

	// fetch SR num from command line, if present
	flag.StringVar(&sr_number, "sr-number", "", "SR number to fetch")
	flag.BoolVar(&backfill, "backfill", false, "run in reverse and backfill data")
	flag.StringVar(&backfill_date, "backfill-from", "", "date to start backfilling data from. Use RFC3339 format. Default will be the time of the least recently updated SR in the database.")
}

func main() {
	defer worker.Db.Close()
	flag.Parse()

	if sr_number != "" {
		sr := fetchSingleRequest(sr_number)
		sr.Save()
		return
	}

	start_backfill_from := backfill_date
	for {
		switch {
		case backfill:
			requests := backFillRequests(start_backfill_from)
			for _, request := range requests {
				request.Save()
			}

			start_backfill_from = requests[len(requests)-1].Updated_datetime // FIXME: is it safe to assume the items are sorted?

		case time.Since(worker.LastRunAt) > (30 * time.Second):
			// load requests from open311
			for _, request := range fetchRequests() {
				request.Save()
			}
			worker.LastRunAt = time.Now()
		default:
			log.Print("sleeping for 10 seconds")
			time.Sleep(10 * time.Second)
		}
	}
}

func (req Open311Request) String() string {
	// pretty print SR information
	return fmt.Sprintf("%s: %s at %s %f,%f, last update %s", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long, req.Updated_datetime)
}

func (req Open311Request) Save() (persisted bool) {
	// create or update a SR

	// open311 says we should always ignore a SR that does not have a SR# assigned
	if req.Service_request_id == "" {
		log.Printf("cowardly refusing to create a new SR record because of empty SR#. Request type is %s", req.Service_name)
		return false
	}

	persisted = false

	// find existing record if exists
	var existing_id int
	err := worker.Db.QueryRow("SELECT id FROM service_requests WHERE service_request_id = $1", req.Service_request_id).Scan(&existing_id)
	switch {
	case err == sql.ErrNoRows:
		// log.Printf("did not find existing record %s", req.Service_request_id)
	case err != nil:
		// log.Print("error searching for existing SR", err)
	default:
		persisted = true
		// log.Printf("found existing sr %s", req.Service_request_id)
	}

	var stmt *sql.Stmt

	if !persisted {
		// create new record
		stmt, err = worker.Db.Prepare("INSERT INTO service_requests(service_request_id," +
			"status, service_name, service_code, agency_responsible, " +
			"address, requested_datetime, updated_datetime, lat, long," +
			"ward, police_district, media_url, channel, duplicate, parent_service_request_id, closed_datetime) " +
			"VALUES ($1::varchar, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17); ")

		// "WHERE NOT EXISTS (SELECT 1 FROM service_requests WHERE service_request_id = $1);")

		if err != nil {
			log.Fatal("error preparing database insert statement", err)
		}

	} else {
		// update existing record
		stmt, err = worker.Db.Prepare("UPDATE service_requests SET " +
			"status = $2, service_name = $3, service_code = $4, agency_responsible = $5, " +
			"address = $6, requested_datetime = $7, updated_datetime = $8, lat = $9, long = $10," +
			"ward = $11, police_district = $12, media_url = $13, channel = $14, duplicate = $15, " +
			"parent_service_request_id = $16, updated_at = NOW(), closed_datetime = $17 WHERE service_request_id = $1;")

		if err != nil {
			log.Fatal("error preparing database update statement", err)
		}
	}

	tx, err := worker.Db.Begin()

	if err != nil {
		log.Fatal("error beginning transaction", err)
	}

	t := req.ExtractClosedDatetime()
	closed_time := pq.NullTime{Time: t, Valid: !t.IsZero()}

	_, err = tx.Stmt(stmt).Exec(req.Service_request_id,
		req.Status,
		req.Service_name,
		req.Service_code,
		req.Agency_responsible,
		req.Address,
		req.Requested_datetime,
		req.Updated_datetime,
		req.Lat,
		req.Long,
		req.Extended_attributes["ward"],
		req.Extended_attributes["police_district"],
		req.Media_url,
		req.Extended_attributes["channel"],
		req.Extended_attributes["duplicate"],
		req.Extended_attributes["parent_service_request_id"],
		closed_time,
	)

	if err != nil {
		log.Fatalf("could not update %s because %s", req.Service_request_id, err)
	} else {
		var verb string
		switch {
		case !persisted && closed_time.Time.IsZero():
			verb = "CREATED"
		case !persisted && !closed_time.Time.IsZero():
			verb = "CREATED/CLOSED"
		case persisted && closed_time.Time.IsZero():
			verb = "UPDATED"
		case persisted && !closed_time.Time.IsZero():
			verb = "UPDATED/CLOSED"
		}

		log.Printf("[%s] %s", verb, req)
		persisted = true
	}

	err = tx.Commit()
	if err != nil {
		log.Fatal("error closing transaction", err)
	}

	return persisted

	// calculate closed time if necessary

}

func (req Open311Request) ExtractClosedDatetime() time.Time {
	// given an extended_attributes JSON blob, pluck out the closed time, if present
	// req.PrintNotes()

	var closed_at time.Time
	for _, note := range req.Notes {
		if note["type"] == "closed" {
			parsed_date, err := time.Parse("2006-01-02T15:04:05-07:00", note["datetime"].(string))
			if err != nil {
				log.Print("error parsing date", err)
			}
			log.Printf("SR %s closed at: %s", req, parsed_date)
			closed_at = parsed_date
		}
	}

	return closed_at
}

func (req Open311Request) PrintNotes() {
	fmt.Printf("Notes for SR %s:\n", req.Service_request_id)

	for _, note := range req.Notes {
		fmt.Printf("%+v\n", note)
	}
}

func fetchSingleRequest(sr_number string) (request Open311Request) {
	// given an SR, fetch the record
	log.Printf("fetching single SR %s", sr_number)
	open311_api_endpoint := fmt.Sprintf("http://311api.cityofchicago.org/open311/v2/requests/%s.json?extensions=true", sr_number)

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	var requests []Open311Request

	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests[0]
}

func fetchRequests() (requests []Open311Request) {
	// find the most recent SR that we know about in the database
	rows, err := worker.Db.Query("SELECT MAX(updated_datetime) FROM service_requests;")
	if err != nil {
		log.Fatal("error finding most recent service request", err)
	}

	last_updated_at := time.Now()
	for rows.Next() {
		if err := rows.Scan(&last_updated_at); err != nil {
			log.Print("error finding most recent SR", err)
		}

		log.Printf("most recent SR timestamp %s", last_updated_at)
	}

	// janky hack to transform the last updated timestamp into
	// a format that plays nicely with the Open311 API
	// FIXME: there HAS to be a better way to handle this.
	formatted_date_string := last_updated_at.Format(time.RFC3339)
	formatted_date_string_with_tz := formatted_date_string[0:len(formatted_date_string)-1] + "-0500" // trunc the trailing 'Z' and tack on timezone

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_after=" + formatted_date_string_with_tz

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests
}

func backFillRequests(start_from string) (requests []Open311Request) {
	var fetch_from time.Time

	if start_from == "" {
		err := worker.Db.QueryRow("SELECT updated_datetime FROM service_requests ORDER BY updated_datetime ASC LIMIT 1").Scan(&fetch_from)
		if err != nil {
			log.Println("error fetching oldest SR:", err)
		}
		log.Printf("no start_from value provided, so falling back to oldest (by last update) SR in the database: %s", fetch_from)
	} else {
		t, err := time.Parse(time.RFC3339, start_from)
		if err != nil {
			log.Fatal("[backfill] error parsing date to start from", err)
		}
		fetch_from = t
	}

	formatted_date_string_with_tz := fetch_from.Format(time.RFC3339)

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_before=" + formatted_date_string_with_tz

	log.Printf("[backfill] fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("[backfill] error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("[backfill] error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("[backfill] error parsing JSON:", err)
	}

	log.Printf("[backfill] received %d requests from Open311", len(requests))

	return requests
}
<|endoftext|>


<commit_before>// vagrant implements the packer.PostProcessor interface and adds a
// post-processor that turns artifacts of known builders into Vagrant
// boxes.
package vagrant

import (
	"fmt"
	"github.com/mitchellh/mapstructure"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
	"log"
)

var builtins = map[string]string{
	"mitchellh.amazonebs":  "aws",
	"mitchellh.virtualbox": "virtualbox",
	"mitchellh.vmware":     "vmware",
}

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	OutputPath string `mapstructure:"output"`
}

type PostProcessor struct {
	config      Config
	premade     map[string]packer.PostProcessor
	extraConfig map[string]interface{}
}

func (p *PostProcessor) Configure(raws ...interface{}) error {
	_, err := common.DecodeConfig(&p.config, raws...)
	if err != nil {
		return err
	}

	tpl, err := packer.NewConfigTemplate()
	if err != nil {
		return err
	}
	tpl.UserVars = p.config.PackerUserVars

	// Defaults
	if p.config.OutputPath == "" {
		p.config.OutputPath = "packer_{{ .BuildName }}_{{.Provider}}.box"
	}

	// Accumulate any errors
	errs := new(packer.MultiError)
	if err := tpl.Validate(p.config.OutputPath); err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Error parsing output template: %s", err))
	}

	// Store extra configuration we'll send to each post-processor type
	p.extraConfig = make(map[string]interface{})
	p.extraConfig["output"] = p.config.OutputPath
	p.extraConfig["packer_build_name"] = p.config.PackerBuildName
	p.extraConfig["packer_builder_type"] = p.config.PackerBuilderType
	p.extraConfig["packer_debug"] = p.config.PackerDebug
	p.extraConfig["packer_force"] = p.config.PackerForce
	p.extraConfig["packer_user_variables"] = p.config.PackerUserVars

	// TODO(mitchellh): Properly handle multiple raw configs. This isn't
	// very pressing at the moment because at the time of this comment
	// only the first member of raws can contain the actual type-overrides.
	var mapConfig map[string]interface{}
	if err := mapstructure.Decode(raws[0], &mapConfig); err != nil {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("Failed to decode config: %s", err))
		return errs
	}

	p.premade = make(map[string]packer.PostProcessor)
	for k, raw := range mapConfig {
		pp, err := p.subPostProcessor(k, raw, p.extraConfig)
		if err != nil {
			errs = packer.MultiErrorAppend(errs, err)
			continue
		}

		if pp == nil {
			continue
		}

		p.premade[k] = pp
	}

	if len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *PostProcessor) PostProcess(ui packer.Ui, artifact packer.Artifact) (packer.Artifact, bool, error) {
	ppName, ok := builtins[artifact.BuilderId()]
	if !ok {
		return nil, false, fmt.Errorf("Unknown artifact type, can't build box: %s", artifact.BuilderId())
	}

	// Use the premade PostProcessor if we have one. Otherwise, we
	// create it and configure it here.
	pp, ok := p.premade[ppName]
	if !ok {
		log.Printf("Premade post-processor for '%s' not found. Creating.", ppName)

		var err error
		pp, err = p.subPostProcessor(ppName, nil, p.extraConfig)
		if err != nil {
			return nil, false, err
		}

		if pp == nil {
			return nil, false, fmt.Errorf("Vagrant box post-processor not found: %s", ppName)
		}
	}

	ui.Say(fmt.Sprintf("Creating Vagrant box for '%s' provider", ppName))
	return pp.PostProcess(ui, artifact)
}

func (p *PostProcessor) subPostProcessor(key string, specific interface{}, extra map[string]interface{}) (packer.PostProcessor, error) {
	pp := keyToPostProcessor(key)
	if pp == nil {
		return nil, nil
	}

	if err := pp.Configure(extra, specific); err != nil {
		return nil, err
	}

	return pp, nil
}

// keyToPostProcessor maps a configuration key to the actual post-processor
// it will be configuring. This returns a new instance of that post-processor.
func keyToPostProcessor(key string) packer.PostProcessor {
	switch key {
	case "aws":
		return new(AWSBoxPostProcessor)
	case "virtualbox":
		return new(VBoxBoxPostProcessor)
	case "vmware":
		return new(VMwareBoxPostProcessor)
	default:
		return nil
	}
}
<commit_msg>Use the same vagrant post-processor for amazon instances than EBS. Fixes #502<commit_after>// vagrant implements the packer.PostProcessor interface and adds a
// post-processor that turns artifacts of known builders into Vagrant
// boxes.
package vagrant

import (
	"fmt"
	"github.com/mitchellh/mapstructure"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
	"log"
)

var builtins = map[string]string{
	"mitchellh.amazonebs":       "aws",
	"mitchellh.amazon.instance": "aws",
	"mitchellh.virtualbox":      "virtualbox",
	"mitchellh.vmware":          "vmware",
}

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	OutputPath string `mapstructure:"output"`
}

type PostProcessor struct {
	config      Config
	premade     map[string]packer.PostProcessor
	extraConfig map[string]interface{}
}

func (p *PostProcessor) Configure(raws ...interface{}) error {
	_, err := common.DecodeConfig(&p.config, raws...)
	if err != nil {
		return err
	}

	tpl, err := packer.NewConfigTemplate()
	if err != nil {
		return err
	}
	tpl.UserVars = p.config.PackerUserVars

	// Defaults
	if p.config.OutputPath == "" {
		p.config.OutputPath = "packer_{{ .BuildName }}_{{.Provider}}.box"
	}

	// Accumulate any errors
	errs := new(packer.MultiError)
	if err := tpl.Validate(p.config.OutputPath); err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Error parsing output template: %s", err))
	}

	// Store extra configuration we'll send to each post-processor type
	p.extraConfig = make(map[string]interface{})
	p.extraConfig["output"] = p.config.OutputPath
	p.extraConfig["packer_build_name"] = p.config.PackerBuildName
	p.extraConfig["packer_builder_type"] = p.config.PackerBuilderType
	p.extraConfig["packer_debug"] = p.config.PackerDebug
	p.extraConfig["packer_force"] = p.config.PackerForce
	p.extraConfig["packer_user_variables"] = p.config.PackerUserVars

	// TODO(mitchellh): Properly handle multiple raw configs. This isn't
	// very pressing at the moment because at the time of this comment
	// only the first member of raws can contain the actual type-overrides.
	var mapConfig map[string]interface{}
	if err := mapstructure.Decode(raws[0], &mapConfig); err != nil {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("Failed to decode config: %s", err))
		return errs
	}

	p.premade = make(map[string]packer.PostProcessor)
	for k, raw := range mapConfig {
		pp, err := p.subPostProcessor(k, raw, p.extraConfig)
		if err != nil {
			errs = packer.MultiErrorAppend(errs, err)
			continue
		}

		if pp == nil {
			continue
		}

		p.premade[k] = pp
	}

	if len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *PostProcessor) PostProcess(ui packer.Ui, artifact packer.Artifact) (packer.Artifact, bool, error) {
	ppName, ok := builtins[artifact.BuilderId()]
	if !ok {
		return nil, false, fmt.Errorf("Unknown artifact type, can't build box: %s", artifact.BuilderId())
	}

	// Use the premade PostProcessor if we have one. Otherwise, we
	// create it and configure it here.
	pp, ok := p.premade[ppName]
	if !ok {
		log.Printf("Premade post-processor for '%s' not found. Creating.", ppName)

		var err error
		pp, err = p.subPostProcessor(ppName, nil, p.extraConfig)
		if err != nil {
			return nil, false, err
		}

		if pp == nil {
			return nil, false, fmt.Errorf("Vagrant box post-processor not found: %s", ppName)
		}
	}

	ui.Say(fmt.Sprintf("Creating Vagrant box for '%s' provider", ppName))
	return pp.PostProcess(ui, artifact)
}

func (p *PostProcessor) subPostProcessor(key string, specific interface{}, extra map[string]interface{}) (packer.PostProcessor, error) {
	pp := keyToPostProcessor(key)
	if pp == nil {
		return nil, nil
	}

	if err := pp.Configure(extra, specific); err != nil {
		return nil, err
	}

	return pp, nil
}

// keyToPostProcessor maps a configuration key to the actual post-processor
// it will be configuring. This returns a new instance of that post-processor.
func keyToPostProcessor(key string) packer.PostProcessor {
	switch key {
	case "aws":
		return new(AWSBoxPostProcessor)
	case "virtualbox":
		return new(VBoxBoxPostProcessor)
	case "vmware":
		return new(VMwareBoxPostProcessor)
	default:
		return nil
	}
}
<|endoftext|>


<commit_before>// Package gpio implements a timer/counter unit.
// Tested compatibility:
//   ATmega48/88/168 (timer 0 only)
// Untested compatability:
//   ATmega48/88/168 (timers 1 and 2)
//   ATtiny4/5/9/10
package timer

import (
    "fmt"
    "github.com/kierdavis/avr/emulator"
    "github.com/kierdavis/avr/hardware/gpio"
    "log"
)

// TODO: for a OSCx output, require corresponding DDR bit to be set to output

// TODO: thread-safety!!

// TODO: OCFy/TOV flags should not be cleared upon interrupt execution if the
// emulator's global interrupt enable flag is cleared

type Timer struct {
    em *emulator.Emulator
    digit uint
    controlA uint8
    controlB uint8
    count uint8
    compareValA uint8
    compareValB uint8
    interruptMask uint8
    interruptFlags uint8
    downwards bool // count direction
    ocPinStates [2]bool
    ocPinCallbacks [2]func(bool)
    logging bool
    inhibitCompareMatch bool // set when TCNT is written to prevent a compare match on the next clock
    excessTicks uint
}

func New(digit uint) (t *Timer) {
    return &Timer{
        digit: digit,
    }
}

func (t *Timer) SetLogging(enabled bool) {
    t.logging = enabled
}

func (t *Timer) AddTo(em *emulator.Emulator) {
    t.em = em

    em.RegisterPortByName(fmt.Sprintf("TCCR%dA", t.digit), tccra{t})
    em.RegisterPortByName(fmt.Sprintf("TCCR%dB", t.digit), tccrb{t})
    em.RegisterPortByName(fmt.Sprintf("TCNT%d", t.digit), tcnt{t})
    em.RegisterPortByName(fmt.Sprintf("OCR%dA", t.digit), ocra{t})
    em.RegisterPortByName(fmt.Sprintf("OCR%dB", t.digit), ocrb{t})
    em.RegisterPortByName(fmt.Sprintf("TIMSK%d", t.digit), timsk{t})
    em.RegisterPortByName(fmt.Sprintf("TIFR%d", t.digit), tifr{t})
}

// Connect an output-compare pin to a GPIO port by calling the GPIO's
// OverrideOutput method.
func (t *Timer) OverrideOCPin(ocPinNum uint, gpioPinNum uint, g *gpio.GPIO) {
    t.ocPinCallbacks[ocPinNum] = g.OverrideOutput(gpioPinNum)
}

// Note: in PWM modes, OCRA/OCRB do not exhibit a newly written value until the count overflows

// TODO: in PC-PWM mode, an OCRx bit may transition without a compare match for two reasons (see datasheet page 98)

func (t *Timer) Run(ticks uint) {
    var ticksIncr uint

    switch t.controlB & 0x07 {
    case 0: // disabled
        t.excessTicks = 0
        return
    case 1: // divider = 1
        ticksIncr = 1
    case 2: // divider = 8
        ticksIncr = 8
    case 3: // divider = 64
        ticksIncr = 64
    case 4: // divider = 256
        ticksIncr = 256
    case 5: // divider = 1024
        ticksIncr = 1024
    case 6, 7:
        panic("(*Timer).Run: external clock sources not implemented")
    }
    ticksExecuted := t.excessTicks

    for ticksExecuted < ticks {
        t.Tick()
        ticksExecuted += ticksIncr
    }

    t.excessTicks = ticksExecuted - ticks
}

// Tick the timer.
func (t *Timer) Tick() {
    // Get the waveform generation mode bits
    wgmA := t.controlA & 0x03 // bits 1 and 0
    wgmB := (t.controlB & 0x08) >> 1 // bit 3 (shifted to bit 2)
    wgm := wgmB | wgmA

    // Handle match-compare interrupts
    if t.inhibitCompareMatch {
        t.inhibitCompareMatch = false
    } else {
        if t.count == t.compareValA {
            t.setOCF(0)
        }
        if t.count == t.compareValB {
            t.setOCF(1)
        }
    }

    // Prepare to tick counter
    switch wgm {
    case 0: // Normal
        t.tickNormalMode()
    case 1: // Phase-correct PWM (TOP = 0xFF)
        t.tickPCPWMMode(0xFF)
    case 2: // Clear timer on compare
        t.tickCTCMode()
    case 3: // Fast PWM (TOP = 0xFF)
        t.tickFastPWMMode(0xFF)
    case 5: // Phase-correct PWM (TOP = OCRA)
        t.tickPCPWMMode(t.compareValA)
    case 7: // Fast PWM (TOP = OCRA)
        t.tickPCPWMMode(t.compareValA)
    }

    // Actually tick the counter
    if t.downwards {
        t.count--
    } else {
        t.count++
    }

    // Trigger interrupts, if possible
    if t.em != nil && t.em.InterruptsEnabled() {
        intName := ""
        if t.interruptFlags & 0x01 != 0 && t.interruptMask & 0x01 != 0 {
            t.interruptFlags &= 0xFE
            intName = fmt.Sprintf("TIMER%d_OVF", t.digit)
        } else if t.interruptFlags & 0x02 != 0 && t.interruptMask & 0x02 != 0 {
            t.interruptFlags &= 0xFD
            intName = fmt.Sprintf("TIMER%d_COMPA", t.digit)
        } else if t.interruptFlags & 0x04 != 0 && t.interruptMask & 0x04 != 0 {
            t.interruptFlags &= 0xFB
            intName = fmt.Sprintf("TIMER%d_COMPB", t.digit)
        }

        if intName != "" {
            ok := t.em.InterruptByName(intName)
            if !ok && t.logging {
                log.Printf("[avr/hardware/timer:(*Timer).Tick] failed to trigger interrupt %s", intName)
            }
        }
    }
}

// Tick the timer in Normal mode.
func (t *Timer) tickNormalMode() {
    t.downwards = false
    if t.count == 0xFF { // Overflow
        t.setTOV()
    }

    t.checkOCPinNormalMode(0, t.compareValA)
    t.checkOCPinNormalMode(1, t.compareValB)
}

// Check for output-compare in normal mode.
func (t *Timer) checkOCPinNormalMode(ocPinNum uint, compareVal uint8) {
    if t.count == compareVal {
        // Get COMxy bits
        shiftAmt := 6 - 2*ocPinNum // 0 => 6, 1 => 4
        com := (t.controlA >> shiftAmt) & 0x03
        switch com {
        case 0: // OCy disabled
            // do nothing
        case 1: // toggle OCy
            t.toggleOCPin(ocPinNum)
        case 2: // clear OCy
            t.clearOCPin(ocPinNum)
        case 3: // set OCy
            t.setOCPin(ocPinNum)
        }
    }
}

// Tick the timer in phase-corrected PWM mode.
func (t *Timer) tickPCPWMMode(top uint8) {
    if t.downwards {
        if t.count == 0x00 { // Reached BOTTOM
            t.downwards = false // Begin counting upwards
            t.setTOV()
        }
    } else {
        if t.count == top { // Reached TOP
            t.downwards = true // Begin counting downwards
        }
    }

    t.checkOCPinPCPWMMode(0, t.compareValA)
    t.checkOCPinPCPWMMode(1, t.compareValB)
}

// TODO: there is a special case in the OC pin checking for both PWM modes when OCRy == TOP and COMxy1 is set (see datasheet page 101)

// Check for output-compare in phase-corrected PWM mode.
func (t *Timer) checkOCPinPCPWMMode(ocPinNum uint, compareVal uint8) {
    if t.count == compareVal {
        // Get COMxy bits
        shiftAmt := 6 - 2*ocPinNum // 0 => 6, 1 => 4
        com := (t.controlA >> shiftAmt) & 0x03
        switch com {
        case 0: // OCy disabled
            // do nothing
        case 1: // Toggle OCy (only on OC pin 0 with WGM2 bit set)
            if ocPinNum == 0 && (t.controlB & 0x80) != 0 {
                t.toggleOCPin(ocPinNum)
            }
        case 2: // Clear OCy if counting upwards or set OCy if counting downwards
            if t.downwards {
                t.setOCPin(ocPinNum)
            } else {
                t.clearOCPin(ocPinNum)
            }
        case 3: // Set OCy if counting upwards or clear OCy if counting downwards
            if t.downwards {
                t.clearOCPin(ocPinNum)
            } else {
                t.setOCPin(ocPinNum)
            }
        }
    }
}

// Tick the timer in clear-timer-on-compare mode.
func (t *Timer) tickCTCMode() {
    t.downwards = false

    if t.count == 0xFF { // Overflow
        t.setTOV()
    }

    if t.count == t.compareValA {
        // this tick should set counter to 0
        t.count = 0xFF
    }
}

// Tick the timer in fast PWM mode.
func (t *Timer) tickFastPWMMode(top uint8) {
    t.downwards = false

    if t.count == top {
        // this tick should set counter to 0
        t.count = 0xFF
        t.setTOV()
    }

    t.checkOCPinFastPWMMode(0, t.compareValA)
    t.checkOCPinFastPWMMode(1, t.compareValB)
}

// Check for output-compare in fast PWM mode.
func (t *Timer) checkOCPinFastPWMMode(ocPinNum uint, compareVal uint8) {
    // BOTTOM
    if t.count == 0x00 {
        // Get COMxy bits
        shiftAmt := 6 - 2*ocPinNum // 0 => 6, 1 => 4
        com := (t.controlA >> shiftAmt) & 0x03
        switch com {
        case 0: // OCy disabled
            // do nothing
        case 1: // Toggle OCy on compare match (only on OC pin 0 with WGM2 bit set)
            // do nothing
        case 2: // Clear OCy on compare match, set OCy at BOTTOM
            t.setOCPin(ocPinNum)
        case 3: // Set OCy on compare match, clear OCy at BOTTOM
            t.clearOCPin(ocPinNum)
        }
    }

    // Compare match
    if t.count == compareVal {
        // Get COMxy bits
        shiftAmt := 6 - 2*ocPinNum // 0 => 6, 1 => 4
        com := (t.controlA >> shiftAmt) & 0x03
        switch com {
        case 0: // OCy disabled
            // do nothing
        case 1: // Toggle OCy on compare match (only on OC pin 0 with WGM2 bit set)
            if ocPinNum == 0 && (t.controlB & 0x80) != 0 {
                t.toggleOCPin(ocPinNum)
            }
        case 2: // Clear OCy on compare match, set OCy at BOTTOM
            t.clearOCPin(ocPinNum)
        case 3: // Set OCy on compare match, clear OCy at BOTTOM
            t.setOCPin(ocPinNum)
        }
    }
}

// Toggle an output-compare pin.
func (t *Timer) toggleOCPin(ocPinNum uint) {
    t.ocPinStates[ocPinNum] = !t.ocPinStates[ocPinNum]
    t.updateOCPin(ocPinNum)
}

// Set an output-compare pin to low.
func (t *Timer) clearOCPin(ocPinNum uint) {
    if t.ocPinStates[ocPinNum] {
        t.ocPinStates[ocPinNum] = false
        t.updateOCPin(ocPinNum)
    }
}

// Set an output-compare pin to high.
func (t *Timer) setOCPin(ocPinNum uint) {
    if !t.ocPinStates[ocPinNum] {
        t.ocPinStates[ocPinNum] = true
        t.updateOCPin(ocPinNum)
    }
}

// Push the new status of an output-compare pin to the GPIO layer.
func (t *Timer) updateOCPin(ocPinNum uint) {
    callback := t.ocPinCallbacks[ocPinNum]
    if callback != nil {
        callback(t.ocPinStates[ocPinNum])
    }
}

// Set an output compare match (OCFy) flag.
func (t *Timer) setOCF(ocPinNum uint) {
    if ocPinNum == 0 { // A
        t.interruptFlags |= 0x02
    } else { // B
        t.interruptFlags |= 0x04
    }
}

// Set the timer overflow (TOV) flag.
func (t *Timer) setTOV() {
    t.interruptFlags |= 0x01
}
<commit_msg>Factor out updating of OC pin in normal mode into a separate function<commit_after>// Package gpio implements a timer/counter unit.
// Tested compatibility:
//   ATmega48/88/168 (timer 0 only)
// Untested compatability:
//   ATmega48/88/168 (timers 1 and 2)
//   ATtiny4/5/9/10
package timer

import (
    "fmt"
    "github.com/kierdavis/avr/emulator"
    "github.com/kierdavis/avr/hardware/gpio"
    "log"
)

// TODO: for a OSCx output, require corresponding DDR bit to be set to output

// TODO: thread-safety!!

// TODO: OCFy/TOV flags should not be cleared upon interrupt execution if the
// emulator's global interrupt enable flag is cleared

type Timer struct {
    em *emulator.Emulator
    digit uint
    controlA uint8
    controlB uint8
    count uint8
    compareValA uint8
    compareValB uint8
    interruptMask uint8
    interruptFlags uint8
    downwards bool // count direction
    ocPinStates [2]bool
    ocPinCallbacks [2]func(bool)
    logging bool
    inhibitCompareMatch bool // set when TCNT is written to prevent a compare match on the next clock
    excessTicks uint
}

func New(digit uint) (t *Timer) {
    return &Timer{
        digit: digit,
    }
}

func (t *Timer) SetLogging(enabled bool) {
    t.logging = enabled
}

func (t *Timer) AddTo(em *emulator.Emulator) {
    t.em = em

    em.RegisterPortByName(fmt.Sprintf("TCCR%dA", t.digit), tccra{t})
    em.RegisterPortByName(fmt.Sprintf("TCCR%dB", t.digit), tccrb{t})
    em.RegisterPortByName(fmt.Sprintf("TCNT%d", t.digit), tcnt{t})
    em.RegisterPortByName(fmt.Sprintf("OCR%dA", t.digit), ocra{t})
    em.RegisterPortByName(fmt.Sprintf("OCR%dB", t.digit), ocrb{t})
    em.RegisterPortByName(fmt.Sprintf("TIMSK%d", t.digit), timsk{t})
    em.RegisterPortByName(fmt.Sprintf("TIFR%d", t.digit), tifr{t})
}

// Connect an output-compare pin to a GPIO port by calling the GPIO's
// OverrideOutput method.
func (t *Timer) OverrideOCPin(ocPinNum uint, gpioPinNum uint, g *gpio.GPIO) {
    t.ocPinCallbacks[ocPinNum] = g.OverrideOutput(gpioPinNum)
}

// Note: in PWM modes, OCRA/OCRB do not exhibit a newly written value until the count overflows

// TODO: in PC-PWM mode, an OCRx bit may transition without a compare match for two reasons (see datasheet page 98)

func (t *Timer) Run(ticks uint) {
    var ticksIncr uint

    switch t.controlB & 0x07 {
    case 0: // disabled
        t.excessTicks = 0
        return
    case 1: // divider = 1
        ticksIncr = 1
    case 2: // divider = 8
        ticksIncr = 8
    case 3: // divider = 64
        ticksIncr = 64
    case 4: // divider = 256
        ticksIncr = 256
    case 5: // divider = 1024
        ticksIncr = 1024
    case 6, 7:
        panic("(*Timer).Run: external clock sources not implemented")
    }
    ticksExecuted := t.excessTicks

    for ticksExecuted < ticks {
        t.Tick()
        ticksExecuted += ticksIncr
    }

    t.excessTicks = ticksExecuted - ticks
}

// Tick the timer.
func (t *Timer) Tick() {
    // Get the waveform generation mode bits
    wgmA := t.controlA & 0x03 // bits 1 and 0
    wgmB := (t.controlB & 0x08) >> 1 // bit 3 (shifted to bit 2)
    wgm := wgmB | wgmA

    // Handle match-compare interrupts
    if t.inhibitCompareMatch {
        t.inhibitCompareMatch = false
    } else {
        if t.count == t.compareValA {
            t.setOCF(0)
        }
        if t.count == t.compareValB {
            t.setOCF(1)
        }
    }

    // Prepare to tick counter
    switch wgm {
    case 0: // Normal
        t.tickNormalMode()
    case 1: // Phase-correct PWM (TOP = 0xFF)
        t.tickPCPWMMode(0xFF)
    case 2: // Clear timer on compare
        t.tickCTCMode()
    case 3: // Fast PWM (TOP = 0xFF)
        t.tickFastPWMMode(0xFF)
    case 5: // Phase-correct PWM (TOP = OCRA)
        t.tickPCPWMMode(t.compareValA)
    case 7: // Fast PWM (TOP = OCRA)
        t.tickPCPWMMode(t.compareValA)
    }

    // Actually tick the counter
    if t.downwards {
        t.count--
    } else {
        t.count++
    }

    // Trigger interrupts, if possible
    if t.em != nil && t.em.InterruptsEnabled() {
        intName := ""
        if t.interruptFlags & 0x01 != 0 && t.interruptMask & 0x01 != 0 {
            t.interruptFlags &= 0xFE
            intName = fmt.Sprintf("TIMER%d_OVF", t.digit)
        } else if t.interruptFlags & 0x02 != 0 && t.interruptMask & 0x02 != 0 {
            t.interruptFlags &= 0xFD
            intName = fmt.Sprintf("TIMER%d_COMPA", t.digit)
        } else if t.interruptFlags & 0x04 != 0 && t.interruptMask & 0x04 != 0 {
            t.interruptFlags &= 0xFB
            intName = fmt.Sprintf("TIMER%d_COMPB", t.digit)
        }

        if intName != "" {
            ok := t.em.InterruptByName(intName)
            if !ok && t.logging {
                log.Printf("[avr/hardware/timer:(*Timer).Tick] failed to trigger interrupt %s", intName)
            }
        }
    }
}

// Tick the timer in Normal mode.
func (t *Timer) tickNormalMode() {
    t.downwards = false
    if t.count == 0xFF { // Overflow
        t.setTOV()
    }

    t.checkOCPinNormalMode(0, t.compareValA)
    t.checkOCPinNormalMode(1, t.compareValB)
}

// Check for output-compare in normal mode.
func (t *Timer) checkOCPinNormalMode(ocPinNum uint, compareVal uint8) {
    if t.count == compareVal {
        t.changeOCPinNormalMode(ocPinNum)
    }
}

func (t *Timer) changeOCPinNormalMode(ocPinNum uint) {
    // Get COMxy bits
    shiftAmt := 6 - 2*ocPinNum // 0 => 6, 1 => 4
    com := (t.controlA >> shiftAmt) & 0x03
    switch com {
    case 0: // OCy disabled
        // do nothing
    case 1: // toggle OCy
        t.toggleOCPin(ocPinNum)
    case 2: // clear OCy
        t.clearOCPin(ocPinNum)
    case 3: // set OCy
        t.setOCPin(ocPinNum)
    }
}

// Tick the timer in phase-corrected PWM mode.
func (t *Timer) tickPCPWMMode(top uint8) {
    if t.downwards {
        if t.count == 0x00 { // Reached BOTTOM
            t.downwards = false // Begin counting upwards
            t.setTOV()
        }
    } else {
        if t.count == top { // Reached TOP
            t.downwards = true // Begin counting downwards
        }
    }

    t.checkOCPinPCPWMMode(0, t.compareValA)
    t.checkOCPinPCPWMMode(1, t.compareValB)
}

// TODO: there is a special case in the OC pin checking for both PWM modes when OCRy == TOP and COMxy1 is set (see datasheet page 101)

// Check for output-compare in phase-corrected PWM mode.
func (t *Timer) checkOCPinPCPWMMode(ocPinNum uint, compareVal uint8) {
    if t.count == compareVal {
        // Get COMxy bits
        shiftAmt := 6 - 2*ocPinNum // 0 => 6, 1 => 4
        com := (t.controlA >> shiftAmt) & 0x03
        switch com {
        case 0: // OCy disabled
            // do nothing
        case 1: // Toggle OCy (only on OC pin 0 with WGM2 bit set)
            if ocPinNum == 0 && (t.controlB & 0x80) != 0 {
                t.toggleOCPin(ocPinNum)
            }
        case 2: // Clear OCy if counting upwards or set OCy if counting downwards
            if t.downwards {
                t.setOCPin(ocPinNum)
            } else {
                t.clearOCPin(ocPinNum)
            }
        case 3: // Set OCy if counting upwards or clear OCy if counting downwards
            if t.downwards {
                t.clearOCPin(ocPinNum)
            } else {
                t.setOCPin(ocPinNum)
            }
        }
    }
}

// Tick the timer in clear-timer-on-compare mode.
func (t *Timer) tickCTCMode() {
    t.downwards = false

    if t.count == 0xFF { // Overflow
        t.setTOV()
    }

    if t.count == t.compareValA {
        // this tick should set counter to 0
        t.count = 0xFF
    }
}

// Tick the timer in fast PWM mode.
func (t *Timer) tickFastPWMMode(top uint8) {
    t.downwards = false

    if t.count == top {
        // this tick should set counter to 0
        t.count = 0xFF
        t.setTOV()
    }

    t.checkOCPinFastPWMMode(0, t.compareValA)
    t.checkOCPinFastPWMMode(1, t.compareValB)
}

// Check for output-compare in fast PWM mode.
func (t *Timer) checkOCPinFastPWMMode(ocPinNum uint, compareVal uint8) {
    // BOTTOM
    if t.count == 0x00 {
        // Get COMxy bits
        shiftAmt := 6 - 2*ocPinNum // 0 => 6, 1 => 4
        com := (t.controlA >> shiftAmt) & 0x03
        switch com {
        case 0: // OCy disabled
            // do nothing
        case 1: // Toggle OCy on compare match (only on OC pin 0 with WGM2 bit set)
            // do nothing
        case 2: // Clear OCy on compare match, set OCy at BOTTOM
            t.setOCPin(ocPinNum)
        case 3: // Set OCy on compare match, clear OCy at BOTTOM
            t.clearOCPin(ocPinNum)
        }
    }

    // Compare match
    if t.count == compareVal {
        // Get COMxy bits
        shiftAmt := 6 - 2*ocPinNum // 0 => 6, 1 => 4
        com := (t.controlA >> shiftAmt) & 0x03
        switch com {
        case 0: // OCy disabled
            // do nothing
        case 1: // Toggle OCy on compare match (only on OC pin 0 with WGM2 bit set)
            if ocPinNum == 0 && (t.controlB & 0x80) != 0 {
                t.toggleOCPin(ocPinNum)
            }
        case 2: // Clear OCy on compare match, set OCy at BOTTOM
            t.clearOCPin(ocPinNum)
        case 3: // Set OCy on compare match, clear OCy at BOTTOM
            t.setOCPin(ocPinNum)
        }
    }
}

// Toggle an output-compare pin.
func (t *Timer) toggleOCPin(ocPinNum uint) {
    t.ocPinStates[ocPinNum] = !t.ocPinStates[ocPinNum]
    t.updateOCPin(ocPinNum)
}

// Set an output-compare pin to low.
func (t *Timer) clearOCPin(ocPinNum uint) {
    if t.ocPinStates[ocPinNum] {
        t.ocPinStates[ocPinNum] = false
        t.updateOCPin(ocPinNum)
    }
}

// Set an output-compare pin to high.
func (t *Timer) setOCPin(ocPinNum uint) {
    if !t.ocPinStates[ocPinNum] {
        t.ocPinStates[ocPinNum] = true
        t.updateOCPin(ocPinNum)
    }
}

// Push the new status of an output-compare pin to the GPIO layer.
func (t *Timer) updateOCPin(ocPinNum uint) {
    callback := t.ocPinCallbacks[ocPinNum]
    if callback != nil {
        callback(t.ocPinStates[ocPinNum])
    }
}

// Set an output compare match (OCFy) flag.
func (t *Timer) setOCF(ocPinNum uint) {
    if ocPinNum == 0 { // A
        t.interruptFlags |= 0x02
    } else { // B
        t.interruptFlags |= 0x04
    }
}

// Set the timer overflow (TOV) flag.
func (t *Timer) setTOV() {
    t.interruptFlags |= 0x01
}
<|endoftext|>