<commit_before>package raft

import (
	"bufio"
	"encoding/json"
	"fmt"
	"hash"
	"hash/crc64"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"
)

const (
	snapPath      = "snapshots"
	metaFilePath  = "meta.json"
	stateFilePath = "state.bin"
	tmpSuffix     = ".tmp"
)

// FileSnapshotStore implements the SnapshotStore interface and allows
// snapshots to be made on the local disk.
type FileSnapshotStore struct {
	path   string
	retain int
}

// Implements the SnapshotSink
type FileSnapshotSink struct {
	dir  string
	meta fileSnapshotMeta

	stateFile *os.File
	stateHash hash.Hash64
	buffered  *bufio.Writer
}

type fileSnapshotMeta struct {
	SnapshotMeta
	CRC []byte
}

// NewFileSnapshotStore creates a new FileSnapshotStore based
// on a base directory. The `retain` parameter controls how many
// snapshots are retained. Must be at least 1.
func NewFileSnapshotStore(base string, retain int) (*FileSnapshotStore, error) {
	if retain < 1 {
		return nil, fmt.Errorf("must retain at least one snapshot")
	}

	path := filepath.Join(base, snapPath)
	store := &FileSnapshotStore{
		path:   path,
		retain: retain,
	}
	return store, nil
}

// Create is used to start a new snapshot
func (f *FileSnapshotStore) Create(index, term uint64, peers []byte) (SnapshotSink, error) {
	// Create a new path
	name := time.Now().Format(time.RFC3339) + tmpSuffix
	path := filepath.Join(f.path, name)
	log.Printf("[INFO] Creating new snapshot at %s", path)

	// Make the directory
	if err := os.Mkdir(path, 0755); err != nil {
		log.Printf("[ERR] Failed to make snapshot directory: %v", err)
		return nil, err
	}

	// Create the sink
	sink := &FileSnapshotSink{
		dir: path,
		meta: fileSnapshotMeta{
			SnapshotMeta: SnapshotMeta{
				ID:    path,
				Index: index,
				Term:  term,
				Peers: peers,
			},
			CRC: nil,
		},
	}

	// Write out the meta data
	if err := sink.writeMeta(); err != nil {
		log.Printf("[ERR] Failed to write metadata: %v", err)
		return nil, err
	}

	// Open the state file
	statePath := filepath.Join(path, stateFilePath)
	fh, err := os.Create(statePath)
	if err != nil {
		log.Printf("[ERR] Failed to create state file: %v", err)
		return nil, err
	}
	sink.stateFile = fh

	// Create a CRC64 hash
	sink.stateHash = crc64.New(crc64.MakeTable(crc64.ECMA))

	// Wrap both the hash and file in a MultiWriter with buffering
	multi := io.MultiWriter(sink.stateFile, sink.stateHash)
	sink.buffered = bufio.NewWriter(multi)

	// Done
	return sink, nil
}

func (f *FileSnapshotStore) List() []*SnapshotMeta {
	return nil
}

func (f *FileSnapshotStore) Open(id string) (io.ReadCloser, error) {
	return nil, nil
}

// Write is used to append to the state file. We write to the
// buffered IO object to reduce the amount of context switches
func (s *FileSnapshotSink) Write(b []byte) (int, error) {
	return s.buffered.Write(b)
}

// Close is used to indicate a successful end
func (s *FileSnapshotSink) Close() error {
	// Close the open handles
	if err := s.finalize(); err != nil {
		log.Printf("[ERR] Failed to finalize snapshot: %v", err)
		return err
	}

	// Write out the meta data
	if err := s.writeMeta(); err != nil {
		log.Printf("[ERR] Failed to write metadata: %v", err)
		return err
	}

	// Move the directory into place
	newPath := strings.TrimSuffix(s.dir, tmpSuffix)
	if err := os.Rename(s.dir, newPath); err != nil {
		log.Printf("[ERR] Failed to move snapshot into place: %v", err)
		return err
	}
	return nil
}

// Cancel is used to indicate an unsuccessful end
func (s *FileSnapshotSink) Cancel() error {
	// Close the open handles
	if err := s.finalize(); err != nil {
		log.Printf("[ERR] Failed to finalize snapshot: %v", err)
		return err
	}

	// Attempt to remove all artifacts
	return os.RemoveAll(s.dir)
}

// finalize is used to close all of our resources
func (s *FileSnapshotSink) finalize() error {
	if err := s.buffered.Flush(); err != nil {
		return err
	}
	if err := s.stateFile.Close(); err != nil {
		return err
	}
	s.meta.CRC = s.stateHash.Sum(nil)
	return nil
}

// writeMeta is used to write out the metadata we have
func (s *FileSnapshotSink) writeMeta() error {
	// Open the meta file
	metaPath := filepath.Join(s.dir, metaFilePath)
	fh, err := os.Create(metaPath)
	if err != nil {
		return err
	}
	defer fh.Close()

	// Buffer the file IO
	buffered := bufio.NewWriter(fh)
	defer buffered.Flush()

	// Write out as JSON
	enc := json.NewEncoder(fh)
	if err := enc.Encode(&s.meta); err != nil {
		return err
	}
	return nil
}
<commit_msg>First pass at List<commit_after>package raft

import (
	"bufio"
	"encoding/json"
	"fmt"
	"hash"
	"hash/crc64"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"
)

const (
	snapPath      = "snapshots"
	metaFilePath  = "meta.json"
	stateFilePath = "state.bin"
	tmpSuffix     = ".tmp"
)

// FileSnapshotStore implements the SnapshotStore interface and allows
// snapshots to be made on the local disk.
type FileSnapshotStore struct {
	path   string
	retain int
}

// Implements the SnapshotSink
type FileSnapshotSink struct {
	dir  string
	meta fileSnapshotMeta

	stateFile *os.File
	stateHash hash.Hash64
	buffered  *bufio.Writer
}

type fileSnapshotMeta struct {
	SnapshotMeta
	CRC []byte
}

// NewFileSnapshotStore creates a new FileSnapshotStore based
// on a base directory. The `retain` parameter controls how many
// snapshots are retained. Must be at least 1.
func NewFileSnapshotStore(base string, retain int) (*FileSnapshotStore, error) {
	if retain < 1 {
		return nil, fmt.Errorf("must retain at least one snapshot")
	}

	path := filepath.Join(base, snapPath)
	store := &FileSnapshotStore{
		path:   path,
		retain: retain,
	}
	return store, nil
}

// Create is used to start a new snapshot
func (f *FileSnapshotStore) Create(index, term uint64, peers []byte) (SnapshotSink, error) {
	// Create a new path
	name := time.Now().Format(time.RFC3339) + tmpSuffix
	path := filepath.Join(f.path, name)
	log.Printf("[INFO] Creating new snapshot at %s", path)

	// Make the directory
	if err := os.Mkdir(path, 0755); err != nil {
		log.Printf("[ERR] Failed to make snapshot directory: %v", err)
		return nil, err
	}

	// Create the sink
	sink := &FileSnapshotSink{
		dir: path,
		meta: fileSnapshotMeta{
			SnapshotMeta: SnapshotMeta{
				ID:    path,
				Index: index,
				Term:  term,
				Peers: peers,
			},
			CRC: nil,
		},
	}

	// Write out the meta data
	if err := sink.writeMeta(); err != nil {
		log.Printf("[ERR] Failed to write metadata: %v", err)
		return nil, err
	}

	// Open the state file
	statePath := filepath.Join(path, stateFilePath)
	fh, err := os.Create(statePath)
	if err != nil {
		log.Printf("[ERR] Failed to create state file: %v", err)
		return nil, err
	}
	sink.stateFile = fh

	// Create a CRC64 hash
	sink.stateHash = crc64.New(crc64.MakeTable(crc64.ECMA))

	// Wrap both the hash and file in a MultiWriter with buffering
	multi := io.MultiWriter(sink.stateFile, sink.stateHash)
	sink.buffered = bufio.NewWriter(multi)

	// Done
	return sink, nil
}

func (f *FileSnapshotStore) List() ([]*SnapshotMeta, error) {
	// Get the eligible snapshots
	snapshots, err := ioutil.ReadDir(f.path)
	if err != nil {
		log.Printf("[ERR] Failed to scan snapshot dir: %v", err)
		return nil, err
	}

	// Populate the metadata, reverse order (newest first)
	var snapMeta []*SnapshotMeta
	for i := len(snapshots) - 1; i >= 0; i++ {
		// Ignore any files
		if !snapshots[i].IsDir() {
			continue
		}

		// Ignore any temporary snapshots
		dirName := snapshots[i].Name()
		if strings.HasSuffix(dirName, tmpSuffix) {
			log.Printf("[WARN] Found temporary snapshot: %v", dirName)
			continue
		}

		// Try to read the meta data
		meta, err := f.readMeta(dirName)
		if err != nil {
			log.Printf("[WARN] Failed to read metadata for %v: %v", dirName, err)
			continue
		}

		snapMeta = append(snapMeta, meta)
	}

	return snapMeta, nil
}

// readMeta is used to read the meta data for a given named backup
func (f *FileSnapshotStore) readMeta(name string) (*SnapshotMeta, error) {
	// Open the meta file
	metaPath := filepath.Join(f.path, name, metaFilePath)
	fh, err := os.Open(metaPath)
	if err != nil {
		return nil, err
	}
	defer fh.Close()

	// Buffer the file IO
	buffered := bufio.NewReader(fh)

	// Read in the JSON
	meta := &SnapshotMeta{}
	dec := json.NewDecoder(buffered)
	if err := dec.Decode(meta); err != nil {
		return nil, err
	}
	return meta, nil
}

func (f *FileSnapshotStore) Open(id string) (io.ReadCloser, error) {
	return nil, nil
}

// Write is used to append to the state file. We write to the
// buffered IO object to reduce the amount of context switches
func (s *FileSnapshotSink) Write(b []byte) (int, error) {
	return s.buffered.Write(b)
}

// Close is used to indicate a successful end
func (s *FileSnapshotSink) Close() error {
	// Close the open handles
	if err := s.finalize(); err != nil {
		log.Printf("[ERR] Failed to finalize snapshot: %v", err)
		return err
	}

	// Write out the meta data
	if err := s.writeMeta(); err != nil {
		log.Printf("[ERR] Failed to write metadata: %v", err)
		return err
	}

	// Move the directory into place
	newPath := strings.TrimSuffix(s.dir, tmpSuffix)
	if err := os.Rename(s.dir, newPath); err != nil {
		log.Printf("[ERR] Failed to move snapshot into place: %v", err)
		return err
	}

	// TODO: Retain count
	return nil
}

// Cancel is used to indicate an unsuccessful end
func (s *FileSnapshotSink) Cancel() error {
	// Close the open handles
	if err := s.finalize(); err != nil {
		log.Printf("[ERR] Failed to finalize snapshot: %v", err)
		return err
	}

	// Attempt to remove all artifacts
	return os.RemoveAll(s.dir)
}

// finalize is used to close all of our resources
func (s *FileSnapshotSink) finalize() error {
	if err := s.buffered.Flush(); err != nil {
		return err
	}
	if err := s.stateFile.Close(); err != nil {
		return err
	}
	s.meta.CRC = s.stateHash.Sum(nil)
	return nil
}

// writeMeta is used to write out the metadata we have
func (s *FileSnapshotSink) writeMeta() error {
	// Open the meta file
	metaPath := filepath.Join(s.dir, metaFilePath)
	fh, err := os.Create(metaPath)
	if err != nil {
		return err
	}
	defer fh.Close()

	// Buffer the file IO
	buffered := bufio.NewWriter(fh)
	defer buffered.Flush()

	// Write out as JSON
	enc := json.NewEncoder(buffered)
	if err := enc.Encode(&s.meta); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

// Package fileutil contains code to work with shell files, also known
// as shell scripts.
package fileutil // import "mvdan.cc/sh/fileutil"

import (
	"os"
	"regexp"
	"strings"
)

var (
	shebangRe = regexp.MustCompile(`^#!\s?/(usr/)?bin/(env\s+)?(sh|bash)\s`)
	extRe     = regexp.MustCompile(`\.(sh|bash)$`)
)

// HasShebang reports whether bs begins with a valid sh or bash shebang.
// It supports variations with /usr and env.
func HasShebang(bs []byte) bool {
	return shebangRe.Match(bs)
}

// ScriptConfidence defines how likely a file is to be a shell script,
// from complete certainty that it is not one to complete certainty that
// it is one.
type ScriptConfidence int

const (
	ConfNotScript ScriptConfidence = iota
	ConfIfShebang
	ConfIsScript
)

// CouldBeScript reports how likely a file is to be a shell script. It
// discards directories, hidden files and files with non-shell
// extensions.
func CouldBeScript(info os.FileInfo) ScriptConfidence {
	name := info.Name()
	switch {
	case info.IsDir(), name[0] == '.':
		return ConfNotScript
	case info.Mode()&os.ModeSymlink != 0:
		return ConfNotScript
	case extRe.MatchString(name):
		return ConfIsScript
	case strings.IndexByte(name, '.') > 0:
		return ConfNotScript // different extension
	case info.Size() < int64(len("#/bin/sh\n")):
		return ConfNotScript // cannot possibly hold valid shebang
	default:
		return ConfIfShebang
	}
}
<commit_msg>fileutil: clarify CouldBeScript with symlinks<commit_after>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

// Package fileutil contains code to work with shell files, also known
// as shell scripts.
package fileutil // import "mvdan.cc/sh/fileutil"

import (
	"os"
	"regexp"
	"strings"
)

var (
	shebangRe = regexp.MustCompile(`^#!\s?/(usr/)?bin/(env\s+)?(sh|bash)\s`)
	extRe     = regexp.MustCompile(`\.(sh|bash)$`)
)

// HasShebang reports whether bs begins with a valid sh or bash shebang.
// It supports variations with /usr and env.
func HasShebang(bs []byte) bool {
	return shebangRe.Match(bs)
}

// ScriptConfidence defines how likely a file is to be a shell script,
// from complete certainty that it is not one to complete certainty that
// it is one.
type ScriptConfidence int

const (
	ConfNotScript ScriptConfidence = iota
	ConfIfShebang
	ConfIsScript
)

// CouldBeScript reports how likely a file is to be a shell script. It
// discards directories, symlinks, hidden files and files with non-shell
// extensions.
func CouldBeScript(info os.FileInfo) ScriptConfidence {
	name := info.Name()
	switch {
	case info.IsDir(), name[0] == '.':
		return ConfNotScript
	case info.Mode()&os.ModeSymlink != 0:
		return ConfNotScript
	case extRe.MatchString(name):
		return ConfIsScript
	case strings.IndexByte(name, '.') > 0:
		return ConfNotScript // different extension
	case info.Size() < int64(len("#/bin/sh\n")):
		return ConfNotScript // cannot possibly hold valid shebang
	default:
		return ConfIfShebang
	}
}
<|endoftext|>


<commit_before>package tsize

import (
	"testing"
	"time"
)

func TestGetSize(t *testing.T) {
	defSize := Size{10, 20}
	fakeSize(defSize)
	defer unFakeSize()

	s, err := GetSize()

	if err != nil {
		t.Fatal("Failed with", err)
	}

	if s.Width != defSize.Width || s.Height != defSize.Height {
		t.Fatal("Terminal size should not be", s.Width, s.Height)
	}
}

func TestFgetSize(t *testing.T) {
	_, err := FgetSize(nil)

	if err != ErrNotATerminal {
		t.Fatal("Should fail with NotATerminal")
	}
}

func TestSizeListener(t *testing.T) {
	defSize := Size{10, 20}
	fakeSize(defSize)
	defer unFakeSize()

	sc, err := NewSizeListener()

	if err != nil {
		t.Fatal("Creating SizeChanger failed with", err)
	}

	triggerFakeResize()
	select {
	case s := <-sc.Change:
		if s.Width != defSize.Width || s.Height != defSize.Height {
			t.Fatal("Terminal size should not be", s.Width, s.Height)
		}
	case <-time.After(1 * time.Second):
		t.Fatal("Resize didn't trigger")
	}

	sc.Close()
	if sc.Change != nil {
		t.Fatal("Closing should nil the Change channel")
	}
}
<commit_msg>test: print out the proper error message<commit_after>package tsize

import (
	"testing"
	"time"
)

func TestGetSize(t *testing.T) {
	defSize := Size{10, 20}
	fakeSize(defSize)
	defer unFakeSize()

	s, err := GetSize()

	if err != nil {
		t.Fatal("Failed with", err)
	}

	if s.Width != defSize.Width || s.Height != defSize.Height {
		t.Fatal("Terminal size should not be", s.Width, s.Height)
	}
}

func TestFgetSize(t *testing.T) {
	_, err := FgetSize(nil)

	if err != ErrNotATerminal {
		t.Fatal("Should fail with", ErrNotATerminal)
	}
}

func TestSizeListener(t *testing.T) {
	defSize := Size{10, 20}
	fakeSize(defSize)
	defer unFakeSize()

	sc, err := NewSizeListener()

	if err != nil {
		t.Fatal("Creating SizeChanger failed with", err)
	}

	triggerFakeResize()
	select {
	case s := <-sc.Change:
		if s.Width != defSize.Width || s.Height != defSize.Height {
			t.Fatal("Terminal size should not be", s.Width, s.Height)
		}
	case <-time.After(1 * time.Second):
		t.Fatal("Resize didn't trigger")
	}

	sc.Close()
	if sc.Change != nil {
		t.Fatal("Closing should nil the Change channel")
	}
}
<|endoftext|>


<commit_before>package flickr

import (
	"bytes"
	flickErr "github.com/masci/flickr.go/flickr/error"
	"io"
	"mime/multipart"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

// Encode the file and request parameters in a multipart body
func getUploadBody(client *FlickrClient, file *os.File) (*bytes.Buffer, string, error) {
	// instance an empty request body
	body := &bytes.Buffer{}
	// multipart writer to fill the body
	writer := multipart.NewWriter(body)
	// dump the file in the "photo" field
	part, err := writer.CreateFormFile("photo", filepath.Base(file.Name()))
	if err != nil {
		return nil, "", err
	}
	_, err = io.Copy(part, file)
	// dump other params
	for key, val := range client.Args {
		_ = writer.WriteField(key, val[0])
	}
	err = writer.Close()
	if err != nil {
		return nil, "", err
	}
	// evaluate the content type and the boundary
	contentType := writer.FormDataContentType()

	return body, contentType, nil
}

// A convenience struct wrapping all optional upload parameters
type UploadParams struct {
	Title, Description           string
	Tags                         []string
	IsPublic, IsFamily, IsFriend bool
	ContentType                  int
	Hidden                       int
	SafetyLevel                  int
}

// Provide meaningful default values
func NewUploadParams() *UploadParams {
	ret := &UploadParams{}
	ret.ContentType = 1 // photo
	ret.Hidden = 2      // hidden from public searchesi
	ret.SafetyLevel = 1 // safe
	return ret
}

// Type representing a successful upload response from the api
type UploadResponse struct {
	FlickrResponse
	Id int `xml:"photoid"`
}

// Set client query arguments based on the contents of the UploadParams struct
func fillArgsWithParams(client *FlickrClient, params *UploadParams) {
	if params.Title != "" {
		client.Args.Set("title", params.Title)
	}

	if params.Description != "" {
		client.Args.Set("description", params.Description)
	}

	if len(params.Tags) > 0 {
		client.Args.Set("tags", strings.Join(params.Tags, " "))
	}

	var boolString = func(b bool) string {
		if b {
			return "1"
		}
		return "0"
	}
	client.Args.Set("is_public", boolString(params.IsPublic))
	client.Args.Set("is_friend", boolString(params.IsFriend))
	client.Args.Set("is_family", boolString(params.IsFamily))

	if params.ContentType >= 1 && params.ContentType <= 3 {
		client.Args.Set("content_type", strconv.Itoa(params.ContentType))
	}

	if params.Hidden >= 1 && params.Hidden <= 2 {
		client.Args.Set("hidden", strconv.Itoa(params.Hidden))
	}

	if params.SafetyLevel >= 1 && params.SafetyLevel <= 3 {
		client.Args.Set("safety_level", strconv.Itoa(params.SafetyLevel))
	}
}

// Perform a file upload using the Flickr API. If optionalParams is nil,
// no parameters will be added to the request and Flickr will set User's
// default preferences.
// This call must be signed with write permissions
func UploadPhoto(client *FlickrClient, path string, optionalParams *UploadParams) (*UploadResponse, error) {
	client.EndpointUrl = UPLOAD_ENDPOINT
	client.HTTPVerb = "POST"
	client.SetDefaultArgs()
	client.Args.Set("oauth_token", client.OAuthToken)
	client.Args.Set("oauth_consumer_key", client.ApiKey)

	if optionalParams != nil {
		fillArgsWithParams(client, optionalParams)
	}

	client.Sign(client.OAuthTokenSecret)

	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	body, ctype, err := getUploadBody(client, file)
	if err != nil {
		return nil, err
	}

	resp := &UploadResponse{}
	err = DoPost(client, body, ctype, resp)

	if err == nil {
		return nil, err
	}

	if resp.HasErrors() {
		return resp, flickErr.NewError(10)
	}

	return resp, nil
}
<commit_msg>fixed huge bug<commit_after>package flickr

import (
	"bytes"
	flickErr "github.com/masci/flickr.go/flickr/error"
	"io"
	"mime/multipart"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

// Encode the file and request parameters in a multipart body
func getUploadBody(client *FlickrClient, file *os.File) (*bytes.Buffer, string, error) {
	// instance an empty request body
	body := &bytes.Buffer{}
	// multipart writer to fill the body
	writer := multipart.NewWriter(body)
	// dump the file in the "photo" field
	part, err := writer.CreateFormFile("photo", filepath.Base(file.Name()))
	if err != nil {
		return nil, "", err
	}
	_, err = io.Copy(part, file)
	// dump other params
	for key, val := range client.Args {
		_ = writer.WriteField(key, val[0])
	}
	err = writer.Close()
	if err != nil {
		return nil, "", err
	}
	// evaluate the content type and the boundary
	contentType := writer.FormDataContentType()

	return body, contentType, nil
}

// A convenience struct wrapping all optional upload parameters
type UploadParams struct {
	Title, Description           string
	Tags                         []string
	IsPublic, IsFamily, IsFriend bool
	ContentType                  int
	Hidden                       int
	SafetyLevel                  int
}

// Provide meaningful default values
func NewUploadParams() *UploadParams {
	ret := &UploadParams{}
	ret.ContentType = 1 // photo
	ret.Hidden = 2      // hidden from public searchesi
	ret.SafetyLevel = 1 // safe
	return ret
}

// Type representing a successful upload response from the api
type UploadResponse struct {
	FlickrResponse
	Id int `xml:"photoid"`
}

// Set client query arguments based on the contents of the UploadParams struct
func fillArgsWithParams(client *FlickrClient, params *UploadParams) {
	if params.Title != "" {
		client.Args.Set("title", params.Title)
	}

	if params.Description != "" {
		client.Args.Set("description", params.Description)
	}

	if len(params.Tags) > 0 {
		client.Args.Set("tags", strings.Join(params.Tags, " "))
	}

	var boolString = func(b bool) string {
		if b {
			return "1"
		}
		return "0"
	}
	client.Args.Set("is_public", boolString(params.IsPublic))
	client.Args.Set("is_friend", boolString(params.IsFriend))
	client.Args.Set("is_family", boolString(params.IsFamily))

	if params.ContentType >= 1 && params.ContentType <= 3 {
		client.Args.Set("content_type", strconv.Itoa(params.ContentType))
	}

	if params.Hidden >= 1 && params.Hidden <= 2 {
		client.Args.Set("hidden", strconv.Itoa(params.Hidden))
	}

	if params.SafetyLevel >= 1 && params.SafetyLevel <= 3 {
		client.Args.Set("safety_level", strconv.Itoa(params.SafetyLevel))
	}
}

// Perform a file upload using the Flickr API. If optionalParams is nil,
// no parameters will be added to the request and Flickr will set User's
// default preferences.
// This call must be signed with write permissions
func UploadPhoto(client *FlickrClient, path string, optionalParams *UploadParams) (*UploadResponse, error) {
	client.EndpointUrl = UPLOAD_ENDPOINT
	client.HTTPVerb = "POST"
	client.SetDefaultArgs()
	client.Args.Set("oauth_token", client.OAuthToken)
	client.Args.Set("oauth_consumer_key", client.ApiKey)

	if optionalParams != nil {
		fillArgsWithParams(client, optionalParams)
	}

	client.Sign(client.OAuthTokenSecret)

	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	body, ctype, err := getUploadBody(client, file)
	if err != nil {
		return nil, err
	}

	resp := &UploadResponse{}
	err = DoPost(client, body, ctype, resp)

	if err != nil {
		return nil, err
	}

	if resp.HasErrors() {
		return resp, flickErr.NewError(10)
	}

	return resp, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"strings"

	"github.com/flynn/flynn-controller/client"
	ct "github.com/flynn/flynn-controller/types"
	"github.com/flynn/go-discoverd"
	"github.com/flynn/go-flynn/cluster"
	"github.com/flynn/go-flynn/exec"
)

var clusterc *cluster.Client

func init() {
	log.SetFlags(0)

	var err error
	clusterc, err = cluster.NewClient()
	if err != nil {
		log.Fatalln("Error connecting to cluster leader:", err)
	}
}

var typesPattern = regexp.MustCompile(`types.* -> (.+)$`)

func main() {
	client, err := controller.NewClient("", os.Getenv("CONTROLLER_AUTH_KEY"))
	if err != nil {
		log.Fatalln("Unable to connect to controller:", err)
	}
	// TODO: use discoverd http dialer here?
	services, err := discoverd.Services("shelf", discoverd.DefaultTimeout)
	if err != nil || len(services) < 1 {
		log.Fatalf("Unable to discover shelf %q", err)
	}
	shelfHost := services[0].Addr

	app := os.Args[1]
	commit := os.Args[2]

	_, err = client.GetApp(app)
	if err == controller.ErrNotFound {
		log.Fatalf("Unknown app %q", app)
	} else if err != nil {
		log.Fatalln("Error retrieving app:", err)
	}

	fmt.Printf("-----> Building %s...\n", app)

	var output bytes.Buffer
	cmd := exec.Command("flynn/slugbuilder", fmt.Sprintf("http://%s/%s/tgz", shelfHost, commit))
	cmd.Stdout = io.MultiWriter(os.Stdout, &output)
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin

	if err := cmd.Run(); err != nil {
		log.Fatalln("Build failed:", err)
	}

	var types []string
	if match := typesPattern.FindSubmatch(output.Bytes()); match != nil {
		types = strings.Split(string(match[1]), ", ")
	}

	fmt.Printf("-----> Creating release...\n")

	prevRelease, err := client.GetAppRelease(app)
	if err == controller.ErrNotFound {
		prevRelease = &ct.Release{}
	} else if err != nil {
		log.Fatalln("Error creating getting current app release:", err)
	}
	artifact := &ct.Artifact{URI: "docker://flynn/slugrunner"}
	if err := client.CreateArtifact(artifact); err != nil {
		log.Fatalln("Error creating artifact:", err)
	}

	release := &ct.Release{
		ArtifactID: artifact.ID,
		Env:        prevRelease.Env,
	}
	procs := make(map[string]ct.ProcessType)
	for _, t := range types {
		proc := prevRelease.Processes[t]
		proc.Cmd = []string{"start", t}
		if t == "web" {
			proc.Ports.TCP = 1
			if proc.Env == nil {
				proc.Env = make(map[string]string)
			}
			proc.Env["SD_NAME"] = app + "-web"
		}
		procs[t] = proc
	}
	release.Processes = procs
	if release.Env == nil {
		release.Env = make(map[string]string)
	}
	release.Env["SLUG_URL"] = "http://" + shelfHost + "/" + commit + ".tgz"

	if err := client.CreateRelease(release); err != nil {
		log.Fatalln("Error creating release:", err)
	}
	if err := client.SetAppRelease(app, release.ID); err != nil {
		log.Fatalln("Error setting app release:", err)
	}

	fmt.Println("=====> Application deployed")
}
<commit_msg>receiver: Fix typo<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"strings"

	"github.com/flynn/flynn-controller/client"
	ct "github.com/flynn/flynn-controller/types"
	"github.com/flynn/go-discoverd"
	"github.com/flynn/go-flynn/cluster"
	"github.com/flynn/go-flynn/exec"
)

var clusterc *cluster.Client

func init() {
	log.SetFlags(0)

	var err error
	clusterc, err = cluster.NewClient()
	if err != nil {
		log.Fatalln("Error connecting to cluster leader:", err)
	}
}

var typesPattern = regexp.MustCompile(`types.* -> (.+)$`)

func main() {
	client, err := controller.NewClient("", os.Getenv("CONTROLLER_AUTH_KEY"))
	if err != nil {
		log.Fatalln("Unable to connect to controller:", err)
	}
	// TODO: use discoverd http dialer here?
	services, err := discoverd.Services("shelf", discoverd.DefaultTimeout)
	if err != nil || len(services) < 1 {
		log.Fatalf("Unable to discover shelf %q", err)
	}
	shelfHost := services[0].Addr

	app := os.Args[1]
	commit := os.Args[2]

	_, err = client.GetApp(app)
	if err == controller.ErrNotFound {
		log.Fatalf("Unknown app %q", app)
	} else if err != nil {
		log.Fatalln("Error retrieving app:", err)
	}

	fmt.Printf("-----> Building %s...\n", app)

	var output bytes.Buffer
	slugURL := fmt.Sprintf("http://%s/%s.tgz", shelfHost, commit)
	cmd := exec.Command("flynn/slugbuilder", slugURL)
	cmd.Stdout = io.MultiWriter(os.Stdout, &output)
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin

	if err := cmd.Run(); err != nil {
		log.Fatalln("Build failed:", err)
	}

	var types []string
	if match := typesPattern.FindSubmatch(output.Bytes()); match != nil {
		types = strings.Split(string(match[1]), ", ")
	}

	fmt.Printf("-----> Creating release...\n")

	prevRelease, err := client.GetAppRelease(app)
	if err == controller.ErrNotFound {
		prevRelease = &ct.Release{}
	} else if err != nil {
		log.Fatalln("Error creating getting current app release:", err)
	}
	artifact := &ct.Artifact{URI: "docker://flynn/slugrunner"}
	if err := client.CreateArtifact(artifact); err != nil {
		log.Fatalln("Error creating artifact:", err)
	}

	release := &ct.Release{
		ArtifactID: artifact.ID,
		Env:        prevRelease.Env,
	}
	procs := make(map[string]ct.ProcessType)
	for _, t := range types {
		proc := prevRelease.Processes[t]
		proc.Cmd = []string{"start", t}
		if t == "web" {
			proc.Ports.TCP = 1
			if proc.Env == nil {
				proc.Env = make(map[string]string)
			}
			proc.Env["SD_NAME"] = app + "-web"
		}
		procs[t] = proc
	}
	release.Processes = procs
	if release.Env == nil {
		release.Env = make(map[string]string)
	}
	release.Env["SLUG_URL"] = slugURL

	if err := client.CreateRelease(release); err != nil {
		log.Fatalln("Error creating release:", err)
	}
	if err := client.SetAppRelease(app, release.ID); err != nil {
		log.Fatalln("Error setting app release:", err)
	}

	fmt.Println("=====> Application deployed")
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime/pprof"
	"sort"
	"text/tabwriter"
)

const VERSION = `0.3.1`

var languages = []Language{
	Language{"Thrift", mExt(".thrift"), cComments},

	Language{"C", mExt(".c", ".h"), cComments},
	Language{"C++", mExt(".cc", ".cpp", ".cxx", ".hh", ".hpp", ".hxx"), cComments},
	Language{"Go", mExt(".go"), cComments},
	Language{"Rust", mExt(".rs", ".rc"), cComments},
	Language{"Scala", mExt(".scala"), cComments},
	Language{"Java", mExt(".java"), cComments},

	Language{"YACC", mExt(".y"), cComments},
	Language{"Lex", mExt(".l"), cComments},

	Language{"Lua", mExt(".lua"), luaComments},

	Language{"SQL", mExt(".sql"), sqlComments},

	Language{"Haskell", mExt(".hs", ".lhs"), hsComments},
	Language{"ML", mExt(".ml", ".mli"), mlComments},

	Language{"Perl", mExt(".pl", ".pm"), perlComments},
	Language{"PHP", mExt(".php"), cComments},

	Language{"Shell", mExt(".sh"), shComments},
	Language{"Bash", mExt(".bash"), shComments},
	Language{"R", mExt(".r", ".R"), shComments},
	Language{"Tcl", mExt(".tcl"), shComments},

	Language{"MATLAB", mExt(".m"), matlabComments},

	Language{"Ruby", mExt(".rb"), rubyComments},
	Language{"Python", mExt(".py"), pyComments},
	Language{"Assembly", mExt(".asm", ".s"), semiComments},
	Language{"Lisp", mExt(".lsp", ".lisp"), semiComments},
	Language{"Scheme", mExt(".scm", ".scheme"), semiComments},

	Language{"Make", mName("makefile", "Makefile", "MAKEFILE"), shComments},
	Language{"CMake", mName("CMakeLists.txt"), shComments},
	Language{"Jam", mName("Jamfile", "Jamrules"), shComments},

	Language{"Markdown", mExt(".md"), noComments},

	Language{"HAML", mExt(".haml"), noComments},
	Language{"SASS", mExt(".sass"), cssComments},
	Language{"SCSS", mExt(".scss"), cssComments},

	Language{"HTML", mExt(".htm", ".html", ".xhtml"), xmlComments},
	Language{"XML", mExt(".xml"), xmlComments},
	Language{"CSS", mExt(".css"), cssComments},
	Language{"JavaScript", mExt(".js"), cComments},
	Language{"CoffeeScript", mExt(".coffee"), coffeeComments},
	Language{"JSON", mExt(".json"), noComments},

	Language{"Erlang", mExt(".erl"), erlangComments},
}

type Commenter struct {
	LineComment  string
	StartComment string
	EndComment   string
	Nesting      bool
}

var (
	noComments     = Commenter{"\000", "\000", "\000", false}
	xmlComments    = Commenter{"\000", `<!--`, `-->`, false}
	cComments      = Commenter{`//`, `/*`, `*/`, false}
	cssComments    = Commenter{"\000", `/*`, `*/`, false}
	shComments     = Commenter{`#`, "\000", "\000", false}
	semiComments   = Commenter{`;`, "\000", "\000", false}
	hsComments     = Commenter{`--`, `{-`, `-}`, true}
	mlComments     = Commenter{`\000`, `(*`, `*)`, false}
	sqlComments    = Commenter{`--`, `/*`, `*/`, false}
	luaComments    = Commenter{`--`, `--[[`, `]]`, false}
	pyComments     = Commenter{`#`, `"""`, `"""`, false}
	matlabComments = Commenter{`%`, `%{`, `%}`, false}
	erlangComments = Commenter{`%`, "\000", "\000", false}
	rubyComments   = Commenter{`#`, "=begin", "=end", false}
	coffeeComments = Commenter{`#`, "###", "###", false}

	// TODO support POD and __END__
	perlComments = Commenter{`#`, "\000", "\000", false}
)

type Language struct {
	Namer
	Matcher
	Commenter
}

// TODO work properly with unicode
func (l Language) Update(c []byte, s *Stats) {
	s.FileCount++

	commentLen := 0
	codeLen := 0
	inComment := 0 // this is an int for nesting
	inLComment := false
	lc := []byte(l.LineComment)
	sc := []byte(l.StartComment)
	ec := []byte(l.EndComment)
	lp, sp, ep := 0, 0, 0

	for _, b := range c {
		if b != byte(' ') && b != byte('\t') && b != byte('\n') && b != byte('\r') {
			if !inLComment && inComment == 0 {
				codeLen++
			} else {
				commentLen++
			}
		}
		if inComment == 0 && b == lc[lp] {
			lp++
			if lp == len(lc) {
				if !inLComment {
					codeLen -= lp
				}
				inLComment = true
				lp = 0
			}
		} else {
			lp = 0
		}
		if !inLComment && b == sc[sp] {
			sp++
			if sp == len(sc) {
				if inComment == 0 {
					codeLen -= sp
				}
				inComment++
				if inComment > 1 && !l.Nesting {
					inComment = 1
				}
				sp = 0
			}
		} else {
			sp = 0
		}
		if !inLComment && inComment > 0 && b == ec[ep] {
			ep++
			if ep == len(ec) {
				if inComment > 0 {
					inComment--
				}
				if inComment == 0 {
					commentLen -= ep
				}
				ep = 0
			}
		} else {
			ep = 0
		}

		if b == byte('\n') {
			s.TotalLines++
			if commentLen > 0 {
				s.CommentLines++
			}
			if codeLen > 0 {
				s.CodeLines++
			}
			if commentLen == 0 && codeLen == 0 {
				s.BlankLines++
			}
			inLComment = false
			codeLen = 0
			commentLen = 0
			continue
		}
	}
}

type Namer string

func (l Namer) Name() string { return string(l) }

type Matcher func(string) bool

func (m Matcher) Match(fname string) bool { return m(fname) }

func mExt(exts ...string) Matcher {
	return func(fname string) bool {
		for _, ext := range exts {
			if ext == path.Ext(fname) {
				return true
			}
		}
		return false
	}
}

func mName(names ...string) Matcher {
	return func(fname string) bool {
		for _, name := range names {
			if name == path.Base(fname) {
				return true
			}
		}
		return false
	}
}

type Stats struct {
	FileCount    int
	TotalLines   int
	CodeLines    int
	BlankLines   int
	CommentLines int
}

var info = map[string]*Stats{}

func handleFileLang(fname string, l Language) {
	i, ok := info[l.Name()]
	if !ok {
		i = &Stats{}
		info[l.Name()] = i
	}
	c, err := ioutil.ReadFile(fname)
	if err != nil {
		fmt.Fprintf(os.Stderr, "  ! %s\n", fname)
		return
	}
	l.Update(c, i)
}

func handleFile(fname string) {
	for _, lang := range languages {
		if lang.Match(fname) {
			handleFileLang(fname, lang)
			return
		}
	}
	// TODO No recognized extension - check for hashbang
}

var files []string

func add(n string) {
	fi, err := os.Stat(n)
	if err != nil {
		goto invalid
	}
	if fi.IsDir() {
		fs, err := ioutil.ReadDir(n)
		if err != nil {
			goto invalid
		}
		for _, f := range fs {
			if f.Name() == ".nosloc" {
				return
			}
		}
		for _, f := range fs {
			if f.Name()[0] != '.' {
				add(path.Join(n, f.Name()))
			}
		}
		return
	}
	if fi.Mode()&os.ModeType == 0 {
		files = append(files, n)
		return
	}

	println(fi.Mode())

invalid:
	fmt.Fprintf(os.Stderr, "  ! %s\n", n)
}

type LData []LResult

func (d LData) Len() int { return len(d) }

func (d LData) Less(i, j int) bool {
	if d[i].CodeLines == d[j].CodeLines {
		return d[i].Name > d[j].Name
	}
	return d[i].CodeLines > d[j].CodeLines
}

func (d LData) Swap(i, j int) {
	d[i], d[j] = d[j], d[i]
}

type LResult struct {
	Name         string
	FileCount    int
	CodeLines    int
	CommentLines int
	BlankLines   int
	TotalLines   int
}

func (r *LResult) Add(a LResult) {
	r.FileCount += a.FileCount
	r.CodeLines += a.CodeLines
	r.CommentLines += a.CommentLines
	r.BlankLines += a.BlankLines
	r.TotalLines += a.TotalLines
}

func printJSON() {
	bs, err := json.MarshalIndent(info, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(bs))
}

func printInfo() {
	w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', tabwriter.AlignRight)
	fmt.Fprintln(w, "Language\tFiles\tCode\tComment\tBlank\tTotal\t")
	d := LData([]LResult{})
	total := &LResult{}
	total.Name = "Total"
	for n, i := range info {
		r := LResult{n, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines}
		d = append(d, r)
		total.Add(r)
	}
	d = append(d, *total)
	sort.Sort(d)
	for _, i := range d {
		fmt.Fprintf(w, "%s\t%d\t%d\t%d\t%d\t%d\t\n", i.Name, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines)
	}

	w.Flush()
}

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	useJson    = flag.Bool("json", false, "JSON-format output")
	version    = flag.Bool("V", false, "display version info and exit")
)

func main() {
	flag.Parse()
	if *version {
		fmt.Printf("sloc %s\n", VERSION)
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %s\n", err.Error())
			return
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	args := flag.Args()
	if len(args) == 0 {
		args = append(args, `.`)
	}

	for _, n := range args {
		add(n)
	}

	for _, f := range files {
		handleFile(f)
	}

	if *useJson {
		printJSON()
	} else {
		printInfo()
	}
}
<commit_msg>objc + swift support<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime/pprof"
	"sort"
	"text/tabwriter"
)

const VERSION = `1.0.0`

var languages = []Language{
	Language{"Thrift", mExt(".thrift"), cComments},

	Language{"C", mExt(".c"), cComments},
	Language{"C++", mExt(".cc", ".cpp", ".cxx", ".hh", ".hpp", ".hxx"), cComments},
	Language{"Go", mExt(".go"), cComments},
	Language{"Rust", mExt(".rs", ".rc"), cComments},
	Language{"Scala", mExt(".scala"), cComments},
	Language{"Java", mExt(".java"), cComments},
	Language{"Objc", mExt(".m, .h"), cComments},
	Language{"Swift", mExt(".swift"), cComments},

	Language{"YACC", mExt(".y"), cComments},
	Language{"Lex", mExt(".l"), cComments},

	Language{"Lua", mExt(".lua"), luaComments},

	Language{"SQL", mExt(".sql"), sqlComments},

	Language{"Haskell", mExt(".hs", ".lhs"), hsComments},
	Language{"ML", mExt(".ml", ".mli"), mlComments},

	Language{"Perl", mExt(".pl", ".pm"), perlComments},
	Language{"PHP", mExt(".php"), cComments},

	Language{"Shell", mExt(".sh"), shComments},
	Language{"Bash", mExt(".bash"), shComments},
	Language{"R", mExt(".r", ".R"), shComments},
	Language{"Tcl", mExt(".tcl"), shComments},

	Language{"MATLAB", mExt(".m"), matlabComments},

	Language{"Ruby", mExt(".rb"), rubyComments},
	Language{"Python", mExt(".py"), pyComments},
	Language{"Assembly", mExt(".asm", ".s"), semiComments},
	Language{"Lisp", mExt(".lsp", ".lisp"), semiComments},
	Language{"Scheme", mExt(".scm", ".scheme"), semiComments},

	Language{"Make", mName("makefile", "Makefile", "MAKEFILE"), shComments},
	Language{"CMake", mName("CMakeLists.txt"), shComments},
	Language{"Jam", mName("Jamfile", "Jamrules"), shComments},

	Language{"Markdown", mExt(".md"), noComments},

	Language{"HAML", mExt(".haml"), noComments},
	Language{"SASS", mExt(".sass"), cssComments},
	Language{"SCSS", mExt(".scss"), cssComments},

	Language{"HTML", mExt(".htm", ".html", ".xhtml"), xmlComments},
	Language{"XML", mExt(".xml"), xmlComments},
	Language{"CSS", mExt(".css"), cssComments},
	Language{"JavaScript", mExt(".js"), cComments},
	Language{"CoffeeScript", mExt(".coffee"), coffeeComments},
	Language{"JSON", mExt(".json"), noComments},

	Language{"Erlang", mExt(".erl"), erlangComments},
}

type Commenter struct {
	LineComment  string
	StartComment string
	EndComment   string
	Nesting      bool
}

var (
	noComments     = Commenter{"\000", "\000", "\000", false}
	xmlComments    = Commenter{"\000", `<!--`, `-->`, false}
	cComments      = Commenter{`//`, `/*`, `*/`, false}
	cssComments    = Commenter{"\000", `/*`, `*/`, false}
	shComments     = Commenter{`#`, "\000", "\000", false}
	semiComments   = Commenter{`;`, "\000", "\000", false}
	hsComments     = Commenter{`--`, `{-`, `-}`, true}
	mlComments     = Commenter{`\000`, `(*`, `*)`, false}
	sqlComments    = Commenter{`--`, `/*`, `*/`, false}
	luaComments    = Commenter{`--`, `--[[`, `]]`, false}
	pyComments     = Commenter{`#`, `"""`, `"""`, false}
	matlabComments = Commenter{`%`, `%{`, `%}`, false}
	erlangComments = Commenter{`%`, "\000", "\000", false}
	rubyComments   = Commenter{`#`, "=begin", "=end", false}
	coffeeComments = Commenter{`#`, "###", "###", false}

	// TODO support POD and __END__
	perlComments = Commenter{`#`, "\000", "\000", false}
)

type Language struct {
	Namer
	Matcher
	Commenter
}

// TODO work properly with unicode
func (l Language) Update(c []byte, s *Stats) {
	s.FileCount++

	commentLen := 0
	codeLen := 0
	inComment := 0 // this is an int for nesting
	inLComment := false
	lc := []byte(l.LineComment)
	sc := []byte(l.StartComment)
	ec := []byte(l.EndComment)
	lp, sp, ep := 0, 0, 0

	for _, b := range c {
		if b != byte(' ') && b != byte('\t') && b != byte('\n') && b != byte('\r') {
			if !inLComment && inComment == 0 {
				codeLen++
			} else {
				commentLen++
			}
		}
		if inComment == 0 && b == lc[lp] {
			lp++
			if lp == len(lc) {
				if !inLComment {
					codeLen -= lp
				}
				inLComment = true
				lp = 0
			}
		} else {
			lp = 0
		}
		if !inLComment && b == sc[sp] {
			sp++
			if sp == len(sc) {
				if inComment == 0 {
					codeLen -= sp
				}
				inComment++
				if inComment > 1 && !l.Nesting {
					inComment = 1
				}
				sp = 0
			}
		} else {
			sp = 0
		}
		if !inLComment && inComment > 0 && b == ec[ep] {
			ep++
			if ep == len(ec) {
				if inComment > 0 {
					inComment--
				}
				if inComment == 0 {
					commentLen -= ep
				}
				ep = 0
			}
		} else {
			ep = 0
		}

		if b == byte('\n') {
			s.TotalLines++
			if commentLen > 0 {
				s.CommentLines++
			}
			if codeLen > 0 {
				s.CodeLines++
			}
			if commentLen == 0 && codeLen == 0 {
				s.BlankLines++
			}
			inLComment = false
			codeLen = 0
			commentLen = 0
			continue
		}
	}
}

type Namer string

func (l Namer) Name() string { return string(l) }

type Matcher func(string) bool

func (m Matcher) Match(fname string) bool { return m(fname) }

func mExt(exts ...string) Matcher {
	return func(fname string) bool {
		for _, ext := range exts {
			if ext == path.Ext(fname) {
				return true
			}
		}
		return false
	}
}

func mName(names ...string) Matcher {
	return func(fname string) bool {
		for _, name := range names {
			if name == path.Base(fname) {
				return true
			}
		}
		return false
	}
}

type Stats struct {
	FileCount    int
	TotalLines   int
	CodeLines    int
	BlankLines   int
	CommentLines int
}

var info = map[string]*Stats{}

func handleFileLang(fname string, l Language) {
	i, ok := info[l.Name()]
	if !ok {
		i = &Stats{}
		info[l.Name()] = i
	}
	c, err := ioutil.ReadFile(fname)
	if err != nil {
		fmt.Fprintf(os.Stderr, "  ! %s\n", fname)
		return
	}
	l.Update(c, i)
}

func handleFile(fname string) {
	for _, lang := range languages {
		if lang.Match(fname) {
			handleFileLang(fname, lang)
			return
		}
	}
	// TODO No recognized extension - check for hashbang
}

var files []string

func add(n string) {
	fi, err := os.Stat(n)
	if err != nil {
		goto invalid
	}
	if fi.IsDir() {
		fs, err := ioutil.ReadDir(n)
		if err != nil {
			goto invalid
		}
		for _, f := range fs {
			if f.Name() == ".nosloc" {
				return
			}
		}
		for _, f := range fs {
			if f.Name()[0] != '.' {
				add(path.Join(n, f.Name()))
			}
		}
		return
	}
	if fi.Mode()&os.ModeType == 0 {
		files = append(files, n)
		return
	}

	println(fi.Mode())

invalid:
	fmt.Fprintf(os.Stderr, "  ! %s\n", n)
}

type LData []LResult

func (d LData) Len() int { return len(d) }

func (d LData) Less(i, j int) bool {
	if d[i].CodeLines == d[j].CodeLines {
		return d[i].Name > d[j].Name
	}
	return d[i].CodeLines > d[j].CodeLines
}

func (d LData) Swap(i, j int) {
	d[i], d[j] = d[j], d[i]
}

type LResult struct {
	Name         string
	FileCount    int
	CodeLines    int
	CommentLines int
	BlankLines   int
	TotalLines   int
}

func (r *LResult) Add(a LResult) {
	r.FileCount += a.FileCount
	r.CodeLines += a.CodeLines
	r.CommentLines += a.CommentLines
	r.BlankLines += a.BlankLines
	r.TotalLines += a.TotalLines
}

func printJSON() {
	bs, err := json.MarshalIndent(info, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(bs))
}

func printInfo() {
	w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', tabwriter.AlignRight)
	fmt.Fprintln(w, "Language\tFiles\tCode\tComment\tBlank\tTotal\t")
	d := LData([]LResult{})
	total := &LResult{}
	total.Name = "Total"
	for n, i := range info {
		r := LResult{n, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines}
		d = append(d, r)
		total.Add(r)
	}
	d = append(d, *total)
	sort.Sort(d)
	for _, i := range d {
		fmt.Fprintf(w, "%s\t%d\t%d\t%d\t%d\t%d\t\n", i.Name, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines)
	}

	w.Flush()
}

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	useJson    = flag.Bool("json", false, "JSON-format output")
	version    = flag.Bool("V", false, "display version info and exit")
)

func main() {
	flag.Parse()
	if *version {
		fmt.Printf("sloc %s\n", VERSION)
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %s\n", err.Error())
			return
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	args := flag.Args()
	if len(args) == 0 {
		args = append(args, `.`)
	}

	for _, n := range args {
		add(n)
	}

	for _, f := range files {
		handleFile(f)
	}

	if *useJson {
		printJSON()
	} else {
		printInfo()
	}
}
<|endoftext|>


<commit_before>package solr

import (
	"github.com/samuel/go-zookeeper/zk"
)

func (s *solrZkInstance) Listen() error {
	err := s.zookeeper.Connect()
	s.currentNode = 0
	if err != nil {
		return err
	}
	s.clusterState = ClusterState{}

	collectionsEvents, err := s.initCollectionsListener()
	if err != nil {
		return err
	}
	liveNodeEvents, err := s.initLiveNodesListener()
	if err != nil {
		return err
	}
	leaderEvents, err := s.initLeaderElectListener()
	if err != nil {
		return err
	}
	//loop forever
	go func() {
		for {
			select {
			case cEvent := <-collectionsEvents:
				// do something if its not a session or disconnect
				if cEvent.Type > zk.EventSession {
					collections, version, err := s.zookeeper.GetClusterState()
					if err != nil {
						continue
					}
					s.setCollections(collections, version)
				}
				if cEvent.State < zk.StateConnected {
					s.logger.Printf("[Error] solr cluster zk disconnected  %v", cEvent)
				} else {
					s.logger.Printf("go-solr: solr cluster zk state changed zkType: %d zkState: %d", cEvent.Type, cEvent.State)
				}
			case lEvent := <-leaderEvents:
				if lEvent.Type == zk.EventNodeChildrenChanged || lEvent.Type == zk.EventNodeDataChanged {
					// s.Logger().Printf("Leader changed pausing")
				}
			case nEvent := <-liveNodeEvents:
				// do something if its not a session or disconnect
				if nEvent.Type > zk.EventSession {
					liveNodes, err := s.zookeeper.GetLiveNodes()
					if err != nil {
						continue
					}
					s.setLiveNodes(liveNodes)
				}
				if nEvent.State < zk.StateConnected {
					s.logger.Printf("[Error] solr cluster zk live nodes disconnected zkType: %v ", nEvent)
				} else {
					s.logger.Printf("go-solr: solr cluster zk live nodes state changed zkType: %d zkState: %d", nEvent.Type, nEvent.State)
				}
			}
		}
	}()
	s.listening = true
	return nil
}

func (s *solrZkInstance) initCollectionsListener() (<-chan zk.Event, error) {
	s.clusterState = ClusterState{}
	collections, version, collectionsEvents, err := s.zookeeper.GetClusterStateW()
	if err != nil {
		return nil, err
	}
	s.setCollections(collections, version)
	return collectionsEvents, nil
}

func (s *solrZkInstance) initLiveNodesListener() (<-chan zk.Event, error) {
	liveNodes, liveNodeEvents, err := s.zookeeper.GetLiveNodesW()
	if err != nil {
		return nil, err
	}
	s.setLiveNodes(liveNodes)
	return liveNodeEvents, nil
}

func (s *solrZkInstance) initLeaderElectListener() (<-chan zk.Event, error) {
	leaderEvents, err := s.zookeeper.GetLeaderElectW()
	if err != nil {
		return nil, err
	}
	return leaderEvents, nil
}

// GetClusterState Intentionally return a copy vs a pointer want to be thread safe
func (s *solrZkInstance) GetClusterState() (ClusterState, error) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	return s.clusterState, nil
}

func (s *solrZkInstance) setLiveNodes(nodes []string) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	s.clusterState.LiveNodes = nodes
	s.logger.Printf("go-solr: zk livenodes updated %v ", s.clusterState.LiveNodes)
}

func (s *solrZkInstance) setCollections(collections map[string]Collection, version int) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	s.clusterState.Collections = collections
	s.clusterState.Version = version
	s.logger.Printf("go-solr: zk collections updated %v ", s.clusterState.Collections)
}
<commit_msg>incremental backoff<commit_after>package solr

import (
	"github.com/samuel/go-zookeeper/zk"
	"time"
)

func (s *solrZkInstance) Listen() error {
	err := s.zookeeper.Connect()
	s.currentNode = 0
	if err != nil {
		return err
	}
	s.clusterState = ClusterState{}

	collectionsEvents, err := s.initCollectionsListener()
	if err != nil {
		return err
	}
	liveNodeEvents, err := s.initLiveNodesListener()
	if err != nil {
		return err
	}
	leaderEvents, err := s.initLeaderElectListener()
	if err != nil {
		return err
	}
	//loop forever
	go func() {
		errCount := 0
		for {
			select {
			case cEvent := <-collectionsEvents:
				// do something if its not a session or disconnect
				if cEvent.Type == zk.EventNodeDataChanged {
					collections, version, err := s.zookeeper.GetClusterState()
					if err != nil {
						errCount++
						time.Sleep(time.Duration(errCount*500) * time.Millisecond)
						continue
					}
					errCount = 0
					s.setCollections(collections, version)
				}
				if cEvent.State < zk.StateConnected {
					s.logger.Printf("[Error] solr cluster zk disconnected  %v", cEvent)
				} else {
					s.logger.Printf("go-solr: solr cluster zk state changed zkType: %d zkState: %d", cEvent.Type, cEvent.State)
				}
			case lEvent := <-leaderEvents:
				if lEvent.Type == zk.EventNodeChildrenChanged || lEvent.Type == zk.EventNodeDataChanged {
					// s.Logger().Printf("Leader changed pausing")
				}
			case nEvent := <-liveNodeEvents:
				// do something if its not a session or disconnect
				if nEvent.Type == zk.EventNodeDataChanged || nEvent.Type == zk.EventNodeChildrenChanged {
					liveNodes, err := s.zookeeper.GetLiveNodes()
					if err != nil {
						errCount++
						time.Sleep(time.Duration(errCount*500) * time.Millisecond)
						continue
					}
					errCount = 0
					s.setLiveNodes(liveNodes)
				}
				if nEvent.State < zk.StateConnected {
					s.logger.Printf("[Error] solr cluster zk live nodes disconnected zkType: %v ", nEvent)
				} else {
					s.logger.Printf("go-solr: solr cluster zk live nodes state changed zkType: %d zkState: %d", nEvent.Type, nEvent.State)
				}
			}
		}
	}()
	s.listening = true
	return nil
}

func (s *solrZkInstance) initCollectionsListener() (<-chan zk.Event, error) {
	s.clusterState = ClusterState{}
	collections, version, collectionsEvents, err := s.zookeeper.GetClusterStateW()
	if err != nil {
		return nil, err
	}
	s.setCollections(collections, version)
	return collectionsEvents, nil
}

func (s *solrZkInstance) initLiveNodesListener() (<-chan zk.Event, error) {
	liveNodes, liveNodeEvents, err := s.zookeeper.GetLiveNodesW()
	if err != nil {
		return nil, err
	}
	s.setLiveNodes(liveNodes)
	return liveNodeEvents, nil
}

func (s *solrZkInstance) initLeaderElectListener() (<-chan zk.Event, error) {
	leaderEvents, err := s.zookeeper.GetLeaderElectW()
	if err != nil {
		return nil, err
	}
	return leaderEvents, nil
}

// GetClusterState Intentionally return a copy vs a pointer want to be thread safe
func (s *solrZkInstance) GetClusterState() (ClusterState, error) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	return s.clusterState, nil
}

func (s *solrZkInstance) setLiveNodes(nodes []string) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	s.clusterState.LiveNodes = nodes
	s.logger.Printf("go-solr: zk livenodes updated %v ", s.clusterState.LiveNodes)
}

func (s *solrZkInstance) setCollections(collections map[string]Collection, version int) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	s.clusterState.Collections = collections
	s.clusterState.Version = version
	s.logger.Printf("go-solr: zk collections updated %v ", s.clusterState.Collections)
}
<|endoftext|>


<commit_before>package solr_test

import (
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/sendgrid/go-solr"
)

var _ = Describe("Solr Client", func() {
	var solrClient solr.SolrZK
	var solrHttp solr.SolrHTTP
	var solrHttpRetrier solr.SolrHTTP
	var locator solr.SolrLocator
	solrClient = solr.NewSolrZK("zk:2181", "solr", "solrtest")
	locator = solrClient.GetSolrLocator()

	var err error
	err = solrClient.Listen()
	BeforeEach(func() {
		Expect(err).To(BeNil())
		https, _ := solrClient.UseHTTPS()
		solrHttp, err = solr.NewSolrHTTP(https, "solrtest", solr.User("solr"), solr.Password("admin"), solr.MinRF(2))
		Expect(err).To(BeNil())
		solrHttpRetrier = solr.NewSolrHttpRetrier(solrHttp, 5, 100*time.Millisecond)
	})
	It("construct", func() {
		solrClient := solr.NewSolrZK("test", "solr", "solrtest")
		Expect(solrClient).To(Not(BeNil()))
		err := solrClient.Listen()
		Expect(err).To(Not(BeNil()))

	})

	Describe("Test Connection", func() {

		It("can get clusterstate", func() {
			state, err := solrClient.GetClusterState()
			Expect(err).To(BeNil())
			Expect(state).To(Not(BeNil()))
			Expect(state.Version > 0).To(BeTrue())
			Expect(len(state.Collections)).To(Equal(1))
		})

		It("can find a leader", func() {
			state, err := solrClient.GetClusterState()
			Expect(err).To(BeNil())
			Expect(state).To(Not(BeNil()))
			Expect(len(state.Collections)).To(Equal(1))
			leaders, err := locator.GetLeaders("mycrazyshardkey1!test.1@test.com")
			Expect(err).To(BeNil())
			leader := leaders[0]
			Expect(leader).To(Not(BeNil()))
			Expect(leader).To(ContainSubstring(":8983/solr"))
			Expect(leader).To(ContainSubstring("http://"))
		})

		It("can find a replica", func() {
			state, err := solrClient.GetClusterState()
			Expect(err).To(BeNil())
			Expect(state).To(Not(BeNil()))
			Expect(len(state.Collections)).To(Equal(1))
			replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey1!")
			Expect(err).To(BeNil())
			Expect(len(replicas)).To(Not(BeZero()))

			Expect(replicas[0]).To(ContainSubstring(":8983/solr"))
			Expect(replicas[0]).To(ContainSubstring("http://"))
		})

		Describe("Test Requests", func() {
			It("can get requests", func() {
				replicas, err := locator.GetReplicaUris()
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.FilterQuery("*:*"), solr.Rows(10))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
			})
			It("can update requests", func() {
				uuid, _ := newUUID()

				doc := map[string]interface{}{
					"id":         "mycrazyshardkey1!" + uuid,
					"email":      uuid + "feldman1@sendgrid.com",
					"first_name": "shawn1" + uuid,
					"last_name":  uuid + "feldman1",
				}
				leader, err := locator.GetLeaders("mycrazyshardkey1!" + uuid)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, true, doc, solr.Commit(true))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey1!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("first_name:shawn1"+uuid), solr.Rows(10))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(1))
			})

			It("can update requests with no doc id", func() {
				uuid, _ := newUUID()

				doc := map[string]interface{}{
					"id":         "mycrazyshardkey1!" + uuid,
					"email":      uuid + "feldman1@sendgrid.com",
					"first_name": "shawn1" + uuid,
					"last_name":  uuid + "feldman1",
				}
				leader, err := locator.GetLeaders("mycrazyshardkey1!" + uuid)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, true, doc, solr.Commit(true))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey1!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("first_name:shawn1"+uuid), solr.Rows(10))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(1))
			})

			It("can update requests with route", func() {
				uuid, _ := newUUID()

				doc := map[string]interface{}{
					"id":         "mycrazyshardkey1!" + uuid,
					"email":      uuid + "feldman1@sendgrid.com",
					"first_name": "shawn1" + uuid,
					"last_name":  uuid + "feldman1",
				}
				leader, err := locator.GetLeaders("mycrazyshardkey1!" + uuid)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, true, doc, solr.Commit(true), solr.Route("mycrazyshardkey1!"))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey1!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("first_name:shawn1"+uuid), solr.Rows(10))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(1))
			})

			It("can update requests with route with version", func() {
				uuid, _ := newUUID()

				doc := map[string]interface{}{
					"id":         "mycrazyshardkey1!" + uuid,
					"email":      uuid + "feldman1@sendgrid.com",
					"first_name": "shawn1" + uuid,
					"last_name":  uuid + "feldman1",
				}
				state, err := solrClient.GetClusterState()
				Expect(err).To(BeNil())
				leader, err := locator.GetLeaders("mycrazyshardkey1!" + uuid)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, true, doc, solr.Commit(true), solr.Route("mycrazyshardkey1!"), solr.ClusterStateVersion(state.Version, "goseg"))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey1!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("first_name:shawn1"+uuid), solr.Rows(10), solr.ClusterStateVersion(state.Version, "goseg"))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(1))
			})

			It("can update requests and read with route", func() {
				uuid, _ := newUUID()

				doc := map[string]interface{}{
					"id":         "mycrazyshardkey3!" + uuid,
					"email":      uuid + "feldman@sendgrid.com",
					"first_name": "shawn3" + uuid,
					"last_name":  uuid,
				}
				leader, err := locator.GetLeaders("mycrazyshardkey3!" + uuid)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, true, doc, solr.Commit(true))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey3!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("last_name:"+uuid), solr.Rows(10))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(1))
			})

			It("can update requests and read with route many times", func() {
				const limit int = 10
				uuid, _ := newUUID()
				for i := 0; i < limit; i++ {
					iterationId, _ := newUUID()
					doc := map[string]interface{}{
						"id":         "mycrazyshardkey4!rando" + iterationId,
						"email":      "rando" + iterationId + "@sendgrid.com",
						"first_name": "tester" + iterationId,
						"last_name":  uuid,
					}
					if i < limit-1 {
						leader, err := locator.GetLeaders(doc["id"].(string))
						Expect(err).To(BeNil())
						err = solrHttp.Update(leader, true, doc, solr.Commit(false))
						Expect(err).To(BeNil())
					} else {
						leader, err := locator.GetLeaders(doc["id"].(string))
						Expect(err).To(BeNil())
						err = solrHttp.Update(leader, true, doc, solr.Commit(true))
						Expect(err).To(BeNil())
					}

				}
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey4!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("last_name:"+uuid), solr.Rows(1000))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(limit))
			})

			It("can test the retrier requests and read with route many times", func() {
				const limit int = 100
				uuid, _ := newUUID()
				for i := 0; i < limit; i++ {
					iterationId, _ := newUUID()
					doc := map[string]interface{}{
						"id":         "mycrazyshardkey4!rando" + iterationId,
						"email":      "rando" + iterationId + "@sendgrid.com",
						"first_name": "tester" + iterationId,
						"last_name":  uuid,
					}
					if i < limit-1 {
						leader, err := locator.GetLeaders(doc["id"].(string))
						Expect(err).To(BeNil())
						err = solrHttpRetrier.Update(leader, true, doc, solr.Commit(false))
						Expect(err).To(BeNil())
					} else {
						leader, err := locator.GetLeaders(doc["id"].(string))
						Expect(err).To(BeNil())
						err = solrHttpRetrier.Update(leader, true, doc, solr.Commit(true))
						Expect(err).To(BeNil())
					}

				}
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey4!")
				Expect(err).To(BeNil())
				r, err := solrHttpRetrier.Read(replicas, solr.Query("*:*"), solr.FilterQuery("last_name:"+uuid), solr.Rows(1000))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(limit))
			})

			It("can delete all", func() {
				lastId := ""
				const limit int = 10
				uuid, _ := newUUID()
				shardKey := "mycrazysha" + uuid
				for i := 0; i < limit; i++ {
					iterationId, _ := newUUID()
					lastId := shardKey + "!rando" + iterationId
					doc := map[string]interface{}{
						"id":         lastId,
						"email":      "rando" + iterationId + "@sendgrid.com",
						"first_name": "tester" + iterationId,
						"last_name":  uuid,
					}
					leader, err := locator.GetLeaders(doc["id"].(string))
					Expect(err).To(BeNil())

					if i < limit-1 {
						err := solrHttp.Update(leader, true, doc, solr.Commit(false))
						Expect(err).To(BeNil())
					} else {
						err := solrHttp.Update(leader, true, doc, solr.Commit(true))
						Expect(err).To(BeNil())
					}

				}
				leader, err := locator.GetLeaders(lastId)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, false, nil, solr.Commit(true), solr.DeleteStreamBody("last_name:*"))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute(shardKey + "!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Route(shardKey), solr.Query("*:*"), solr.FilterQuery("last_name:"+uuid), solr.Rows(1000))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(0))
			})

			It("can get the shard for a route", func() {
				shard, err := locator.GetShardFromRoute("mycrazyshardkey3!")
				Expect(err).To(BeNil())
				Expect(shard).To(Equal("shard1"))
			})
		})
	})
	Describe("Basic Auth Fails", func() {
		It("can get requests", func() {
			solrNoAuthClient := solr.NewSolrZK("zk:2181", "solr", "solrtest")
			err := solrNoAuthClient.Listen()
			Expect(err).To(BeNil())
			https, _ := solrClient.UseHTTPS()
			solrNoAuthHttp, err := solr.NewSolrHTTP(https, "solrtest")
			Expect(err).To(BeNil())
			err = solrNoAuthClient.Listen()
			Expect(err).To(BeNil())
			replicas, err := locator.GetReplicaUris()
			r, err := solrNoAuthHttp.Read(replicas, solr.FilterQuery("*:*"), solr.Rows(10))
			Expect(err).To(Not(BeNil()))
			Expect(strings.Contains(err.Error(), "401")).To(BeTrue())
			Expect(r.Status).To(BeEquivalentTo(401))
		})

	})

})
<commit_msg>shard test<commit_after>package solr_test

import (
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/sendgrid/go-solr"
)

var _ = Describe("Solr Client", func() {
	var solrClient solr.SolrZK
	var solrHttp solr.SolrHTTP
	var solrHttpRetrier solr.SolrHTTP
	var locator solr.SolrLocator
	solrClient = solr.NewSolrZK("zk:2181", "solr", "solrtest")
	locator = solrClient.GetSolrLocator()

	var err error
	err = solrClient.Listen()
	BeforeEach(func() {
		Expect(err).To(BeNil())
		https, _ := solrClient.UseHTTPS()
		solrHttp, err = solr.NewSolrHTTP(https, "solrtest", solr.User("solr"), solr.Password("admin"), solr.MinRF(2))
		Expect(err).To(BeNil())
		solrHttpRetrier = solr.NewSolrHttpRetrier(solrHttp, 5, 100*time.Millisecond)
	})
	It("construct", func() {
		solrClient := solr.NewSolrZK("test", "solr", "solrtest")
		Expect(solrClient).To(Not(BeNil()))
		err := solrClient.Listen()
		Expect(err).To(Not(BeNil()))

	})

	Describe("Test Connection", func() {

		It("can get clusterstate", func() {
			state, err := solrClient.GetClusterState()
			Expect(err).To(BeNil())
			Expect(state).To(Not(BeNil()))
			Expect(state.Version > 0).To(BeTrue())
			Expect(len(state.Collections)).To(Equal(1))
		})

		It("can find a leader", func() {
			state, err := solrClient.GetClusterState()
			Expect(err).To(BeNil())
			Expect(state).To(Not(BeNil()))
			Expect(len(state.Collections)).To(Equal(1))
			leaders, err := locator.GetLeaders("mycrazyshardkey1!test.1@test.com")
			Expect(err).To(BeNil())
			leader := leaders[0]
			Expect(leader).To(Not(BeNil()))
			Expect(leader).To(ContainSubstring(":8983/solr"))
			Expect(leader).To(ContainSubstring("http://"))
		})

		It("can find a replica", func() {
			state, err := solrClient.GetClusterState()
			Expect(err).To(BeNil())
			Expect(state).To(Not(BeNil()))
			Expect(len(state.Collections)).To(Equal(1))
			replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey1!")
			Expect(err).To(BeNil())
			Expect(len(replicas)).To(Not(BeZero()))

			Expect(replicas[0]).To(ContainSubstring(":8983/solr"))
			Expect(replicas[0]).To(ContainSubstring("http://"))
		})

		Describe("Test Requests", func() {
			It("can get requests", func() {
				replicas, err := locator.GetReplicaUris()
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.FilterQuery("*:*"), solr.Rows(10))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
			})
			It("can update requests", func() {
				uuid, _ := newUUID()

				doc := map[string]interface{}{
					"id":         "mycrazyshardkey1!" + uuid,
					"email":      uuid + "feldman1@sendgrid.com",
					"first_name": "shawn1" + uuid,
					"last_name":  uuid + "feldman1",
				}
				leader, err := locator.GetLeaders("mycrazyshardkey1!" + uuid)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, true, doc, solr.Commit(true))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey1!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("first_name:shawn1"+uuid), solr.Rows(10))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(1))
			})

			It("can update requests with no doc id", func() {
				uuid, _ := newUUID()

				doc := map[string]interface{}{
					"id":         "mycrazyshardkey1!" + uuid,
					"email":      uuid + "feldman1@sendgrid.com",
					"first_name": "shawn1" + uuid,
					"last_name":  uuid + "feldman1",
				}
				leader, err := locator.GetLeaders("mycrazyshardkey1!" + uuid)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, true, doc, solr.Commit(true))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey1!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("first_name:shawn1"+uuid), solr.Rows(10))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(1))
			})

			It("can update requests with route", func() {
				uuid, _ := newUUID()

				doc := map[string]interface{}{
					"id":         "mycrazyshardkey1!" + uuid,
					"email":      uuid + "feldman1@sendgrid.com",
					"first_name": "shawn1" + uuid,
					"last_name":  uuid + "feldman1",
				}
				leader, err := locator.GetLeaders("mycrazyshardkey1!" + uuid)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, true, doc, solr.Commit(true), solr.Route("mycrazyshardkey1!"))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey1!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("first_name:shawn1"+uuid), solr.Rows(10))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(1))
			})

			It("can update requests with route with version", func() {
				uuid, _ := newUUID()

				doc := map[string]interface{}{
					"id":         "mycrazyshardkey1!" + uuid,
					"email":      uuid + "feldman1@sendgrid.com",
					"first_name": "shawn1" + uuid,
					"last_name":  uuid + "feldman1",
				}
				state, err := solrClient.GetClusterState()
				Expect(err).To(BeNil())
				leader, err := locator.GetLeaders("mycrazyshardkey1!" + uuid)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, true, doc, solr.Commit(true), solr.Route("mycrazyshardkey1!"), solr.ClusterStateVersion(state.Version, "goseg"))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey1!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("first_name:shawn1"+uuid), solr.Rows(10), solr.ClusterStateVersion(state.Version, "goseg"))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(1))
			})

			It("can update requests and read with route", func() {
				uuid, _ := newUUID()

				doc := map[string]interface{}{
					"id":         "mycrazyshardkey3!" + uuid,
					"email":      uuid + "feldman@sendgrid.com",
					"first_name": "shawn3" + uuid,
					"last_name":  uuid,
				}
				leader, err := locator.GetLeaders("mycrazyshardkey3!" + uuid)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, true, doc, solr.Commit(true))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey3!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("last_name:"+uuid), solr.Rows(10))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(1))
			})

			It("can update requests and read with route many times", func() {
				const limit int = 10
				uuid, _ := newUUID()
				for i := 0; i < limit; i++ {
					iterationId, _ := newUUID()
					doc := map[string]interface{}{
						"id":         "mycrazyshardkey4!rando" + iterationId,
						"email":      "rando" + iterationId + "@sendgrid.com",
						"first_name": "tester" + iterationId,
						"last_name":  uuid,
					}
					if i < limit-1 {
						leader, err := locator.GetLeaders(doc["id"].(string))
						Expect(err).To(BeNil())
						err = solrHttp.Update(leader, true, doc, solr.Commit(false))
						Expect(err).To(BeNil())
					} else {
						leader, err := locator.GetLeaders(doc["id"].(string))
						Expect(err).To(BeNil())
						err = solrHttp.Update(leader, true, doc, solr.Commit(true))
						Expect(err).To(BeNil())
					}

				}
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey4!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Query("*:*"), solr.FilterQuery("last_name:"+uuid), solr.Rows(1000))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(limit))
			})

			It("can test the retrier requests and read with route many times", func() {
				const limit int = 100
				uuid, _ := newUUID()
				for i := 0; i < limit; i++ {
					iterationId, _ := newUUID()
					doc := map[string]interface{}{
						"id":         "mycrazyshardkey4!rando" + iterationId,
						"email":      "rando" + iterationId + "@sendgrid.com",
						"first_name": "tester" + iterationId,
						"last_name":  uuid,
					}
					if i < limit-1 {
						leader, err := locator.GetLeaders(doc["id"].(string))
						Expect(err).To(BeNil())
						err = solrHttpRetrier.Update(leader, true, doc, solr.Commit(false))
						Expect(err).To(BeNil())
					} else {
						leader, err := locator.GetLeaders(doc["id"].(string))
						Expect(err).To(BeNil())
						err = solrHttpRetrier.Update(leader, true, doc, solr.Commit(true))
						Expect(err).To(BeNil())
					}

				}
				replicas, err := locator.GetReplicasFromRoute("mycrazyshardkey4!")
				Expect(err).To(BeNil())
				r, err := solrHttpRetrier.Read(replicas, solr.Query("*:*"), solr.FilterQuery("last_name:"+uuid), solr.Rows(1000))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(limit))
			})

			It("can delete all", func() {
				lastId := ""
				const limit int = 10
				uuid, _ := newUUID()
				shardKey := "mycrazysha" + uuid
				for i := 0; i < limit; i++ {
					iterationId, _ := newUUID()
					lastId := shardKey + "!rando" + iterationId
					doc := map[string]interface{}{
						"id":         lastId,
						"email":      "rando" + iterationId + "@sendgrid.com",
						"first_name": "tester" + iterationId,
						"last_name":  uuid,
					}
					leader, err := locator.GetLeaders(doc["id"].(string))
					Expect(err).To(BeNil())

					if i < limit-1 {
						err := solrHttp.Update(leader, true, doc, solr.Commit(false))
						Expect(err).To(BeNil())
					} else {
						err := solrHttp.Update(leader, true, doc, solr.Commit(true))
						Expect(err).To(BeNil())
					}

				}
				leader, err := locator.GetLeaders(lastId)
				Expect(err).To(BeNil())
				err = solrHttp.Update(leader, false, nil, solr.Commit(true), solr.DeleteStreamBody("last_name:*"))
				Expect(err).To(BeNil())
				replicas, err := locator.GetReplicasFromRoute(shardKey + "!")
				Expect(err).To(BeNil())
				r, err := solrHttp.Read(replicas, solr.Route(shardKey), solr.Query("*:*"), solr.FilterQuery("last_name:"+uuid), solr.Rows(1000))
				Expect(err).To(BeNil())
				Expect(r).To(Not(BeNil()))
				Expect(r.Response.NumFound).To(BeEquivalentTo(0))
			})

			It("can get the shard for a route", func() {
				shard, err := locator.GetShardFromRoute("mycrazyshardkey3!")
				Expect(err).To(BeNil())
				Expect(shard).To(Not(BeNil()))
			})
		})
	})
	Describe("Basic Auth Fails", func() {
		It("can get requests", func() {
			solrNoAuthClient := solr.NewSolrZK("zk:2181", "solr", "solrtest")
			err := solrNoAuthClient.Listen()
			Expect(err).To(BeNil())
			https, _ := solrClient.UseHTTPS()
			solrNoAuthHttp, err := solr.NewSolrHTTP(https, "solrtest")
			Expect(err).To(BeNil())
			err = solrNoAuthClient.Listen()
			Expect(err).To(BeNil())
			replicas, err := locator.GetReplicaUris()
			r, err := solrNoAuthHttp.Read(replicas, solr.FilterQuery("*:*"), solr.Rows(10))
			Expect(err).To(Not(BeNil()))
			Expect(strings.Contains(err.Error(), "401")).To(BeTrue())
			Expect(r.Status).To(BeEquivalentTo(401))
		})

	})

})
<|endoftext|>


<commit_before>package sql

import (
	"database/sql"
	"errors"
	"fmt"
	"reflect"
	"strconv"
	"strings"
)

type shard struct {
	db *sql.DB
	shardConn
}

type sqlConn interface {
	Exec(query string, args ...interface{}) (sql.Result, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
}

type shardConn struct {
	sqlConn
}

func (s *shard) Autocommit(acFun TxFunc) (err error) {
	tx, err := s.db.Begin()
	if err != nil {
		return
	}

	return acFun(&shardConn{tx})
}

func (s *shard) Transact(txFun TxFunc) (err error) {
	tx, err := s.db.Begin()
	if err != nil {
		return
	}
	_, err = tx.Exec("START TRANSACTION")
	if err != nil {
		return
	}

	err = txFun(&shardConn{tx})
	if err != nil {
		rbErr := tx.Rollback()
		if rbErr != nil {
			return errors.New("Rollback error: " + rbErr.Error() + " Query error: " + err.Error())
		}

	} else {
		err = tx.Commit()
		if err != nil {
			return
		}
	}

	return
}

// Query with fixed args
func (s *shardConn) Query(query string, args ...interface{}) (rows *sql.Rows, err error) {
	fixArgs(args)
	rows, err = s.sqlConn.Query(query, args...)
	if err != nil {
		err = errors.New("sql.Query Error: " + err.Error() + ". Query: " + query + " Args: " + fmt.Sprint(args))
	}
	return
}

// Execute with fixed args
func (s *shardConn) Exec(query string, args ...interface{}) (res sql.Result, err error) {
	fixArgs(args)
	res, err = s.sqlConn.Exec(query, args...)
	if err != nil {
		err = errors.New("sql.Exec Error: " + err.Error() + ". Query: " + query + " Args: " + fmt.Sprint(args))
	}
	return
}

/*
Fix args by converting them to values of their underlying kind.
This avoids problems in database/sql with e.g custom string types.
Without fixArgs, the following code:

	type Foo string
	...
	pool.Query("SELECT * WHERE Foo=?", Foo("bar"))

would give you the error:

	sql: converting Exec argument #1's type: unsupported type Foo, a string
*/
func fixArgs(args []interface{}) {
	for i, arg := range args {
		vArg := reflect.ValueOf(arg)
		switch vArg.Kind() {
		case reflect.String:
			args[i] = vArg.String()
			if args[i] == "" {
				args[i] = nil
			}
		}
	}
}

func (s *shardConn) SelectInt(query string, args ...interface{}) (num int, err error) {
	found, err := s.queryOne(query, args, &num)
	if !found {
		err = errors.New("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectString(query string, args ...interface{}) (str string, err error) {
	found, err := s.queryOne(query, args, &str)
	if !found {
		err = errors.New("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectUint(query string, args ...interface{}) (num uint, err error) {
	found, err := s.queryOne(query, args, &num)
	if !found {
		err = errors.New("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectIntForce(query string, args ...interface{}) (num int, err error) {
	found, err := s.queryOne(query, args, &num)
	if !found {
		panic("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectStringForce(query string, args ...interface{}) (str string, err error) {
	found, err := s.queryOne(query, args, &str)
	if !found {
		panic("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectUintForce(query string, args ...interface{}) (num uint, err error) {
	found, err := s.queryOne(query, args, &num)
	if !found {
		panic("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectIntMaybe(query string, args ...interface{}) (num int, found bool, err error) {
	found, err = s.queryOne(query, args, &num)
	return
}

func (s *shardConn) SelectStringMaybe(query string, args ...interface{}) (str string, found bool, err error) {
	found, err = s.queryOne(query, args, &str)
	return
}

func (s *shardConn) SelectUintMaybe(query string, args ...interface{}) (num uint, found bool, err error) {
	found, err = s.queryOne(query, args, &num)
	return
}

func (s *shardConn) queryOne(query string, args []interface{}, out interface{}) (found bool, err error) {
	rows, err := s.Query(query, args...)
	if err != nil {
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.Scan(out)
		if err != nil {
			return
		}
		if rows.Next() {
			err = errors.New("Query returned too many rows")
			return
		}
		found = true
	} else {
		found = false
	}

	err = rows.Err()
	if err != nil {
		return
	}

	return
}

func (s *shardConn) UpdateOne(query string, args ...interface{}) (err error) {
	return s.UpdateNum(1, query, args...)
}

func (s *shardConn) UpdateNum(num int64, query string, args ...interface{}) (err error) {
	rowsAffected, err := s.Update(query, args...)
	if err != nil {
		return
	}
	if rowsAffected != num {
		return errors.New(fmt.Sprintf("UpdateOne affected %d rows. Query: %q Args: %q", rowsAffected, query, args))
	}
	return
}

func (s *shardConn) Update(query string, args ...interface{}) (rowsAffected int64, err error) {
	res, err := s.Exec(query, args...)
	if err != nil {
		return
	}

	rowsAffected, err = res.RowsAffected()
	return
}

func (s *shardConn) InsertIgnoreId(query string, args ...interface{}) (err error) {
	_, err = s.Insert(query, args...)
	return
}

func (s *shardConn) InsertIgnoreDuplicates(query string, args ...interface{}) (err error) {
	_, err = s.Insert(query, args...)
	if err != nil && strings.Contains(err.Error(), "Duplicate entry") {
		err = nil
	}
	return
}

func (s *shardConn) Insert(query string, args ...interface{}) (id int64, err error) {
	res, err := s.Exec(query, args...)
	if err != nil {
		return
	}
	id, err = res.LastInsertId()
	return
}

func (s *shardConn) Select(output interface{}, sql string, args ...interface{}) error {
	// Check types
	var outputPtr = reflect.ValueOf(output)
	if outputPtr.Kind() != reflect.Ptr {
		return errors.New("fun/sql.Select: expects a pointer to a slice of items")
	}
	var outputReflection = reflect.Indirect(outputPtr)
	if outputReflection.Kind() != reflect.Slice {
		return errors.New("fun/sql.Select: expects items to be a slice")
	}
	if outputReflection.Len() != 0 {
		return errors.New("fun/sql.Select: expects items to be empty")
	}
	outputItemType := outputReflection.Type().Elem().Elem()
	if outputItemType.Kind() != reflect.Struct {
		return errors.New("fun/sql.Select: expects items to be a slice of structs")
	}

	// Query DB
	var rows, err = s.Query(sql, args...)
	if err != nil {
		return err
	}
	defer rows.Close()

	outputReflection.Set(reflect.MakeSlice(outputReflection.Type(), 0, 0))
	// Reflect onto structs
	columns, err := rows.Columns()
	if err != nil {
		return err
	}

	structType := outputReflection.Type().Elem()
	for rows.Next() {
		structPtrVal := reflect.New(structType.Elem())
		outputItemStructVal := structPtrVal.Elem()
		err = structFromRow(outputItemStructVal, columns, rows)
		if err != nil {
			return err
		}

		outputReflection.Set(reflect.Append(outputReflection, structPtrVal))
	}

	err = rows.Err()
	if err != nil {
		return err
	}

	return nil
}

const scanOneTypeError = "fun/sql.SelectOne: expects a **struct, e.g var person *Person; c.SelectOne(&person, sql)"

func (s *shardConn) SelectOne(output interface{}, query string, args ...interface{}) error {
	return s.scanOne(output, query, true, args...)
}
func (s *shardConn) SelectMaybe(output interface{}, query string, args ...interface{}) error {
	return s.scanOne(output, query, false, args...)
}
func (s *shardConn) scanOne(output interface{}, query string, required bool, args ...interface{}) error {
	// Check types
	var outputReflectionPtr = reflect.ValueOf(output)
	if !outputReflectionPtr.IsValid() {
		panic(scanOneTypeError)
	}
	if outputReflectionPtr.Kind() != reflect.Ptr {
		panic(scanOneTypeError)
	}
	var outputReflection = outputReflectionPtr.Elem()
	if outputReflection.Kind() != reflect.Ptr {
		panic(scanOneTypeError)
	}

	// Query DB
	var rows, err = s.Query(query, args...)
	if err != nil {
		return err
	}
	defer rows.Close()

	// Reflect onto struct
	columns, err := rows.Columns()
	if err != nil {
		return err
	}
	if !rows.Next() {
		if required {
			return errors.New("SelectOne got 0 rows. Query: " + query + " Args: " + fmt.Sprint(args))
		} else {
			return nil
		}
	}

	var vStruct reflect.Value
	if outputReflection.IsNil() {
		structPtrVal := reflect.New(outputReflection.Type().Elem())
		outputReflection.Set(structPtrVal)
		vStruct = structPtrVal.Elem()
	} else {
		vStruct = outputReflection.Elem()
	}

	err = structFromRow(vStruct, columns, rows)
	if err != nil {
		return err
	}

	if rows.Next() {
		return errors.New("fun/sql.SelectOne: got multiple rows. Query: " + query + " Args: " + fmt.Sprint(args))
	}

	err = rows.Err()
	if err != nil {
		return err
	}

	return nil
}

func structFromRow(outputItemStructVal reflect.Value, columns []string, rows *sql.Rows) (err error) {
	vals := make([]interface{}, len(columns))
	for i, _ := range columns {
		vals[i] = &sql.RawBytes{}
		// vals[i] = &[]byte{}
	}
	err = rows.Scan(vals...)
	if err != nil {
		return
	}

	var uintVal uint64
	var intVal int64
	for i, column := range columns {
		bytes := []byte(*vals[i].(*sql.RawBytes))
		// bytes := []byte(*vals[i].(*[]byte))
		if bytes == nil {
			continue // Leave struct field empty
		}
		var outputItemField = outputItemStructVal.FieldByName(column)
		switch outputItemField.Kind() {
		case reflect.String:
			outputItemField.SetString(string(bytes))
		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
			uintVal, err = strconv.ParseUint(string(bytes), 10, 64)
			if err != nil {
				return
			}
			outputItemField.SetUint(reflect.ValueOf(uintVal).Uint())
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			intVal, err = strconv.ParseInt(string(bytes), 10, 64)
			if err != nil {
				return
			}
			outputItemField.SetInt(reflect.ValueOf(intVal).Int())
		default:
			err = errors.New("fun/sql: Bad row value for column " + column + ": " + outputItemField.Kind().String())
			return
		}
	}

	return
}
<commit_msg>Begin already calls START TRANSACTION<commit_after>package sql

import (
	"database/sql"
	"errors"
	"fmt"
	"reflect"
	"strconv"
	"strings"
)

type shard struct {
	db *sql.DB
	shardConn
}

type sqlConn interface {
	Exec(query string, args ...interface{}) (sql.Result, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
}

type shardConn struct {
	sqlConn
}

func (s *shard) Autocommit(acFun TxFunc) (err error) {
	tx, err := s.db.Begin()
	if err != nil {
		return
	}

	return acFun(&shardConn{tx})
}

func (s *shard) Transact(txFun TxFunc) (err error) {
	tx, err := s.db.Begin()
	if err != nil {
		return
	}

	err = txFun(&shardConn{tx})
	if err != nil {
		rbErr := tx.Rollback()
		if rbErr != nil {
			return errors.New("Rollback error: " + rbErr.Error() + " Query error: " + err.Error())
		}

	} else {
		err = tx.Commit()
		if err != nil {
			return
		}
	}

	return
}

// Query with fixed args
func (s *shardConn) Query(query string, args ...interface{}) (rows *sql.Rows, err error) {
	fixArgs(args)
	rows, err = s.sqlConn.Query(query, args...)
	if err != nil {
		err = errors.New("sql.Query Error: " + err.Error() + ". Query: " + query + " Args: " + fmt.Sprint(args))
	}
	return
}

// Execute with fixed args
func (s *shardConn) Exec(query string, args ...interface{}) (res sql.Result, err error) {
	fixArgs(args)
	res, err = s.sqlConn.Exec(query, args...)
	if err != nil {
		err = errors.New("sql.Exec Error: " + err.Error() + ". Query: " + query + " Args: " + fmt.Sprint(args))
	}
	return
}

/*
Fix args by converting them to values of their underlying kind.
This avoids problems in database/sql with e.g custom string types.
Without fixArgs, the following code:

	type Foo string
	...
	pool.Query("SELECT * WHERE Foo=?", Foo("bar"))

would give you the error:

	sql: converting Exec argument #1's type: unsupported type Foo, a string
*/
func fixArgs(args []interface{}) {
	for i, arg := range args {
		vArg := reflect.ValueOf(arg)
		switch vArg.Kind() {
		case reflect.String:
			args[i] = vArg.String()
			if args[i] == "" {
				args[i] = nil
			}
		}
	}
}

func (s *shardConn) SelectInt(query string, args ...interface{}) (num int, err error) {
	found, err := s.queryOne(query, args, &num)
	if !found {
		err = errors.New("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectString(query string, args ...interface{}) (str string, err error) {
	found, err := s.queryOne(query, args, &str)
	if !found {
		err = errors.New("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectUint(query string, args ...interface{}) (num uint, err error) {
	found, err := s.queryOne(query, args, &num)
	if !found {
		err = errors.New("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectIntForce(query string, args ...interface{}) (num int, err error) {
	found, err := s.queryOne(query, args, &num)
	if !found {
		panic("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectStringForce(query string, args ...interface{}) (str string, err error) {
	found, err := s.queryOne(query, args, &str)
	if !found {
		panic("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectUintForce(query string, args ...interface{}) (num uint, err error) {
	found, err := s.queryOne(query, args, &num)
	if !found {
		panic("Query returned no rows")
	}
	return
}

func (s *shardConn) SelectIntMaybe(query string, args ...interface{}) (num int, found bool, err error) {
	found, err = s.queryOne(query, args, &num)
	return
}

func (s *shardConn) SelectStringMaybe(query string, args ...interface{}) (str string, found bool, err error) {
	found, err = s.queryOne(query, args, &str)
	return
}

func (s *shardConn) SelectUintMaybe(query string, args ...interface{}) (num uint, found bool, err error) {
	found, err = s.queryOne(query, args, &num)
	return
}

func (s *shardConn) queryOne(query string, args []interface{}, out interface{}) (found bool, err error) {
	rows, err := s.Query(query, args...)
	if err != nil {
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.Scan(out)
		if err != nil {
			return
		}
		if rows.Next() {
			err = errors.New("Query returned too many rows")
			return
		}
		found = true
	} else {
		found = false
	}

	err = rows.Err()
	if err != nil {
		return
	}

	return
}

func (s *shardConn) UpdateOne(query string, args ...interface{}) (err error) {
	return s.UpdateNum(1, query, args...)
}

func (s *shardConn) UpdateNum(num int64, query string, args ...interface{}) (err error) {
	rowsAffected, err := s.Update(query, args...)
	if err != nil {
		return
	}
	if rowsAffected != num {
		return errors.New(fmt.Sprintf("UpdateOne affected %d rows. Query: %q Args: %q", rowsAffected, query, args))
	}
	return
}

func (s *shardConn) Update(query string, args ...interface{}) (rowsAffected int64, err error) {
	res, err := s.Exec(query, args...)
	if err != nil {
		return
	}

	rowsAffected, err = res.RowsAffected()
	return
}

func (s *shardConn) InsertIgnoreId(query string, args ...interface{}) (err error) {
	_, err = s.Insert(query, args...)
	return
}

func (s *shardConn) InsertIgnoreDuplicates(query string, args ...interface{}) (err error) {
	_, err = s.Insert(query, args...)
	if err != nil && strings.Contains(err.Error(), "Duplicate entry") {
		err = nil
	}
	return
}

func (s *shardConn) Insert(query string, args ...interface{}) (id int64, err error) {
	res, err := s.Exec(query, args...)
	if err != nil {
		return
	}
	id, err = res.LastInsertId()
	return
}

func (s *shardConn) Select(output interface{}, sql string, args ...interface{}) error {
	// Check types
	var outputPtr = reflect.ValueOf(output)
	if outputPtr.Kind() != reflect.Ptr {
		return errors.New("fun/sql.Select: expects a pointer to a slice of items")
	}
	var outputReflection = reflect.Indirect(outputPtr)
	if outputReflection.Kind() != reflect.Slice {
		return errors.New("fun/sql.Select: expects items to be a slice")
	}
	if outputReflection.Len() != 0 {
		return errors.New("fun/sql.Select: expects items to be empty")
	}
	outputItemType := outputReflection.Type().Elem().Elem()
	if outputItemType.Kind() != reflect.Struct {
		return errors.New("fun/sql.Select: expects items to be a slice of structs")
	}

	// Query DB
	var rows, err = s.Query(sql, args...)
	if err != nil {
		return err
	}
	defer rows.Close()

	outputReflection.Set(reflect.MakeSlice(outputReflection.Type(), 0, 0))
	// Reflect onto structs
	columns, err := rows.Columns()
	if err != nil {
		return err
	}

	structType := outputReflection.Type().Elem()
	for rows.Next() {
		structPtrVal := reflect.New(structType.Elem())
		outputItemStructVal := structPtrVal.Elem()
		err = structFromRow(outputItemStructVal, columns, rows)
		if err != nil {
			return err
		}

		outputReflection.Set(reflect.Append(outputReflection, structPtrVal))
	}

	err = rows.Err()
	if err != nil {
		return err
	}

	return nil
}

const scanOneTypeError = "fun/sql.SelectOne: expects a **struct, e.g var person *Person; c.SelectOne(&person, sql)"

func (s *shardConn) SelectOne(output interface{}, query string, args ...interface{}) error {
	return s.scanOne(output, query, true, args...)
}
func (s *shardConn) SelectMaybe(output interface{}, query string, args ...interface{}) error {
	return s.scanOne(output, query, false, args...)
}
func (s *shardConn) scanOne(output interface{}, query string, required bool, args ...interface{}) error {
	// Check types
	var outputReflectionPtr = reflect.ValueOf(output)
	if !outputReflectionPtr.IsValid() {
		panic(scanOneTypeError)
	}
	if outputReflectionPtr.Kind() != reflect.Ptr {
		panic(scanOneTypeError)
	}
	var outputReflection = outputReflectionPtr.Elem()
	if outputReflection.Kind() != reflect.Ptr {
		panic(scanOneTypeError)
	}

	// Query DB
	var rows, err = s.Query(query, args...)
	if err != nil {
		return err
	}
	defer rows.Close()

	// Reflect onto struct
	columns, err := rows.Columns()
	if err != nil {
		return err
	}
	if !rows.Next() {
		if required {
			return errors.New("SelectOne got 0 rows. Query: " + query + " Args: " + fmt.Sprint(args))
		} else {
			return nil
		}
	}

	var vStruct reflect.Value
	if outputReflection.IsNil() {
		structPtrVal := reflect.New(outputReflection.Type().Elem())
		outputReflection.Set(structPtrVal)
		vStruct = structPtrVal.Elem()
	} else {
		vStruct = outputReflection.Elem()
	}

	err = structFromRow(vStruct, columns, rows)
	if err != nil {
		return err
	}

	if rows.Next() {
		return errors.New("fun/sql.SelectOne: got multiple rows. Query: " + query + " Args: " + fmt.Sprint(args))
	}

	err = rows.Err()
	if err != nil {
		return err
	}

	return nil
}

func structFromRow(outputItemStructVal reflect.Value, columns []string, rows *sql.Rows) (err error) {
	vals := make([]interface{}, len(columns))
	for i, _ := range columns {
		vals[i] = &sql.RawBytes{}
		// vals[i] = &[]byte{}
	}
	err = rows.Scan(vals...)
	if err != nil {
		return
	}

	var uintVal uint64
	var intVal int64
	for i, column := range columns {
		bytes := []byte(*vals[i].(*sql.RawBytes))
		// bytes := []byte(*vals[i].(*[]byte))
		if bytes == nil {
			continue // Leave struct field empty
		}
		var outputItemField = outputItemStructVal.FieldByName(column)
		switch outputItemField.Kind() {
		case reflect.String:
			outputItemField.SetString(string(bytes))
		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
			uintVal, err = strconv.ParseUint(string(bytes), 10, 64)
			if err != nil {
				return
			}
			outputItemField.SetUint(reflect.ValueOf(uintVal).Uint())
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			intVal, err = strconv.ParseInt(string(bytes), 10, 64)
			if err != nil {
				return
			}
			outputItemField.SetInt(reflect.ValueOf(intVal).Int())
		default:
			err = errors.New("fun/sql: Bad row value for column " + column + ": " + outputItemField.Kind().String())
			return
		}
	}

	return
}
<|endoftext|>