<commit_before>package resource

import (
	"errors"
	"net/http"

	"github.com/LewisWatson/carshare-back/model"
	"github.com/LewisWatson/carshare-back/storage"
	"github.com/benbjohnson/clock"
	"github.com/manyminds/api2go"
)

// TripResource for api2go routes
type TripResource struct {
	TripStorage     storage.TripStorage
	UserStorage     storage.UserStorage
	CarShareStorage storage.CarShareStorage
	Clock           clock.Clock
}

// FindAll trips
func (t TripResource) FindAll(r api2go.Request) (api2go.Responder, error) {
	var result []model.Trip

	trips, err := t.TripStorage.GetAll()
	if err != nil {
		return &Response{}, err
	}

	for _, trip := range trips {

		if trip.CarShareID != "" {
			carShare, err := t.CarShareStorage.GetOne(trip.CarShareID)
			if err != nil {
				return &Response{}, err
			}
			trip.CarShare = &carShare
		}

		if trip.DriverID != "" {
			driver, err := t.UserStorage.GetOne(trip.DriverID)
			if err != nil {
				return &Response{}, err
			}
			trip.Driver = &driver
		}

		for _, passenger := range trip.Passengers {
			passenger, err := t.UserStorage.GetOne(passenger.GetID())
			if err != nil {
				return &Response{}, err
			}
			trip.Passengers = append(trip.Passengers, &passenger)
		}

		result = append(result, trip)
	}

	return &Response{Res: result}, nil
}

// FindOne trip
func (t TripResource) FindOne(ID string, r api2go.Request) (api2go.Responder, error) {
	trip, err := t.TripStorage.GetOne(ID)
	if err != nil {
		return &Response{}, err
	}

	if trip.CarShareID != "" {
		carShare, err := t.CarShareStorage.GetOne(trip.CarShareID)
		if err != nil {
			return &Response{}, err
		}
		trip.CarShare = &carShare
	}

	if trip.DriverID != "" {
		driver, err := t.UserStorage.GetOne(trip.DriverID)
		if err != nil {
			return &Response{}, err
		}
		trip.Driver = &driver
	}

	for _, passenger := range trip.Passengers {
		passenger, err := t.UserStorage.GetOne(passenger.GetID())
		if err != nil {
			return &Response{}, err
		}
		trip.Passengers = append(trip.Passengers, &passenger)
	}

	return &Response{Res: trip}, err
}

// Create a new trip
func (t TripResource) Create(obj interface{}, r api2go.Request) (api2go.Responder, error) {
	trip, ok := obj.(model.Trip)
	if !ok {
		return &Response{}, api2go.NewHTTPError(errors.New("Invalid instance given"), "Invalid instance given", http.StatusBadRequest)
	}

	if trip.CarShareID != "" {
		carShare, err := t.CarShareStorage.GetOne(trip.CarShareID)
		if err != nil {
			return &Response{}, err
		}
		trip.CarShare = &carShare
	}

	if trip.DriverID != "" {
		driver, err := t.UserStorage.GetOne(trip.DriverID)
		if err != nil {
			return &Response{}, err
		}
		trip.Driver = &driver
	}

	for _, passengerID := range trip.PassengerIDs {
		passenger, err := t.UserStorage.GetOne(passengerID)
		if err != nil {
			return &Response{}, err
		}
		trip.Passengers = append(trip.Passengers, &passenger)
	}

	trip.Scores = make(map[string]model.Score)
	if trip.CarShareID != "" {
		latestTrip, err := t.TripStorage.GetLatest(trip.CarShareID)
		if err != nil {
			return &Response{}, err
		}
		trip.CalculateScores(latestTrip.Scores)
	}

	trip.TimeStamp = t.Clock.Now().UTC()

	id, err := t.TripStorage.Insert(trip)
	if err != nil {
		return &Response{}, err
	}

	trip.SetID(id)

	return &Response{Res: trip, Code: http.StatusCreated}, nil
}

// Delete a trip :(
func (t TripResource) Delete(id string, r api2go.Request) (api2go.Responder, error) {
	err := t.TripStorage.Delete(id)
	return &Response{Code: http.StatusOK}, err
}

// Update a trip
func (t TripResource) Update(obj interface{}, r api2go.Request) (api2go.Responder, error) {
	trip, ok := obj.(model.Trip)
	if !ok {
		return &Response{}, api2go.NewHTTPError(errors.New("Invalid instance given"), "Invalid instance given", http.StatusBadRequest)
	}

	if trip.CarShareID != "" {
		latestTrip, err := t.TripStorage.GetLatest(trip.CarShareID)
		if err != nil {
			return &Response{}, err
		}
		trip.CalculateScores(latestTrip.Scores)
	}

	err := t.TripStorage.Update(trip)
	return &Response{Res: trip, Code: http.StatusNoContent}, err
}
<commit_msg>Fix compiler warnings and trip passenger population issue<commit_after>package resource

import (
	"errors"
	"net/http"

	"github.com/LewisWatson/carshare-back/model"
	"github.com/LewisWatson/carshare-back/storage"
	"github.com/benbjohnson/clock"
	"github.com/manyminds/api2go"
)

// TripResource for api2go routes
type TripResource struct {
	TripStorage     storage.TripStorage
	UserStorage     storage.UserStorage
	CarShareStorage storage.CarShareStorage
	Clock           clock.Clock
}

// FindAll trips
func (t TripResource) FindAll(r api2go.Request) (api2go.Responder, error) {
	var result []model.Trip

	trips, err := t.TripStorage.GetAll()
	if err != nil {
		return &Response{}, err
	}

	for _, trip := range trips {

		if trip.CarShareID != "" {
			var carShare model.CarShare
			carShare, err = t.CarShareStorage.GetOne(trip.CarShareID)
			if err != nil {
				return &Response{}, err
			}
			trip.CarShare = &carShare
		}

		if trip.DriverID != "" {
			var driver model.User
			driver, err = t.UserStorage.GetOne(trip.DriverID)
			if err != nil {
				return &Response{}, err
			}
			trip.Driver = &driver
		}

		for _, passengerID := range trip.PassengerIDs {
			var passenger model.User
			passenger, err = t.UserStorage.GetOne(passengerID)
			if err != nil {
				return &Response{}, err
			}
			trip.Passengers = append(trip.Passengers, &passenger)
		}

		result = append(result, trip)
	}

	return &Response{Res: result}, nil
}

// FindOne trip
func (t TripResource) FindOne(ID string, r api2go.Request) (api2go.Responder, error) {

	trip, err := t.TripStorage.GetOne(ID)
	if err != nil {
		return &Response{}, err
	}

	if trip.CarShareID != "" {
		var carShare model.CarShare
		carShare, err = t.CarShareStorage.GetOne(trip.CarShareID)
		if err != nil {
			return &Response{}, err
		}
		trip.CarShare = &carShare
	}

	if trip.DriverID != "" {
		var driver model.User
		driver, err = t.UserStorage.GetOne(trip.DriverID)
		if err != nil {
			return &Response{}, err
		}
		trip.Driver = &driver
	}

	for _, passengerID := range trip.PassengerIDs {
		var passenger model.User
		passenger, err = t.UserStorage.GetOne(passengerID)
		if err != nil {
			return &Response{}, err
		}
		trip.Passengers = append(trip.Passengers, &passenger)
	}

	return &Response{Res: trip}, err
}

// Create a new trip
func (t TripResource) Create(obj interface{}, r api2go.Request) (api2go.Responder, error) {
	trip, ok := obj.(model.Trip)
	if !ok {
		return &Response{}, api2go.NewHTTPError(errors.New("Invalid instance given"), "Invalid instance given", http.StatusBadRequest)
	}

	if trip.CarShareID != "" {
		carShare, err := t.CarShareStorage.GetOne(trip.CarShareID)
		if err != nil {
			return &Response{}, err
		}
		trip.CarShare = &carShare
	}

	if trip.DriverID != "" {
		driver, err := t.UserStorage.GetOne(trip.DriverID)
		if err != nil {
			return &Response{}, err
		}
		trip.Driver = &driver
	}

	for _, passengerID := range trip.PassengerIDs {
		passenger, err := t.UserStorage.GetOne(passengerID)
		if err != nil {
			return &Response{}, err
		}
		trip.Passengers = append(trip.Passengers, &passenger)
	}

	trip.Scores = make(map[string]model.Score)
	if trip.CarShareID != "" {
		latestTrip, err := t.TripStorage.GetLatest(trip.CarShareID)
		if err != nil {
			return &Response{}, err
		}
		trip.CalculateScores(latestTrip.Scores)
	}

	trip.TimeStamp = t.Clock.Now().UTC()

	id, err := t.TripStorage.Insert(trip)
	if err != nil {
		return &Response{}, err
	}

	trip.SetID(id)

	return &Response{Res: trip, Code: http.StatusCreated}, nil
}

// Delete a trip :(
func (t TripResource) Delete(id string, r api2go.Request) (api2go.Responder, error) {
	err := t.TripStorage.Delete(id)
	return &Response{Code: http.StatusOK}, err
}

// Update a trip
func (t TripResource) Update(obj interface{}, r api2go.Request) (api2go.Responder, error) {
	trip, ok := obj.(model.Trip)
	if !ok {
		return &Response{}, api2go.NewHTTPError(errors.New("Invalid instance given"), "Invalid instance given", http.StatusBadRequest)
	}

	if trip.CarShareID != "" {
		latestTrip, err := t.TripStorage.GetLatest(trip.CarShareID)
		if err != nil {
			return &Response{}, err
		}
		trip.CalculateScores(latestTrip.Scores)
	}

	err := t.TripStorage.Update(trip)
	return &Response{Res: trip, Code: http.StatusNoContent}, err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/DimensionDataResearch/go-dd-cloud-compute/compute"
	"github.com/hashicorp/terraform/helper/schema"
	"log"
	"time"
)

const (
	resourceKeyNetworkDomainName        = "name"
	resourceKeyNetworkDomainDescription = "description"
	resourceKeyNetworkDomainPlan        = "plan"
	resourceKeyNetworkDomainDataCenter  = "datacenter"
)

func resourceNetworkDomain() *schema.Resource {
	return &schema.Resource{
		Create: resourceNetworkDomainCreate,
		Read:   resourceNetworkDomainRead,
		Update: resourceNetworkDomainUpdate,
		Delete: resourceNetworkDomainDelete,

		Schema: map[string]*schema.Schema{
			resourceKeyNetworkDomainName: &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			resourceKeyNetworkDomainDescription: &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "",
			},
			resourceKeyNetworkDomainPlan: &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "ESSENTIALS",
			},
			resourceKeyNetworkDomainDataCenter: &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
		},
	}
}

// Create a network domain resource.
func resourceNetworkDomainCreate(data *schema.ResourceData, provider interface{}) error {
	var name, description, plan, dataCenterID string

	name = data.Get(resourceKeyNetworkDomainName).(string)
	description = data.Get(resourceKeyNetworkDomainDataCenter).(string)
	plan = data.Get(resourceKeyNetworkDomainPlan).(string)
	dataCenterID = data.Get(resourceKeyNetworkDomainDataCenter).(string)

	log.Printf("Create network domain '%s' in data center '%s' (plan = '%s', description = '%s').", name, dataCenterID, plan, description)

	providerClient := provider.(*compute.Client)

	networkDomainID, err := providerClient.DeployNetworkDomain(name, description, plan, dataCenterID)
	if err != nil {
		return err
	}

	data.SetId(networkDomainID)

	log.Printf("Network domain '%s' is being provisioned...", networkDomainID)

	timeout := time.NewTimer(60 * time.Second)
	defer timeout.Stop()

	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-timeout.C:
			return fmt.Errorf("Timed out after waiting %d seconds for provisioning of network domain '%s' to complete.", 60, networkDomainID)

		case <-ticker.C:
			log.Printf("Polling status for network domain '%s'...", networkDomainID)
			networkDomain, err := providerClient.GetNetworkDomain(networkDomainID)
			if err != nil {
				return err
			}

			if networkDomain == nil {
				return fmt.Errorf("Newly-created network domain was not found with Id '%s'.", networkDomainID)
			}

			switch networkDomain.State {
			case "PENDING_ADD":
				log.Printf("Network domain '%s' is still being provisioned...", networkDomainID)

				continue
			case "NORMAL":
				log.Printf("Network domain '%s' has been successfully provisioned.", networkDomainID)

				return nil
			default:
				log.Printf("Unexpected status for network domain '%s' ('%s').", networkDomainID, networkDomain.State)

				return fmt.Errorf("Failed to provision network domain '%s' ('%s'): encountered unexpected state '%s'.", networkDomainID, name, networkDomain.State)
			}
		}
	}
}

// Read a network domain resource.
func resourceNetworkDomainRead(data *schema.ResourceData, provider interface{}) error {
	var name, description, plan, dataCenterID string

	id := data.Id()
	name = data.Get(resourceKeyNetworkDomainName).(string)
	description = data.Get(resourceKeyNetworkDomainDescription).(string)
	plan = data.Get(resourceKeyNetworkDomainPlan).(string)
	dataCenterID = data.Get(resourceKeyNetworkDomainDataCenter).(string)

	log.Printf("Read network domain '%s' (Id = '%s') in data center '%s' (plan = '%s', description = '%s').", name, id, dataCenterID, plan, description)

	providerClient := provider.(*compute.Client)

	networkDomain, err := providerClient.GetNetworkDomain(id)
	if err != nil {
		return err
	}

	if networkDomain != nil {
		data.Set(resourceKeyNetworkDomainName, networkDomain.Name)
		data.Set(resourceKeyNetworkDomainDescription, networkDomain.Description)
		data.Set(resourceKeyNetworkDomainPlan, networkDomain.Type)
		data.Set(resourceKeyNetworkDomainDataCenter, networkDomain.DatacenterID)
	} else {
		data.SetId("") // Mark resource as deleted.
	}

	return nil
}

// Update a network domain resource.
func resourceNetworkDomainUpdate(data *schema.ResourceData, provider interface{}) error {
	var id, name, description, plan string

	id = data.Id()

	if data.HasChange(resourceKeyNetworkDomainName) {
		name = data.Get(resourceKeyNetworkDomainName).(string)
	}

	if data.HasChange(resourceKeyNetworkDomainDescription) {
		description = data.Get(resourceKeyNetworkDomainDescription).(string)
	}

	if data.HasChange(resourceKeyNetworkDomainPlan) {
		plan = data.Get(resourceKeyNetworkDomainPlan).(string)
	}

	log.Printf("Update network domain '%s' (Name = '%s', Description = '%s', Plan = '%s').", data.Id(), name, description, plan)

	providerClient := provider.(*compute.Client)

	return providerClient.EditNetworkDomain(id, name, description, plan)
}

// Delete a network domain resource.
func resourceNetworkDomainDelete(data *schema.ResourceData, provider interface{}) error {
	id := data.Id()
	name := data.Get(resourceKeyNetworkDomainName).(string)
	dataCenterID := data.Get(resourceKeyNetworkDomainDataCenter).(string)

	log.Printf("Delete network domain '%s' ('%s') in data center '%s'.", id, name, dataCenterID)

	providerClient := provider.(*compute.Client)
	err := providerClient.DeleteNetworkDomain(id)

	return err
}
<commit_msg>Add calculated 'nat_ipv4_address' property to 'ddcloud_networkdomain' resource type.<commit_after>package main

import (
	"fmt"
	"github.com/DimensionDataResearch/go-dd-cloud-compute/compute"
	"github.com/hashicorp/terraform/helper/schema"
	"log"
	"time"
)

const (
	resourceKeyNetworkDomainName           = "name"
	resourceKeyNetworkDomainDescription    = "description"
	resourceKeyNetworkDomainPlan           = "plan"
	resourceKeyNetworkDomainDataCenter     = "datacenter"
	resourceKeyNetworkDomainNatIPv4Address = "nat_ipv4_address"
)

func resourceNetworkDomain() *schema.Resource {
	return &schema.Resource{
		Create: resourceNetworkDomainCreate,
		Read:   resourceNetworkDomainRead,
		Update: resourceNetworkDomainUpdate,
		Delete: resourceNetworkDomainDelete,

		Schema: map[string]*schema.Schema{
			resourceKeyNetworkDomainName: &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			resourceKeyNetworkDomainDescription: &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "",
			},
			resourceKeyNetworkDomainPlan: &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "ESSENTIALS",
			},
			resourceKeyNetworkDomainDataCenter: &schema.Schema{
				Type:     schema.TypeString,
				ForceNew: true,
				Required: true,
			},
			resourceKeyNetworkDomainNatIPv4Address: &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

// Create a network domain resource.
func resourceNetworkDomainCreate(data *schema.ResourceData, provider interface{}) error {
	var name, description, plan, dataCenterID string

	name = data.Get(resourceKeyNetworkDomainName).(string)
	description = data.Get(resourceKeyNetworkDomainDataCenter).(string)
	plan = data.Get(resourceKeyNetworkDomainPlan).(string)
	dataCenterID = data.Get(resourceKeyNetworkDomainDataCenter).(string)

	log.Printf("Create network domain '%s' in data center '%s' (plan = '%s', description = '%s').", name, dataCenterID, plan, description)

	providerClient := provider.(*compute.Client)

	networkDomainID, err := providerClient.DeployNetworkDomain(name, description, plan, dataCenterID)
	if err != nil {
		return err
	}

	data.SetId(networkDomainID)

	log.Printf("Network domain '%s' is being provisioned...", networkDomainID)

	timeout := time.NewTimer(60 * time.Second)
	defer timeout.Stop()

	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-timeout.C:
			return fmt.Errorf("Timed out after waiting %d seconds for provisioning of network domain '%s' to complete.", 60, networkDomainID)

		case <-ticker.C:
			log.Printf("Polling status for network domain '%s'...", networkDomainID)
			networkDomain, err := providerClient.GetNetworkDomain(networkDomainID)
			if err != nil {
				return err
			}

			if networkDomain == nil {
				return fmt.Errorf("Newly-created network domain was not found with Id '%s'.", networkDomainID)
			}

			switch networkDomain.State {
			case "PENDING_ADD":
				log.Printf("Network domain '%s' is still being provisioned...", networkDomainID)

				continue
			case "NORMAL":
				log.Printf("Network domain '%s' has been successfully provisioned.", networkDomainID)

				// Capture IPv4 NAT address.
				data.Set(resourceKeyNetworkDomainNatIPv4Address, networkDomain.NatIPv4Address)

				return nil
			default:
				log.Printf("Unexpected status for network domain '%s' ('%s').", networkDomainID, networkDomain.State)

				return fmt.Errorf("Failed to provision network domain '%s' ('%s'): encountered unexpected state '%s'.", networkDomainID, name, networkDomain.State)
			}
		}
	}
}

// Read a network domain resource.
func resourceNetworkDomainRead(data *schema.ResourceData, provider interface{}) error {
	var name, description, plan, dataCenterID string

	id := data.Id()
	name = data.Get(resourceKeyNetworkDomainName).(string)
	description = data.Get(resourceKeyNetworkDomainDescription).(string)
	plan = data.Get(resourceKeyNetworkDomainPlan).(string)
	dataCenterID = data.Get(resourceKeyNetworkDomainDataCenter).(string)

	log.Printf("Read network domain '%s' (Id = '%s') in data center '%s' (plan = '%s', description = '%s').", name, id, dataCenterID, plan, description)

	providerClient := provider.(*compute.Client)

	networkDomain, err := providerClient.GetNetworkDomain(id)
	if err != nil {
		return err
	}

	if networkDomain != nil {
		data.Set(resourceKeyNetworkDomainName, networkDomain.Name)
		data.Set(resourceKeyNetworkDomainDescription, networkDomain.Description)
		data.Set(resourceKeyNetworkDomainPlan, networkDomain.Type)
		data.Set(resourceKeyNetworkDomainDataCenter, networkDomain.DatacenterID)
		data.Set(resourceKeyNetworkDomainNatIPv4Address, networkDomain.NatIPv4Address)
	} else {
		data.SetId("") // Mark resource as deleted.
	}

	return nil
}

// Update a network domain resource.
func resourceNetworkDomainUpdate(data *schema.ResourceData, provider interface{}) error {
	var id, name, description, plan string

	id = data.Id()

	if data.HasChange(resourceKeyNetworkDomainName) {
		name = data.Get(resourceKeyNetworkDomainName).(string)
	}

	if data.HasChange(resourceKeyNetworkDomainDescription) {
		description = data.Get(resourceKeyNetworkDomainDescription).(string)
	}

	if data.HasChange(resourceKeyNetworkDomainPlan) {
		plan = data.Get(resourceKeyNetworkDomainPlan).(string)
	}

	log.Printf("Update network domain '%s' (Name = '%s', Description = '%s', Plan = '%s').", data.Id(), name, description, plan)

	providerClient := provider.(*compute.Client)

	return providerClient.EditNetworkDomain(id, name, description, plan)
}

// Delete a network domain resource.
func resourceNetworkDomainDelete(data *schema.ResourceData, provider interface{}) error {
	id := data.Id()
	name := data.Get(resourceKeyNetworkDomainName).(string)
	dataCenterID := data.Get(resourceKeyNetworkDomainDataCenter).(string)

	log.Printf("Delete network domain '%s' ('%s') in data center '%s'.", id, name, dataCenterID)

	providerClient := provider.(*compute.Client)
	err := providerClient.DeleteNetworkDomain(id)

	return err
}
<|endoftext|>


<commit_before>package model

import (
	"database/sql"
	"fmt"
	"log"

	"github.com/ViBiOh/funds/db"
)

const alertsOpenedQuery = `
SELECT
  isin,
  type,
  score
FROM
  alerts
WHERE
  isin IN (
    SELECT
      isin
    FROM
      alerts
    GROUP BY
      isin
    HAVING
      MOD(COUNT(type), 2) = 1
  )
ORDER BY
  isin          ASC,
  creation_date DESC
`

const alertsCreateQuery = `
INSERT INTO
  alerts
(
  isin,
  score,
  type
) VALUES (
  $1,
  $2,
  $3
)
`

// ReadAlertsOpened retrieves current Alerts (only one mail sent)
func ReadAlertsOpened() (alerts []Alert, err error) {
	rows, err := db.Query(alertsOpenedQuery)
	if err != nil {
		err = fmt.Errorf(`Error while querying opened alerts: %v`, err)
		return
	}

	defer func() {
		if endErr := rows.Close(); endErr != nil {
			if err == nil {
				err = endErr
			} else {
				log.Printf(`Error while closing opened alerts: %v`, endErr)
			}
		}
	}()

	var (
		isin      string
		alertType string
		score     float64
	)

	for rows.Next() {
		if err = rows.Scan(&isin, &alertType, &score); err != nil {
			err = fmt.Errorf(`Error while scanning opened alerts: %v`, err)
			return
		}

		alerts = append(alerts, Alert{Isin: isin, AlertType: alertType, Score: score})
	}

	return
}

// SaveAlert saves Alert
func SaveAlert(alert *Alert, tx *sql.Tx) (err error) {
	if alert == nil {
		return fmt.Errorf(`Unable to save nil Alert`)
	}

	var usedTx *sql.Tx

	if usedTx, err = db.GetTx(tx); err != nil {
		err = fmt.Errorf(`Error while getting transaction for creating alert: %v`, err)
		return
	}

	if usedTx != tx {
		defer func() {
			err = db.EndTx(usedTx, err)
		}()
	}

	if _, err = usedTx.Exec(alertsCreateQuery, alert.Isin, alert.Score, alert.AlertType); err != nil {
		err = fmt.Errorf(`Error while creating alert for isin=%s: %v`, alert.Isin, err)
	}

	return
}
<commit_msg>Update alert_db.go<commit_after>package model

import (
	"database/sql"
	"fmt"
	"log"

	"github.com/ViBiOh/funds/db"
)

const alertsOpenedQuery = `
SELECT
  isin,
  type,
  score
FROM
  alerts
WHERE
  isin IN (
    SELECT
      isin
    FROM
      alerts
    GROUP BY
      isin
    HAVING
      MOD(COUNT(type), 2) = 1
  )
ORDER BY
  isin          ASC,
  creation_date DESC
`

const alertsCreateQuery = `
INSERT INTO
  alerts
(
  isin,
  score,
  type
) VALUES (
  $1,
  $2,
  $3
)
`

// ReadAlertsOpened retrieves current Alerts (only one mail sent)
func ReadAlertsOpened() (alerts []Alert, err error) {
	rows, err := db.Query(alertsOpenedQuery)
	if err != nil {
		err = fmt.Errorf(`Error while querying opened alerts: %v`, err)
		return
	}

	defer func() {
		err = db.RowsClose(rows, err, `opened alerts`)
	}()

	var (
		isin      string
		alertType string
		score     float64
	)

	for rows.Next() {
		if err = rows.Scan(&isin, &alertType, &score); err != nil {
			err = fmt.Errorf(`Error while scanning opened alerts: %v`, err)
			return
		}

		alerts = append(alerts, Alert{Isin: isin, AlertType: alertType, Score: score})
	}

	return
}

// SaveAlert saves Alert
func SaveAlert(alert *Alert, tx *sql.Tx) (err error) {
	if alert == nil {
		return fmt.Errorf(`Unable to save nil Alert`)
	}

	var usedTx *sql.Tx

	if usedTx, err = db.GetTx(tx); err != nil {
		err = fmt.Errorf(`Error while getting transaction for creating alert: %v`, err)
		return
	}

	if usedTx != tx {
		defer func() {
			err = db.EndTx(usedTx, err)
		}()
	}

	if _, err = usedTx.Exec(alertsCreateQuery, alert.Isin, alert.Score, alert.AlertType); err != nil {
		err = fmt.Errorf(`Error while creating alert for isin=%s: %v`, alert.Isin, err)
	}

	return
}
<|endoftext|>


<commit_before>package cache

import (
	"fmt"
	"sort"
	"time"

	"devroom.ru/lomik/carbon/points"

	"github.com/Sirupsen/logrus"
)

type queue []*queueItem

func (v queue) Len() int           { return len(v) }
func (v queue) Swap(i, j int)      { v[i], v[j] = v[j], v[i] }
func (v queue) Less(i, j int) bool { return v[i].count < v[j].count }

// Cache stores and aggregate metrics in memory
type Cache struct {
	data        map[string]*points.Points
	size        int
	maxSize     int
	inputChan   chan *points.Points // from receivers
	outputChan  chan *points.Points // to persisters
	queryChan   chan *Query         // from carbonlink
	exitChan    chan bool           // close for stop worker
	graphPrefix string
	queryCnt    int
	oversizeCnt int // drop packages if cache full
	queue       queue
}

// New create Cache instance and run in/out goroutine
func New() *Cache {
	cache := &Cache{
		data:        make(map[string]*points.Points, 0),
		size:        0,
		maxSize:     1000000,
		inputChan:   make(chan *points.Points, 1024),
		exitChan:    make(chan bool),
		queryChan:   make(chan *Query, 16),
		graphPrefix: "carbon.",
		queryCnt:    0,
		queue:       make(queue, 0),
	}
	return cache
}

// Get any key/values pair from Cache
func (c *Cache) Get() *points.Points {
	for {
		size := len(c.queue)
		if size == 0 {
			break
		}
		cacheRecord := c.queue[size-1]
		c.queue = c.queue[:size-1]

		if values, ok := c.data[cacheRecord.metric]; ok {
			return values
		}
	}
	for _, values := range c.data {
		return values
	}
	return nil
}

// Remove key from cache
func (c *Cache) Remove(key string) {
	if value, exists := c.data[key]; exists {
		c.size -= len(value.Data)
		delete(c.data, key)
	}
}

// Pop return and remove next for save point from cache
func (c *Cache) Pop() *points.Points {
	v := c.Get()
	if v != nil {
		c.Remove(v.Metric)
	}
	return v
}

// Add points to cache
func (c *Cache) Add(p *points.Points) {
	if values, exists := c.data[p.Metric]; exists {
		values.Data = append(values.Data, p.Data...)
	} else {
		c.data[p.Metric] = p
	}
	c.size += len(p.Data)
}

// SetGraphPrefix for internal cache metrics
func (c *Cache) SetGraphPrefix(prefix string) {
	c.graphPrefix = prefix
}

// SetMaxSize of cache
func (c *Cache) SetMaxSize(maxSize int) {
	c.maxSize = maxSize
}

// Size returns size
func (c *Cache) Size() int {
	return c.size
}

type queueItem struct {
	metric string
	count  int
}

// stat send internal statistics of cache
func (c *Cache) stat(metric string, value float64) {
	key := fmt.Sprintf("%scache.%s", c.graphPrefix, metric)
	c.Add(points.OnePoint(key, value, time.Now().Unix()))
	c.queue = append(c.queue, &queueItem{key, 1})
}

// doCheckpoint reorder save queue, add carbon metrics to queue
func (c *Cache) doCheckpoint() {
	start := time.Now()

	newQueue := make(queue, 0)

	for key, values := range c.data {
		newQueue = append(newQueue, &queueItem{key, len(values.Data)})
	}

	sort.Sort(newQueue)

	c.queue = newQueue

	worktime := time.Now().Sub(start)

	c.stat("size", float64(c.size))
	c.stat("metrics", float64(len(c.data)))
	c.stat("queries", float64(c.queryCnt))
	c.stat("oversize", float64(c.oversizeCnt))
	c.stat("checkpoint_time", worktime.Seconds())

	logrus.WithFields(logrus.Fields{
		"time":     worktime.String(),
		"size":     c.size,
		"metrics":  len(c.data),
		"queries":  c.queryCnt,
		"oversize": c.oversizeCnt,
	}).Info("doCheckpoint()")

	c.queryCnt = 0
	c.oversizeCnt = 0
}

func (c *Cache) worker() {
	var values *points.Points
	var sendTo chan *points.Points

	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for {
		if values == nil {
			values = c.Pop()

			if values != nil {
				sendTo = c.outputChan
			} else {
				sendTo = nil
			}
		}

		select {
		case <-ticker.C: // checkpoint
			c.doCheckpoint()
		case query := <-c.queryChan: // carbonlink
			c.queryCnt++
			reply := NewReply()

			if values != nil && values.Metric == query.Metric {
				reply.Points = values.Copy()
			} else if v, ok := c.data[query.Metric]; ok {
				reply.Points = v.Copy()
			}

			query.ReplyChan <- reply
		case sendTo <- values: // to persister
			values = nil
		case msg := <-c.inputChan: // from receiver
			if c.maxSize == 0 || c.size < c.maxSize {
				c.Add(msg)
			} else {
				c.oversizeCnt++
			}
		case <-c.exitChan: // exit
			break
		}
	}

}

// In returns input channel
func (c *Cache) In() chan *points.Points {
	return c.inputChan
}

// Out returns output channel
func (c *Cache) Out() chan *points.Points {
	return c.outputChan
}

// Query returns carbonlink query channel
func (c *Cache) Query() chan *Query {
	return c.queryChan
}

// SetOutputChanSize ...
func (c *Cache) SetOutputChanSize(size int) {
	c.outputChan = make(chan *points.Points, size)
}

// Start worker
func (c *Cache) Start() {
	if c.outputChan == nil {
		c.outputChan = make(chan *points.Points, 1024)
	}
	go c.worker()
}

// Stop worker
func (c *Cache) Stop() {
	close(c.exitChan)
}
<commit_msg>rename metrics<commit_after>package cache

import (
	"fmt"
	"sort"
	"time"

	"devroom.ru/lomik/carbon/points"

	"github.com/Sirupsen/logrus"
)

type queue []*queueItem

func (v queue) Len() int           { return len(v) }
func (v queue) Swap(i, j int)      { v[i], v[j] = v[j], v[i] }
func (v queue) Less(i, j int) bool { return v[i].count < v[j].count }

// Cache stores and aggregate metrics in memory
type Cache struct {
	data        map[string]*points.Points
	size        int
	maxSize     int
	inputChan   chan *points.Points // from receivers
	outputChan  chan *points.Points // to persisters
	queryChan   chan *Query         // from carbonlink
	exitChan    chan bool           // close for stop worker
	graphPrefix string
	queryCnt    int
	oversizeCnt int // drop packages if cache full
	queue       queue
}

// New create Cache instance and run in/out goroutine
func New() *Cache {
	cache := &Cache{
		data:        make(map[string]*points.Points, 0),
		size:        0,
		maxSize:     1000000,
		inputChan:   make(chan *points.Points, 1024),
		exitChan:    make(chan bool),
		queryChan:   make(chan *Query, 16),
		graphPrefix: "carbon.",
		queryCnt:    0,
		queue:       make(queue, 0),
	}
	return cache
}

// Get any key/values pair from Cache
func (c *Cache) Get() *points.Points {
	for {
		size := len(c.queue)
		if size == 0 {
			break
		}
		cacheRecord := c.queue[size-1]
		c.queue = c.queue[:size-1]

		if values, ok := c.data[cacheRecord.metric]; ok {
			return values
		}
	}
	for _, values := range c.data {
		return values
	}
	return nil
}

// Remove key from cache
func (c *Cache) Remove(key string) {
	if value, exists := c.data[key]; exists {
		c.size -= len(value.Data)
		delete(c.data, key)
	}
}

// Pop return and remove next for save point from cache
func (c *Cache) Pop() *points.Points {
	v := c.Get()
	if v != nil {
		c.Remove(v.Metric)
	}
	return v
}

// Add points to cache
func (c *Cache) Add(p *points.Points) {
	if values, exists := c.data[p.Metric]; exists {
		values.Data = append(values.Data, p.Data...)
	} else {
		c.data[p.Metric] = p
	}
	c.size += len(p.Data)
}

// SetGraphPrefix for internal cache metrics
func (c *Cache) SetGraphPrefix(prefix string) {
	c.graphPrefix = prefix
}

// SetMaxSize of cache
func (c *Cache) SetMaxSize(maxSize int) {
	c.maxSize = maxSize
}

// Size returns size
func (c *Cache) Size() int {
	return c.size
}

type queueItem struct {
	metric string
	count  int
}

// stat send internal statistics of cache
func (c *Cache) stat(metric string, value float64) {
	key := fmt.Sprintf("%scache.%s", c.graphPrefix, metric)
	c.Add(points.OnePoint(key, value, time.Now().Unix()))
	c.queue = append(c.queue, &queueItem{key, 1})
}

// doCheckpoint reorder save queue, add carbon metrics to queue
func (c *Cache) doCheckpoint() {
	start := time.Now()

	newQueue := make(queue, 0)

	for key, values := range c.data {
		newQueue = append(newQueue, &queueItem{key, len(values.Data)})
	}

	sort.Sort(newQueue)

	c.queue = newQueue

	worktime := time.Now().Sub(start)

	c.stat("size", float64(c.size))
	c.stat("metrics", float64(len(c.data)))
	c.stat("queries", float64(c.queryCnt))
	c.stat("oversizeDrops", float64(c.oversizeCnt))
	c.stat("checkpointTime", worktime.Seconds())

	logrus.WithFields(logrus.Fields{
		"time":     worktime.String(),
		"size":     c.size,
		"metrics":  len(c.data),
		"queries":  c.queryCnt,
		"oversize": c.oversizeCnt,
	}).Info("doCheckpoint()")

	c.queryCnt = 0
	c.oversizeCnt = 0
}

func (c *Cache) worker() {
	var values *points.Points
	var sendTo chan *points.Points

	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for {
		if values == nil {
			values = c.Pop()

			if values != nil {
				sendTo = c.outputChan
			} else {
				sendTo = nil
			}
		}

		select {
		case <-ticker.C: // checkpoint
			c.doCheckpoint()
		case query := <-c.queryChan: // carbonlink
			c.queryCnt++
			reply := NewReply()

			if values != nil && values.Metric == query.Metric {
				reply.Points = values.Copy()
			} else if v, ok := c.data[query.Metric]; ok {
				reply.Points = v.Copy()
			}

			query.ReplyChan <- reply
		case sendTo <- values: // to persister
			values = nil
		case msg := <-c.inputChan: // from receiver
			if c.maxSize == 0 || c.size < c.maxSize {
				c.Add(msg)
			} else {
				c.oversizeCnt++
			}
		case <-c.exitChan: // exit
			break
		}
	}

}

// In returns input channel
func (c *Cache) In() chan *points.Points {
	return c.inputChan
}

// Out returns output channel
func (c *Cache) Out() chan *points.Points {
	return c.outputChan
}

// Query returns carbonlink query channel
func (c *Cache) Query() chan *Query {
	return c.queryChan
}

// SetOutputChanSize ...
func (c *Cache) SetOutputChanSize(size int) {
	c.outputChan = make(chan *points.Points, size)
}

// Start worker
func (c *Cache) Start() {
	if c.outputChan == nil {
		c.outputChan = make(chan *points.Points, 1024)
	}
	go c.worker()
}

// Stop worker
func (c *Cache) Stop() {
	close(c.exitChan)
}
<|endoftext|>


<commit_before>package cache

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/asjoyner/shade"
	"github.com/asjoyner/shade/drive"
)

var (
	defaultCacheDir = path.Join(shade.ConfigDir(), "cache")
	cacheDir        = flag.String("cache", defaultCacheDir, "Where to store the drive data cache")
	cacheDebug      = flag.Bool("cacheDebug", false, "Print cache debugging traces")
)

// A very compact representation of a file
type Node struct {
	Filename     string
	Filesize     uint64
	ModifiedTime time.Time
	Sha256sum    []byte // the sha of the associated File
	// TODO(asjoyner): use a struct{} here for efficiency?
	Children map[string]bool
}

func (n *Node) Synthetic() bool {
	if n.Sha256sum == nil {
		return true
	}
	return false
}

// Reader is a wrapper around a slice of cloud storage backends.  It presents an
// interface to query for the union of the set of known files by an integer ID,
// which will be stable across single processes invoking this cache, a node
// representing that file, or a single chunk of that file.  It can also cache a
// configurable quantity of chunks to disk.
//
// TODO(asjoyner): implement disk caching of data blocks.
type Reader struct {
	clients []drive.Client
	nodes   map[string]Node // full path to node
	sync.RWMutex
}

func NewReader(clients []drive.Client, t *time.Ticker) (*Reader, error) {
	c := &Reader{
		clients: clients,
		nodes: map[string]Node{
			"/": {
				Filename: "/",
				Children: make(map[string]bool),
			}},
	}
	if err := c.refresh(); err != nil {
		return nil, fmt.Errorf("initializing cache: %s", err)
	}
	go c.periodicRefresh(t)
	return c, nil
}

// NodeByPath returns the current file object for a given path.
func (c *Reader) NodeByPath(p string) (Node, error) {
	c.RLock()
	defer c.RUnlock()
	if n, ok := c.nodes[p]; ok {
		return n, nil
	}
	log.Printf("%+v\n", c.nodes)
	return Node{}, fmt.Errorf("no such node: %q", p)
}

// FileByNode returns the full shade.File object for a given node.
func (c *Reader) FileByNode(n Node) (*shade.File, error) {
	if n.Synthetic() {
		return nil, errors.New("no shade.File defined")
	}
	var fj []byte
	var err error
	for _, client := range c.clients {
		fj, err = client.GetChunk(n.Sha256sum)
		if err != nil {
			log.Printf("Failed to fetch %s: %s", n.Sha256sum, err)
			continue
		}
	}
	if fj == nil || len(fj) == 0 {
		return nil, fmt.Errorf("Could not find JSON for node: %q", n.Filename)
	}
	file := &shade.File{}
	if err := json.Unmarshal(fj, file); err != nil {
		return nil, fmt.Errorf("Failed to unmarshal sha256sum %x: %s", n.Sha256sum, err)
	}
	return file, nil
}

func (c *Reader) HasChild(parent, child string) bool {
	c.RLock()
	defer c.RUnlock()
	return c.nodes[parent].Children[child]
}

// NumNodes returns the number of nodes (files + synthetic directories) in the
// system.
func (c *Reader) NumNodes() int {
	c.RLock()
	defer c.RUnlock()
	return len(c.nodes)
}

func (c *Reader) GetChunk(sha256sum []byte) {
}

// refresh updates the cache
func (c *Reader) refresh() error {
	debug("Begining cache refresh cycle.")
	// key is a string([]byte) representation of the file's SHA2
	knownNodes := make(map[string]bool)
	for _, client := range c.clients {
		lfm, err := client.ListFiles()
		if err != nil {
			return fmt.Errorf("%q ListFiles(): %s", client.GetConfig().Provider, err)
		}
		debug(fmt.Sprintf("Found %d file(s) via %s", len(lfm), client.GetConfig().Provider))
		// fetch all those files into the local disk cache
		for _, sha256sum := range lfm {
			// check if we have already processed this Node
			if knownNodes[string(sha256sum)] {
				continue // we've already processed this file
			}

			// fetch the file Chunk
			f, err := client.GetChunk(sha256sum)
			if err != nil {
				// TODO(asjoyner): if !client.Local()... retry?
				log.Printf("Failed to fetch file %x: %s", sha256sum, err)
				continue
			}
			// ensure this file is known to all the writable clients
			for _, lc := range c.clients {
				if lc.GetConfig().Write {
					if err := lc.PutFile(sha256sum, f); err != nil {
						log.Printf("Failed to store checksum %x in %s: %s", sha256sum, client.GetConfig().Provider, err)
					}
				}
			}

			// unmarshal and populate c.nodes as the shade.files go by
			file := &shade.File{}
			if err := json.Unmarshal(f, file); err != nil {
				log.Printf("Failed to unmarshal file %x: %s", sha256sum, err)
				continue
			}
			node := Node{
				Filename:     file.Filename,
				Filesize:     uint64(file.Filesize),
				ModifiedTime: file.ModifiedTime,
				Sha256sum:    sha256sum,
				Children:     nil,
			}
			c.Lock()
			existing, ok := c.nodes[node.Filename]
			// TODO(asjoyner): handle file + directory collisions
			if ok && existing.ModifiedTime.After(node.ModifiedTime) {
				c.Unlock()
				continue
			}
			c.nodes[node.Filename] = node
			c.addParents(node.Filename)
			c.Unlock()
			knownNodes[string(sha256sum)] = true
		}
	}
	debug(fmt.Sprintf("Refresh complete with %d file(s).", len(knownNodes)))
	return nil
}

// recursive function to update parent dirs
func (c *Reader) addParents(filepath string) {
	dir, f := path.Split(filepath)
	if dir == "" {
		dir = "/"
	} else {
		dir = strings.TrimSuffix(dir, "/")
	}
	debug(fmt.Sprintf("adding %q as a child of %q", f, dir))
	// TODO(asjoyner): handle file + directory collisions
	parent, ok := c.nodes[dir]
	if !ok {
		// if the parent node doesn't yet exist, initialize it
		c.nodes[dir] = Node{
			Filename: dir,
			Children: map[string]bool{f: true},
		}
	} else {
		parent.Children[f] = true
	}
	if dir == "/" {
		return
	} else {
		c.addParents(dir)
	}
}

func (c *Reader) periodicRefresh(t *time.Ticker) {
	for {
		<-t.C
		c.refresh()
	}
}

func debug(args interface{}) {
	if !*cacheDebug {
		return
	}
	log.Printf("CACHE: %s\n", args)
}
<commit_msg>factored out some duplicate code<commit_after>package cache

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/asjoyner/shade"
	"github.com/asjoyner/shade/drive"
)

var (
	defaultCacheDir = path.Join(shade.ConfigDir(), "cache")
	cacheDir        = flag.String("cache", defaultCacheDir, "Where to store the drive data cache")
	cacheDebug      = flag.Bool("cacheDebug", false, "Print cache debugging traces")
)

// A very compact representation of a file
type Node struct {
	Filename     string
	Filesize     uint64
	ModifiedTime time.Time
	Sha256sum    []byte // the sha of the associated File
	// TODO(asjoyner): use a struct{} here for efficiency?
	Children map[string]bool
}

func (n *Node) Synthetic() bool {
	if n.Sha256sum == nil {
		return true
	}
	return false
}

// Reader is a wrapper around a slice of cloud storage backends.  It presents an
// interface to query for the union of the set of known files by an integer ID,
// which will be stable across single processes invoking this cache, a node
// representing that file, or a single chunk of that file.  It can also cache a
// configurable quantity of chunks to disk.
//
// TODO(asjoyner): implement disk caching of data blocks.
type Reader struct {
	clients []drive.Client
	nodes   map[string]Node // full path to node
	sync.RWMutex
}

func NewReader(clients []drive.Client, t *time.Ticker) (*Reader, error) {
	c := &Reader{
		clients: clients,
		nodes: map[string]Node{
			"/": {
				Filename: "/",
				Children: make(map[string]bool),
			}},
	}
	if err := c.refresh(); err != nil {
		return nil, fmt.Errorf("initializing cache: %s", err)
	}
	go c.periodicRefresh(t)
	return c, nil
}

// NodeByPath returns the current file object for a given path.
func (c *Reader) NodeByPath(p string) (Node, error) {
	c.RLock()
	defer c.RUnlock()
	if n, ok := c.nodes[p]; ok {
		return n, nil
	}
	log.Printf("%+v\n", c.nodes)
	return Node{}, fmt.Errorf("no such node: %q", p)
}

func unmarshalChunk(fj, sha []byte) (*shade.File, error) {
	file := &shade.File{}
	if err := json.Unmarshal(fj, file); err != nil {
		return nil, fmt.Errorf("Failed to unmarshal sha256sum %x: %s", sha, err)
	}
	return file, nil
}

// FileByNode returns the full shade.File object for a given node.
func (c *Reader) FileByNode(n Node) (*shade.File, error) {
	if n.Synthetic() {
		return nil, errors.New("no shade.File defined")
	}
	var fj []byte
	var err error
	for _, client := range c.clients {
		fj, err = client.GetChunk(n.Sha256sum)
		if err != nil {
			log.Printf("Failed to fetch %s: %s", n.Sha256sum, err)
			continue
		}
	}
	if fj == nil || len(fj) == 0 {
		return nil, fmt.Errorf("Could not find JSON for node: %q", n.Filename)
	}
	return unmarshalChunk(fj, n.Sha256sum)
}

func (c *Reader) HasChild(parent, child string) bool {
	c.RLock()
	defer c.RUnlock()
	return c.nodes[parent].Children[child]
}

// NumNodes returns the number of nodes (files + synthetic directories) in the
// system.
func (c *Reader) NumNodes() int {
	c.RLock()
	defer c.RUnlock()
	return len(c.nodes)
}

func (c *Reader) GetChunk(sha256sum []byte) {
}

// refresh updates the cache
func (c *Reader) refresh() error {
	debug("Begining cache refresh cycle.")
	// key is a string([]byte) representation of the file's SHA2
	knownNodes := make(map[string]bool)
	for _, client := range c.clients {
		lfm, err := client.ListFiles()
		if err != nil {
			return fmt.Errorf("%q ListFiles(): %s", client.GetConfig().Provider, err)
		}
		debug(fmt.Sprintf("Found %d file(s) via %s", len(lfm), client.GetConfig().Provider))
		// fetch all those files into the local disk cache
		for _, sha256sum := range lfm {
			// check if we have already processed this Node
			if knownNodes[string(sha256sum)] {
				continue // we've already processed this file
			}

			// fetch the file Chunk
			f, err := client.GetChunk(sha256sum)
			if err != nil {
				// TODO(asjoyner): if !client.Local()... retry?
				log.Printf("Failed to fetch file %x: %s", sha256sum, err)
				continue
			}
			// ensure this file is known to all the writable clients
			for _, lc := range c.clients {
				if lc.GetConfig().Write {
					if err := lc.PutFile(sha256sum, f); err != nil {
						log.Printf("Failed to store checksum %x in %s: %s", sha256sum, client.GetConfig().Provider, err)
					}
				}
			}

			// unmarshal and populate c.nodes as the shade.files go by
			file, err := unmarshalChunk(f, sha256sum)
			if err != nil {
				log.Printf("%v", err)
				continue
			}
			node := Node{
				Filename:     file.Filename,
				Filesize:     uint64(file.Filesize),
				ModifiedTime: file.ModifiedTime,
				Sha256sum:    sha256sum,
				Children:     nil,
			}
			c.Lock()
			existing, ok := c.nodes[node.Filename]
			// TODO(asjoyner): handle file + directory collisions
			if ok && existing.ModifiedTime.After(node.ModifiedTime) {
				c.Unlock()
				continue
			}
			c.nodes[node.Filename] = node
			c.addParents(node.Filename)
			c.Unlock()
			knownNodes[string(sha256sum)] = true
		}
	}
	debug(fmt.Sprintf("Refresh complete with %d file(s).", len(knownNodes)))
	return nil
}

// recursive function to update parent dirs
func (c *Reader) addParents(filepath string) {
	dir, f := path.Split(filepath)
	if dir == "" {
		dir = "/"
	} else {
		dir = strings.TrimSuffix(dir, "/")
	}
	debug(fmt.Sprintf("adding %q as a child of %q", f, dir))
	// TODO(asjoyner): handle file + directory collisions
	parent, ok := c.nodes[dir]
	if !ok {
		// if the parent node doesn't yet exist, initialize it
		c.nodes[dir] = Node{
			Filename: dir,
			Children: map[string]bool{f: true},
		}
	} else {
		parent.Children[f] = true
	}
	if dir == "/" {
		return
	} else {
		c.addParents(dir)
	}
}

func (c *Reader) periodicRefresh(t *time.Ticker) {
	for {
		<-t.C
		c.refresh()
	}
}

func debug(args interface{}) {
	if !*cacheDebug {
		return
	}
	log.Printf("CACHE: %s\n", args)
}
<|endoftext|>


<commit_before>package cache

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/asjoyner/shade"
	"github.com/asjoyner/shade/drive"
)

var (
	defaultCacheDir = path.Join(shade.ConfigDir(), "cache")
	cacheDir        = flag.String("cache", defaultCacheDir, "Where to store the drive data cache")
	cacheDebug      = flag.Bool("cacheDebug", false, "Print cache debugging traces")
)

// A very compact representation of a file
type Node struct {
	Filename     string
	Filesize     uint64
	ModifiedTime time.Time
	Sha256sum    []byte // the sha of the associated File
	// TODO(asjoyner): use a struct{} here for efficiency?
	Children map[string]bool
}

func (n *Node) Synthetic() bool {
	if n.Sha256sum == nil {
		return true
	}
	return false
}

// Reader is a wrapper around a slice of cloud storage backends.  It presents an
// interface to query for the union of the set of known files by an integer ID,
// which will be stable across single processes invoking this cache, a node
// representing that file, or a single chunk of that file.  It can also cache a
// configurable quantity of chunks to disk.
//
// TODO(asjoyner): implement disk caching of data blocks.
type Reader struct {
	clients []drive.Client
	nodes   map[string]Node // full path to node
	sync.RWMutex
}

func NewReader(clients []drive.Client, t *time.Ticker) (*Reader, error) {
	c := &Reader{
		clients: clients,
		nodes: map[string]Node{
			"/": {
				Filename: "/",
				Children: make(map[string]bool),
			}},
	}
	if err := c.refresh(); err != nil {
		return nil, fmt.Errorf("initializing cache: %s", err)
	}
	go c.periodicRefresh(t)
	return c, nil
}

// NodeByPath returns the current file object for a given path.
func (c *Reader) NodeByPath(p string) (Node, error) {
	c.RLock()
	defer c.RUnlock()
	if n, ok := c.nodes[p]; ok {
		return n, nil
	}
	// TODO(shanel): Should this be debug?
	log.Printf("%+v\n", c.nodes)
	return Node{}, fmt.Errorf("no such node: %q", p)
}

func unmarshalChunk(fj, sha []byte) (*shade.File, error) {
	file := &shade.File{}
	if err := json.Unmarshal(fj, file); err != nil {
		return nil, fmt.Errorf("Failed to unmarshal sha256sum %x: %s", sha, err)
	}
	return file, nil
}

// FileByNode returns the full shade.File object for a given node.
func (c *Reader) FileByNode(n Node) (*shade.File, error) {
	if n.Synthetic() {
		return nil, errors.New("no shade.File defined")
	}
	var fj []byte
	var err error
	for _, client := range c.clients {
		fj, err = client.GetChunk(n.Sha256sum)
		if err != nil {
			log.Printf("Failed to fetch %s: %s", n.Sha256sum, err)
			continue
		}
	}
	if fj == nil || len(fj) == 0 {
		return nil, fmt.Errorf("Could not find JSON for node: %q", n.Filename)
	}
	return unmarshalChunk(fj, n.Sha256sum)
}

func (c *Reader) HasChild(parent, child string) bool {
	c.RLock()
	defer c.RUnlock()
	return c.nodes[parent].Children[child]
}

// NumNodes returns the number of nodes (files + synthetic directories) in the
// system.
func (c *Reader) NumNodes() int {
	c.RLock()
	defer c.RUnlock()
	return len(c.nodes)
}

func (c *Reader) GetChunk(sha256sum []byte) {
}

// refresh updates the cache
func (c *Reader) refresh() error {
	debug("Begining cache refresh cycle.")
	// key is a string([]byte) representation of the file's SHA2
	knownNodes := make(map[string]bool)
	for _, client := range c.clients {
		lfm, err := client.ListFiles()
		if err != nil {
			return fmt.Errorf("%q ListFiles(): %s", client.GetConfig().Provider, err)
		}
		debug(fmt.Sprintf("Found %d file(s) via %s", len(lfm), client.GetConfig().Provider))
		// fetch all those files into the local disk cache
		for _, sha256sum := range lfm {
			// check if we have already processed this Node
			if knownNodes[string(sha256sum)] {
				continue // we've already processed this file
			}

			// fetch the file Chunk
			f, err := client.GetChunk(sha256sum)
			if err != nil {
				// TODO(asjoyner): if !client.Local()... retry?
				log.Printf("Failed to fetch file %x: %s", sha256sum, err)
				continue
			}
			// ensure this file is known to all the writable clients
			for _, lc := range c.clients {
				if lc.GetConfig().Write {
					if err := lc.PutFile(sha256sum, f); err != nil {
						log.Printf("Failed to store checksum %x in %s: %s", sha256sum, client.GetConfig().Provider, err)
					}
				}
			}

			// unmarshal and populate c.nodes as the shade.files go by
			file, err := unmarshalChunk(f, sha256sum)
			if err != nil {
				log.Printf("%v", err)
				continue
			}
			node := Node{
				Filename:     file.Filename,
				Filesize:     uint64(file.Filesize),
				ModifiedTime: file.ModifiedTime,
				Sha256sum:    sha256sum,
				Children:     nil,
			}
			c.Lock()
			existing, ok := c.nodes[node.Filename]
			// TODO(asjoyner): handle file + directory collisions
			if ok && existing.ModifiedTime.After(node.ModifiedTime) {
				c.Unlock()
				continue
			}
			c.nodes[node.Filename] = node
			c.addParents(node.Filename)
			c.Unlock()
			knownNodes[string(sha256sum)] = true
		}
	}
	debug(fmt.Sprintf("Refresh complete with %d file(s).", len(knownNodes)))
	return nil
}

// recursive function to update parent dirs
func (c *Reader) addParents(filepath string) {
	dir, f := path.Split(filepath)
	if dir == "" {
		dir = "/"
	} else {
		dir = strings.TrimSuffix(dir, "/")
	}
	debug(fmt.Sprintf("adding %q as a child of %q", f, dir))
	// TODO(asjoyner): handle file + directory collisions
	if parent, ok := c.nodes[dir]; !ok {
		// if the parent node doesn't yet exist, initialize it
		c.nodes[dir] = Node{
			Filename: dir,
			Children: map[string]bool{f: true},
		}
	} else {
		parent.Children[f] = true
	}
	if dir != "/" {
		c.addParents(dir)
	}
}

func (c *Reader) periodicRefresh(t *time.Ticker) {
	for {
		<-t.C
		c.refresh()
	}
}

func debug(args interface{}) {
	if *cacheDebug {
		log.Printf("CACHE: %s\n", args)
	}
}
<commit_msg>more cleanup<commit_after>package cache

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/asjoyner/shade"
	"github.com/asjoyner/shade/drive"
)

var (
	defaultCacheDir = path.Join(shade.ConfigDir(), "cache")
	cacheDir        = flag.String("cache", defaultCacheDir, "Where to store the drive data cache")
	cacheDebug      = flag.Bool("cacheDebug", false, "Print cache debugging traces")
)

// A very compact representation of a file
type Node struct {
	Filename     string
	Filesize     uint64
	ModifiedTime time.Time
	Sha256sum    []byte // the sha of the associated File
	// TODO(asjoyner): use a struct{} here for efficiency?
	Children map[string]bool
}

func (n *Node) Synthetic() bool {
	if n.Sha256sum == nil {
		return true
	}
	return false
}

// Reader is a wrapper around a slice of cloud storage backends.  It presents an
// interface to query for the union of the set of known files by an integer ID,
// which will be stable across single processes invoking this cache, a node
// representing that file, or a single chunk of that file.  It can also cache a
// configurable quantity of chunks to disk.
//
// TODO(asjoyner): implement disk caching of data blocks.
type Reader struct {
	clients []drive.Client
	nodes   map[string]Node // full path to node
	sync.RWMutex
}

func NewReader(clients []drive.Client, t *time.Ticker) (*Reader, error) {
	c := &Reader{
		clients: clients,
		nodes: map[string]Node{
			"/": {
				Filename: "/",
				Children: make(map[string]bool),
			}},
	}
	if err := c.refresh(); err != nil {
		return nil, fmt.Errorf("initializing cache: %s", err)
	}
	go c.periodicRefresh(t)
	return c, nil
}

// NodeByPath returns the current file object for a given path.
func (c *Reader) NodeByPath(p string) (Node, error) {
	c.RLock()
	defer c.RUnlock()
	if n, ok := c.nodes[p]; ok {
		return n, nil
	}
	// TODO(shanel): Should this be debug?
	log.Printf("%+v\n", c.nodes)
	return Node{}, fmt.Errorf("no such node: %q", p)
}

func unmarshalChunk(fj, sha []byte) (*shade.File, error) {
	file := &shade.File{}
	if err := json.Unmarshal(fj, file); err != nil {
		return nil, fmt.Errorf("Failed to unmarshal sha256sum %x: %s", sha, err)
	}
	return file, nil
}

// FileByNode returns the full shade.File object for a given node.
func (c *Reader) FileByNode(n Node) (*shade.File, error) {
	if n.Synthetic() {
		return nil, errors.New("no shade.File defined")
	}
	var fj []byte
	var err error
	for _, client := range c.clients {
		fj, err = client.GetChunk(n.Sha256sum)
		if err != nil {
			log.Printf("Failed to fetch %s: %s", n.Sha256sum, err)
			continue
		}
	}
	if fj == nil || len(fj) == 0 {
		return nil, fmt.Errorf("Could not find JSON for node: %q", n.Filename)
	}
	return unmarshalChunk(fj, n.Sha256sum)
}

func (c *Reader) HasChild(parent, child string) bool {
	c.RLock()
	defer c.RUnlock()
	return c.nodes[parent].Children[child]
}

// NumNodes returns the number of nodes (files + synthetic directories) in the
// system.
func (c *Reader) NumNodes() int {
	c.RLock()
	defer c.RUnlock()
	return len(c.nodes)
}

func (c *Reader) GetChunk(sha256sum []byte) {
}

// refresh updates the cache
func (c *Reader) refresh() error {
	debug("Begining cache refresh cycle.")
	// key is a string([]byte) representation of the file's SHA2
	knownNodes := make(map[string]bool)
	for _, client := range c.clients {
		lfm, err := client.ListFiles()
		if err != nil {
			return fmt.Errorf("%q ListFiles(): %s", client.GetConfig().Provider, err)
		}
		debug(fmt.Sprintf("Found %d file(s) via %s", len(lfm), client.GetConfig().Provider))
		// fetch all those files into the local disk cache
		for _, sha256sum := range lfm {
			// check if we have already processed this Node
			if knownNodes[string(sha256sum)] {
				continue // we've already processed this file
			}

			// fetch the file Chunk
			f, err := client.GetChunk(sha256sum)
			if err != nil {
				// TODO(asjoyner): if !client.Local()... retry?
				log.Printf("Failed to fetch file %x: %s", sha256sum, err)
				continue
			}
			// ensure this file is known to all the writable clients
			for _, lc := range c.clients {
				if lc.GetConfig().Write {
					if err := lc.PutFile(sha256sum, f); err != nil {
						log.Printf("Failed to store checksum %x in %s: %s", sha256sum, client.GetConfig().Provider, err)
					}
				}
			}

			// unmarshal and populate c.nodes as the shade.files go by
			file, err := unmarshalChunk(f, sha256sum)
			if err != nil {
				log.Printf("%v", err)
				continue
			}
			node := Node{
				Filename:     file.Filename,
				Filesize:     uint64(file.Filesize),
				ModifiedTime: file.ModifiedTime,
				Sha256sum:    sha256sum,
				Children:     nil,
			}
			c.Lock()
			// TODO(asjoyner): handle file + directory collisions
			if existing, ok := c.nodes[node.Filename]; ok && existing.ModifiedTime.After(node.ModifiedTime) {
				c.Unlock()
				continue
			}
			c.nodes[node.Filename] = node
			c.addParents(node.Filename)
			c.Unlock()
			knownNodes[string(sha256sum)] = true
		}
	}
	debug(fmt.Sprintf("Refresh complete with %d file(s).", len(knownNodes)))
	return nil
}

// recursive function to update parent dirs
func (c *Reader) addParents(filepath string) {
	dir, f := path.Split(filepath)
	if dir == "" {
		dir = "/"
	} else {
		dir = strings.TrimSuffix(dir, "/")
	}
	debug(fmt.Sprintf("adding %q as a child of %q", f, dir))
	// TODO(asjoyner): handle file + directory collisions
	if parent, ok := c.nodes[dir]; !ok {
		// if the parent node doesn't yet exist, initialize it
		c.nodes[dir] = Node{
			Filename: dir,
			Children: map[string]bool{f: true},
		}
	} else {
		parent.Children[f] = true
	}
	if dir != "/" {
		c.addParents(dir)
	}
}

func (c *Reader) periodicRefresh(t *time.Ticker) {
	for {
		<-t.C
		c.refresh()
	}
}

func debug(args interface{}) {
	if *cacheDebug {
		log.Printf("CACHE: %s\n", args)
	}
}
<|endoftext|>


<commit_before>package s3

import (
	"bytes"
	"errors"
	"io"
	"os"
	"path/filepath"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"

	"github.com/go-spatial/tegola"
	"github.com/go-spatial/tegola/cache"
	"github.com/go-spatial/tegola/util/dict"
)

var (
	ErrMissingBucket = errors.New("s3cache: missing required param 'bucket'")
)

const CacheType = "s3"

const (
	// required
	ConfigKeyBucket = "bucket"
	// optional
	ConfigKeyBasepath       = "basepath"
	ConfigKeyMaxZoom        = "max_zoom"
	ConfigKeyRegion         = "region" // defaults to "us-east-1"
	ConfigKeyAWSAccessKeyID = "aws_access_key_id"
	ConfigKeyAWSSecretKey   = "aws_secret_access_key"
)

const (
	DefaultRegion = "us-east-1"
)

func init() {
	cache.Register(CacheType, New)
}

// New instantiates a S3 cache. The config expects the following params:
//
// 	required:
// 		bucket (string): the name of the s3 bucket to write to
//
// 	optional:
// 		region (string): the AWS region the bucket is located. defaults to 'us-east-1'
// 		aws_access_key_id (string): an AWS access key id
// 		aws_secret_access_key (string): an AWS secret access key
// 		basepath (string): a path prefix added to all cache operations inside of the S3 bucket
// 		max_zoom (int): max zoom to use the cache. beyond this zoom cache Set() calls will be ignored

func New(config map[string]interface{}) (cache.Interface, error) {
	var err error

	s3cache := Cache{}

	// parse the config
	c := dict.M(config)

	// the config map's underlying value is int
	defaultMaxZoom := uint(tegola.MaxZ)
	maxZoom, err := c.Uint(ConfigKeyMaxZoom, &defaultMaxZoom)
	if err != nil {
		return nil, err
	}

	s3cache.MaxZoom = maxZoom

	s3cache.Bucket, err = c.String(ConfigKeyBucket, nil)
	if err != nil {
		return nil, ErrMissingBucket
	}
	if s3cache.Bucket == "" {
		return nil, ErrMissingBucket
	}

	// basepath
	basepath := ""
	s3cache.Basepath, err = c.String(ConfigKeyBasepath, &basepath)
	if err != nil {
		return nil, err
	}

	// check for region env var
	region := os.Getenv("AWS_REGION")
	if region == "" {
		region = DefaultRegion
	}
	region, err = c.String(ConfigKeyRegion, &region)
	if err != nil {
		return nil, err
	}

	accessKey := ""
	accessKey, err = c.String(ConfigKeyAWSAccessKeyID, &accessKey)
	if err != nil {
		return nil, err
	}
	secretKey := ""
	secretKey, err = c.String(ConfigKeyAWSSecretKey, &secretKey)
	if err != nil {
		return nil, err
	}

	awsConfig := aws.Config{
		Region: aws.String(region),
	}

	// support for static credentials, this is not recommended by AWS but
	// necessary for some environments
	if accessKey != "" && secretKey != "" {
		awsConfig.Credentials = credentials.NewStaticCredentials(accessKey, secretKey, "")
	}

	// setup the s3 session.
	// if the accessKey and secreteKey are not provided (static creds) then the provider chain is used
	// http://docs.aws.amazon.com/sdk-for-go/v1/developer-guide/configuring-sdk.html
	s3cache.Client = s3.New(
		session.New(&awsConfig),
	)

	// in order to confirm we have the correct permissions on the bucket create a small file
	// and test a PUT, GET and DELETE to the bucket
	key := cache.Key{
		MapName:   "tegola-test-map",
		LayerName: "test-layer",
		Z:         0,
		X:         0,
		Y:         0,
	}
	// write a test file
	if err := s3cache.Set(&key, []byte("\x53\x69\x6c\x61\x73")); err != nil {
		e := cache.ErrSettingToCache{
			CacheType: CacheType,
			Err:       err,
		}

		return nil, e
	}

	// read the test file
	_, hit, err := s3cache.Get(&key)
	if err != nil {
		e := cache.ErrGettingFromCache{
			CacheType: CacheType,
			Err:       err,
		}

		return nil, e
	}
	if !hit {
		// return an error?
	}

	// purge the test file
	if err := s3cache.Purge(&key); err != nil {
		e := cache.ErrPurgingCache{
			CacheType: CacheType,
			Err:       err,
		}

		return nil, e
	}

	return &s3cache, nil
}

type Cache struct {
	// Bucket is the name of the s3 bucket to operate on
	Bucket string

	// Basepath is a path prefix added to all cache operations inside of the S3 bucket
	// helpful so a bucket does not need to be dedicated to only this cache
	Basepath string

	// MaxZoom determines the max zoom the cache to persist. Beyond this
	// zoom, cache Set() calls will be ignored. This is useful if the cache
	// should not be leveraged for higher zooms when data changes often.
	MaxZoom uint

	// client holds a reference to the s3 client. it's expected the client
	// has an active session and read, write, delete permissions have been checked
	Client *s3.S3
}

func (s3c *Cache) Set(key *cache.Key, val []byte) error {
	var err error

	// check for maxzoom
	if key.Z > s3c.MaxZoom {
		return nil
	}

	// add our basepath
	k := filepath.Join(s3c.Basepath, key.String())

	input := s3.PutObjectInput{
		Body:   aws.ReadSeekCloser(bytes.NewReader(val)),
		Bucket: aws.String(s3c.Bucket),
		Key:    aws.String(k),
	}

	_, err = s3c.Client.PutObject(&input)
	if err != nil {
		return err
	}

	return nil
}

func (s3c *Cache) Get(key *cache.Key) ([]byte, bool, error) {
	var err error

	// add our basepath
	k := filepath.Join(s3c.Basepath, key.String())

	input := s3.GetObjectInput{
		Bucket: aws.String(s3c.Bucket),
		Key:    aws.String(k),
	}

	result, err := s3c.Client.GetObject(&input)
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok {
			switch aerr.Code() {
			case s3.ErrCodeNoSuchKey:
				return nil, false, nil
			default:
				return nil, false, aerr
			}
		}
		return nil, false, err
	}

	var buf bytes.Buffer
	_, err = io.Copy(&buf, result.Body)
	if err != nil {
		return nil, false, err
	}

	return buf.Bytes(), true, nil
}

func (s3c *Cache) Purge(key *cache.Key) error {
	var err error

	// add our basepath
	k := filepath.Join(s3c.Basepath, key.String())

	input := s3.DeleteObjectInput{
		Bucket: aws.String(s3c.Bucket),
		Key:    aws.String(k),
	}

	_, err = s3c.Client.DeleteObject(&input)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>#413 add support for non AWS S3 endpoints (#414)<commit_after>package s3

import (
	"bytes"
	"errors"
	"io"
	"os"
	"path/filepath"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"

	"github.com/go-spatial/tegola"
	"github.com/go-spatial/tegola/cache"
	"github.com/go-spatial/tegola/util/dict"
)

var (
	ErrMissingBucket = errors.New("s3cache: missing required param 'bucket'")
)

const CacheType = "s3"

const (
	// required
	ConfigKeyBucket = "bucket"
	// optional
	ConfigKeyBasepath       = "basepath"
	ConfigKeyMaxZoom        = "max_zoom"
	ConfigKeyRegion         = "region" // defaults to "us-east-1"
	ConfigKeyEndpoint       = "endpoint" //	defaults to ""
	ConfigKeyAWSAccessKeyID = "aws_access_key_id"
	ConfigKeyAWSSecretKey   = "aws_secret_access_key"
	ConfigKeyACL      		= "access_control_list" //	defaults to ""
)

const (
	DefaultRegion = "us-east-1"
)

const (
	DefaultEndpoint = ""
)

func init() {
	cache.Register(CacheType, New)
}

// New instantiates a S3 cache. The config expects the following params:
//
// 	required:
// 		bucket (string): the name of the s3 bucket to write to
//
// 	optional:
// 		region (string): the AWS region the bucket is located. defaults to 'us-east-1'
// 		aws_access_key_id (string): an AWS access key id
// 		aws_secret_access_key (string): an AWS secret access key
// 		basepath (string): a path prefix added to all cache operations inside of the S3 bucket
// 		max_zoom (int): max zoom to use the cache. beyond this zoom cache Set() calls will be ignored
// 		endpoint (string): the S3 endpoint the bucket is located. defaults to '' and only needed for non-AWS endpoints
//  	access_control_list (string) : the S3 access control to set on the file when putting the file. Empty is the default for the bucket.

func New(config map[string]interface{}) (cache.Interface, error) {
	var err error

	s3cache := Cache{}

	// parse the config
	c := dict.M(config)

	// the config map's underlying value is int
	defaultMaxZoom := uint(tegola.MaxZ)
	maxZoom, err := c.Uint(ConfigKeyMaxZoom, &defaultMaxZoom)
	if err != nil {
		return nil, err
	}

	s3cache.MaxZoom = maxZoom

	s3cache.Bucket, err = c.String(ConfigKeyBucket, nil)
	if err != nil {
		return nil, ErrMissingBucket
	}
	if s3cache.Bucket == "" {
		return nil, ErrMissingBucket
	}

	// basepath
	basepath := ""
	s3cache.Basepath, err = c.String(ConfigKeyBasepath, &basepath)
	if err != nil {
		return nil, err
	}

	// check for region env var
	region := os.Getenv("AWS_REGION")
	if region == "" {
		region = DefaultRegion
	}
	region, err = c.String(ConfigKeyRegion, &region)
	if err != nil {
		return nil, err
	}

	accessKey := ""
	accessKey, err = c.String(ConfigKeyAWSAccessKeyID, &accessKey)
	if err != nil {
		return nil, err
	}
	secretKey := ""
	secretKey, err = c.String(ConfigKeyAWSSecretKey, &secretKey)
	if err != nil {
		return nil, err
	}

	awsConfig := aws.Config{
		Region: aws.String(region),
	}
	// check for endpoint env var
	endpoint := os.Getenv("AWS_ENDPOINT")
	if endpoint == "" {
		endpoint = DefaultEndpoint
	}
	endpoint, err = c.String(ConfigKeyEndpoint, &endpoint)
	if err != nil {
		return nil, err
	}

	// support for static credentials, this is not recommended by AWS but
	// necessary for some environments
	if accessKey != "" && secretKey != "" {
		awsConfig.Credentials = credentials.NewStaticCredentials(accessKey, secretKey, "")
	}

	// if an endpoint is set, add it to the awsConfig
	// otherwise do not set it and it will automatically use the correct aws-s3 endpoint
	if endpoint != "" {
		awsConfig.Endpoint = aws.String(endpoint)
	}

	// setup the s3 session.
	// if the accessKey and secreteKey are not provided (static creds) then the provider chain is used
	// http://docs.aws.amazon.com/sdk-for-go/v1/developer-guide/configuring-sdk.html
	s3cache.Client = s3.New(
		session.New(&awsConfig),
	)

	// check for control_access_list env var
	acl := os.Getenv("AWS_ACL")
	acl, err = c.String(ConfigKeyACL, &acl)
	if err != nil {
		return nil, err
	}
	s3cache.ACL = acl

	// in order to confirm we have the correct permissions on the bucket create a small file
	// and test a PUT, GET and DELETE to the bucket
	key := cache.Key{
		MapName:   "tegola-test-map",
		LayerName: "test-layer",
		Z:         0,
		X:         0,
		Y:         0,
	}
	// write a test file
	if err := s3cache.Set(&key, []byte("\x53\x69\x6c\x61\x73")); err != nil {
		e := cache.ErrSettingToCache{
			CacheType: CacheType,
			Err:       err,
		}

		return nil, e
	}

	// read the test file
	_, hit, err := s3cache.Get(&key)
	if err != nil {
		e := cache.ErrGettingFromCache{
			CacheType: CacheType,
			Err:       err,
		}

		return nil, e
	}
	if !hit {
		// return an error?
	}

	// purge the test file
	if err := s3cache.Purge(&key); err != nil {
		e := cache.ErrPurgingCache{
			CacheType: CacheType,
			Err:       err,
		}

		return nil, e
	}

	return &s3cache, nil
}

type Cache struct {
	// Bucket is the name of the s3 bucket to operate on
	Bucket string

	// Basepath is a path prefix added to all cache operations inside of the S3 bucket
	// helpful so a bucket does not need to be dedicated to only this cache
	Basepath string

	// MaxZoom determines the max zoom the cache to persist. Beyond this
	// zoom, cache Set() calls will be ignored. This is useful if the cache
	// should not be leveraged for higher zooms when data changes often.
	MaxZoom uint

	// client holds a reference to the s3 client. it's expected the client
	// has an active session and read, write, delete permissions have been checked
	Client *s3.S3

	// ACL is the aws ACL, if the not set it will use the default value for aws.
	ACL string
}

func (s3c *Cache) Set(key *cache.Key, val []byte) error {
	var err error

	// check for maxzoom
	if key.Z > s3c.MaxZoom {
		return nil
	}

	// add our basepath
	k := filepath.Join(s3c.Basepath, key.String())

	input := s3.PutObjectInput{
		Body:   aws.ReadSeekCloser(bytes.NewReader(val)),
		Bucket: aws.String(s3c.Bucket),
		Key:    aws.String(k),
	}
	if s3c.ACL != "" {
		input.ACL = aws.String(s3c.ACL)
   	}

	_, err = s3c.Client.PutObject(&input)
	if err != nil {
		return err
	}

	return nil
}

func (s3c *Cache) Get(key *cache.Key) ([]byte, bool, error) {
	var err error

	// add our basepath
	k := filepath.Join(s3c.Basepath, key.String())

	input := s3.GetObjectInput{
		Bucket: aws.String(s3c.Bucket),
		Key:    aws.String(k),
	}

	result, err := s3c.Client.GetObject(&input)
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok {
			switch aerr.Code() {
			case s3.ErrCodeNoSuchKey:
				return nil, false, nil
			default:
				return nil, false, aerr
			}
		}
		return nil, false, err
	}

	var buf bytes.Buffer
	_, err = io.Copy(&buf, result.Body)
	if err != nil {
		return nil, false, err
	}

	return buf.Bytes(), true, nil
}

func (s3c *Cache) Purge(key *cache.Key) error {
	var err error

	// add our basepath
	k := filepath.Join(s3c.Basepath, key.String())

	input := s3.DeleteObjectInput{
		Bucket: aws.String(s3c.Bucket),
		Key:    aws.String(k),
	}

	_, err = s3c.Client.DeleteObject(&input)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"strings"
	"text/template"
)

var pac struct {
	template        *template.Template
	topLevelDomain  string
	proxyServerAddr string
}

func init() {
	const pacRawTmpl = `var direct = 'DIRECT';
var httpProxy = '{{.ProxyAddr}}';

var directList = [
"localhost",
"0.1"{{.DirectDomains}}
];

var directAcc = {};
for (var i = 0; i < directList.length; i += 1) {
	directAcc[directList[i]] = true;
}

var topLevel = {
{{.TopLevel}}
};

function host2domain(host) {
	var lastDot = host.lastIndexOf(".");
	if (lastDot === -1)
		return host;
	// Find the second last dot
	dot2ndLast = host.lastIndexOf(".", lastDot-1);
	if (dot2ndLast === -1)
		return host;

	var part = host.substring(dot2ndLast+1, lastDot)
	if (topLevel[part]) {
		var dot3rdLast = host.lastIndexOf(".", dot2ndLast-1)
		if (dot3rdLast === -1) {
			return host
		}
		return host.substring(dot3rdLast+1)
	}
	return host.substring(dot2ndLast+1);
};

function FindProxyForURL(url, host) {
	return directAcc[host2domain(host)] ? direct : httpProxy;
};
`
	var err error
	pac.template, err = template.New("pac").Parse(pacRawTmpl)
	if err != nil {
		fmt.Println("Internal error on generating pac file template")
		os.Exit(1)
	}

	var buf bytes.Buffer
	for k, _ := range topLevelDomain {
		buf.WriteString(fmt.Sprintf("\t\"%s\": true,\n", k))
	}
	pac.topLevelDomain = buf.String()[:buf.Len()-2] // remove the final comma
}

func initProxyServerAddr() {
	listen, port := splitHostPort(config.ListenAddr)
	if listen == "0.0.0.0" {
		addrs, err := hostIP()
		if err != nil {
			errl.Println("Either change listen address to specific IP, or correct your host network settings.")
			os.Exit(1)
		}

		for _, ip := range addrs {
			pac.proxyServerAddr += fmt.Sprintf("PROXY %s:%s; ", ip, port)
		}
		pac.proxyServerAddr += "DIRECT"
		info.Printf("proxy listen address is %s, PAC will have proxy address: %s\n",
			config.ListenAddr, pac.proxyServerAddr)
	} else {
		pac.proxyServerAddr = fmt.Sprintf("PROXY %s; DIRECT", config.ListenAddr)
	}
}

// No need for content-length as we are closing connection
var pacHeader = []byte("HTTP/1.1 200 OK\r\nServer: cow-proxy\r\n" +
	"Content-Type: application/x-ns-proxy-autoconfig\r\nConnection: close\r\n\r\n")
var pacDirect = []byte("function FindProxyForURL(url, host) { return 'DIRECT'; };")

func sendPAC(w io.Writer) {
	// domains in PAC file needs double quote
	ds1 := strings.Join(domainSet.alwaysDirect.toSlice(), "\",\n\"")
	ds2 := strings.Join(domainSet.direct.toSlice(), "\",\n\"")
	var ds string
	if ds1 == "" {
		ds = ds2
	} else if ds2 == "" {
		ds = ds1
	} else {
		ds = ds1 + "\",\n\"" + ds2
	}
	if ds == "" {
		// Empty direct domain list
		w.Write(pacHeader)
		w.Write(pacDirect)
		return
	}

	data := struct {
		ProxyAddr     string
		DirectDomains string
		TopLevel      string
	}{
		pac.proxyServerAddr,
		",\n\"" + ds + "\"",
		pac.topLevelDomain,
	}

	if _, err := w.Write(pacHeader); err != nil {
		debug.Println("Error writing pac header")
		return
	}
	// debug.Println("direct:", data.DirectDomains)
	buf := new(bytes.Buffer)
	if err := pac.template.Execute(buf, data); err != nil {
		errl.Println("Error generating pac file:", err)
	}
	if _, err := w.Write(buf.Bytes()); err != nil {
		debug.Println("Error writing pac content:", err)
	}
}
<commit_msg>Tell browser to use proxy if direct list is empty.<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"strings"
	"text/template"
)

var pac struct {
	template        *template.Template
	topLevelDomain  string
	proxyServerAddr string
}

func init() {
	const pacRawTmpl = `var direct = 'DIRECT';
var httpProxy = '{{.ProxyAddr}}';

var directList = [
"localhost",
"0.1"{{.DirectDomains}}
];

var directAcc = {};
for (var i = 0; i < directList.length; i += 1) {
	directAcc[directList[i]] = true;
}

var topLevel = {
{{.TopLevel}}
};

function host2domain(host) {
	var lastDot = host.lastIndexOf(".");
	if (lastDot === -1)
		return host;
	// Find the second last dot
	dot2ndLast = host.lastIndexOf(".", lastDot-1);
	if (dot2ndLast === -1)
		return host;

	var part = host.substring(dot2ndLast+1, lastDot)
	if (topLevel[part]) {
		var dot3rdLast = host.lastIndexOf(".", dot2ndLast-1)
		if (dot3rdLast === -1) {
			return host
		}
		return host.substring(dot3rdLast+1)
	}
	return host.substring(dot2ndLast+1);
};

function FindProxyForURL(url, host) {
	return directAcc[host2domain(host)] ? direct : httpProxy;
};
`
	var err error
	pac.template, err = template.New("pac").Parse(pacRawTmpl)
	if err != nil {
		fmt.Println("Internal error on generating pac file template")
		os.Exit(1)
	}

	var buf bytes.Buffer
	for k, _ := range topLevelDomain {
		buf.WriteString(fmt.Sprintf("\t\"%s\": true,\n", k))
	}
	pac.topLevelDomain = buf.String()[:buf.Len()-2] // remove the final comma
}

func initProxyServerAddr() {
	listen, port := splitHostPort(config.ListenAddr)
	if listen == "0.0.0.0" {
		addrs, err := hostIP()
		if err != nil {
			errl.Println("Either change listen address to specific IP, or correct your host network settings.")
			os.Exit(1)
		}

		for _, ip := range addrs {
			pac.proxyServerAddr += fmt.Sprintf("PROXY %s:%s; ", ip, port)
		}
		pac.proxyServerAddr += "DIRECT"
		info.Printf("proxy listen address is %s, PAC will have proxy address: %s\n",
			config.ListenAddr, pac.proxyServerAddr)
	} else {
		pac.proxyServerAddr = fmt.Sprintf("PROXY %s; DIRECT", config.ListenAddr)
	}
}

// No need for content-length as we are closing connection
var pacHeader = []byte("HTTP/1.1 200 OK\r\nServer: cow-proxy\r\n" +
	"Content-Type: application/x-ns-proxy-autoconfig\r\nConnection: close\r\n\r\n")

func sendPAC(w io.Writer) {
	// domains in PAC file needs double quote
	ds1 := strings.Join(domainSet.alwaysDirect.toSlice(), "\",\n\"")
	ds2 := strings.Join(domainSet.direct.toSlice(), "\",\n\"")
	var ds string
	if ds1 == "" {
		ds = ds2
	} else if ds2 == "" {
		ds = ds1
	} else {
		ds = ds1 + "\",\n\"" + ds2
	}
	if ds == "" {
		// Empty direct domain list
		w.Write(pacHeader)
		pacproxy := fmt.Sprintf("function FindProxyForURL(url, host) { return '%s'; };",
			pac.proxyServerAddr)
		w.Write([]byte(pacproxy))
		return
	}

	data := struct {
		ProxyAddr     string
		DirectDomains string
		TopLevel      string
	}{
		pac.proxyServerAddr,
		",\n\"" + ds + "\"",
		pac.topLevelDomain,
	}

	if _, err := w.Write(pacHeader); err != nil {
		debug.Println("Error writing pac header")
		return
	}
	// debug.Println("direct:", data.DirectDomains)
	buf := new(bytes.Buffer)
	if err := pac.template.Execute(buf, data); err != nil {
		errl.Println("Error generating pac file:", err)
	}
	if _, err := w.Write(buf.Bytes()); err != nil {
		debug.Println("Error writing pac content:", err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Matthew Baird
// 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 elastigo

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strconv"
	"sync"
	"time"
)

const (
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer = 16384
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs = 100
	// Max delay before forcing a flush to Elasticearch
	BulkDelaySeconds = 5
	// maximum wait shutdown seconds
	MAX_SHUTDOWN_SECS = 5
)

type ErrorBuffer struct {
	Err error
	Buf *bytes.Buffer
}

// A bulk indexer creates goroutines, and channels for connecting and sending data
// to elasticsearch in bulk, using buffers.
type BulkIndexer struct {
	conn *Conn

	// We are creating a variable defining the func responsible for sending
	// to allow a mock sendor for test purposes
	Sender func(*bytes.Buffer) error

	// If we encounter an error in sending, we are going to retry for this long
	// before returning an error
	// if 0 it will not retry
	RetryForSeconds int

	// channel for getting errors
	ErrorChannel chan *ErrorBuffer

	// channel for sending to background indexer
	bulkChannel chan []byte

	// numErrors is a running total of errors seen
	numErrors uint64

	// shutdown channel
	shutdownChan chan chan struct{}
	// Channel to shutdown http send go-routines
	httpDoneChan chan bool
	// channel to shutdown timer
	timerDoneChan chan bool
	// channel to shutdown doc go-routines
	docDoneChan chan bool

	// Channel to send a complete byte.Buffer to the http sendor
	sendBuf chan *bytes.Buffer
	// byte buffer for docs that have been converted to bytes, but not yet sent
	buf *bytes.Buffer
	// Buffer for Max number of time before forcing flush
	BufferDelayMax time.Duration
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer int // 1048576
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs int // 100

	// Number of documents we have send through so far on this session
	docCt int
	// Max number of http conns in flight at one time
	maxConns int
	// If we are indexing enough docs per bufferdelaymax, we won't need to do time
	// based eviction, else we do.
	needsTimeBasedFlush bool
	// Lock for document writes/operations
	mu sync.Mutex
	// Wait Group for the http sends
	sendWg *sync.WaitGroup
}

func (b *BulkIndexer) NumErrors() uint64 {
	return b.numErrors
}

func (c *Conn) NewBulkIndexer(maxConns int) *BulkIndexer {
	b := BulkIndexer{conn: c, sendBuf: make(chan *bytes.Buffer, maxConns)}
	b.needsTimeBasedFlush = true
	b.buf = new(bytes.Buffer)
	b.maxConns = maxConns
	b.BulkMaxBuffer = BulkMaxBuffer
	b.BulkMaxDocs = BulkMaxDocs
	b.BufferDelayMax = time.Duration(BulkDelaySeconds) * time.Second
	b.bulkChannel = make(chan []byte, 100)
	b.sendWg = new(sync.WaitGroup)
	b.docDoneChan = make(chan bool)
	b.timerDoneChan = make(chan bool)
	b.httpDoneChan = make(chan bool)
	return &b
}

// A bulk indexer with more control over error handling
//    @maxConns is the max number of in flight http requests
//    @retrySeconds is # of seconds to wait before retrying falied requests
//
//   done := make(chan bool)
//   BulkIndexerGlobalRun(100, done)
func (c *Conn) NewBulkIndexerErrors(maxConns, retrySeconds int) *BulkIndexer {
	b := c.NewBulkIndexer(maxConns)
	b.RetryForSeconds = retrySeconds
	b.ErrorChannel = make(chan *ErrorBuffer, 20)
	return b
}

// Starts this bulk Indexer running, this Run opens a go routine so is
// Non blocking
func (b *BulkIndexer) Start() {
	b.shutdownChan = make(chan chan struct{})

	go func() {
		// XXX(j): Refactor this stuff to use an interface.
		if b.Sender == nil {
			b.Sender = b.Send
		}
		// Backwards compatibility
		b.startHttpSender()
		b.startDocChannel()
		b.startTimer()
		ch := <-b.shutdownChan
		b.Flush()
		b.shutdown()
		ch <- struct{}{}
		close(ch)
	}()
}

// Stop stops the bulk indexer, blocking the caller until it is complete.
func (b *BulkIndexer) Stop() {
	ch := make(chan struct{})
	b.shutdownChan <- ch
	<-ch
	close(b.shutdownChan)
}

// Make a channel that will close when the given WaitGroup is done.
func wgChan(wg *sync.WaitGroup) <-chan interface{} {
	ch := make(chan interface{})
	go func() {
		wg.Wait()
		close(ch)
	}()
	return ch
}

func (b *BulkIndexer) PendingDocuments() int {
	return b.docCt
}

// Flush all current documents to ElasticSearch
func (b *BulkIndexer) Flush() {
	b.mu.Lock()
	if b.docCt > 0 {
		b.send(b.buf)
	}
	b.mu.Unlock()
	for {
		select {
		case <-wgChan(b.sendWg):
			// done
			return
		case <-time.After(time.Second * time.Duration(MAX_SHUTDOWN_SECS)):
			// timeout!
			return
		}
	}
}

func (b *BulkIndexer) startHttpSender() {

	// this sends http requests to elasticsearch it uses maxConns to open up that
	// many goroutines, each of which will synchronously call ElasticSearch
	// in theory, the whole set will cause a backup all the way to IndexBulk if
	// we have consumed all maxConns
	for i := 0; i < b.maxConns; i++ {
		go func() {
			for {
				select {
				case buf := <-b.sendBuf:
					b.sendWg.Add(1)
					err := b.Sender(buf)

					// Perhaps a b.FailureStrategy(err)  ??  with different types of strategies
					//  1.  Retry, then panic
					//  2.  Retry then return error and let runner decide
					//  3.  Retry, then log to disk?   retry later?
					if err != nil {
						if b.RetryForSeconds > 0 {
							time.Sleep(time.Second * time.Duration(b.RetryForSeconds))
							err = b.Sender(buf)
							if err == nil {
								// Successfully re-sent with no error
								b.sendWg.Done()
								continue
							}
						}
						if b.ErrorChannel != nil {
							b.ErrorChannel <- &ErrorBuffer{err, buf}
						}
					}
					b.sendWg.Done()
				case <-b.httpDoneChan:
					// shutdown this go routine
					return
				}

			}
		}()
	}
}

// start a timer for checking back and forcing flush ever BulkDelaySeconds seconds
// even if we haven't hit max messages/size
func (b *BulkIndexer) startTimer() {
	ticker := time.NewTicker(b.BufferDelayMax)
	go func() {
		for {
			select {
			case <-ticker.C:
				b.mu.Lock()
				// don't send unless last sendor was the time,
				// otherwise an indication of other thresholds being hit
				// where time isn't needed
				if b.buf.Len() > 0 && b.needsTimeBasedFlush {
					b.needsTimeBasedFlush = true
					b.send(b.buf)
				} else if b.buf.Len() > 0 {
					b.needsTimeBasedFlush = true
				}
				b.mu.Unlock()
			case <-b.timerDoneChan:
				// shutdown this go routine
				return
			}

		}
	}()
}

func (b *BulkIndexer) startDocChannel() {
	// This goroutine accepts incoming byte arrays from the IndexBulk function and
	// writes to buffer
	go func() {
		for {
			select {
			case docBytes := <-b.bulkChannel:
				b.mu.Lock()
				b.docCt += 1
				b.buf.Write(docBytes)
				if b.buf.Len() >= b.BulkMaxBuffer || b.docCt >= b.BulkMaxDocs {
					b.needsTimeBasedFlush = false
					//log.Printf("Send due to size:  docs=%d  bufsize=%d", b.docCt, b.buf.Len())
					b.send(b.buf)
				}
				b.mu.Unlock()
			case <-b.docDoneChan:
				// shutdown this go routine
				return
			}
		}
	}()
}

func (b *BulkIndexer) send(buf *bytes.Buffer) {
	//b2 := *b.buf
	b.sendBuf <- buf
	b.buf = new(bytes.Buffer)
	//	b.buf.Reset()
	b.docCt = 0
}

func (b *BulkIndexer) shutdown() {
	// This must be called After flush
	b.docDoneChan <- true
	b.timerDoneChan <- true
	for i := 0; i < b.maxConns; i++ {
		b.httpDoneChan <- true
	}
}

// The index bulk API adds or updates a typed JSON document to a specific index, making it searchable.
// it operates by buffering requests, and ocassionally flushing to elasticsearch
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func (b *BulkIndexer) Index(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("index", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

func (b *BulkIndexer) Update(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("update", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

// This does the actual send of a buffer, which has already been formatted
// into bytes of ES formatted bulk data
func (b *BulkIndexer) Send(buf *bytes.Buffer) error {
	type responseStruct struct {
		Took   int64                    `json:"took"`
		Errors bool                     `json:"errors"`
		Items  []map[string]interface{} `json:"items"`
	}

	response := responseStruct{}

	body, err := b.conn.DoCommand("POST", "/_bulk", nil, buf)

	if err != nil {
		b.numErrors += 1
		return err
	}
	// check for response errors, bulk insert will give 200 OK but then include errors in response
	jsonErr := json.Unmarshal(body, &response)
	if jsonErr == nil {
		if response.Errors {
			b.numErrors += uint64(len(response.Items))
			return fmt.Errorf("Bulk Insertion Error. Failed item count [%d]", len(response.Items))
		}
	}
	return nil
}

// Given a set of arguments for index, type, id, data create a set of bytes that is formatted for bulkd index
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func WriteBulkBytes(op string, index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) ([]byte, error) {
	// only index and update are currently supported
	if op != "index" && op != "update" {
		return nil, errors.New(fmt.Sprintf("Operation '%s' is not yet supported", op))
	}

	// First line
	buf := bytes.Buffer{}
	buf.WriteString(fmt.Sprintf(`{"%s":{"_index":"`, op))
	buf.WriteString(index)
	buf.WriteString(`","_type":"`)
	buf.WriteString(_type)
	buf.WriteString(`"`)
	if len(id) > 0 {
		buf.WriteString(`,"_id":"`)
		buf.WriteString(id)
		buf.WriteString(`"`)
	}

	if op == "update" {
		buf.WriteString(`,"retry_on_conflict":3`)
	}

	if len(ttl) > 0 {
		buf.WriteString(`,"ttl":"`)
		buf.WriteString(ttl)
		buf.WriteString(`"`)
	}
	if date != nil {
		buf.WriteString(`,"_timestamp":"`)
		buf.WriteString(strconv.FormatInt(date.UnixNano()/1e6, 10))
		buf.WriteString(`"`)
	}
	if refresh {
		buf.WriteString(`,"refresh":true`)
	}
	buf.WriteString(`}}`)
	buf.WriteRune('\n')
	//buf.WriteByte('\n')
	switch v := data.(type) {
	case *bytes.Buffer:
		io.Copy(&buf, v)
	case []byte:
		buf.Write(v)
	case string:
		buf.WriteString(v)
	default:
		body, jsonErr := json.Marshal(data)
		if jsonErr != nil {
			return nil, jsonErr
		}
		buf.Write(body)
	}
	buf.WriteRune('\n')
	return buf.Bytes(), nil
}
<commit_msg>Need to save the bulk for the retry.<commit_after>// Copyright 2013 Matthew Baird
// 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 elastigo

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strconv"
	"sync"
	"time"
)

const (
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer = 16384
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs = 100
	// Max delay before forcing a flush to Elasticearch
	BulkDelaySeconds = 5
	// maximum wait shutdown seconds
	MAX_SHUTDOWN_SECS = 5
)

type ErrorBuffer struct {
	Err error
	Buf *bytes.Buffer
}

// A bulk indexer creates goroutines, and channels for connecting and sending data
// to elasticsearch in bulk, using buffers.
type BulkIndexer struct {
	conn *Conn

	// We are creating a variable defining the func responsible for sending
	// to allow a mock sendor for test purposes
	Sender func(*bytes.Buffer) error

	// If we encounter an error in sending, we are going to retry for this long
	// before returning an error
	// if 0 it will not retry
	RetryForSeconds int

	// channel for getting errors
	ErrorChannel chan *ErrorBuffer

	// channel for sending to background indexer
	bulkChannel chan []byte

	// numErrors is a running total of errors seen
	numErrors uint64

	// shutdown channel
	shutdownChan chan chan struct{}
	// Channel to shutdown http send go-routines
	httpDoneChan chan bool
	// channel to shutdown timer
	timerDoneChan chan bool
	// channel to shutdown doc go-routines
	docDoneChan chan bool

	// Channel to send a complete byte.Buffer to the http sendor
	sendBuf chan *bytes.Buffer
	// byte buffer for docs that have been converted to bytes, but not yet sent
	buf *bytes.Buffer
	// Buffer for Max number of time before forcing flush
	BufferDelayMax time.Duration
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer int // 1048576
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs int // 100

	// Number of documents we have send through so far on this session
	docCt int
	// Max number of http conns in flight at one time
	maxConns int
	// If we are indexing enough docs per bufferdelaymax, we won't need to do time
	// based eviction, else we do.
	needsTimeBasedFlush bool
	// Lock for document writes/operations
	mu sync.Mutex
	// Wait Group for the http sends
	sendWg *sync.WaitGroup
}

func (b *BulkIndexer) NumErrors() uint64 {
	return b.numErrors
}

func (c *Conn) NewBulkIndexer(maxConns int) *BulkIndexer {
	b := BulkIndexer{conn: c, sendBuf: make(chan *bytes.Buffer, maxConns)}
	b.needsTimeBasedFlush = true
	b.buf = new(bytes.Buffer)
	b.maxConns = maxConns
	b.BulkMaxBuffer = BulkMaxBuffer
	b.BulkMaxDocs = BulkMaxDocs
	b.BufferDelayMax = time.Duration(BulkDelaySeconds) * time.Second
	b.bulkChannel = make(chan []byte, 100)
	b.sendWg = new(sync.WaitGroup)
	b.docDoneChan = make(chan bool)
	b.timerDoneChan = make(chan bool)
	b.httpDoneChan = make(chan bool)
	return &b
}

// A bulk indexer with more control over error handling
//    @maxConns is the max number of in flight http requests
//    @retrySeconds is # of seconds to wait before retrying falied requests
//
//   done := make(chan bool)
//   BulkIndexerGlobalRun(100, done)
func (c *Conn) NewBulkIndexerErrors(maxConns, retrySeconds int) *BulkIndexer {
	b := c.NewBulkIndexer(maxConns)
	b.RetryForSeconds = retrySeconds
	b.ErrorChannel = make(chan *ErrorBuffer, 20)
	return b
}

// Starts this bulk Indexer running, this Run opens a go routine so is
// Non blocking
func (b *BulkIndexer) Start() {
	b.shutdownChan = make(chan chan struct{})

	go func() {
		// XXX(j): Refactor this stuff to use an interface.
		if b.Sender == nil {
			b.Sender = b.Send
		}
		// Backwards compatibility
		b.startHttpSender()
		b.startDocChannel()
		b.startTimer()
		ch := <-b.shutdownChan
		b.Flush()
		b.shutdown()
		ch <- struct{}{}
		close(ch)
	}()
}

// Stop stops the bulk indexer, blocking the caller until it is complete.
func (b *BulkIndexer) Stop() {
	ch := make(chan struct{})
	b.shutdownChan <- ch
	<-ch
	close(b.shutdownChan)
}

// Make a channel that will close when the given WaitGroup is done.
func wgChan(wg *sync.WaitGroup) <-chan interface{} {
	ch := make(chan interface{})
	go func() {
		wg.Wait()
		close(ch)
	}()
	return ch
}

func (b *BulkIndexer) PendingDocuments() int {
	return b.docCt
}

// Flush all current documents to ElasticSearch
func (b *BulkIndexer) Flush() {
	b.mu.Lock()
	if b.docCt > 0 {
		b.send(b.buf)
	}
	b.mu.Unlock()
	for {
		select {
		case <-wgChan(b.sendWg):
			// done
			return
		case <-time.After(time.Second * time.Duration(MAX_SHUTDOWN_SECS)):
			// timeout!
			return
		}
	}
}

func (b *BulkIndexer) startHttpSender() {

	// this sends http requests to elasticsearch it uses maxConns to open up that
	// many goroutines, each of which will synchronously call ElasticSearch
	// in theory, the whole set will cause a backup all the way to IndexBulk if
	// we have consumed all maxConns
	for i := 0; i < b.maxConns; i++ {
		go func() {
			for {
				select {
				case buf := <-b.sendBuf:
					b.sendWg.Add(1)
					// Copy for the potential re-send.
					bufCopy := bytes.NewBuffer(buf.Bytes())
					err := b.Sender(buf)

					// Perhaps a b.FailureStrategy(err)  ??  with different types of strategies
					//  1.  Retry, then panic
					//  2.  Retry then return error and let runner decide
					//  3.  Retry, then log to disk?   retry later?
					if err != nil {
						if b.RetryForSeconds > 0 {
							time.Sleep(time.Second * time.Duration(b.RetryForSeconds))
							err = b.Sender(bufCopy)
							if err == nil {
								// Successfully re-sent with no error
								b.sendWg.Done()
								continue
							}
						}
						if b.ErrorChannel != nil {
							b.ErrorChannel <- &ErrorBuffer{err, buf}
						}
					}
					b.sendWg.Done()
				case <-b.httpDoneChan:
					// shutdown this go routine
					return
				}

			}
		}()
	}
}

// start a timer for checking back and forcing flush ever BulkDelaySeconds seconds
// even if we haven't hit max messages/size
func (b *BulkIndexer) startTimer() {
	ticker := time.NewTicker(b.BufferDelayMax)
	go func() {
		for {
			select {
			case <-ticker.C:
				b.mu.Lock()
				// don't send unless last sendor was the time,
				// otherwise an indication of other thresholds being hit
				// where time isn't needed
				if b.buf.Len() > 0 && b.needsTimeBasedFlush {
					b.needsTimeBasedFlush = true
					b.send(b.buf)
				} else if b.buf.Len() > 0 {
					b.needsTimeBasedFlush = true
				}
				b.mu.Unlock()
			case <-b.timerDoneChan:
				// shutdown this go routine
				return
			}

		}
	}()
}

func (b *BulkIndexer) startDocChannel() {
	// This goroutine accepts incoming byte arrays from the IndexBulk function and
	// writes to buffer
	go func() {
		for {
			select {
			case docBytes := <-b.bulkChannel:
				b.mu.Lock()
				b.docCt += 1
				b.buf.Write(docBytes)
				if b.buf.Len() >= b.BulkMaxBuffer || b.docCt >= b.BulkMaxDocs {
					b.needsTimeBasedFlush = false
					//log.Printf("Send due to size:  docs=%d  bufsize=%d", b.docCt, b.buf.Len())
					b.send(b.buf)
				}
				b.mu.Unlock()
			case <-b.docDoneChan:
				// shutdown this go routine
				return
			}
		}
	}()
}

func (b *BulkIndexer) send(buf *bytes.Buffer) {
	//b2 := *b.buf
	b.sendBuf <- buf
	b.buf = new(bytes.Buffer)
	//	b.buf.Reset()
	b.docCt = 0
}

func (b *BulkIndexer) shutdown() {
	// This must be called After flush
	b.docDoneChan <- true
	b.timerDoneChan <- true
	for i := 0; i < b.maxConns; i++ {
		b.httpDoneChan <- true
	}
}

// The index bulk API adds or updates a typed JSON document to a specific index, making it searchable.
// it operates by buffering requests, and ocassionally flushing to elasticsearch
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func (b *BulkIndexer) Index(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("index", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

func (b *BulkIndexer) Update(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("update", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

// This does the actual send of a buffer, which has already been formatted
// into bytes of ES formatted bulk data
func (b *BulkIndexer) Send(buf *bytes.Buffer) error {
	type responseStruct struct {
		Took   int64                    `json:"took"`
		Errors bool                     `json:"errors"`
		Items  []map[string]interface{} `json:"items"`
	}

	response := responseStruct{}

	body, err := b.conn.DoCommand("POST", "/_bulk", nil, buf)

	if err != nil {
		b.numErrors += 1
		return err
	}
	// check for response errors, bulk insert will give 200 OK but then include errors in response
	jsonErr := json.Unmarshal(body, &response)
	if jsonErr == nil {
		if response.Errors {
			b.numErrors += uint64(len(response.Items))
			return fmt.Errorf("Bulk Insertion Error. Failed item count [%d]", len(response.Items))
		}
	}
	return nil
}

// Given a set of arguments for index, type, id, data create a set of bytes that is formatted for bulkd index
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func WriteBulkBytes(op string, index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) ([]byte, error) {
	// only index and update are currently supported
	if op != "index" && op != "update" {
		return nil, errors.New(fmt.Sprintf("Operation '%s' is not yet supported", op))
	}

	// First line
	buf := bytes.Buffer{}
	buf.WriteString(fmt.Sprintf(`{"%s":{"_index":"`, op))
	buf.WriteString(index)
	buf.WriteString(`","_type":"`)
	buf.WriteString(_type)
	buf.WriteString(`"`)
	if len(id) > 0 {
		buf.WriteString(`,"_id":"`)
		buf.WriteString(id)
		buf.WriteString(`"`)
	}

	if op == "update" {
		buf.WriteString(`,"retry_on_conflict":3`)
	}

	if len(ttl) > 0 {
		buf.WriteString(`,"ttl":"`)
		buf.WriteString(ttl)
		buf.WriteString(`"`)
	}
	if date != nil {
		buf.WriteString(`,"_timestamp":"`)
		buf.WriteString(strconv.FormatInt(date.UnixNano()/1e6, 10))
		buf.WriteString(`"`)
	}
	if refresh {
		buf.WriteString(`,"refresh":true`)
	}
	buf.WriteString(`}}`)
	buf.WriteRune('\n')
	//buf.WriteByte('\n')
	switch v := data.(type) {
	case *bytes.Buffer:
		io.Copy(&buf, v)
	case []byte:
		buf.Write(v)
	case string:
		buf.WriteString(v)
	default:
		body, jsonErr := json.Marshal(data)
		if jsonErr != nil {
			return nil, jsonErr
		}
		buf.Write(body)
	}
	buf.WriteRune('\n')
	return buf.Bytes(), nil
}
<|endoftext|>