<commit_before>package api

import (
	"encoding/json"
	"log"
	"net/http"
	"sort"
	"strconv"
	"time"

	"github.com/willf/bloom"

	"github.com/brnstz/bus/internal/etc"
	"github.com/brnstz/bus/internal/models"
	"github.com/brnstz/bus/internal/partners"
)

var (
	// stopChan is a channel for receiving requests to get live departure
	// data
	stopChan chan *stopLiveRequest

	// workers is the number of workers processing requestChan concurrently
	stopWorkers = 10

	// Formula for determining m and k values: http://hur.st/bloomfilter
	// n = approx number of items to insert
	// p = desired false positive rate (between 0 and 1)
	// m = ceil((n * log(p)) / log(1.0 / (pow(2.0, log(2.0)))))
	// k = round(log(2.0) * m / n)
	// with n = 300 and p = 0.001
	bloomM uint = 4314
	bloomK uint = 10

	minFirstDepartureTime = time.Duration(2) * time.Hour
)

type stopLiveRequest struct {
	route    *models.Route
	stop     *models.Stop
	partner  partners.P
	response chan error
}

func init() {
	stopChan = make(chan *stopLiveRequest, 100000)

	for i := 0; i < stopWorkers; i++ {
		go stopWorker()
	}
}

// stop worker calls the partner's live departure API and sets
// req.stop.Live
func stopWorker() {
	for req := range stopChan {
		liveDepartures, liveVehicles, err := req.partner.Live(req.route.AgencyID, req.route.RouteID, req.stop.StopID, req.stop.DirectionID)
		if err != nil {
			req.response <- err
			continue
		}

		if len(liveVehicles) > 0 {
			req.stop.Vehicles = liveVehicles
		}

		sd := models.SortableDepartures(liveDepartures)
		sort.Sort(sd)
		liveDepartures = []*models.Departure(sd)

		if len(liveDepartures) > 0 {
			liveTripIDs := map[string]bool{}

			// Remove any of the same trip ids that appear in scheduled
			// departures. Live info is better for that trip, but there
			// might still be scheduled departures later we want to use.
			for _, d := range liveDepartures {
				liveTripIDs[d.TripID] = true
			}

			// If there are less than max departures, then add scheduled
			// departures that are after our last live departure and
			// don't have dupe trip IDs
			count := len(liveDepartures)
			lastLiveDeparture := liveDepartures[count-1]

			i := -1
			for {
				i++

				// Stop once we have enough departures
				if count >= models.MaxDepartures {
					break
				}

				// Stop if we reach the end of the scheduled departures
				if i >= len(req.stop.Departures) {
					break
				}

				// Ignore departures with trip IDs that we know of
				if liveTripIDs[req.stop.Departures[i].TripID] {
					continue
				}

				if req.stop.Departures[i].Time.After(lastLiveDeparture.Time) {
					liveDepartures = append(liveDepartures, req.stop.Departures[i])
				}

			}

			if len(liveDepartures) > 5 {
				req.stop.Departures = liveDepartures[0:5]
			} else {
				req.stop.Departures = liveDepartures
			}

		}

		req.response <- nil
	}
}

// hereResponse is the value returned by getHere
type hereResponse struct {
	Stops  []*models.Stop     `json:"stops"`
	Routes []*models.Route    `json:"routes"`
	Trips  []*models.Trip     `json:"trips"`
	Filter *bloom.BloomFilter `json:"filter"`
}

func getHere(w http.ResponseWriter, r *http.Request) {
	var err error
	var resp hereResponse
	var routes []*models.Route
	var now time.Time

	if len(r.FormValue("now")) > 0 {
		now, err = time.ParseInLocation("2006-01-02 15:04:05", r.FormValue("now"), time.Local)
		if err != nil {
			log.Println("can't parse time", err)
			apiErr(w, errBadRequest)
			return
		}
	} else {
		now = time.Now()
	}

	// Read values incoming from http request
	lat, err := floatOrDie(r.FormValue("lat"))
	if err != nil {
		apiErr(w, err)
		return
	}

	lon, err := floatOrDie(r.FormValue("lon"))
	if err != nil {
		apiErr(w, err)
		return
	}

	SWLat, err := floatOrDie(r.FormValue("sw_lat"))
	if err != nil {
		apiErr(w, err)
		return
	}

	SWLon, err := floatOrDie(r.FormValue("sw_lon"))
	if err != nil {
		apiErr(w, err)
		return
	}

	NELat, err := floatOrDie(r.FormValue("ne_lat"))
	if err != nil {
		apiErr(w, err)
		return
	}

	NELon, err := floatOrDie(r.FormValue("ne_lon"))
	if err != nil {
		apiErr(w, err)
		return
	}

	// Initialize or read incoming bloom filter
	filter := r.FormValue("filter")

	var routeTypes []int
	for _, v := range r.Form["route_type"] {
		intv, err := strconv.Atoi(v)
		if err != nil {
			apiErr(w, errBadRequest)
		}

		routeTypes = append(routeTypes, intv)
	}

	if len(filter) < 1 {
		// If there is no filter, then create a new one
		resp.Filter = bloom.New(bloomM, bloomK)

	} else {
		resp.Filter = &bloom.BloomFilter{}
		// Otherwise read the passed value as JSON string
		err = json.Unmarshal([]byte(filter), resp.Filter)
		if err != nil {
			log.Println("can't read incoming bloom filter JSON", err)
			apiErr(w, errBadRequest)
			return
		}
	}

	t1 := time.Now()
	hq, err := models.NewHereQuery(
		lat, lon, SWLat, SWLon, NELat, NELon, routeTypes, now,
	)
	if err != nil {
		log.Println("can't create here query", err)
		apiErr(w, err)
		return
	}

	stops, stopRoutes, err := models.GetHereResults(etc.DBConn, hq)
	if err != nil {
		log.Println("can't get here results", err)
		apiErr(w, err)
		return
	}
	log.Println("here query:        ", time.Now().Sub(t1))

	t2 := time.Now()
	// Create a channel for receiving responses to stopLiveRequest values
	respch := make(chan error, len(stops))
	count := 0

	// save the first scheduled departure of each stop, so that we can
	// use it in case the live tripID cannot be found
	firstDepart := map[string]*models.Departure{}

	for _, s := range stops {

		firstDepart[s.UniqueID] = s.Departures[0]
		route := stopRoutes[s.UniqueID]
		routes = append(routes, route)

		// Get a live partner or skip it
		partner, err := partners.Find(*route)
		if err != nil {
			log.Println(err)
			continue
		}

		// Create a request to get live info and send it on the channel
		req := &stopLiveRequest{
			route:    route,
			stop:     s,
			partner:  partner,
			response: respch,
		}
		stopChan <- req
		count++
	}

	// Wait for all responses
	for i := 0; i < count; i++ {
		err = <-respch
		if err != nil {
			log.Println(err)
		}
	}
	log.Println("partner info:      ", time.Now().Sub(t2))

	t3 := time.Now()
	// Add any routes to the response that the bloom filter says we don't have
	for _, route := range routes {
		exists := resp.Filter.TestString(route.UniqueID)
		// If the route doesn't exist in our filter, then we want to pull
		// the shapes and also append it to our response list.
		if !exists {
			route.RouteShapes, err = models.GetSavedRouteShapes(
				etc.DBConn, route.AgencyID, route.RouteID,
			)
			if err != nil {
				// This is a fatal error because the front end code
				// assumes the route will be there
				log.Println("can't get route shapes", route, err)
				apiErr(w, err)
				return
			}

			resp.Filter.AddString(route.UniqueID)
			resp.Routes = append(resp.Routes, route)
		}
	}
	log.Println("route shape:       ", time.Now().Sub(t3))

	t4 := time.Now()

	// Set stop value of the response
	resp.Stops = stops

	// Add the first trip of each stop response that is not already in our
	// bloom filter
	for i, stop := range resp.Stops {
		var trip models.Trip

		if len(stop.Departures) < 1 {
			continue
		}

		if now.Add(minFirstDepartureTime).Before(stop.Departures[0].Time) {
			continue
		}

		// Get info for the trip
		tripID := stop.Departures[0].TripID
		uniqueID := stop.AgencyID + "|" + tripID

		// Check if the trip already exists
		exists := resp.Filter.TestString(uniqueID)

		// If it exists, skip it
		if exists {
			continue
		}

		// Get the full trip with stop and shape details. If we succeed, we can
		// move onto next trip
		trip, err = models.GetTrip(etc.DBConn, stop.AgencyID, stop.RouteID, tripID)
		if err == nil {
			resp.Filter.AddString(uniqueID)
			resp.Trips = append(resp.Trips, &trip)
			continue
		}

		// If the error is unexpected, we should error out immediately
		if err != models.ErrNotFound {
			log.Println("can't get trip", err)
			apiErr(w, err)
			return
		}

		// Here we weren't able to find the trip ID in the database. This is
		// typically due to a response from a realtime source which gives us
		// TripIDs that are not in the static feed or are partial matches.
		// Let's first look for a partial match. If that fails, let's just get
		// the use the first scheduled departure instead.

		// Checking for partial match.
		tripID, err = models.GetPartialTripIDMatch(
			etc.DBConn, stop.AgencyID, stop.RouteID, tripID,
		)

		// If we get one, then update the uniqueID and the relevant stop /
		// departure's ID, adding it to our filter.
		if err == nil {
			uniqueID = stop.AgencyID + "|" + tripID
			resp.Stops[i].Departures[0].TripID = tripID
			resp.Stops[i].Initialize()

			// Re-get the trip with update ID
			trip, err = models.GetTrip(etc.DBConn, stop.AgencyID, stop.RouteID,
				tripID)
			if err != nil {
				log.Println("can't get trip", err)
				apiErr(w, err)
				return
			}

			resp.Filter.AddString(uniqueID)
			resp.Trips = append(resp.Trips, &trip)

			continue
		}

		// If the error is unexpected, we should error out immediately
		if err != models.ErrNotFound {
			log.Println("can't get trip", err)
			apiErr(w, err)
			return
		}

		// Our last hope is take the first scheduled departure
		tripID = firstDepart[stop.UniqueID].TripID

		uniqueID = stop.AgencyID + "|" + tripID
		resp.Stops[i].Departures[0].TripID = tripID
		resp.Stops[i].Initialize()

		// Re-get the trip with update ID
		trip, err = models.GetTrip(etc.DBConn, stop.AgencyID, stop.RouteID,
			tripID)
		if err != nil {
			log.Println("can't get trip", err)
			apiErr(w, err)
			return
		}

		resp.Filter.AddString(uniqueID)
		resp.Trips = append(resp.Trips, &trip)
	}
	log.Println("trips:             ", time.Now().Sub(t4))

	b, err := json.Marshal(resp)
	if err != nil {
		log.Println("can't marshal to json", err)
		apiErr(w, err)
		return
	}

	w.Write(b)
}
<commit_msg>fixing live departure compass dir<commit_after>package api

import (
	"encoding/json"
	"log"
	"net/http"
	"sort"
	"strconv"
	"time"

	"github.com/willf/bloom"

	"github.com/brnstz/bus/internal/etc"
	"github.com/brnstz/bus/internal/models"
	"github.com/brnstz/bus/internal/partners"
)

var (
	// stopChan is a channel for receiving requests to get live departure
	// data
	stopChan chan *stopLiveRequest

	// workers is the number of workers processing requestChan concurrently
	stopWorkers = 10

	// Formula for determining m and k values: http://hur.st/bloomfilter
	// n = approx number of items to insert
	// p = desired false positive rate (between 0 and 1)
	// m = ceil((n * log(p)) / log(1.0 / (pow(2.0, log(2.0)))))
	// k = round(log(2.0) * m / n)
	// with n = 300 and p = 0.001
	bloomM uint = 4314
	bloomK uint = 10

	minFirstDepartureTime = time.Duration(2) * time.Hour
)

type stopLiveRequest struct {
	route    *models.Route
	stop     *models.Stop
	partner  partners.P
	response chan error
}

func init() {
	stopChan = make(chan *stopLiveRequest, 100000)

	for i := 0; i < stopWorkers; i++ {
		go stopWorker()
	}
}

// stop worker calls the partner's live departure API and sets
// req.stop.Live
func stopWorker() {
	for req := range stopChan {
		liveDepartures, liveVehicles, err := req.partner.Live(req.route.AgencyID, req.route.RouteID, req.stop.StopID, req.stop.DirectionID)
		if err != nil {
			req.response <- err
			continue
		}

		if len(liveVehicles) > 0 {
			req.stop.Vehicles = liveVehicles
		}

		// FIXME: assume compass dir for live departures is
		// the first scheduled departure's dir
		compassDir := req.stop.Departures[0].CompassDir

		sd := models.SortableDepartures(liveDepartures)
		sort.Sort(sd)
		liveDepartures = []*models.Departure(sd)

		if len(liveDepartures) > 0 {
			liveTripIDs := map[string]bool{}

			// Remove any of the same trip ids that appear in scheduled
			// departures. Live info is better for that trip, but there
			// might still be scheduled departures later we want to use.
			for _, d := range liveDepartures {
				liveTripIDs[d.TripID] = true
			}

			// If there are less than max departures, then add scheduled
			// departures that are after our last live departure and
			// don't have dupe trip IDs
			count := len(liveDepartures)
			lastLiveDeparture := liveDepartures[count-1]

			i := -1
			for {
				i++

				// Stop once we have enough departures
				if count >= models.MaxDepartures {
					break
				}

				// Stop if we reach the end of the scheduled departures
				if i >= len(req.stop.Departures) {
					break
				}

				// Ignore departures with trip IDs that we know of
				if liveTripIDs[req.stop.Departures[i].TripID] {
					continue
				}

				if req.stop.Departures[i].Time.After(lastLiveDeparture.Time) {
					liveDepartures = append(liveDepartures, req.stop.Departures[i])
				}

			}

			for i := 0; i < models.MaxDepartures && i < len(liveDepartures); i++ {
				liveDepartures[i].CompassDir = compassDir
				req.stop.Departures[i] = liveDepartures[i]

			}

		}

		req.response <- nil
	}
}

// hereResponse is the value returned by getHere
type hereResponse struct {
	Stops  []*models.Stop     `json:"stops"`
	Routes []*models.Route    `json:"routes"`
	Trips  []*models.Trip     `json:"trips"`
	Filter *bloom.BloomFilter `json:"filter"`
}

func getHere(w http.ResponseWriter, r *http.Request) {
	var err error
	var resp hereResponse
	var routes []*models.Route
	var now time.Time

	if len(r.FormValue("now")) > 0 {
		now, err = time.ParseInLocation("2006-01-02 15:04:05", r.FormValue("now"), time.Local)
		if err != nil {
			log.Println("can't parse time", err)
			apiErr(w, errBadRequest)
			return
		}
	} else {
		now = time.Now()
	}

	// Read values incoming from http request
	lat, err := floatOrDie(r.FormValue("lat"))
	if err != nil {
		apiErr(w, err)
		return
	}

	lon, err := floatOrDie(r.FormValue("lon"))
	if err != nil {
		apiErr(w, err)
		return
	}

	SWLat, err := floatOrDie(r.FormValue("sw_lat"))
	if err != nil {
		apiErr(w, err)
		return
	}

	SWLon, err := floatOrDie(r.FormValue("sw_lon"))
	if err != nil {
		apiErr(w, err)
		return
	}

	NELat, err := floatOrDie(r.FormValue("ne_lat"))
	if err != nil {
		apiErr(w, err)
		return
	}

	NELon, err := floatOrDie(r.FormValue("ne_lon"))
	if err != nil {
		apiErr(w, err)
		return
	}

	// Initialize or read incoming bloom filter
	filter := r.FormValue("filter")

	var routeTypes []int
	for _, v := range r.Form["route_type"] {
		intv, err := strconv.Atoi(v)
		if err != nil {
			apiErr(w, errBadRequest)
		}

		routeTypes = append(routeTypes, intv)
	}

	if len(filter) < 1 {
		// If there is no filter, then create a new one
		resp.Filter = bloom.New(bloomM, bloomK)

	} else {
		resp.Filter = &bloom.BloomFilter{}
		// Otherwise read the passed value as JSON string
		err = json.Unmarshal([]byte(filter), resp.Filter)
		if err != nil {
			log.Println("can't read incoming bloom filter JSON", err)
			apiErr(w, errBadRequest)
			return
		}
	}

	t1 := time.Now()
	hq, err := models.NewHereQuery(
		lat, lon, SWLat, SWLon, NELat, NELon, routeTypes, now,
	)
	if err != nil {
		log.Println("can't create here query", err)
		apiErr(w, err)
		return
	}

	stops, stopRoutes, err := models.GetHereResults(etc.DBConn, hq)
	if err != nil {
		log.Println("can't get here results", err)
		apiErr(w, err)
		return
	}
	log.Println("here query:        ", time.Now().Sub(t1))

	t2 := time.Now()
	// Create a channel for receiving responses to stopLiveRequest values
	respch := make(chan error, len(stops))
	count := 0

	// save the first scheduled departure of each stop, so that we can
	// use it in case the live tripID cannot be found
	firstDepart := map[string]*models.Departure{}

	for _, s := range stops {

		firstDepart[s.UniqueID] = s.Departures[0]
		route := stopRoutes[s.UniqueID]
		routes = append(routes, route)

		// Get a live partner or skip it
		partner, err := partners.Find(*route)
		if err != nil {
			log.Println(err)
			continue
		}

		// Create a request to get live info and send it on the channel
		req := &stopLiveRequest{
			route:    route,
			stop:     s,
			partner:  partner,
			response: respch,
		}
		stopChan <- req
		count++
	}

	// Wait for all responses
	for i := 0; i < count; i++ {
		err = <-respch
		if err != nil {
			log.Println(err)
		}
	}
	log.Println("partner info:      ", time.Now().Sub(t2))

	t3 := time.Now()
	// Add any routes to the response that the bloom filter says we don't have
	for _, route := range routes {
		exists := resp.Filter.TestString(route.UniqueID)
		// If the route doesn't exist in our filter, then we want to pull
		// the shapes and also append it to our response list.
		if !exists {
			route.RouteShapes, err = models.GetSavedRouteShapes(
				etc.DBConn, route.AgencyID, route.RouteID,
			)
			if err != nil {
				// This is a fatal error because the front end code
				// assumes the route will be there
				log.Println("can't get route shapes", route, err)
				apiErr(w, err)
				return
			}

			resp.Filter.AddString(route.UniqueID)
			resp.Routes = append(resp.Routes, route)
		}
	}
	log.Println("route shape:       ", time.Now().Sub(t3))

	t4 := time.Now()

	// Set stop value of the response
	resp.Stops = stops

	// Add the first trip of each stop response that is not already in our
	// bloom filter
	for i, stop := range resp.Stops {
		var trip models.Trip

		if len(stop.Departures) < 1 {
			continue
		}

		if now.Add(minFirstDepartureTime).Before(stop.Departures[0].Time) {
			continue
		}

		// Get info for the trip
		tripID := stop.Departures[0].TripID
		uniqueID := stop.AgencyID + "|" + tripID

		// Check if the trip already exists
		exists := resp.Filter.TestString(uniqueID)

		// If it exists, skip it
		if exists {
			continue
		}

		// Get the full trip with stop and shape details. If we succeed, we can
		// move onto next trip
		trip, err = models.GetTrip(etc.DBConn, stop.AgencyID, stop.RouteID, tripID)
		if err == nil {
			resp.Filter.AddString(uniqueID)
			resp.Trips = append(resp.Trips, &trip)
			continue
		}

		// If the error is unexpected, we should error out immediately
		if err != models.ErrNotFound {
			log.Println("can't get trip", err)
			apiErr(w, err)
			return
		}

		// Here we weren't able to find the trip ID in the database. This is
		// typically due to a response from a realtime source which gives us
		// TripIDs that are not in the static feed or are partial matches.
		// Let's first look for a partial match. If that fails, let's just get
		// the use the first scheduled departure instead.

		// Checking for partial match.
		tripID, err = models.GetPartialTripIDMatch(
			etc.DBConn, stop.AgencyID, stop.RouteID, tripID,
		)

		// If we get one, then update the uniqueID and the relevant stop /
		// departure's ID, adding it to our filter.
		if err == nil {
			uniqueID = stop.AgencyID + "|" + tripID
			resp.Stops[i].Departures[0].TripID = tripID
			resp.Stops[i].Initialize()

			// Re-get the trip with update ID
			trip, err = models.GetTrip(etc.DBConn, stop.AgencyID, stop.RouteID,
				tripID)
			if err != nil {
				log.Println("can't get trip", err)
				apiErr(w, err)
				return
			}

			resp.Filter.AddString(uniqueID)
			resp.Trips = append(resp.Trips, &trip)

			continue
		}

		// If the error is unexpected, we should error out immediately
		if err != models.ErrNotFound {
			log.Println("can't get trip", err)
			apiErr(w, err)
			return
		}

		// Our last hope is take the first scheduled departure
		tripID = firstDepart[stop.UniqueID].TripID

		uniqueID = stop.AgencyID + "|" + tripID
		resp.Stops[i].Departures[0].TripID = tripID
		resp.Stops[i].Initialize()

		// Re-get the trip with update ID
		trip, err = models.GetTrip(etc.DBConn, stop.AgencyID, stop.RouteID,
			tripID)
		if err != nil {
			log.Println("can't get trip", err)
			apiErr(w, err)
			return
		}

		resp.Filter.AddString(uniqueID)
		resp.Trips = append(resp.Trips, &trip)
	}
	log.Println("trips:             ", time.Now().Sub(t4))

	b, err := json.Marshal(resp)
	if err != nil {
		log.Println("can't marshal to json", err)
		apiErr(w, err)
		return
	}

	w.Write(b)
}
<|endoftext|>


<commit_before>// Copyright 2016 Circonus, Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package api

import (
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
)

func callServer() *httptest.Server {
	f := func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintln(w, r.Method)
	}

	return httptest.NewServer(http.HandlerFunc(f))
}

// func retryServer() *httptest.Server {
// 	f := func(w http.ResponseWriter, r *http.Request) {
// 		w.WriteHeader(500)
// 		w.Header().Set("Content-Type", "application/json")
// 		fmt.Fprintln(w, "blah blah blah")
// 	}
//
// 	return httptest.NewServer(http.HandlerFunc(f))
// }

func TestNew(t *testing.T) {
	var expectedError error

	t.Log("Testing correct error return when no API config supplied")
	expectedError = errors.New("Invalid API configuration (nil)")
	_, err := New(nil)
	if err == nil {
		t.Error("Expected an error")
	}
	if err.Error() != expectedError.Error() {
		t.Errorf("Expected an '%#v' error, got '%#v'", expectedError, err)
	}

	t.Log("Testing correct error return when no API Token supplied")
	expectedError = errors.New("API Token is required")
	ac := &Config{}
	_, err = New(ac)
	if err == nil {
		t.Error("Expected an error")
	}
	if err.Error() != expectedError.Error() {
		t.Errorf("Expected an '%#v' error, got '%#v'", expectedError, err)
	}

	t.Log("Testing correct return when an API Token is supplied")
	ac = &Config{
		TokenKey: "abc123",
	}
	_, err = New(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	t.Log("Testing correct return when an API Token and App are supplied")
	ac = &Config{
		TokenKey: "abc123",
		TokenApp: "someapp",
	}
	_, err = New(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	t.Log("Testing correct return when an API Token, App, and URL (host) are supplied")
	ac = &Config{
		TokenKey: "abc123",
		TokenApp: "someapp",
		URL:      "something.somewhere.com",
	}
	_, err = New(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	t.Log("Testing correct return when an API Token, App, and URL (w/trailing '/') are supplied")
	ac = &Config{
		TokenKey: "abc123",
		TokenApp: "someapp",
		URL:      "something.somewhere.com/somepath/",
	}
	_, err = New(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	t.Log("Testing correct return when an API Token, App, and [invalid] URL are supplied")
	expectedError = errors.New("parse http://something.somewhere.com\\somepath$: invalid character \"\\\\\" in host name")
	ac = &Config{
		TokenKey: "abc123",
		TokenApp: "someapp",
		URL:      "http://something.somewhere.com\\somepath$",
	}
	_, err = New(ac)
	if err == nil {
		t.Error("Expected an error")
	}
	if err.Error() != expectedError.Error() {
		t.Errorf("Expected an '%#v' error, got '%#v'", expectedError, err)
	}

	t.Log("Testing correct return when an Debug true but no log.Logger are supplied")
	ac = &Config{
		TokenKey: "abc123",
		Debug:    true,
	}
	_, err = New(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}
}

func TestApiCall(t *testing.T) {
	server := callServer()
	defer server.Close()

	ac := &Config{
		TokenKey: "foo",
		TokenApp: "bar",
		URL:      server.URL,
	}

	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%+v'", err)
	}

	t.Log("Testing invalid URL path")
	{
		_, err := apih.apiCall("GET", "", nil)
		expectedError := errors.New("Invalid URL path")
		if err == nil {
			t.Errorf("Expected error")
		}
		if err.Error() != expectedError.Error() {
			t.Errorf("Expected %+v go '%+v'", expectedError, err)
		}
	}

	t.Log("Testing URL path fixup, prefix '/'")
	{
		call := "GET"
		resp, err := apih.apiCall(call, "nothing", nil)
		if err != nil {
			t.Errorf("Expected no error, got '%+v'", resp)
		}
		expected := fmt.Sprintf("%s\n", call)
		if string(resp) != expected {
			t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
		}
	}

	t.Log("Testing URL path fixup, remove '/v2' prefix")
	{
		call := "GET"
		resp, err := apih.apiCall(call, "/v2/nothing", nil)
		if err != nil {
			t.Errorf("Expected no error, got '%+v'", resp)
		}
		expected := fmt.Sprintf("%s\n", call)
		if string(resp) != expected {
			t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
		}
	}

	calls := []string{"GET", "PUT", "POST", "DELETE"}
	for _, call := range calls {
		t.Logf("Testing %s call", call)
		resp, err := apih.apiCall(call, "/", nil)
		if err != nil {
			t.Errorf("Expected no error, got '%+v'", resp)
		}

		expected := fmt.Sprintf("%s\n", call)
		if string(resp) != expected {
			t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
		}
	}

}

func TestApiGet(t *testing.T) {
	server := callServer()
	defer server.Close()

	ac := &Config{
		TokenKey: "foo",
		TokenApp: "bar",
		URL:      server.URL,
	}

	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%+v'", err)
	}

	resp, err := apih.Get("/")

	if err != nil {
		t.Errorf("Expected no error, got '%+v'", resp)
	}

	expected := "GET\n"
	if string(resp) != expected {
		t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
	}

}

func TestApiPut(t *testing.T) {
	server := callServer()
	defer server.Close()

	ac := &Config{
		TokenKey: "foo",
		TokenApp: "bar",
		URL:      server.URL,
	}

	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%+v'", err)
	}

	resp, err := apih.Put("/", nil)

	if err != nil {
		t.Errorf("Expected no error, got '%+v'", resp)
	}

	expected := "PUT\n"
	if string(resp) != expected {
		t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
	}

}

func TestApiPost(t *testing.T) {
	server := callServer()
	defer server.Close()

	ac := &Config{
		TokenKey: "foo",
		TokenApp: "bar",
		URL:      server.URL,
	}

	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%+v'", err)
	}

	resp, err := apih.Post("/", nil)

	if err != nil {
		t.Errorf("Expected no error, got '%+v'", resp)
	}

	expected := "POST\n"
	if string(resp) != expected {
		t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
	}

}

func TestApiDelete(t *testing.T) {
	server := callServer()
	defer server.Close()

	ac := &Config{
		TokenKey: "foo",
		TokenApp: "bar",
		URL:      server.URL,
	}

	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%+v'", err)
	}

	resp, err := apih.Delete("/")

	if err != nil {
		t.Errorf("Expected no error, got '%+v'", resp)
	}

	expected := "DELETE\n"
	if string(resp) != expected {
		t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
	}

}
<commit_msg>upd: test log messages<commit_after>// Copyright 2016 Circonus, Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package api

import (
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
)

func callServer() *httptest.Server {
	f := func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintln(w, r.Method)
	}

	return httptest.NewServer(http.HandlerFunc(f))
}

func TestNew(t *testing.T) {
	t.Log("invalid config [nil]")
	{
		expectedError := errors.New("Invalid API configuration (nil)")
		_, err := New(nil)
		if err == nil {
			t.Error("Expected an error")
		}
		if err.Error() != expectedError.Error() {
			t.Errorf("Expected an '%#v' error, got '%#v'", expectedError, err)
		}
	}

	t.Log("invalid config [blank]")
	{
		expectedError := errors.New("API Token is required")
		ac := &Config{}
		_, err := New(ac)
		if err == nil {
			t.Error("Expected an error")
		}
		if err.Error() != expectedError.Error() {
			t.Errorf("Expected an '%#v' error, got '%#v'", expectedError, err)
		}
	}

	t.Log("API Token, no API App, no API URL")
	{
		ac := &Config{
			TokenKey: "abc123",
		}
		_, err := New(ac)
		if err != nil {
			t.Errorf("Expected no error, got '%v'", err)
		}
	}

	t.Log("API Token, API App, no API URL")
	{
		ac := &Config{
			TokenKey: "abc123",
			TokenApp: "someapp",
		}
		_, err := New(ac)
		if err != nil {
			t.Errorf("Expected no error, got '%v'", err)
		}
	}

	t.Log("API Token, API App, API URL [host]")
	{
		ac := &Config{
			TokenKey: "abc123",
			TokenApp: "someapp",
			URL:      "something.somewhere.com",
		}
		_, err := New(ac)
		if err != nil {
			t.Errorf("Expected no error, got '%v'", err)
		}
	}

	t.Log("API Token, API App, API URL [trailing '/']")
	{
		ac := &Config{
			TokenKey: "abc123",
			TokenApp: "someapp",
			URL:      "something.somewhere.com/somepath/",
		}
		_, err := New(ac)
		if err != nil {
			t.Errorf("Expected no error, got '%v'", err)
		}
	}

	t.Log("API Token, API App, API URL [w/o trailing '/']")
	{
		ac := &Config{
			TokenKey: "abc123",
			TokenApp: "someapp",
			URL:      "something.somewhere.com/somepath",
		}
		_, err := New(ac)
		if err != nil {
			t.Errorf("Expected no error, got '%v'", err)
		}
	}

	t.Log("API Token, API App, API URL [invalid]")
	{
		expectedError := errors.New("parse http://something.somewhere.com\\somepath$: invalid character \"\\\\\" in host name")
		ac := &Config{
			TokenKey: "abc123",
			TokenApp: "someapp",
			URL:      "http://something.somewhere.com\\somepath$",
		}
		_, err := New(ac)
		if err == nil {
			t.Error("Expected an error")
		}
		if err.Error() != expectedError.Error() {
			t.Errorf("Expected an '%#v' error, got '%#v'", expectedError, err)
		}
	}

	t.Log("Debug true, no log.Logger")
	{
		ac := &Config{
			TokenKey: "abc123",
			Debug:    true,
		}
		_, err := New(ac)
		if err != nil {
			t.Errorf("Expected no error, got '%v'", err)
		}
	}
}

func TestApiCall(t *testing.T) {
	server := callServer()
	defer server.Close()

	ac := &Config{
		TokenKey: "foo",
		TokenApp: "bar",
		URL:      server.URL,
	}

	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%+v'", err)
	}

	t.Log("invalid URL path")
	{
		_, err := apih.apiCall("GET", "", nil)
		expectedError := errors.New("Invalid URL path")
		if err == nil {
			t.Errorf("Expected error")
		}
		if err.Error() != expectedError.Error() {
			t.Errorf("Expected %+v go '%+v'", expectedError, err)
		}
	}

	t.Log("URL path fixup, prefix '/'")
	{
		call := "GET"
		resp, err := apih.apiCall(call, "nothing", nil)
		if err != nil {
			t.Errorf("Expected no error, got '%+v'", resp)
		}
		expected := fmt.Sprintf("%s\n", call)
		if string(resp) != expected {
			t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
		}
	}

	t.Log("URL path fixup, remove '/v2' prefix")
	{
		call := "GET"
		resp, err := apih.apiCall(call, "/v2/nothing", nil)
		if err != nil {
			t.Errorf("Expected no error, got '%+v'", resp)
		}
		expected := fmt.Sprintf("%s\n", call)
		if string(resp) != expected {
			t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
		}
	}

	calls := []string{"GET", "PUT", "POST", "DELETE"}
	for _, call := range calls {
		t.Logf("Testing %s call", call)
		resp, err := apih.apiCall(call, "/", nil)
		if err != nil {
			t.Errorf("Expected no error, got '%+v'", resp)
		}

		expected := fmt.Sprintf("%s\n", call)
		if string(resp) != expected {
			t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
		}
	}

}

func TestApiGet(t *testing.T) {
	server := callServer()
	defer server.Close()

	ac := &Config{
		TokenKey: "foo",
		TokenApp: "bar",
		URL:      server.URL,
	}

	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%+v'", err)
	}

	resp, err := apih.Get("/")

	if err != nil {
		t.Errorf("Expected no error, got '%+v'", resp)
	}

	expected := "GET\n"
	if string(resp) != expected {
		t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
	}

}

func TestApiPut(t *testing.T) {
	server := callServer()
	defer server.Close()

	ac := &Config{
		TokenKey: "foo",
		TokenApp: "bar",
		URL:      server.URL,
	}

	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%+v'", err)
	}

	resp, err := apih.Put("/", nil)

	if err != nil {
		t.Errorf("Expected no error, got '%+v'", resp)
	}

	expected := "PUT\n"
	if string(resp) != expected {
		t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
	}

}

func TestApiPost(t *testing.T) {
	server := callServer()
	defer server.Close()

	ac := &Config{
		TokenKey: "foo",
		TokenApp: "bar",
		URL:      server.URL,
	}

	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%+v'", err)
	}

	resp, err := apih.Post("/", nil)

	if err != nil {
		t.Errorf("Expected no error, got '%+v'", resp)
	}

	expected := "POST\n"
	if string(resp) != expected {
		t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
	}

}

func TestApiDelete(t *testing.T) {
	server := callServer()
	defer server.Close()

	ac := &Config{
		TokenKey: "foo",
		TokenApp: "bar",
		URL:      server.URL,
	}

	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%+v'", err)
	}

	resp, err := apih.Delete("/")

	if err != nil {
		t.Errorf("Expected no error, got '%+v'", resp)
	}

	expected := "DELETE\n"
	if string(resp) != expected {
		t.Errorf("Expected\n'%s'\ngot\n'%s'\n", expected, resp)
	}

}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp 2016 All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

		 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"errors"
	"fmt"
	"encoding/json"
	"github.com/hyperledger/fabric/core/chaincode/shim"
)

// SimpleChaincode example simple Chaincode implementation
type SimpleChaincode struct {

}

// Structure of Patient Details : Custom block

type Patient struct {
	Id    				string  `json:"Id"`
	Name         			string  `json:"Name"`
	Ailment 			string  `json:"Ailment"`
	DateOfBirth 			string  `json:"DateOfBirth"`
	NameOfLab       		string  `json:"NameOfLab"`
	ReportType 			string  `json:"ReportType"`
	Date				string  `json:"Date"`
	Impression 			string  `json:"Impression"`
	Finding 			string  `json:"Finding"`
	Disease 			string  `json:"Disease"`
	OnGoingMedication 		string  `json:"OnGoingMedication"`
	Duration 			string  `json:"Duration"`
	CurrentProblemDescription	string  `json:"CurrentProblemDescription"`
}

func main() {
	err := shim.Start(new(SimpleChaincode))
	if err != nil {
		fmt.Printf("Error starting Simple chaincode: %s", err)
	}
}

// Init resets all the things
func (t *SimpleChaincode) Init(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	if len(args) != 1 {
		return nil, errors.New("Incorrect number of arguments. Expecting 1")
	}

	err := stub.PutState("hello_blockchain", []byte(args[0]))
	if err != nil {
		return nil, err
	}

	return nil, nil
}

// Invoke isur entry point to invoke a chaincode function
func (t *SimpleChaincode) Invoke(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	fmt.Println("invoke is running " + function)

	// Handle different functions
	if function == "init" {
		return t.Init(stub, "init", args)
	} else if function == "write" {
		return t.write(stub, args)
	}
	fmt.Println("invoke did not find func: " + function)

	return nil, errors.New("Received unknown function invocation: " + function)
}

// Query is our entry point for queries
func (t *SimpleChaincode) Query(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	fmt.Println("query is running " + function)

	// Handle different functions
	if function == "read" { //read a variable
		return t.read(stub, args)
	}
	fmt.Println("query did not find func: " + function)

	return nil, errors.New("Received unknown function query: " + function)
}

// write - invoke function to write key/value pair
func (t *SimpleChaincode) write(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	var err error
	fmt.Println("running write()")

	if len(args) != 13 {
		return nil, errors.New("Incorrect number of arguments. Expecting 2. name of the key and value to set")
	}

	m_patient := &Patient{}

	m_patient.Id 				= args[0]
	m_patient.Name 				= args[1]
	m_patient.Ailment			= args[2]
	m_patient.DateOfBirth			= args[3]
	m_patient.NameOfLab			= args[4]
	m_patient.ReportType			= args[5]
	m_patient.Date				= args[6]
	m_patient.Impression			= args[7]
	m_patient.Finding			= args[8]
	m_patient.Disease			= args[9]
	m_patient.OnGoingMedication		= args[10]
	m_patient.Duration			= args[11]
	m_patient.CurrentProblemDescription	= args[12]

	var key = args[0] //rename for funsies

	value, err := json.Marshal(&m_patient)

	if err != nil {
		return nil, err
	}

	err = stub.PutState(key, []byte(value)) //write the variable into the chaincode state
	if err != nil {
		return nil, err
	}
	return nil, nil
}

// read - query function to read key/value pair
func (t *SimpleChaincode) read(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	var key, jsonResp string
	var err error

	if len(args) != 1 {
		return nil, errors.New("Incorrect number of arguments. Expecting name of the key to query")
	}

	key = args[0]
	valAsbytes, err := stub.GetState(key)
	if err != nil {
		jsonResp = "{\"Error\":\"Failed to get state for " + key + "\"}"
		return nil, errors.New(jsonResp)
	}
	return valAsbytes, nil
}<commit_msg>Delete chaincode_healthCare.go<commit_after><|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"database/sql/driver"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"time"
	"unicode/utf8"

	"golang.org/x/crypto/ssh/terminal"

	_ "github.com/go-sql-driver/mysql"
)

const (
	// Amount of CSV write data to buffer between flushes
	flushSize = 26214400 // 25MB

	// Timeout length where ctrl+c is ignored
	// If a second ctrl+c is sent before the timeout the program exits
	signalTimeout = 3 // Seconds

	// Timeout length to wait for a query string sent via stdin
	stdinTimeout = 10 // Milliseconds
)

type (
	// dbInfo contains information necessary to connect to a database
	dbInfo struct {
		user string
		pass string
		host string
		port string
		TLS  bool
	}

	// NullRawBytes represents sql.RawBytes which may be null
	NullRawBytes struct {
		RawBytes sql.RawBytes
		Valid    bool
	}
)

// ShowUsage prints a help screen
func showUsage() {
	fmt.Println(`
	mycsv usage:
	mycsv DB_COMMANDS [CSV OUTPUT FLAGS] [DEBUG FLAGS] [CSV OUTFILE] query

	EXAMPLES:
	mycsv -user=jprunier -p -file=my.csv -query="select * from jjp.example_table where filter in ('1', 'test', 'another')"
	echo "select * from mysql.plugins" | -user=jprunier -pmypass -hremotedb -tls > my.csv
	mycsv -user=jprunier -p -file=my.csv -d="|" -q="'" < queryfile

	DATABASE FLAGS
	==============
	-user: Username (required)
	-pass: Password (interactive prompt if blank)
	-host: Database Host (localhost assumed if blank)
	-port: Port (3306 default)
	-tls:  Use TLS/SSL for database connection (false default)

	CSV FLAGS
	=========
	-file: CSV output filename (Write to stdout if not supplied)
	-query: MySQL query (required, can be sent via stdin redirection)
	-header: Print initial column name header line (true default)
	-d: CSV field delimiter ("," default)
	-q: CSV quote character ("\"" default)
	-e: CSV escape character ("\\" default)
	-t: CSV line terminator ("\n" default)
	-v: Print more information (false default)

	DEBUG FLAGS
	===========
	-debug_cpu: CPU debugging filename
	-debug_mem: Memory debugging filename

	`)
}

func main() {
	start := time.Now()

	// Profiling flags
	var cpuprofile = flag.String("debug_cpu", "", "CPU debugging filename")
	var memprofile = flag.String("debug_mem", "", "Memory debugging filename")

	// Database flags
	dbUser := flag.String("user", "", "Username (required)")
	dbPass := flag.String("pass", "", "Password (interactive prompt if blank)")
	dbHost := flag.String("host", "", "Database Host (localhost assumed if blank)")
	dbPort := flag.String("port", "3306", "Port")
	dbTLS := flag.Bool("tls", false, "Use TLS/SSL for database connection")

	// CSV format flags
	csvDelimiter := flag.String("d", ",", "CSV field delimiter")
	csvQuote := flag.String("q", "\"", "CSV quote character")
	csvEscape := flag.String("e", "\\", "CSV escape character")
	csvTerminator := flag.String("t", "\n", "CSV line terminator")

	csvHeader := flag.Bool("header", true, "Print initial column name header line")
	csvFile := flag.String("file", "", "CSV output filename")
	sqlQuery := flag.String("query", "", "MySQL query")
	verbose := flag.Bool("v", false, "Print more information")

	// Parse flags
	flag.Parse()

	// Print usage
	if flag.NFlag() == 0 {
		showUsage()

		os.Exit(0)
	}

	// Check if an output file was supplied otherwise use standard out
	var writeTo string
	var writerDest io.Writer
	var err error
	if *csvFile == "" {
		writeTo = "standard out"
		writerDest = os.Stdout
	} else {
		f, err := os.Open(*csvFile)
		if err == nil {
			fmt.Fprintln(os.Stderr, *csvFile, "already exists!")
			fmt.Fprintln(os.Stderr, "Please remove it or use a different filename")
			os.Exit(1)
		}
		f.Close()
		writerDest, err = os.Create(*csvFile)
		writeTo = *csvFile
	}

	if *verbose {
		fmt.Println("CSV output will be written to", writeTo)
	}

	// If query not provided read from standard in
	var query string
	queryChan := make(chan string)
	defer close(queryChan)
	if *sqlQuery == "" {
		go func() {
			b, err := ioutil.ReadAll(os.Stdin)
			checkErr(err)

			queryChan <- string(b)
		}()

		select {
		case q := <-queryChan:
			query = q
		case <-time.After(time.Millisecond * stdinTimeout):
			fmt.Fprintln(os.Stderr, "You must supply a query")
			os.Exit(1)
		}
	} else {
		query = *sqlQuery
	}

	// Check if Stdin has been redirected and reset so the user can be prompted for a password
	checkStdin()

	// Listen for SIGINT (ctrl+c)
	catchNotifications()

	// Make sure the query is a select
	if query[0:6] != "select" {
		fmt.Fprintln(os.Stderr, "Query must be a select!")
		os.Exit(1)
	}

	// CPU Profiling
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		checkErr(err)
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// Default to localhost if no host or socket provided
	if *dbHost == "" {
		*dbHost = "127.0.0.1"
	}

	// Need to provide a target
	if *dbUser == "" {
		fmt.Fprintln(os.Stderr, "You must provide a user name!")
		os.Exit(1)
	}

	// If password is blank prompt user
	if *dbPass == "" {
		fmt.Println("Enter password: ")
		pwd, err := terminal.ReadPassword(int(os.Stdin.Fd()))
		if err != nil {
			if err != io.EOF {
				checkErr(err)
			}
		}
		*dbPass = string(pwd)
	}

	// Create a new CSV writer
	CSVWriter := NewWriter(writerDest)
	CSVWriter.Delimiter, _ = utf8.DecodeLastRuneInString(*csvDelimiter)
	CSVWriter.Quote, _ = utf8.DecodeLastRuneInString(*csvQuote)
	CSVWriter.Escape, _ = utf8.DecodeLastRuneInString(*csvEscape)
	CSVWriter.Terminator = *csvTerminator

	// Populate dbInfo struct with cli flags
	dbi := dbInfo{user: *dbUser, pass: *dbPass, host: *dbHost, port: *dbPort, TLS: *dbTLS}

	// Create a *sql.DB connection to the source database
	db, err := dbi.Connect()
	defer db.Close()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Start reading & writing
	dataChan := make(chan []NullRawBytes)
	quitChan := make(chan bool)
	goChan := make(chan bool)
	go readRows(db, query, dataChan, quitChan, goChan, *csvHeader)
	rowCount := writeCSV(CSVWriter, dataChan, goChan)

	<-quitChan
	close(quitChan)
	close(goChan)

	// Memory Profiling
	if *memprofile != "" {
		f, err := os.Create(*memprofile)
		checkErr(err)
		pprof.WriteHeapProfile(f)
		defer f.Close()
	}

	if *verbose {
		fmt.Println()
		fmt.Println(rowCount, "rows written")
		fmt.Println("Total runtime =", time.Since(start))
	}
}

// Pass the buck error catching
func checkErr(e error) {
	if e != nil {
		log.Panic(e)
	}
}

func catchNotifications() {
	state, err := terminal.GetState(int(os.Stdin.Fd()))
	checkErr(err)

	// Trap for SIGINT
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)

	var timer time.Time
	go func() {
		for sig := range sigChan {
			if time.Now().Sub(timer) < time.Second*signalTimeout {
				terminal.Restore(int(os.Stdin.Fd()), state)
				os.Exit(0)
			}

			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, sig, "signal caught!")
			fmt.Fprintf(os.Stderr, "Send signal again within %v seconds to exit\n", signalTimeout)
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")

			timer = time.Now()
		}
	}()
}

// Create a database connection object
func (dbi *dbInfo) Connect() (*sql.DB, error) {
	var db *sql.DB
	var err error
	if dbi.TLS {
		db, err = sql.Open("mysql", dbi.user+":"+dbi.pass+"@tcp("+dbi.host+":"+dbi.port+")/?allowCleartextPasswords=1&tls=skip-verify")
	} else {
		db, err = sql.Open("mysql", dbi.user+":"+dbi.pass+"@tcp("+dbi.host+":"+dbi.port+")/?allowCleartextPasswords=1")
	}
	checkErr(err)

	// Ping database to verify credentials
	err = db.Ping()

	return db, err
}

// Scan impliments the Scanner interface
func (nb *NullRawBytes) Scan(value interface{}) error {
	if value == nil {
		nb.RawBytes, nb.Valid = []byte(string(0x0)), false
		return nil
	}

	switch v := value.(type) {
	case []byte:
		nb.RawBytes = v
	}
	nb.Valid = true

	return nil
}

// Value impliments the sql.driver Valuer interface
func (nb NullRawBytes) Value() (driver.Value, error) {
	if !nb.Valid {
		return nil, nil
	}
	return nb.RawBytes, nil
}

// readRows executes a query and returns the rows
func readRows(db *sql.DB, query string, dataChan chan []NullRawBytes, quitChan chan bool, goChan chan bool, csvHeader bool) {
	rows, err := db.Query(query)
	defer rows.Close()
	if err != nil {
		log.Print(err)
		os.Exit(1)
	}

	cols, err := rows.Columns()
	checkErr(err)

	if csvHeader {
		headers := make([]NullRawBytes, len(cols))
		for i, col := range cols {
			headers[i] = NullRawBytes{RawBytes: []byte(col), Valid: true}
		}
		dataChan <- headers
		<-goChan
	}

	// Need to scan into empty interface since we don't know how many columns or their types
	scanVals := make([]interface{}, len(cols))
	vals := make([]NullRawBytes, len(cols))
	cpy := make([]NullRawBytes, len(cols))
	for i := range vals {
		scanVals[i] = &vals[i]
	}

	for rows.Next() {
		err := rows.Scan(scanVals...)
		checkErr(err)

		copy(cpy, vals)
		dataChan <- cpy

		// Block until writeRows() signals it is safe to proceed
		// This is necessary because sql.RawBytes is a memory pointer and rows.Next() will loop and change the memory address before writeRows can properly process the values
		<-goChan
	}

	err = rows.Err()
	checkErr(err)

	close(dataChan)
	quitChan <- true
}

// writeCSV writes csv output
func writeCSV(w *Writer, dataChan chan []NullRawBytes, goChan chan bool) uint {
	var cnt uint
	// Range over row results from readRows()
	for data := range dataChan {
		// Write CSV
		size, err := w.Write(data)
		checkErr(err)

		cnt++

		// Flush CSV writer contents
		if size > flushSize {
			w.Flush()
			err = w.Error()
			checkErr(err)
		}

		// Allow read function to unblock and loop over rows
		goChan <- true
	}

	// Flush CSV writer contents
	w.Flush()
	err := w.Error()
	checkErr(err)

	return cnt
}
<commit_msg>Fixed select case check in query, write error message and quit for bad CSV delimiter, quote and escape characters<commit_after>package main

import (
	"database/sql"
	"database/sql/driver"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"strings"
	"time"
	"unicode/utf8"

	"golang.org/x/crypto/ssh/terminal"

	_ "github.com/go-sql-driver/mysql"
)

const (
	// Amount of CSV write data to buffer between flushes
	flushSize = 26214400 // 25MB

	// Timeout length where ctrl+c is ignored
	// If a second ctrl+c is sent before the timeout the program exits
	signalTimeout = 3 // Seconds

	// Timeout length to wait for a query string sent via stdin
	stdinTimeout = 10 // Milliseconds
)

type (
	// dbInfo contains information necessary to connect to a database
	dbInfo struct {
		user string
		pass string
		host string
		port string
		TLS  bool
	}

	// NullRawBytes represents sql.RawBytes which may be null
	NullRawBytes struct {
		RawBytes sql.RawBytes
		Valid    bool
	}
)

// ShowUsage prints a help screen
func showUsage() {
	fmt.Println(`
	mycsv usage:
	mycsv DB_COMMANDS [CSV OUTPUT FLAGS] [DEBUG FLAGS] [CSV OUTFILE] query

	EXAMPLES:
	mycsv -user=jprunier -p -file=my.csv -query="select * from jjp.example_table where filter in ('1', 'test', 'another')"
	echo "select * from mysql.plugins" | -user=jprunier -pmypass -hremotedb -tls > my.csv
	mycsv -user=jprunier -p -file=my.csv -d="|" -q="'" < queryfile

	DATABASE FLAGS
	==============
	-user: Username (required)
	-pass: Password (interactive prompt if blank)
	-host: Database Host (localhost assumed if blank)
	-port: Port (3306 default)
	-tls:  Use TLS/SSL for database connection (false default)

	CSV FLAGS
	=========
	-file: CSV output filename (Write to stdout if not supplied)
	-query: MySQL query (required, can be sent via stdin redirection)
	-header: Print initial column name header line (true default)
	-d: CSV field delimiter ("," default)
	-q: CSV quote character ("\"" default)
	-e: CSV escape character ("\\" default)
	-t: CSV line terminator ("\n" default)
	-v: Print more information (false default)

	DEBUG FLAGS
	===========
	-debug_cpu: CPU debugging filename
	-debug_mem: Memory debugging filename

	`)
}

func main() {
	start := time.Now()

	// Profiling flags
	var cpuprofile = flag.String("debug_cpu", "", "CPU debugging filename")
	var memprofile = flag.String("debug_mem", "", "Memory debugging filename")

	// Database flags
	dbUser := flag.String("user", "", "Username (required)")
	dbPass := flag.String("pass", "", "Password (interactive prompt if blank)")
	dbHost := flag.String("host", "", "Database Host (localhost assumed if blank)")
	dbPort := flag.String("port", "3306", "Port")
	dbTLS := flag.Bool("tls", false, "Use TLS/SSL for database connection")

	// CSV format flags
	csvDelimiter := flag.String("d", ",", "CSV field delimiter")
	csvQuote := flag.String("q", "\"", "CSV quote character")
	csvEscape := flag.String("e", "\\", "CSV escape character")
	csvTerminator := flag.String("t", "\n", "CSV line terminator")

	csvHeader := flag.Bool("header", true, "Print initial column name header line")
	csvFile := flag.String("file", "", "CSV output filename")
	sqlQuery := flag.String("query", "", "MySQL query")
	verbose := flag.Bool("v", false, "Print more information")

	// Parse flags
	flag.Parse()

	// Print usage
	if flag.NFlag() == 0 {
		showUsage()

		os.Exit(0)
	}

	// If query not provided read from standard in
	var query string
	queryChan := make(chan string)
	defer close(queryChan)
	if *sqlQuery == "" {
		go func() {
			b, err := ioutil.ReadAll(os.Stdin)
			checkErr(err)

			queryChan <- string(b)
		}()

		select {
		case q := <-queryChan:
			query = q
		case <-time.After(time.Millisecond * stdinTimeout):
			fmt.Fprintln(os.Stderr, "You must supply a query")
			os.Exit(1)
		}
	} else {
		query = *sqlQuery
	}

	// Make sure the query is a select
	if strings.ToLower(query[0:6]) != "select" {
		fmt.Fprintln(os.Stderr, "Query must be a select!")
		os.Exit(1)
	}

	// Check if an output file was supplied otherwise use standard out
	var writeTo string
	var writerDest io.Writer
	var err error
	if *csvFile == "" {
		writeTo = "standard out"
		writerDest = os.Stdout
	} else {
		f, err := os.Open(*csvFile)
		if err == nil {
			fmt.Fprintln(os.Stderr, *csvFile, "already exists!")
			fmt.Fprintln(os.Stderr, "Please remove it or use a different filename")
			os.Exit(1)
		}
		f.Close()
		writerDest, err = os.Create(*csvFile)
		writeTo = *csvFile
	}

	// Create a new CSV writer
	var i int
	CSVWriter := NewWriter(writerDest)
	CSVWriter.Delimiter, i = utf8.DecodeLastRuneInString(*csvDelimiter)
	if i == 0 {
		fmt.Fprintln(os.Stderr, "You must supply a valid delimiter character")
		os.Exit(1)
	}
	CSVWriter.Quote, i = utf8.DecodeLastRuneInString(*csvQuote)
	if i == 0 {
		fmt.Fprintln(os.Stderr, "You must supply a valid quote character")
		os.Exit(1)
	}
	CSVWriter.Escape, i = utf8.DecodeLastRuneInString(*csvEscape)
	if i == 0 {
		fmt.Fprintln(os.Stderr, "You must supply a valid escape character")
		os.Exit(1)
	}
	CSVWriter.Terminator = *csvTerminator

	if *verbose {
		fmt.Println("CSV output will be written to", writeTo)
	}

	// Check if Stdin has been redirected and reset so the user can be prompted for a password
	checkStdin()

	// Listen for SIGINT (ctrl+c)
	catchNotifications()

	// CPU Profiling
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		checkErr(err)
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// Default to localhost if no host or socket provided
	if *dbHost == "" {
		*dbHost = "127.0.0.1"
	}

	// Need to provide a target
	if *dbUser == "" {
		fmt.Fprintln(os.Stderr, "You must provide a user name!")
		os.Exit(1)
	}

	// If password is blank prompt user
	if *dbPass == "" {
		fmt.Println("Enter password: ")
		pwd, err := terminal.ReadPassword(int(os.Stdin.Fd()))
		if err != nil {
			if err != io.EOF {
				checkErr(err)
			}
		}
		*dbPass = string(pwd)
	}

	// Populate dbInfo struct with cli flags
	dbi := dbInfo{user: *dbUser, pass: *dbPass, host: *dbHost, port: *dbPort, TLS: *dbTLS}

	// Create a *sql.DB connection to the source database
	db, err := dbi.Connect()
	defer db.Close()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Start reading & writing
	dataChan := make(chan []NullRawBytes)
	quitChan := make(chan bool)
	goChan := make(chan bool)
	go readRows(db, query, dataChan, quitChan, goChan, *csvHeader)
	rowCount := writeCSV(CSVWriter, dataChan, goChan)

	<-quitChan
	close(quitChan)
	close(goChan)

	// Memory Profiling
	if *memprofile != "" {
		f, err := os.Create(*memprofile)
		checkErr(err)
		pprof.WriteHeapProfile(f)
		defer f.Close()
	}

	if *verbose {
		fmt.Println()
		fmt.Println(rowCount, "rows written")
		fmt.Println("Total runtime =", time.Since(start))
	}
}

// Pass the buck error catching
func checkErr(e error) {
	if e != nil {
		log.Panic(e)
	}
}

func catchNotifications() {
	state, err := terminal.GetState(int(os.Stdin.Fd()))
	checkErr(err)

	// Trap for SIGINT
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)

	var timer time.Time
	go func() {
		for sig := range sigChan {
			if time.Now().Sub(timer) < time.Second*signalTimeout {
				terminal.Restore(int(os.Stdin.Fd()), state)
				os.Exit(0)
			}

			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, sig, "signal caught!")
			fmt.Fprintf(os.Stderr, "Send signal again within %v seconds to exit\n", signalTimeout)
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")

			timer = time.Now()
		}
	}()
}

// Create a database connection object
func (dbi *dbInfo) Connect() (*sql.DB, error) {
	var db *sql.DB
	var err error
	if dbi.TLS {
		db, err = sql.Open("mysql", dbi.user+":"+dbi.pass+"@tcp("+dbi.host+":"+dbi.port+")/?allowCleartextPasswords=1&tls=skip-verify")
	} else {
		db, err = sql.Open("mysql", dbi.user+":"+dbi.pass+"@tcp("+dbi.host+":"+dbi.port+")/?allowCleartextPasswords=1")
	}
	checkErr(err)

	// Ping database to verify credentials
	err = db.Ping()

	return db, err
}

// Scan impliments the Scanner interface
func (nb *NullRawBytes) Scan(value interface{}) error {
	if value == nil {
		nb.RawBytes, nb.Valid = []byte(string(0x0)), false
		return nil
	}

	switch v := value.(type) {
	case []byte:
		nb.RawBytes = v
	}
	nb.Valid = true

	return nil
}

// Value impliments the sql.driver Valuer interface
func (nb NullRawBytes) Value() (driver.Value, error) {
	if !nb.Valid {
		return nil, nil
	}
	return nb.RawBytes, nil
}

// readRows executes a query and returns the rows
func readRows(db *sql.DB, query string, dataChan chan []NullRawBytes, quitChan chan bool, goChan chan bool, csvHeader bool) {
	rows, err := db.Query(query)
	defer rows.Close()
	if err != nil {
		log.Print(err)
		os.Exit(1)
	}

	cols, err := rows.Columns()
	checkErr(err)

	if csvHeader {
		headers := make([]NullRawBytes, len(cols))
		for i, col := range cols {
			headers[i] = NullRawBytes{RawBytes: []byte(col), Valid: true}
		}
		dataChan <- headers
		<-goChan
	}

	// Need to scan into empty interface since we don't know how many columns or their types
	scanVals := make([]interface{}, len(cols))
	vals := make([]NullRawBytes, len(cols))
	cpy := make([]NullRawBytes, len(cols))
	for i := range vals {
		scanVals[i] = &vals[i]
	}

	for rows.Next() {
		err := rows.Scan(scanVals...)
		checkErr(err)

		copy(cpy, vals)
		dataChan <- cpy

		// Block until writeRows() signals it is safe to proceed
		// This is necessary because sql.RawBytes is a memory pointer and rows.Next() will loop and change the memory address before writeRows can properly process the values
		<-goChan
	}

	err = rows.Err()
	checkErr(err)

	close(dataChan)
	quitChan <- true
}

// writeCSV writes csv output
func writeCSV(w *Writer, dataChan chan []NullRawBytes, goChan chan bool) uint {
	var cnt uint
	// Range over row results from readRows()
	for data := range dataChan {
		// Write CSV
		size, err := w.Write(data)
		checkErr(err)

		cnt++

		// Flush CSV writer contents
		if size > flushSize {
			w.Flush()
			err = w.Error()
			checkErr(err)
		}

		// Allow read function to unblock and loop over rows
		goChan <- true
	}

	// Flush CSV writer contents
	w.Flush()
	err := w.Error()
	checkErr(err)

	return cnt
}
<|endoftext|>


<commit_before>package api

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

	"github.com/asdine/storm"
	"github.com/gorilla/mux"
	"github.com/resin-io/edge-node-manager/config"
	"github.com/resin-io/edge-node-manager/device"
	"github.com/resin-io/edge-node-manager/process"
	"github.com/resin-io/edge-node-manager/process/status"

	log "github.com/Sirupsen/logrus"
)

func DependentDeviceUpdate(w http.ResponseWriter, r *http.Request) {
	type dependentDeviceUpdate struct {
		Commit      string      `json:"commit"`
		Environment interface{} `json:"environment"`
	}

	decoder := json.NewDecoder(r.Body)
	var content dependentDeviceUpdate
	if err := decoder.Decode(&content); err != nil {
		log.WithFields(log.Fields{
			"Error": err,
		}).Error("Unable to decode Dependent device update hook")
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	if err := setField(w, r, "TargetCommit", content.Commit); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
	}

	w.WriteHeader(http.StatusAccepted)
}

func DependentDeviceDelete(w http.ResponseWriter, r *http.Request) {
	if err := setField(w, r, "Delete", true); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
	}

	w.WriteHeader(http.StatusOK)
}

func DependentDeviceRestart(w http.ResponseWriter, r *http.Request) {
	if err := setField(w, r, "Restart", true); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
	}

	w.WriteHeader(http.StatusOK)
}

func SetStatus(w http.ResponseWriter, r *http.Request) {
	type s struct {
		TargetStatus status.Status `json:"targetStatus"`
	}

	var content *s
	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(&content); err != nil {
		log.WithFields(log.Fields{
			"Error": err,
		}).Error("Unable to decode status hook")
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	process.TargetStatus = content.TargetStatus

	w.WriteHeader(http.StatusOK)

	log.WithFields(log.Fields{
		"Target status": process.TargetStatus,
	}).Debug("Set status")
}

func GetStatus(w http.ResponseWriter, r *http.Request) {
	type s struct {
		CurrentStatus status.Status `json:"currentStatus"`
		TargetStatus  status.Status `json:"targetStatus"`
	}

	content := &s{
		CurrentStatus: process.CurrentStatus,
		TargetStatus:  process.TargetStatus,
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		log.WithFields(log.Fields{
			"Error": err,
		}).Error("Unable to encode status hook")
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if written, err := w.Write(bytes); (err != nil) || (written != len(bytes)) {
		log.WithFields(log.Fields{
			"Error": err,
		}).Error("Unable to write response")
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	log.WithFields(log.Fields{
		"Target status": process.TargetStatus,
		"Curent status": process.CurrentStatus,
	}).Debug("Get status")
}

func setField(w http.ResponseWriter, r *http.Request, key string, value interface{}) error {
	vars := mux.Vars(r)
	deviceUUID := vars["uuid"]

	db, err := storm.Open(config.GetDbPath())
	if err != nil {
		log.WithFields(log.Fields{
			"Error": err,
		}).Error("Unable to open database")
		return err
	}
	defer db.Close()

	var d device.Device
	if err := db.One("ResinUUID", deviceUUID, &d); err != nil {
		log.WithFields(log.Fields{
			"Error": err,
			"UUID":  deviceUUID,
		}).Error("Unable to find device in database")
		return err
	}

	switch key {
	case "TargetCommit":
		d.TargetCommit = value.(string)
	case "Delete":
		d.DeleteFlag = value.(bool)
	case "Restart":
		d.RestartFlag = value.(bool)
	default:
		log.WithFields(log.Fields{
			"Error": err,
			"UUID":  deviceUUID,
			"Key":   key,
			"value": value,
		}).Error("Unable to set field")
		return err
	}

	if err := db.Update(&d); err != nil {
		log.WithFields(log.Fields{
			"Error": err,
			"UUID":  deviceUUID,
		}).Error("Unable to update device in database")
		return err
	}

	log.WithFields(log.Fields{
		"UUID":  deviceUUID,
		"Key":   key,
		"value": value,
	}).Debug("Dependent device field updated")

	return nil
}
<commit_msg>Ignore db open errors allowing the supervisor to silently retry<commit_after>package api

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

	"github.com/asdine/storm"
	"github.com/gorilla/mux"
	"github.com/resin-io/edge-node-manager/config"
	"github.com/resin-io/edge-node-manager/device"
	"github.com/resin-io/edge-node-manager/process"
	"github.com/resin-io/edge-node-manager/process/status"

	log "github.com/Sirupsen/logrus"
)

func DependentDeviceUpdate(w http.ResponseWriter, r *http.Request) {
	type dependentDeviceUpdate struct {
		Commit      string      `json:"commit"`
		Environment interface{} `json:"environment"`
	}

	decoder := json.NewDecoder(r.Body)
	var content dependentDeviceUpdate
	if err := decoder.Decode(&content); err != nil {
		log.WithFields(log.Fields{
			"Error": err,
		}).Error("Unable to decode Dependent device update hook")
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	if err := setField(r, "TargetCommit", content.Commit); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusAccepted)
}

func DependentDeviceDelete(w http.ResponseWriter, r *http.Request) {
	if err := setField(r, "Delete", true); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
}

func DependentDeviceRestart(w http.ResponseWriter, r *http.Request) {
	if err := setField(r, "Restart", true); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
}

func SetStatus(w http.ResponseWriter, r *http.Request) {
	type s struct {
		TargetStatus status.Status `json:"targetStatus"`
	}

	var content *s
	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(&content); err != nil {
		log.WithFields(log.Fields{
			"Error": err,
		}).Error("Unable to decode status hook")
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	process.TargetStatus = content.TargetStatus

	w.WriteHeader(http.StatusOK)

	log.WithFields(log.Fields{
		"Target status": process.TargetStatus,
	}).Debug("Set status")
}

func GetStatus(w http.ResponseWriter, r *http.Request) {
	type s struct {
		CurrentStatus status.Status `json:"currentStatus"`
		TargetStatus  status.Status `json:"targetStatus"`
	}

	content := &s{
		CurrentStatus: process.CurrentStatus,
		TargetStatus:  process.TargetStatus,
	}

	bytes, err := json.Marshal(content)
	if err != nil {
		log.WithFields(log.Fields{
			"Error": err,
		}).Error("Unable to encode status hook")
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if written, err := w.Write(bytes); (err != nil) || (written != len(bytes)) {
		log.WithFields(log.Fields{
			"Error": err,
		}).Error("Unable to write response")
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	log.WithFields(log.Fields{
		"Target status": process.TargetStatus,
		"Curent status": process.CurrentStatus,
	}).Debug("Get status")
}

func setField(r *http.Request, key string, value interface{}) error {
	vars := mux.Vars(r)
	deviceUUID := vars["uuid"]

	db, err := storm.Open(config.GetDbPath())
	if err != nil {
		return err
	}
	defer db.Close()

	var d device.Device
	if err := db.One("ResinUUID", deviceUUID, &d); err != nil {
		log.WithFields(log.Fields{
			"Error": err,
			"UUID":  deviceUUID,
		}).Error("Unable to find device in database")
		return err
	}

	switch key {
	case "TargetCommit":
		d.TargetCommit = value.(string)
	case "Delete":
		d.DeleteFlag = value.(bool)
	case "Restart":
		d.RestartFlag = value.(bool)
	default:
		log.WithFields(log.Fields{
			"Error": err,
			"UUID":  deviceUUID,
			"Key":   key,
			"value": value,
		}).Error("Unable to set field")
		return err
	}

	if err := db.Update(&d); err != nil {
		log.WithFields(log.Fields{
			"Error": err,
			"UUID":  deviceUUID,
		}).Error("Unable to update device in database")
		return err
	}

	log.WithFields(log.Fields{
		"UUID":  deviceUUID,
		"Key":   key,
		"value": value,
	}).Debug("Dependent device field updated")

	return nil
}
<|endoftext|>


<commit_before>package proxy

import (
	"context"
	"errors"
	"fmt"

	"github.com/digitalocean/godo"
)

type (
	// Domain represents a DNS domain from Digital Ocean.
	Domain struct {
		*godo.Domain
		ZoneName ZoneName
		Records  []*Record
	}
)

var (
	// ErrZoneNotFound will be returned if the zone cannot be found at Digital Ocean.
	ErrZoneNotFound = errors.New("Zone not found")
)

// NewDomain will instantiate a new Domain from the specified zoneName.
func NewDomain(zoneName ZoneName) *Domain {
	return &Domain{
		ZoneName: zoneName,
	}
}

// Find will try to find the matching domain at Digital Ocean.
func (d *Domain) Find(client *godo.Client) error {
	// Check if the domain is registered at Digital Ocean.
	domains, _, err := client.Domains.List(context.TODO(), nil)
	if err != nil {
		return err
	}

	for _, domain := range domains {
		if d.ZoneName == NewZoneName(domain.Name) {
			d.Domain = &domain
			return nil
		}
	}

	return ErrZoneNotFound
}

// Add will add a new zone/domain to Digital Ocean.
func (d *Domain) Add(client *godo.Client) error {
	req := godo.DomainCreateRequest{
		Name:      d.ZoneName.String(""),
		IPAddress: "127.0.0.1", // FIXME: Try to add something that makes more sense
	}

	domain, _, err := client.Domains.Create(context.TODO(), &req)
	if err != nil {
		return err
	}

	d.Domain = domain

	return nil
}

// FindOrAdd will search for the named zone at Digital Ocean - and add it if not
// found.
func (d *Domain) FindOrAdd(client *godo.Client) error {
	err := d.Find(client)
	if err == ErrZoneNotFound {
		return d.Add(client)
	}

	return nil
}

// RefreshRecords retrieves all records from Digital Ocean.
func (d *Domain) RefreshRecords(client *godo.Client) error {
	// create options. initially, these will be blank.
	opt := &godo.ListOptions{
		// DO doesn't support 10.000 entries per page (yet). It will be clamped
		// to 200 (for now), but we set it anyway. Maybe they will some time in
		// the future.
		PerPage: 10000,
	}

	d.Records = nil
	for {
		records, resp, err := client.Domains.Records(context.TODO(), d.ZoneName.String(""), opt)
		if err != nil {
			return err
		}

		for _, dr := range records {
			record := NewRecord(dr, d.ZoneName)
			d.Records = append(d.Records, record)
		}

		fmt.Printf("We got %d records", len(d.Records))

		// if we are at the last page, break out of the for loop.
		if resp.Links == nil || resp.Links.IsLastPage() {
			break
		}

		page, err := resp.Links.CurrentPage()
		if err != nil {
			return err
		}

		// set the page we want for the next request
		opt.Page = page + 1
	}

	// If we arrived here, everything must be good :)
	return nil
}
<commit_msg>Added newline for maximum readability.<commit_after>package proxy

import (
	"context"
	"errors"
	"fmt"

	"github.com/digitalocean/godo"
)

type (
	// Domain represents a DNS domain from Digital Ocean.
	Domain struct {
		*godo.Domain
		ZoneName ZoneName
		Records  []*Record
	}
)

var (
	// ErrZoneNotFound will be returned if the zone cannot be found at Digital Ocean.
	ErrZoneNotFound = errors.New("Zone not found")
)

// NewDomain will instantiate a new Domain from the specified zoneName.
func NewDomain(zoneName ZoneName) *Domain {
	return &Domain{
		ZoneName: zoneName,
	}
}

// Find will try to find the matching domain at Digital Ocean.
func (d *Domain) Find(client *godo.Client) error {
	// Check if the domain is registered at Digital Ocean.
	domains, _, err := client.Domains.List(context.TODO(), nil)
	if err != nil {
		return err
	}

	for _, domain := range domains {
		if d.ZoneName == NewZoneName(domain.Name) {
			d.Domain = &domain
			return nil
		}
	}

	return ErrZoneNotFound
}

// Add will add a new zone/domain to Digital Ocean.
func (d *Domain) Add(client *godo.Client) error {
	req := godo.DomainCreateRequest{
		Name:      d.ZoneName.String(""),
		IPAddress: "127.0.0.1", // FIXME: Try to add something that makes more sense
	}

	domain, _, err := client.Domains.Create(context.TODO(), &req)
	if err != nil {
		return err
	}

	d.Domain = domain

	return nil
}

// FindOrAdd will search for the named zone at Digital Ocean - and add it if not
// found.
func (d *Domain) FindOrAdd(client *godo.Client) error {
	err := d.Find(client)
	if err == ErrZoneNotFound {
		return d.Add(client)
	}

	return nil
}

// RefreshRecords retrieves all records from Digital Ocean.
func (d *Domain) RefreshRecords(client *godo.Client) error {
	// create options. initially, these will be blank.
	opt := &godo.ListOptions{
		// DO doesn't support 10.000 entries per page (yet). It will be clamped
		// to 200 (for now), but we set it anyway. Maybe they will some time in
		// the future.
		PerPage: 10000,
	}

	d.Records = nil
	for {
		records, resp, err := client.Domains.Records(context.TODO(), d.ZoneName.String(""), opt)
		if err != nil {
			return err
		}

		for _, dr := range records {
			record := NewRecord(dr, d.ZoneName)
			d.Records = append(d.Records, record)
		}

		fmt.Printf("We got %d records\n", len(d.Records))

		// if we are at the last page, break out of the for loop.
		if resp.Links == nil || resp.Links.IsLastPage() {
			break
		}

		page, err := resp.Links.CurrentPage()
		if err != nil {
			return err
		}

		// set the page we want for the next request
		opt.Page = page + 1
	}

	// If we arrived here, everything must be good :)
	return nil
}
<|endoftext|>


<commit_before>package impl

import (
	"errors"
	"fmt"
	"time"

	"os"

	"github.com/ellcrys/util"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/postgres" // gorm requires it
	"github.com/kr/pretty"
	"github.com/ncodes/cocoon/core/common"
	"github.com/ncodes/cocoon/core/types"
	logging "github.com/op/go-logging"
)

var log = logging.MustGetLogger("postgres.store")

// ErrChainExist represents an error about an existing chain
var ErrChainExist = errors.New("chain already exists")

// LedgerTableName represents the name of the table where all
// known ledger info are stored.
const LedgerTableName = "ledgers"

// TransactionTableName represents the name of the table where all
// transactions are stored.
const TransactionTableName = "transactions"

// PostgresStore defines a store implementation
// on the postgres database. It implements the Store interface
type PostgresStore struct {
	db         *gorm.DB
	blockchain types.Blockchain
	locker     types.Lock
}

// SetBlockchainImplementation sets sets a reference of the blockchain implementation
func (s *PostgresStore) SetBlockchainImplementation(b types.Blockchain) {
	s.blockchain = b
}

// GetImplementationName returns the name of this store implementation
func (s *PostgresStore) GetImplementationName() string {
	return "postgres.store"
}

// Connect connects to a postgres server and returns a client
// or error if connection failed.
func (s *PostgresStore) Connect(dbAddr string) (interface{}, error) {

	var err error
	s.db, err = gorm.Open("postgres", dbAddr)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to store backend. %s", err)
	}

	s.db.LogMode(false)

	return s.db, nil
}

// MakeLegderHash takes a ledger and computes a hash
func (s *PostgresStore) MakeLegderHash(ledger *types.Ledger) string {
	return util.Sha256(fmt.Sprintf("%s;%t;%d", ledger.Name, ledger.Public, ledger.CreatedAt))
}

// Init initializes the types. Creates the necessary tables such as the
// the tables and public and private system ledgers
func (s *PostgresStore) Init(systemPublicLedgerName, systemPrivateLedgerName string) error {

	// create ledger table if not existing
	if !s.db.HasTable(LedgerTableName) {
		if err := s.db.CreateTable(&types.Ledger{}).Error; err != nil {
			return fmt.Errorf("failed to create `%s` table. %s", LedgerTableName, err)
		}
	}

	// create transaction table if not existing
	if !s.db.HasTable(TransactionTableName) {
		if err := s.db.CreateTable(&types.Transaction{}).
			AddIndex("idx_name_ledger_key_created_at", "ledger", "key", "created_at").Error; err != nil {
			return fmt.Errorf("failed to create `%s` table. %s", TransactionTableName, err)
		}
	}

	// create system ledgers
	var systemLedgers = [][]interface{}{
		[]interface{}{systemPublicLedgerName, true},   // public
		[]interface{}{systemPrivateLedgerName, false}, // private
	}
	for _, ledger := range systemLedgers {
		var c int
		if err := s.db.Model(&types.Ledger{}).Where("name = ?", ledger[0].(string)).Count(&c).Error; err != nil {
			return fmt.Errorf("failed to check existence of ledger named %s: %s", ledger[0].(string), err)
		}
		if c == 0 {
			_, err := s.CreateLedger(ledger[0].(string), true, ledger[1].(bool))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Destroy removes the database tables.
// Will only work in a test environment (Test Only!!!)
func Destroy(dbAddr string) error {

	if os.Getenv("APP_ENV") != "test" {
		return fmt.Errorf("Cowardly refusing to do it! Can only call Destroy() in test environment")
	}

	db, err := gorm.Open("postgres", dbAddr)
	if err != nil {
		return fmt.Errorf("failed to connect to store backend")
	}

	return db.DropTable(types.Ledger{}, types.Transaction{}).Error
}

// Clear the database tables.
// Will only work in a test environment (Test Only!!!)
func Clear(dbAddr string) error {

	if os.Getenv("APP_ENV") != "test" {
		return fmt.Errorf("Cowardly refusing to do it! Can only call Destroy() in test environment")
	}

	db, err := gorm.Open("postgres", dbAddr)
	if err != nil {
		return fmt.Errorf("failed to connect to store backend")
	}

	err = db.Delete(types.Transaction{}).Error
	if err != nil {
		return err
	}

	err = db.Delete(types.Ledger{}).Error
	if err != nil {
		return err
	}

	return nil
}

// CreateLedgerThen creates a new ledger and accepts an additional operation (via the thenFunc) to be
// executed before the ledger creation transaction is committed. If the thenFunc returns an error, the
// transaction is rolled back and error returned
func (s *PostgresStore) CreateLedgerThen(name string, chained, public bool, thenFunc func() error) (*types.Ledger, error) {

	tx := s.db.Begin()

	err := tx.Exec(`SET TRANSACTION isolation level repeatable read`).Error
	if err != nil {
		tx.Rollback()
		return nil, fmt.Errorf("failed to set transaction isolation level. %s", err)
	}

	newLedger := &types.Ledger{
		Name:      name,
		Public:    public,
		Chained:   chained,
		CreatedAt: time.Now().Unix(),
	}

	newLedger.Hash = s.MakeLegderHash(newLedger)

	if err := tx.Create(newLedger).Error; err != nil {
		tx.Rollback()
		if common.IsUniqueConstraintError(err, "name") {
			return nil, fmt.Errorf("ledger with matching name already exists")
		} else if common.IsUniqueConstraintError(err, "hash") {
			return nil, fmt.Errorf("hash is being used by another ledger")
		}
		return nil, err
	}

	// create chain
	if chained {
		_, err := s.blockchain.CreateChain(name, public)
		if err != nil {
			tx.Rollback()
			return nil, err
		}
	}

	// run the companion functions and Rollback
	// the transaction if error was returned
	if thenFunc != nil {
		if err = thenFunc(); err != nil {
			tx.Rollback()
			return nil, err
		}
	}

	tx.Commit()

	return newLedger, nil
}

// CreateLedger creates a new ledger.
func (s *PostgresStore) CreateLedger(name string, chained, public bool) (*types.Ledger, error) {
	return s.CreateLedgerThen(name, chained, public, nil)
}

// GetLedger fetches a ledger meta information
func (s *PostgresStore) GetLedger(name string) (*types.Ledger, error) {

	var l types.Ledger

	err := s.db.Where("name = ?", name).Last(&l).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return nil, fmt.Errorf("failed to get ledger. %s", err)
	} else if err == gorm.ErrRecordNotFound {
		return nil, nil
	}

	return &l, nil
}

// makeTxLockKey constructs a lock key using a transaction key and ledger name
func makeTxLockKey(ledgerName, key string) string {
	return fmt.Sprintf("tx;key;%s;%s", ledgerName, key)
}

// PutThen adds transactions to the store and returns a list of transaction receipts.
// Any transaction that failed to be created will result in an error receipt being created
// and returned along with success receipts of they successfully added transactions.
// However, all transactions will be rolled back if the `thenFunc` returns error. Only the
// transaction that are successfully added will be passed to the thenFunc.
// Future work may allow the caller to determine the behaviour via an additional parameter.
func (s *PostgresStore) PutThen(ledgerName string, txs []*types.Transaction, thenFunc func(validTxss []*types.Transaction) error) ([]*types.TxReceipt, error) {

	var validTxs []*types.Transaction
	txReceipts := []*types.TxReceipt{}

	dbTx := s.db.Begin()
	err := dbTx.Exec(`SET TRANSACTION isolation level repeatable read`).Error
	if err != nil {
		dbTx.Rollback()
		return nil, fmt.Errorf("failed to set transaction isolation level. %s", err)
	}

	// create transactions and add transaction receipts for
	// successfully stored transactions
	for _, tx := range txs {

		var err error

		// acquire lock on the transaction via its key
		lock := common.NewLock(makeTxLockKey(tx.LedgerInternal, tx.Key))
		if err = lock.Acquire(); err != nil {
			if err == types.ErrLockAlreadyAcquired {
				txReceipts = append(txReceipts, &types.TxReceipt{
					ID:  tx.ID,
					Err: "failed to acquire lock. object has been locked by another process",
				})
			} else {
				txReceipts = append(txReceipts, &types.TxReceipt{
					ID:  tx.ID,
					Err: err.Error(),
				})
			}
			continue
		}

		// ensure the current transaction is not a stale transaction
		// when checked with the valid transactions. If this is not done, before we call the tx.Create
		// the entire transaction will not be committed by the postgres driver
		isStale := false
		for _, vTx := range validTxs {
			if len(tx.RevisionTo) > 0 && tx.RevisionTo == vTx.RevisionTo && tx.KeyInternal == vTx.KeyInternal {
				isStale = true
				break
			}
		}
		if isStale {
			log.Info("Stale transaction (%s)", tx.ID)
			txReceipts = append(txReceipts, &types.TxReceipt{
				ID:  tx.ID,
				Err: "stale object",
			})
			lock.Release()
			continue
		}

		tx.Hash = tx.MakeHash()
		tx.Ledger = ledgerName
		txReceipt := &types.TxReceipt{ID: tx.ID}
		err = dbTx.Create(tx).Error
		if err != nil {
			log.Errorf("Failed to create transaction (%s): %s", tx.ID, err)
			txReceipt.Err = err.Error()
			if common.CompareErr(err, fmt.Errorf(`pq: duplicate key value violates unique constraint "idx_name_revision_to"`)) == 0 {
				txReceipt.Err = "stale object"
			}
		} else {
			fmt.Println("Added new valid transaction")
			validTxs = append(validTxs, tx)
		}

		if err = lock.Release(); err != nil {
			fmt.Println("Error releasing lock: ", err)
		}

		txReceipts = append(txReceipts, txReceipt)
	}

	// run the companion functions. Rollback
	// the transactions only if error was returned
	if thenFunc != nil {
		if err = thenFunc(validTxs); err != nil {
			return txReceipts, err
		}
	}

	if err := dbTx.Debug().Commit().Error; err != nil {
		pretty.Println(dbTx.GetErrors())
		return nil, err
	}

	return txReceipts, nil
}

// Put creates one or more transactions associated to a ledger.
// Returns a list of transaction receipts and a general error.
func (s *PostgresStore) Put(ledgerName string, txs []*types.Transaction) ([]*types.TxReceipt, error) {
	return s.PutThen(ledgerName, txs, nil)
}

// Get fetches a transaction by its ledger and key
func (s *PostgresStore) Get(ledger, key string) (*types.Transaction, error) {
	var tx types.Transaction

	// acquire lock on the transaction via its key
	lock := common.NewLock(makeTxLockKey(ledger, key))
	if err := lock.Acquire(); err != nil {
		if err == types.ErrLockAlreadyAcquired {
			return nil, fmt.Errorf("failed to acquire lock. object has been locked by another process")
		}
		return nil, err
	}

	defer lock.Release()

	err := s.db.Where("ledger = ? AND key = ?", ledger, key).Last(&tx).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return nil, fmt.Errorf("failed to get transaction. %s", err)
	} else if err == gorm.ErrRecordNotFound {
		return nil, nil
	}

	return &tx, nil
}

// GetRange fetches transactions with keys included in a specified range.
// No lock is acquired in this operation.
func (s *PostgresStore) GetRange(ledger, startKey, endKey string, inclusive bool, limit, offset int) ([]*types.Transaction, error) {

	var err error
	var txs []*types.Transaction
	var q *gorm.DB

	if len(startKey) > 0 && len(endKey) > 0 {
		if !inclusive {
			q = s.db.Where("ledger = ? AND (key >= ? AND key < ?)", ledger, startKey, endKey)
		} else {
			q = s.db.Where("ledger = ? AND (key >= ? OR key <= ?)", ledger, startKey+"%", endKey+"%")
		}
	} else if len(startKey) > 0 && len(endKey) == 0 {
		q = s.db.Where("ledger = ? AND key like ?", ledger, startKey+"%")
	} else {
		// setting endKey only is a little tricky as the call code may construct
		// through a secondary process or rule, so add the '%' operator will most likely
		// result in wrong query. So we just let the external decide where to put it.
		q = s.db.Where("ledger = ? AND key like ?", ledger, endKey)
	}

	err = q.
		Limit(limit).
		Offset(offset).
		Select("DISTINCT ON (key) *").
		Order("key").
		Order("created_at desc").
		Find(&txs).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return nil, fmt.Errorf("failed to get transactions. %s", err)
	} else if err == gorm.ErrRecordNotFound {
		return nil, nil
	}

	return txs, nil
}

// Close releases any resource held
func (s *PostgresStore) Close() error {
	if s.db != nil {
		return s.db.Close()
	}
	return nil
}
<commit_msg>debug: remove isolation change block<commit_after>package impl

import (
	"errors"
	"fmt"
	"time"

	"os"

	"github.com/ellcrys/util"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/postgres" // gorm requires it
	"github.com/kr/pretty"
	"github.com/ncodes/cocoon/core/common"
	"github.com/ncodes/cocoon/core/types"
	logging "github.com/op/go-logging"
)

var log = logging.MustGetLogger("postgres.store")

// ErrChainExist represents an error about an existing chain
var ErrChainExist = errors.New("chain already exists")

// LedgerTableName represents the name of the table where all
// known ledger info are stored.
const LedgerTableName = "ledgers"

// TransactionTableName represents the name of the table where all
// transactions are stored.
const TransactionTableName = "transactions"

// PostgresStore defines a store implementation
// on the postgres database. It implements the Store interface
type PostgresStore struct {
	db         *gorm.DB
	blockchain types.Blockchain
	locker     types.Lock
}

// SetBlockchainImplementation sets sets a reference of the blockchain implementation
func (s *PostgresStore) SetBlockchainImplementation(b types.Blockchain) {
	s.blockchain = b
}

// GetImplementationName returns the name of this store implementation
func (s *PostgresStore) GetImplementationName() string {
	return "postgres.store"
}

// Connect connects to a postgres server and returns a client
// or error if connection failed.
func (s *PostgresStore) Connect(dbAddr string) (interface{}, error) {

	var err error
	s.db, err = gorm.Open("postgres", dbAddr)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to store backend. %s", err)
	}

	s.db.LogMode(false)

	return s.db, nil
}

// MakeLegderHash takes a ledger and computes a hash
func (s *PostgresStore) MakeLegderHash(ledger *types.Ledger) string {
	return util.Sha256(fmt.Sprintf("%s;%t;%d", ledger.Name, ledger.Public, ledger.CreatedAt))
}

// Init initializes the types. Creates the necessary tables such as the
// the tables and public and private system ledgers
func (s *PostgresStore) Init(systemPublicLedgerName, systemPrivateLedgerName string) error {

	// create ledger table if not existing
	if !s.db.HasTable(LedgerTableName) {
		if err := s.db.CreateTable(&types.Ledger{}).Error; err != nil {
			return fmt.Errorf("failed to create `%s` table. %s", LedgerTableName, err)
		}
	}

	// create transaction table if not existing
	if !s.db.HasTable(TransactionTableName) {
		if err := s.db.CreateTable(&types.Transaction{}).
			AddIndex("idx_name_ledger_key_created_at", "ledger", "key", "created_at").Error; err != nil {
			return fmt.Errorf("failed to create `%s` table. %s", TransactionTableName, err)
		}
	}

	// create system ledgers
	var systemLedgers = [][]interface{}{
		[]interface{}{systemPublicLedgerName, true},   // public
		[]interface{}{systemPrivateLedgerName, false}, // private
	}
	for _, ledger := range systemLedgers {
		var c int
		if err := s.db.Model(&types.Ledger{}).Where("name = ?", ledger[0].(string)).Count(&c).Error; err != nil {
			return fmt.Errorf("failed to check existence of ledger named %s: %s", ledger[0].(string), err)
		}
		if c == 0 {
			_, err := s.CreateLedger(ledger[0].(string), true, ledger[1].(bool))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Destroy removes the database tables.
// Will only work in a test environment (Test Only!!!)
func Destroy(dbAddr string) error {

	if os.Getenv("APP_ENV") != "test" {
		return fmt.Errorf("Cowardly refusing to do it! Can only call Destroy() in test environment")
	}

	db, err := gorm.Open("postgres", dbAddr)
	if err != nil {
		return fmt.Errorf("failed to connect to store backend")
	}

	return db.DropTable(types.Ledger{}, types.Transaction{}).Error
}

// Clear the database tables.
// Will only work in a test environment (Test Only!!!)
func Clear(dbAddr string) error {

	if os.Getenv("APP_ENV") != "test" {
		return fmt.Errorf("Cowardly refusing to do it! Can only call Destroy() in test environment")
	}

	db, err := gorm.Open("postgres", dbAddr)
	if err != nil {
		return fmt.Errorf("failed to connect to store backend")
	}

	err = db.Delete(types.Transaction{}).Error
	if err != nil {
		return err
	}

	err = db.Delete(types.Ledger{}).Error
	if err != nil {
		return err
	}

	return nil
}

// CreateLedgerThen creates a new ledger and accepts an additional operation (via the thenFunc) to be
// executed before the ledger creation transaction is committed. If the thenFunc returns an error, the
// transaction is rolled back and error returned
func (s *PostgresStore) CreateLedgerThen(name string, chained, public bool, thenFunc func() error) (*types.Ledger, error) {

	tx := s.db.Begin()

	err := tx.Exec(`SET TRANSACTION isolation level repeatable read`).Error
	if err != nil {
		tx.Rollback()
		return nil, fmt.Errorf("failed to set transaction isolation level. %s", err)
	}

	newLedger := &types.Ledger{
		Name:      name,
		Public:    public,
		Chained:   chained,
		CreatedAt: time.Now().Unix(),
	}

	newLedger.Hash = s.MakeLegderHash(newLedger)

	if err := tx.Create(newLedger).Error; err != nil {
		tx.Rollback()
		if common.IsUniqueConstraintError(err, "name") {
			return nil, fmt.Errorf("ledger with matching name already exists")
		} else if common.IsUniqueConstraintError(err, "hash") {
			return nil, fmt.Errorf("hash is being used by another ledger")
		}
		return nil, err
	}

	// create chain
	if chained {
		_, err := s.blockchain.CreateChain(name, public)
		if err != nil {
			tx.Rollback()
			return nil, err
		}
	}

	// run the companion functions and Rollback
	// the transaction if error was returned
	if thenFunc != nil {
		if err = thenFunc(); err != nil {
			tx.Rollback()
			return nil, err
		}
	}

	tx.Commit()

	return newLedger, nil
}

// CreateLedger creates a new ledger.
func (s *PostgresStore) CreateLedger(name string, chained, public bool) (*types.Ledger, error) {
	return s.CreateLedgerThen(name, chained, public, nil)
}

// GetLedger fetches a ledger meta information
func (s *PostgresStore) GetLedger(name string) (*types.Ledger, error) {

	var l types.Ledger

	err := s.db.Where("name = ?", name).Last(&l).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return nil, fmt.Errorf("failed to get ledger. %s", err)
	} else if err == gorm.ErrRecordNotFound {
		return nil, nil
	}

	return &l, nil
}

// makeTxLockKey constructs a lock key using a transaction key and ledger name
func makeTxLockKey(ledgerName, key string) string {
	return fmt.Sprintf("tx;key;%s;%s", ledgerName, key)
}

// PutThen adds transactions to the store and returns a list of transaction receipts.
// Any transaction that failed to be created will result in an error receipt being created
// and returned along with success receipts of they successfully added transactions.
// However, all transactions will be rolled back if the `thenFunc` returns error. Only the
// transaction that are successfully added will be passed to the thenFunc.
// Future work may allow the caller to determine the behaviour via an additional parameter.
func (s *PostgresStore) PutThen(ledgerName string, txs []*types.Transaction, thenFunc func(validTxss []*types.Transaction) error) ([]*types.TxReceipt, error) {

	var validTxs []*types.Transaction
	txReceipts := []*types.TxReceipt{}

	dbTx := s.db.Begin()
	// err := dbTx.Exec(`SET TRANSACTION isolation level repeatable read`).Error
	// if err != nil {
	// 	dbTx.Rollback()
	// 	return nil, fmt.Errorf("failed to set transaction isolation level. %s", err)
	// }

	// create transactions and add transaction receipts for
	// successfully stored transactions
	for _, tx := range txs {

		var err error

		// acquire lock on the transaction via its key
		lock := common.NewLock(makeTxLockKey(tx.LedgerInternal, tx.Key))
		if err = lock.Acquire(); err != nil {
			if err == types.ErrLockAlreadyAcquired {
				txReceipts = append(txReceipts, &types.TxReceipt{
					ID:  tx.ID,
					Err: "failed to acquire lock. object has been locked by another process",
				})
			} else {
				txReceipts = append(txReceipts, &types.TxReceipt{
					ID:  tx.ID,
					Err: err.Error(),
				})
			}
			continue
		}

		// ensure the current transaction is not a stale transaction
		// when checked with the valid transactions. If this is not done, before we call the tx.Create
		// the entire transaction will not be committed by the postgres driver
		isStale := false
		for _, vTx := range validTxs {
			if len(tx.RevisionTo) > 0 && tx.RevisionTo == vTx.RevisionTo && tx.KeyInternal == vTx.KeyInternal {
				isStale = true
				break
			}
		}
		if isStale {
			log.Info("Stale transaction (%s)", tx.ID)
			txReceipts = append(txReceipts, &types.TxReceipt{
				ID:  tx.ID,
				Err: "stale object",
			})
			lock.Release()
			continue
		}

		tx.Hash = tx.MakeHash()
		tx.Ledger = ledgerName
		txReceipt := &types.TxReceipt{ID: tx.ID}
		err = dbTx.Create(tx).Error
		if err != nil {
			log.Errorf("Failed to create transaction (%s): %s", tx.ID, err)
			txReceipt.Err = err.Error()
			if common.CompareErr(err, fmt.Errorf(`pq: duplicate key value violates unique constraint "idx_name_revision_to"`)) == 0 {
				txReceipt.Err = "stale object"
			}
		} else {
			fmt.Println("Added new valid transaction")
			validTxs = append(validTxs, tx)
		}

		if err = lock.Release(); err != nil {
			fmt.Println("Error releasing lock: ", err)
		}

		txReceipts = append(txReceipts, txReceipt)
	}

	// run the companion functions. Rollback
	// the transactions only if error was returned
	if thenFunc != nil {
		if err = thenFunc(validTxs); err != nil {
			return txReceipts, err
		}
	}

	if err := dbTx.Debug().Commit().Error; err != nil {
		pretty.Println(dbTx.GetErrors())
		return nil, err
	}

	return txReceipts, nil
}

// Put creates one or more transactions associated to a ledger.
// Returns a list of transaction receipts and a general error.
func (s *PostgresStore) Put(ledgerName string, txs []*types.Transaction) ([]*types.TxReceipt, error) {
	return s.PutThen(ledgerName, txs, nil)
}

// Get fetches a transaction by its ledger and key
func (s *PostgresStore) Get(ledger, key string) (*types.Transaction, error) {
	var tx types.Transaction

	// acquire lock on the transaction via its key
	lock := common.NewLock(makeTxLockKey(ledger, key))
	if err := lock.Acquire(); err != nil {
		if err == types.ErrLockAlreadyAcquired {
			return nil, fmt.Errorf("failed to acquire lock. object has been locked by another process")
		}
		return nil, err
	}

	defer lock.Release()

	err := s.db.Where("ledger = ? AND key = ?", ledger, key).Last(&tx).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return nil, fmt.Errorf("failed to get transaction. %s", err)
	} else if err == gorm.ErrRecordNotFound {
		return nil, nil
	}

	return &tx, nil
}

// GetRange fetches transactions with keys included in a specified range.
// No lock is acquired in this operation.
func (s *PostgresStore) GetRange(ledger, startKey, endKey string, inclusive bool, limit, offset int) ([]*types.Transaction, error) {

	var err error
	var txs []*types.Transaction
	var q *gorm.DB

	if len(startKey) > 0 && len(endKey) > 0 {
		if !inclusive {
			q = s.db.Where("ledger = ? AND (key >= ? AND key < ?)", ledger, startKey, endKey)
		} else {
			q = s.db.Where("ledger = ? AND (key >= ? OR key <= ?)", ledger, startKey+"%", endKey+"%")
		}
	} else if len(startKey) > 0 && len(endKey) == 0 {
		q = s.db.Where("ledger = ? AND key like ?", ledger, startKey+"%")
	} else {
		// setting endKey only is a little tricky as the call code may construct
		// through a secondary process or rule, so add the '%' operator will most likely
		// result in wrong query. So we just let the external decide where to put it.
		q = s.db.Where("ledger = ? AND key like ?", ledger, endKey)
	}

	err = q.
		Limit(limit).
		Offset(offset).
		Select("DISTINCT ON (key) *").
		Order("key").
		Order("created_at desc").
		Find(&txs).Error
	if err != nil && err != gorm.ErrRecordNotFound {
		return nil, fmt.Errorf("failed to get transactions. %s", err)
	} else if err == gorm.ErrRecordNotFound {
		return nil, nil
	}

	return txs, nil
}

// Close releases any resource held
func (s *PostgresStore) Close() error {
	if s.db != nil {
		return s.db.Close()
	}
	return nil
}
<|endoftext|>