<commit_before>package main

import (
	"errors"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	"cloud.google.com/go/spanner"
	"google.golang.org/api/iterator"
	adminpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1"
	spannerpb "google.golang.org/genproto/googleapis/spanner/v1"
)

type Statement interface {
	Execute(session *Session) (*Result, error)
}

type Result struct {
	ColumnNames []string
	Rows        []Row
	QueryStats  QueryStats
	IsMutation  bool
}

type Row struct {
	Columns []string
}

type QueryStats struct {
	Rows        int
	ElapsedTime string
}

var (
	exitRe            = regexp.MustCompile(`(?i)^EXIT$`)
	selectRe          = regexp.MustCompile(`(?i)^SELECT\s.+$`)
	createTableRe     = regexp.MustCompile(`(?i)^CREATE\s+TABLE\s.+$`)
	showDatabasesRe   = regexp.MustCompile(`(?i)^SHOW\s+DATABASES$`)
	showCreateTableRe = regexp.MustCompile(`(?i)^SHOW\s+CREATE\s+TABLE\s+(.*)$`)
	showTablesRe      = regexp.MustCompile(`(?i)^SHOW\s+TABLES$`)
)

var (
	statementExitError = errors.New("exit")
)

func buildStatement(input string) (Statement, error) {
	var stmt Statement

	if exitRe.MatchString(input) {
		return nil, statementExitError
	} else if selectRe.MatchString(input) {
		stmt = &QueryStatement{
			text: input,
		}
	} else if createTableRe.MatchString(input) {
		stmt = &CreateTableStatement{
			text: input,
		}
	} else if showDatabasesRe.MatchString(input) {
		stmt = &ShowDatabasesStatement{}
	} else if showCreateTableRe.MatchString(input) {
		matched := showCreateTableRe.FindStringSubmatch(input)
		stmt = &ShowCreateTableStatement{
			table: matched[1],
		}
	} else if showTablesRe.MatchString(input) {
		stmt = &ShowTablesStatement{}
	}

	if stmt == nil {
		return nil, errors.New("invalid statement")
	}

	return stmt, nil
}

type QueryStatement struct {
	text string
}

func (s *QueryStatement) Execute(session *Session) (*Result, error) {
	stmt := spanner.NewStatement(s.text)
	iter := session.client.Single().QueryWithStats(session.ctx, stmt)

	result := &Result{
		ColumnNames: make([]string, 0),
		Rows:        make([]Row, 0),
		IsMutation:  false,
	}

	defer iter.Stop()
	for {
		row, err := iter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}

		resultRow := Row{
			Columns: make([]string, row.Size()),
		}

		result.ColumnNames = row.ColumnNames() // TODO

		for i := 0; i < row.Size(); i++ {
			var column spanner.GenericColumnValue
			err := row.Column(i, &column)
			if err != nil {
				return nil, err
			}
			// fmt.Println(column.Type.Code)
			switch column.Type.Code {
			case spannerpb.TypeCode_INT64:
				var v int64
				if err := column.Decode(&v); err != nil {
					return nil, err
				}
				resultRow.Columns[i] = fmt.Sprintf("%d", v)
			case spannerpb.TypeCode_STRING:
				var v string
				if err := column.Decode(&v); err != nil {
					return nil, err
				}
				resultRow.Columns[i] = v
			default:
				resultRow.Columns[i] = fmt.Sprintf("%s", column.Value)
			}
		}

		result.Rows = append(result.Rows, resultRow)
	}

	rowsReturned, _ := strconv.Atoi(iter.QueryStats["rows_returned"].(string))
	elapsedTime := iter.QueryStats["elapsed_time"].(string)
	result.QueryStats = QueryStats{
		Rows:        rowsReturned,
		ElapsedTime: elapsedTime,
	}

	return result, nil
}

type CreateTableStatement struct {
	text string
}

func (s *CreateTableStatement) Execute(session *Session) (*Result, error) {
	result := &Result{
		ColumnNames: make([]string, 0),
		Rows:        make([]Row, 0),
		IsMutation:  true,
	}

	t1 := time.Now()
	op, err := session.adminClient.UpdateDatabaseDdl(session.ctx, &adminpb.UpdateDatabaseDdlRequest{
		Database:   session.GetDatabasePath(),
		Statements: []string{s.text},
	})
	if err != nil {
		return nil, err
	}
	if err := op.Wait(session.ctx); err != nil {
		return nil, err
	}
	elapsed := time.Since(t1).String()

	result.QueryStats = QueryStats{
		Rows:        0,
		ElapsedTime: elapsed,
	}

	return result, nil
}

type ShowDatabasesStatement struct {
}

func (s *ShowDatabasesStatement) Execute(session *Session) (*Result, error) {
	result := &Result{
		ColumnNames: []string{"Database"},
		Rows:        make([]Row, 0),
		IsMutation:  false,
	}

	t1 := time.Now()

	dbIter := session.adminClient.ListDatabases(session.ctx, &adminpb.ListDatabasesRequest{
		Parent: session.GetInstancePath(),
	})

	for {
		database, err := dbIter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}

		re := regexp.MustCompile(`projects/[^/]+/instances/[^/]+/databases/(.+)`)
		matched := re.FindStringSubmatch(database.GetName())
		dbname := matched[1]
		resultRow := Row{
			Columns: []string{dbname},
		}
		result.Rows = append(result.Rows, resultRow)
	}

	elapsed := time.Since(t1).String()

	result.QueryStats = QueryStats{
		Rows:        len(result.Rows),
		ElapsedTime: elapsed,
	}

	return result, nil
}

type ShowCreateTableStatement struct {
	table string
}

func (s *ShowCreateTableStatement) Execute(session *Session) (*Result, error) {
	result := &Result{
		ColumnNames: []string{"Table", "Create Table"},
		Rows:        make([]Row, 0),
		IsMutation:  false,
	}

	t1 := time.Now()

	ddlResponse, err := session.adminClient.GetDatabaseDdl(session.ctx, &adminpb.GetDatabaseDdlRequest{
		Database: session.GetDatabasePath(),
	})
	if err != nil {
		return nil, err
	}
	for _, statement := range ddlResponse.Statements {
		if strings.HasPrefix(statement, fmt.Sprintf("CREATE TABLE %s", s.table)) {
			resultRow := Row{
				Columns: []string{s.table, statement},
			}
			result.Rows = append(result.Rows, resultRow)
			break
		}
	}

	elapsed := time.Since(t1).String()

	result.QueryStats = QueryStats{
		Rows:        len(result.Rows),
		ElapsedTime: elapsed,
	}

	return result, nil
}

type ShowTablesStatement struct{}

func (s *ShowTablesStatement) Execute(session *Session) (*Result, error) {
	query := QueryStatement{
		text: `SELECT t.table_name FROM information_schema.tables AS t WHERE t.table_catalog = '' and t.table_schema = ''`,
	}

	result, err := query.Execute(session)
	if err != nil {
		return nil, err
	}

	// rename column name
	if len(result.ColumnNames) == 1 {
		result.ColumnNames[0] = fmt.Sprintf("Tables_in_%s", session.databaseId)
	}

	return result, nil
}
<commit_msg>implement DML<commit_after>package main

import (
	"context"
	"errors"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	"cloud.google.com/go/spanner"
	"google.golang.org/api/iterator"
	adminpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1"
	spannerpb "google.golang.org/genproto/googleapis/spanner/v1"
)

type Statement interface {
	Execute(session *Session) (*Result, error)
}

type Result struct {
	ColumnNames []string
	Rows        []Row
	QueryStats  QueryStats
	IsMutation  bool
}

type Row struct {
	Columns []string
}

type QueryStats struct {
	Rows        int
	ElapsedTime string
}

var (
	exitRe            = regexp.MustCompile(`(?i)^EXIT$`)
	selectRe          = regexp.MustCompile(`(?i)^SELECT\s.+$`)
	createTableRe     = regexp.MustCompile(`(?i)^CREATE\s+TABLE\s.+$`)
	showDatabasesRe   = regexp.MustCompile(`(?i)^SHOW\s+DATABASES$`)
	showCreateTableRe = regexp.MustCompile(`(?i)^SHOW\s+CREATE\s+TABLE\s+(.*)$`)
	showTablesRe      = regexp.MustCompile(`(?i)^SHOW\s+TABLES$`)
	insertRe          = regexp.MustCompile(`(?i)^INSERT\s+INTO.+$`)
	updateRe          = regexp.MustCompile(`(?i)^UPDATE\s+.+$`)
	deleteRe          = regexp.MustCompile(`(?i)^DELETE\s+.+$`)
)

var (
	statementExitError = errors.New("exit")
)

func buildStatement(input string) (Statement, error) {
	var stmt Statement

	if exitRe.MatchString(input) {
		return nil, statementExitError
	} else if selectRe.MatchString(input) {
		stmt = &QueryStatement{
			text: input,
		}
	} else if createTableRe.MatchString(input) {
		stmt = &CreateTableStatement{
			text: input,
		}
	} else if showDatabasesRe.MatchString(input) {
		stmt = &ShowDatabasesStatement{}
	} else if showCreateTableRe.MatchString(input) {
		matched := showCreateTableRe.FindStringSubmatch(input)
		stmt = &ShowCreateTableStatement{
			table: matched[1],
		}
	} else if showTablesRe.MatchString(input) {
		stmt = &ShowTablesStatement{}
	} else if insertRe.MatchString(input) || updateRe.MatchString(input) || deleteRe.MatchString(input) {
		stmt = &DmlStatement{
			text: input,
		}
	}

	if stmt == nil {
		return nil, errors.New("invalid statement")
	}

	return stmt, nil
}

type QueryStatement struct {
	text string
}

func (s *QueryStatement) Execute(session *Session) (*Result, error) {
	stmt := spanner.NewStatement(s.text)
	iter := session.client.Single().QueryWithStats(session.ctx, stmt)

	result := &Result{
		ColumnNames: make([]string, 0),
		Rows:        make([]Row, 0),
		IsMutation:  false,
	}

	defer iter.Stop()
	for {
		row, err := iter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}

		resultRow := Row{
			Columns: make([]string, row.Size()),
		}

		result.ColumnNames = row.ColumnNames() // TODO

		for i := 0; i < row.Size(); i++ {
			var column spanner.GenericColumnValue
			err := row.Column(i, &column)
			if err != nil {
				return nil, err
			}
			// fmt.Println(column.Type.Code)
			switch column.Type.Code {
			case spannerpb.TypeCode_INT64:
				var v int64
				if err := column.Decode(&v); err != nil {
					return nil, err
				}
				resultRow.Columns[i] = fmt.Sprintf("%d", v)
			case spannerpb.TypeCode_STRING:
				var v string
				if err := column.Decode(&v); err != nil {
					return nil, err
				}
				resultRow.Columns[i] = v
			default:
				resultRow.Columns[i] = fmt.Sprintf("%s", column.Value)
			}
		}

		result.Rows = append(result.Rows, resultRow)
	}

	rowsReturned, _ := strconv.Atoi(iter.QueryStats["rows_returned"].(string))
	elapsedTime := iter.QueryStats["elapsed_time"].(string)
	result.QueryStats = QueryStats{
		Rows:        rowsReturned,
		ElapsedTime: elapsedTime,
	}

	return result, nil
}

type CreateTableStatement struct {
	text string
}

func (s *CreateTableStatement) Execute(session *Session) (*Result, error) {
	result := &Result{
		ColumnNames: make([]string, 0),
		Rows:        make([]Row, 0),
		IsMutation:  true,
	}

	t1 := time.Now()
	op, err := session.adminClient.UpdateDatabaseDdl(session.ctx, &adminpb.UpdateDatabaseDdlRequest{
		Database:   session.GetDatabasePath(),
		Statements: []string{s.text},
	})
	if err != nil {
		return nil, err
	}
	if err := op.Wait(session.ctx); err != nil {
		return nil, err
	}
	elapsed := time.Since(t1).String()

	result.QueryStats = QueryStats{
		Rows:        0,
		ElapsedTime: elapsed,
	}

	return result, nil
}

type ShowDatabasesStatement struct {
}

func (s *ShowDatabasesStatement) Execute(session *Session) (*Result, error) {
	result := &Result{
		ColumnNames: []string{"Database"},
		Rows:        make([]Row, 0),
		IsMutation:  false,
	}

	t1 := time.Now()

	dbIter := session.adminClient.ListDatabases(session.ctx, &adminpb.ListDatabasesRequest{
		Parent: session.GetInstancePath(),
	})

	for {
		database, err := dbIter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}

		re := regexp.MustCompile(`projects/[^/]+/instances/[^/]+/databases/(.+)`)
		matched := re.FindStringSubmatch(database.GetName())
		dbname := matched[1]
		resultRow := Row{
			Columns: []string{dbname},
		}
		result.Rows = append(result.Rows, resultRow)
	}

	elapsed := time.Since(t1).String()

	result.QueryStats = QueryStats{
		Rows:        len(result.Rows),
		ElapsedTime: elapsed,
	}

	return result, nil
}

type ShowCreateTableStatement struct {
	table string
}

func (s *ShowCreateTableStatement) Execute(session *Session) (*Result, error) {
	result := &Result{
		ColumnNames: []string{"Table", "Create Table"},
		Rows:        make([]Row, 0),
		IsMutation:  false,
	}

	t1 := time.Now()

	ddlResponse, err := session.adminClient.GetDatabaseDdl(session.ctx, &adminpb.GetDatabaseDdlRequest{
		Database: session.GetDatabasePath(),
	})
	if err != nil {
		return nil, err
	}
	for _, statement := range ddlResponse.Statements {
		if strings.HasPrefix(statement, fmt.Sprintf("CREATE TABLE %s", s.table)) {
			resultRow := Row{
				Columns: []string{s.table, statement},
			}
			result.Rows = append(result.Rows, resultRow)
			break
		}
	}

	elapsed := time.Since(t1).String()

	result.QueryStats = QueryStats{
		Rows:        len(result.Rows),
		ElapsedTime: elapsed,
	}

	return result, nil
}

type ShowTablesStatement struct{}

func (s *ShowTablesStatement) Execute(session *Session) (*Result, error) {
	query := QueryStatement{
		text: `SELECT t.table_name FROM information_schema.tables AS t WHERE t.table_catalog = '' and t.table_schema = ''`,
	}

	result, err := query.Execute(session)
	if err != nil {
		return nil, err
	}

	// rename column name
	if len(result.ColumnNames) == 1 {
		result.ColumnNames[0] = fmt.Sprintf("Tables_in_%s", session.databaseId)
	}

	return result, nil
}

type DmlStatement struct {
	text string
}

func (s *DmlStatement) Execute(session *Session) (*Result, error) {
	stmt := spanner.NewStatement(s.text)

	result := &Result{
		ColumnNames: make([]string, 0),
		Rows:        make([]Row, 0),
		IsMutation:  true,
	}

	t1 := time.Now()
	_, err := session.client.ReadWriteTransaction(session.ctx, func(ctx context.Context, txn *spanner.ReadWriteTransaction) error {
		numRows, err := txn.Update(ctx, stmt)
		if err != nil {
			return err
		}
		result.QueryStats.Rows = int(numRows) // TODO: int64
		return nil
	})
	if err != nil {
		return nil, err
	}

	elapsed := time.Since(t1).String()
	result.QueryStats.ElapsedTime = elapsed

	return result, nil
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package bleve

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/blevesearch/bleve/analysis"
	"github.com/blevesearch/bleve/search"
)

type numericRange struct {
	Name string   `json:"name,omitempty"`
	Min  *float64 `json:"min,omitempty"`
	Max  *float64 `json:"max,omitempty"`
}

type dateTimeRange struct {
	Name        string    `json:"name,omitempty"`
	Start       time.Time `json:"start,omitempty"`
	End         time.Time `json:"end,omitempty"`
	startString *string
	endString   *string
}

func (dr *dateTimeRange) ParseDates(dateTimeParser analysis.DateTimeParser) {
	if dr.Start.IsZero() && dr.startString != nil {
		start, err := dateTimeParser.ParseDateTime(*dr.startString)
		if err == nil {
			dr.Start = start
		}
	}
	if dr.End.IsZero() && dr.endString != nil {
		end, err := dateTimeParser.ParseDateTime(*dr.endString)
		if err == nil {
			dr.End = end
		}
	}
}

func (dr *dateTimeRange) UnmarshalJSON(input []byte) error {
	var temp struct {
		Name  string  `json:"name,omitempty"`
		Start *string `json:"start,omitempty"`
		End   *string `json:"end,omitempty"`
	}

	err := json.Unmarshal(input, &temp)
	if err != nil {
		return err
	}

	dr.Name = temp.Name
	if temp.Start != nil {
		dr.startString = temp.Start
	}
	if temp.End != nil {
		dr.endString = temp.End
	}

	return nil
}

// A FacetRequest describes an facet or aggregation
// of the result document set you would like to be
// built.
type FacetRequest struct {
	Size           int
	Field          string
	NumericRanges  []*numericRange  `json:"numeric_ranges,omitempty"`
	DateTimeRanges []*dateTimeRange `json:"date_ranges,omitempty"`
}

// NewFacetRequest creates a facet on the specified
// field that limits the number of entries to the
// specified size.
func NewFacetRequest(field string, size int) *FacetRequest {
	return &FacetRequest{
		Field: field,
		Size:  size,
	}
}

// AddDateTimeRange adds a bucket to a field
// containing date values.  Documents with a
// date value falling into this range are tabulated
// as part of this bucket/range.
func (fr *FacetRequest) AddDateTimeRange(name string, start, end time.Time) {
	if fr.DateTimeRanges == nil {
		fr.DateTimeRanges = make([]*dateTimeRange, 0, 1)
	}
	fr.DateTimeRanges = append(fr.DateTimeRanges, &dateTimeRange{Name: name, Start: start, End: end})
}

// AddNumericRange adds a bucket to a field
// containing numeric values.  Documents with a
// numeric value falling into this range are
// tabulated as part of this bucket/range.
func (fr *FacetRequest) AddNumericRange(name string, min, max *float64) {
	if fr.NumericRanges == nil {
		fr.NumericRanges = make([]*numericRange, 0, 1)
	}
	fr.NumericRanges = append(fr.NumericRanges, &numericRange{Name: name, Min: min, Max: max})
}

// FacetsRequest groups together all the
// FacetRequest objects for a single query.
type FacetsRequest map[string]*FacetRequest

// HighlightRequest describes how field matches
// should be highlighted.
type HighlightRequest struct {
	Style  *string  `json:"style"`
	Fields []string `json:"fields"`
}

// NewHighlight creates a default
// HighlightRequest.
func NewHighlight() *HighlightRequest {
	return &HighlightRequest{}
}

// NewHighlightWithStyle creates a HighlightRequest
// with an alternate style.
func NewHighlightWithStyle(style string) *HighlightRequest {
	return &HighlightRequest{
		Style: &style,
	}
}

// A SearchRequest describes all the parameters
// needed to search the index.
// Query is required.
// Size/From describe how much and which part of the
// result set to return.
// Highlight describes optional search result
// highlighting.
// Fields desribed a list of field values whcih
// should be retrieved for result documents.
// Facets describe the set of facets to be computed.
// Explain triggers inclusion of additional search
// result score explanations.
type SearchRequest struct {
	Query     Query             `json:"query"`
	Size      int               `json:"size"`
	From      int               `json:"from"`
	Highlight *HighlightRequest `json:"highlight"`
	Fields    []string          `json:"fields"`
	Facets    FacetsRequest     `json:"facets"`
	Explain   bool              `json:"explain"`
}

// AddFacet adds a FacetRequest to this SearchRequest
func (r *SearchRequest) AddFacet(facetName string, f *FacetRequest) {
	if r.Facets == nil {
		r.Facets = make(FacetsRequest, 1)
	}
	r.Facets[facetName] = f
}

// UnmarshalJSON deserializes a JSON representation of
// a SearchRequest
func (r *SearchRequest) UnmarshalJSON(input []byte) error {
	var temp struct {
		Q         json.RawMessage   `json:"query"`
		Size      int               `json:"size"`
		From      int               `json:"from"`
		Highlight *HighlightRequest `json:"highlight"`
		Fields    []string          `json:"fields"`
		Facets    FacetsRequest     `json:"facets"`
		Explain   bool              `json:"explain"`
	}

	err := json.Unmarshal(input, &temp)
	if err != nil {
		return err
	}

	r.Size = temp.Size
	r.From = temp.From
	r.Explain = temp.Explain
	r.Highlight = temp.Highlight
	r.Fields = temp.Fields
	r.Facets = temp.Facets
	r.Query, err = ParseQuery(temp.Q)
	if err != nil {
		return err
	}

	if r.Size <= 0 {
		r.Size = 10
	}
	if r.From <= 0 {
		r.From = 0
	}

	return nil

}

// NewSearchRequest creates a new SearchRequest
// for the Query, using default values for all
// other search parameters.
func NewSearchRequest(q Query) *SearchRequest {
	return NewSearchRequestOptions(q, 10, 0, false)
}

// NewSearchRequestOptions creates a new SearchRequest
// for the Query, with the requested size, from
// and explanation search parameters.
func NewSearchRequestOptions(q Query, size, from int, explain bool) *SearchRequest {
	return &SearchRequest{
		Query:   q,
		Size:    size,
		From:    from,
		Explain: explain,
	}
}

// A SearchResult describes the results of executing
// a SearchRequest.
type SearchResult struct {
	Request  *SearchRequest                 `json:"request"`
	Hits     search.DocumentMatchCollection `json:"hits"`
	Total    uint64                         `json:"total_hits"`
	MaxScore float64                        `json:"max_score"`
	Took     time.Duration                  `json:"took"`
	Facets   search.FacetResults            `json:"facets"`
}

func (sr *SearchResult) String() string {
	rv := ""
	if len(sr.Hits) > 0 {
		rv = fmt.Sprintf("%d matches, showing %d through %d, took %s\n", sr.Total, sr.Request.From+1, sr.Request.From+len(sr.Hits), sr.Took)
		for i, hit := range sr.Hits {
			rv += fmt.Sprintf("%5d. %s (%f)\n", i+sr.Request.From+1, hit.ID, hit.Score)
			for fragmentField, fragments := range hit.Fragments {
				rv += fmt.Sprintf("\t%s\n", fragmentField)
				for _, fragment := range fragments {
					rv += fmt.Sprintf("\t\t%s\n", fragment)
				}
			}
		}
	} else {
		rv = "No matches"
	}
	return rv
}
<commit_msg>added convenience method to add field to highlight request<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package bleve

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/blevesearch/bleve/analysis"
	"github.com/blevesearch/bleve/search"
)

type numericRange struct {
	Name string   `json:"name,omitempty"`
	Min  *float64 `json:"min,omitempty"`
	Max  *float64 `json:"max,omitempty"`
}

type dateTimeRange struct {
	Name        string    `json:"name,omitempty"`
	Start       time.Time `json:"start,omitempty"`
	End         time.Time `json:"end,omitempty"`
	startString *string
	endString   *string
}

func (dr *dateTimeRange) ParseDates(dateTimeParser analysis.DateTimeParser) {
	if dr.Start.IsZero() && dr.startString != nil {
		start, err := dateTimeParser.ParseDateTime(*dr.startString)
		if err == nil {
			dr.Start = start
		}
	}
	if dr.End.IsZero() && dr.endString != nil {
		end, err := dateTimeParser.ParseDateTime(*dr.endString)
		if err == nil {
			dr.End = end
		}
	}
}

func (dr *dateTimeRange) UnmarshalJSON(input []byte) error {
	var temp struct {
		Name  string  `json:"name,omitempty"`
		Start *string `json:"start,omitempty"`
		End   *string `json:"end,omitempty"`
	}

	err := json.Unmarshal(input, &temp)
	if err != nil {
		return err
	}

	dr.Name = temp.Name
	if temp.Start != nil {
		dr.startString = temp.Start
	}
	if temp.End != nil {
		dr.endString = temp.End
	}

	return nil
}

// A FacetRequest describes an facet or aggregation
// of the result document set you would like to be
// built.
type FacetRequest struct {
	Size           int
	Field          string
	NumericRanges  []*numericRange  `json:"numeric_ranges,omitempty"`
	DateTimeRanges []*dateTimeRange `json:"date_ranges,omitempty"`
}

// NewFacetRequest creates a facet on the specified
// field that limits the number of entries to the
// specified size.
func NewFacetRequest(field string, size int) *FacetRequest {
	return &FacetRequest{
		Field: field,
		Size:  size,
	}
}

// AddDateTimeRange adds a bucket to a field
// containing date values.  Documents with a
// date value falling into this range are tabulated
// as part of this bucket/range.
func (fr *FacetRequest) AddDateTimeRange(name string, start, end time.Time) {
	if fr.DateTimeRanges == nil {
		fr.DateTimeRanges = make([]*dateTimeRange, 0, 1)
	}
	fr.DateTimeRanges = append(fr.DateTimeRanges, &dateTimeRange{Name: name, Start: start, End: end})
}

// AddNumericRange adds a bucket to a field
// containing numeric values.  Documents with a
// numeric value falling into this range are
// tabulated as part of this bucket/range.
func (fr *FacetRequest) AddNumericRange(name string, min, max *float64) {
	if fr.NumericRanges == nil {
		fr.NumericRanges = make([]*numericRange, 0, 1)
	}
	fr.NumericRanges = append(fr.NumericRanges, &numericRange{Name: name, Min: min, Max: max})
}

// FacetsRequest groups together all the
// FacetRequest objects for a single query.
type FacetsRequest map[string]*FacetRequest

// HighlightRequest describes how field matches
// should be highlighted.
type HighlightRequest struct {
	Style  *string  `json:"style"`
	Fields []string `json:"fields"`
}

// NewHighlight creates a default
// HighlightRequest.
func NewHighlight() *HighlightRequest {
	return &HighlightRequest{}
}

// NewHighlightWithStyle creates a HighlightRequest
// with an alternate style.
func NewHighlightWithStyle(style string) *HighlightRequest {
	return &HighlightRequest{
		Style: &style,
	}
}

func (h *HighlightRequest) AddField(field string) {
	if h.Fields == nil {
		h.Fields = make([]string, 0, 1)
	}
	h.Fields = append(h.Fields, field)
}

// A SearchRequest describes all the parameters
// needed to search the index.
// Query is required.
// Size/From describe how much and which part of the
// result set to return.
// Highlight describes optional search result
// highlighting.
// Fields desribed a list of field values whcih
// should be retrieved for result documents.
// Facets describe the set of facets to be computed.
// Explain triggers inclusion of additional search
// result score explanations.
type SearchRequest struct {
	Query     Query             `json:"query"`
	Size      int               `json:"size"`
	From      int               `json:"from"`
	Highlight *HighlightRequest `json:"highlight"`
	Fields    []string          `json:"fields"`
	Facets    FacetsRequest     `json:"facets"`
	Explain   bool              `json:"explain"`
}

// AddFacet adds a FacetRequest to this SearchRequest
func (r *SearchRequest) AddFacet(facetName string, f *FacetRequest) {
	if r.Facets == nil {
		r.Facets = make(FacetsRequest, 1)
	}
	r.Facets[facetName] = f
}

// UnmarshalJSON deserializes a JSON representation of
// a SearchRequest
func (r *SearchRequest) UnmarshalJSON(input []byte) error {
	var temp struct {
		Q         json.RawMessage   `json:"query"`
		Size      int               `json:"size"`
		From      int               `json:"from"`
		Highlight *HighlightRequest `json:"highlight"`
		Fields    []string          `json:"fields"`
		Facets    FacetsRequest     `json:"facets"`
		Explain   bool              `json:"explain"`
	}

	err := json.Unmarshal(input, &temp)
	if err != nil {
		return err
	}

	r.Size = temp.Size
	r.From = temp.From
	r.Explain = temp.Explain
	r.Highlight = temp.Highlight
	r.Fields = temp.Fields
	r.Facets = temp.Facets
	r.Query, err = ParseQuery(temp.Q)
	if err != nil {
		return err
	}

	if r.Size <= 0 {
		r.Size = 10
	}
	if r.From <= 0 {
		r.From = 0
	}

	return nil

}

// NewSearchRequest creates a new SearchRequest
// for the Query, using default values for all
// other search parameters.
func NewSearchRequest(q Query) *SearchRequest {
	return NewSearchRequestOptions(q, 10, 0, false)
}

// NewSearchRequestOptions creates a new SearchRequest
// for the Query, with the requested size, from
// and explanation search parameters.
func NewSearchRequestOptions(q Query, size, from int, explain bool) *SearchRequest {
	return &SearchRequest{
		Query:   q,
		Size:    size,
		From:    from,
		Explain: explain,
	}
}

// A SearchResult describes the results of executing
// a SearchRequest.
type SearchResult struct {
	Request  *SearchRequest                 `json:"request"`
	Hits     search.DocumentMatchCollection `json:"hits"`
	Total    uint64                         `json:"total_hits"`
	MaxScore float64                        `json:"max_score"`
	Took     time.Duration                  `json:"took"`
	Facets   search.FacetResults            `json:"facets"`
}

func (sr *SearchResult) String() string {
	rv := ""
	if len(sr.Hits) > 0 {
		rv = fmt.Sprintf("%d matches, showing %d through %d, took %s\n", sr.Total, sr.Request.From+1, sr.Request.From+len(sr.Hits), sr.Took)
		for i, hit := range sr.Hits {
			rv += fmt.Sprintf("%5d. %s (%f)\n", i+sr.Request.From+1, hit.ID, hit.Score)
			for fragmentField, fragments := range hit.Fragments {
				rv += fmt.Sprintf("\t%s\n", fragmentField)
				for _, fragment := range fragments {
					rv += fmt.Sprintf("\t\t%s\n", fragment)
				}
			}
		}
	} else {
		rv = "No matches"
	}
	return rv
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"text/template"
)

// Check for errors
func check(err error) {
	if err != nil {
		panic(err)
	}
}

// Page Templates
var (
	uploadTemplate *template.Template
	errorTemplate  *template.Template
)

// MIME Validator
func validateimage(h *multipart.FileHeader) (ext string, err error) {
	mimeArray := h.Header["Content-Type"]
	mime := mimeArray[0]
	if mime == "image/jpeg" {
		ext = ".jpg"
	} else if mime == "image/png" {
		ext = ".png"
	} else {
		err = errors.New("Unsupported filetype uploaded")
	}
	return
}

// Upload Handler
func upload(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		// Load the upload page if they aren't posting an image
		uploadTemplate.Execute(w, nil)
		return
	}

	// Get image from POST
	f, h, err := r.FormFile("image")
	check(err)
	defer f.Close()

	// Check MIME and get a file extension
	_, err = validateimage(h)
	check(err)

	// Read and write the uploaded file to disk
	filename := "./img/" + h.Filename
	filebytes, err := ioutil.ReadAll(f)
	check(err)
	err = ioutil.WriteFile(filename, filebytes, 0744)
	check(err)

	// Redirect to the view page
	http.Redirect(w, r, "/view?id="+filename[6:], http.StatusFound)
}

// View Handler
func view(w http.ResponseWriter, r *http.Request) {
	var filename string
	filename = "./img/" + r.FormValue("id")

	if filename[len(filename)-4:] == ".jpg" {
		w.Header().Set("Content-Type", "image/jpeg")
	} else if filename[len(filename)-4:] == ".png" {
		w.Header().Set("Content-Type", "image/png")
	}

	http.ServeFile(w, r, filename)
}

// Error page's variables
type ErrorPage struct {
	Error error
}

// One clean error page
func errorHandler(fn http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if e, ok := recover().(error); ok {
				contents := ErrorPage{Error: e}
				w.WriteHeader(http.StatusInternalServerError)
				errorTemplate.Execute(w, contents)
			}
		}()
		fn(w, r)
	}
}

func main() {
	var err error

	// Load up templates and check for errors
	uploadTemplate, err = template.ParseFiles("upload.html")
	check(err)
	errorTemplate, err = template.ParseFiles("error.html")
	check(err)

	http.HandleFunc("/", errorHandler(upload))
	http.HandleFunc("/view", errorHandler(view))
	http.ListenAndServe(":3000", nil)
}
<commit_msg>added imgdir<commit_after>package main

import (
	"errors"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"text/template"
)

// Check for errors
func check(err error) {
	if err != nil {
		panic(err)
	}
}

// Page Templates
var (
	uploadTemplate *template.Template
	errorTemplate  *template.Template
	imgdir         string
)

// MIME Validator
func validateimage(h *multipart.FileHeader) (ext string, err error) {
	mimeArray := h.Header["Content-Type"]
	mime := mimeArray[0]
	if mime == "image/jpeg" {
		ext = ".jpg"
	} else if mime == "image/png" {
		ext = ".png"
	} else {
		err = errors.New("Unsupported filetype uploaded")
	}
	return
}

// Upload Handler
func upload(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		// Load the upload page if they aren't posting an image
		uploadTemplate.Execute(w, nil)
		return
	}

	// Get image from POST
	f, h, err := r.FormFile("image")
	defer f.Close()
	check(err)

	// Check MIME and get a file extension
	_, err = validateimage(h)
	check(err)

	// Read and write the uploaded file to disk
	filename := imgdir + h.Filename
	filebytes, err := ioutil.ReadAll(f)
	check(err)
	err = ioutil.WriteFile(filename, filebytes, 0744)
	check(err)

	// Redirect to the view page
	http.Redirect(w, r, "/view?id="+filename[6:], http.StatusFound)
}

// View Handler
func view(w http.ResponseWriter, r *http.Request) {
	var filename string
	filename = imgdir + r.FormValue("id")

	if filename[len(filename)-4:] == ".jpg" {
		w.Header().Set("Content-Type", "image/jpeg")
	} else if filename[len(filename)-4:] == ".png" {
		w.Header().Set("Content-Type", "image/png")
	}

	http.ServeFile(w, r, filename)
}

// Error page's variables
type ErrorPage struct {
	Error error
}

// One clean error page
func errorHandler(fn http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if e, ok := recover().(error); ok {
				contents := ErrorPage{Error: e}
				w.WriteHeader(http.StatusInternalServerError)
				errorTemplate.Execute(w, contents)
			}
		}()
		fn(w, r)
	}
}

func main() {
	var err error

	// Set imgdir
	imgdir = "./img/"

	// Load up templates and check for errors
	uploadTemplate, err = template.ParseFiles("upload.html")
	check(err)
	errorTemplate, err = template.ParseFiles("error.html")
	check(err)

	http.HandleFunc("/", errorHandler(upload))
	http.HandleFunc("/view", errorHandler(view))
	http.ListenAndServe(":3000", nil)
}
<|endoftext|>


<commit_before>package main

const (
	evalInf = 50000
)

func SearchPosition(pos Position, depth int) (int, Move) {
	var score int
	var recommended Move
	for d := 1; d <= depth; d++ {
		score, recommended = negamax(pos, recommended, -evalInf, evalInf, d, true, negamax)
	}
	return score, recommended
}

type SearchFunc func(Position, Move, int, int, int, bool, SearchFunc) (int, Move)

func negamax(pos Position, recommended Move, alpha int, beta int, depth int, allowCutoff bool, search SearchFunc) (bestScore int, bestMove Move) {
	if IsTerminal(pos) { // checkmate or stalemate
		if IsCheck(pos) {
			bestScore = -evalInf
		}
		return
	}
	if depth == 0 {
		bestScore = Eval(pos) * evalMult(pos.ToMove)
		return
	}

	moves := Candidates(pos) // pseudo-legal
	// Initialize bestMove with a legal Move if no recommended Move is provided
	if recommended != (Move{}) {
		moves = reorder(moves, recommended)
		bestMove = recommended
	} else {
		for i, m := range moves {
			if IsLegal(Make(pos, m)) {
				bestMove = m
				moves = moves[i:]
				break
			}
		}
	}
	bestScore = alpha

	for _, m := range moves {
		newpos := Make(pos, m)
		if !IsLegal(newpos) {
			continue
		}

		score, _ := search(newpos, Move{}, -beta, -alpha, depth-1, allowCutoff, search)
		score *= -1

		if score > bestScore {
			bestScore, bestMove = score, m
		}
		if score > alpha {
			alpha = score
		}
		if alpha > beta && allowCutoff {
			break
		}
	}
	return
}

// IsCheck returns whether the king of the side to move is in check.
func IsCheck(pos Position) bool {
	return IsAttacked(pos, pos.KingSquare[pos.ToMove], pos.Opp())
}

// IsLegal returns whether a Position results from a legal move.
// A position is illegal if the king of the side that just moved is in check.
func IsLegal(pos Position) bool {
	return !IsAttacked(pos, pos.KingSquare[pos.Opp()], pos.ToMove)
}

// IsTerminal returns whether or not a Position is checkmate or stalemate.
// A position is checkmate or stalemate if the side to move has no legal moves.
func IsTerminal(pos Position) bool {
	for _, m := range Candidates(pos) {
		if IsLegal(Make(pos, m)) {
			return false
		}
	}
	return true
}

// reorder returns a reordered slice of Moves with the specified Move first.
// The slice is not modified if it does not contain the specified Move.
func reorder(moves []Move, m Move) []Move {
	for i := range moves {
		if moves[i] == m {
			moves = append(append(moves[i:i+1], moves[:i]...), moves[i+1:]...)
			return moves
		}
	}
	return moves
}
<commit_msg>append reordered moves one by one<commit_after>package main

const (
	evalInf = 50000
)

func SearchPosition(pos Position, depth int) (int, Move) {
	var score int
	var recommended Move
	for d := 1; d <= depth; d++ {
		score, recommended = negamax(pos, recommended, -evalInf, evalInf, d, true, negamax)
	}
	return score, recommended
}

type SearchFunc func(Position, Move, int, int, int, bool, SearchFunc) (int, Move)

func negamax(pos Position, recommended Move, alpha int, beta int, depth int, allowCutoff bool, search SearchFunc) (bestScore int, bestMove Move) {
	if IsTerminal(pos) { // checkmate or stalemate
		if IsCheck(pos) {
			bestScore = -evalInf
		}
		return
	}
	if depth == 0 {
		bestScore = Eval(pos) * evalMult(pos.ToMove)
		return
	}

	moves := Candidates(pos) // pseudo-legal
	// Initialize bestMove with a legal Move if no recommended Move is provided
	if recommended != (Move{}) {
		moves = reorder(moves, recommended)
		bestMove = recommended
	} else {
		for i, m := range moves {
			if IsLegal(Make(pos, m)) {
				bestMove = m
				moves = moves[i:]
				break
			}
		}
	}
	bestScore = alpha

	for _, m := range moves {
		newpos := Make(pos, m)
		if !IsLegal(newpos) {
			continue
		}

		score, _ := search(newpos, Move{}, -beta, -alpha, depth-1, allowCutoff, search)
		score *= -1

		if score > bestScore {
			bestScore, bestMove = score, m
		}
		if score > alpha {
			alpha = score
		}
		if alpha > beta && allowCutoff {
			break
		}
	}
	return
}

// IsCheck returns whether the king of the side to move is in check.
func IsCheck(pos Position) bool {
	return IsAttacked(pos, pos.KingSquare[pos.ToMove], pos.Opp())
}

// IsLegal returns whether a Position results from a legal move.
// A position is illegal if the king of the side that just moved is in check.
func IsLegal(pos Position) bool {
	return !IsAttacked(pos, pos.KingSquare[pos.Opp()], pos.ToMove)
}

// IsTerminal returns whether or not a Position is checkmate or stalemate.
// A position is checkmate or stalemate if the side to move has no legal moves.
func IsTerminal(pos Position) bool {
	for _, m := range Candidates(pos) {
		if IsLegal(Make(pos, m)) {
			return false
		}
	}
	return true
}

// reorder returns a reordered slice of Moves with the specified Move first.
// The slice is not modified if it does not contain the specified Move.
func reorder(moves []Move, m Move) []Move {
	for _, n := range moves {
		if n == m {
			s := []Move{m}
			for _, n := range moves {
				if n != m {
					s = append(s, n)
				}
			}
			return s
		}
	}
	return moves
}
<|endoftext|>


<commit_before>// Copyright © 2016 Aaron Longwell
//
// Use of this source code is governed by an MIT licese.
// Details in the LICENSE file.

package trello

type SearchResult struct {
	Options SearchOptions `json:"options"`
	Actions []*Action     `json:"actions,omitempty"`
	Cards   []*Card       `json:"cards,omitempty"`
	Boards  []*Board      `json:"boards,omitempty"`
	Members []*Member     `json:"members,omitempty"`
}

type SearchOptions struct {
	Terms      []SearchTerm `json:"terms"`
	Modifiers  []string     `json:"modifiers,omitempty"`
	ModelTypes []string     `json:"modelTypes,omitempty"`
	Partial    bool         `json:"partial"`
}

type SearchTerm struct {
	Text    string `json:"text"`
	Negated bool   `json:"negated,omitempty"`
}

func (c *Client) SearchCards(query string, args Arguments) (cards []*Card, err error) {
	args["query"] = query
	args["modelTypes"] = "cards"
	res := SearchResult{}
	err = c.Get("search", args, &res)
	cards = res.Cards
	return
}

func (c *Client) SearchBoards(query string, args Arguments) (boards []*Board, err error) {
	args["query"] = query
	args["modelTypes"] = "boards"
	res := SearchResult{}
	err = c.Get("search", args, &res)
	boards = res.Boards
	return
}

func (c *Client) SearchMembers(query string, args Arguments) (members []*Member, err error) {
	args["query"] = query
	err = c.Get("search/members", args, &members)
	return
}
<commit_msg>update searchOptions Modifiers type<commit_after>// Copyright © 2016 Aaron Longwell
//
// Use of this source code is governed by an MIT licese.
// Details in the LICENSE file.

package trello

type SearchResult struct {
	Options SearchOptions `json:"options"`
	Actions []*Action     `json:"actions,omitempty"`
	Cards   []*Card       `json:"cards,omitempty"`
	Boards  []*Board      `json:"boards,omitempty"`
	Members []*Member     `json:"members,omitempty"`
}

type SearchOptions struct {
	Terms      []SearchTerm     `json:"terms"`
	Modifiers  []SearchModifier `json:"modifiers,omitempty"`
	ModelTypes []string         `json:"modelTypes,omitempty"`
	Partial    bool             `json:"partial"`
}

type SearchModifier struct {
	Text string `json:"text"`
}

type SearchTerm struct {
	Text    string `json:"text"`
	Negated bool   `json:"negated,omitempty"`
}

func (c *Client) SearchCards(query string, args Arguments) (cards []*Card, err error) {
	args["query"] = query
	args["modelTypes"] = "cards"
	res := SearchResult{}
	err = c.Get("search", args, &res)
	cards = res.Cards
	return
}

func (c *Client) SearchBoards(query string, args Arguments) (boards []*Board, err error) {
	args["query"] = query
	args["modelTypes"] = "boards"
	res := SearchResult{}
	err = c.Get("search", args, &res)
	boards = res.Boards
	return
}

func (c *Client) SearchMembers(query string, args Arguments) (members []*Member, err error) {
	args["query"] = query
	err = c.Get("search/members", args, &members)
	return
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"errors"
	"fmt"
	"sort"
)

const (
	evalInf = 50000
)

var (
	errCheckmate checkmateError
	errStalemate = errors.New("stalemate")
	errFiftyMove = errors.New("fifty-move rule")
	// TODO: threefold repetition
)

// A checkmateError indicates the number of plies until a forced checkmate can be delivered.
// The zero value of type checkmateError indicates that the current position is checkmate.
type checkmateError int

func (n checkmateError) Error() string {
	if n&1 == 0 {
		return fmt.Sprintf("-#%d", n/2)
	}
	return fmt.Sprintf("#%d", (n+1)/2)
}

// Prev returns the checkmateError corresponding to the previous ply.
func (n checkmateError) Prev() checkmateError { return n + 1 }

// SearchPosition searches a Position to the specified depth via iterative deepening
// and returns the search results.
func SearchPosition(ctx context.Context, pos Position, depth int) Results {
	var results Results
	s := Search{
		allowCutoff: true,
		counters:    make([]int, depth+1),
	}
	for d := 1; d <= depth; d++ {
		_, results, _ = s.negamax(ctx, pos, results, Window{-evalInf, evalInf}, d)
		if ctx.Err() != nil {
			break
		}
	}
	return results
}

type Search struct {
	allowCutoff bool
	counters    []int
}

// negamax recursively searches a Position to the specified depth and returns the evaluation score
// relative to the side to move and the search results. It employs alpha-beta pruning outside of
// the specified Window. If recommended is zero length, negamax will generate and search all legal
// moves; if recommended moves are provided, they must all be legal, and only they will be searched.
func (s *Search) negamax(
	ctx context.Context,
	pos Position,
	recommended Results,
	w Window,
	depth int,
) (bestScore RelScore, results Results, err error) {
	s.counters[len(s.counters)-1-depth]++

	score, recommended, err := checkDone(pos, recommended)
	if err != nil {
		return score, nil, err
	}
	// Invariant: len(recommended) > 0 and recommended contains only legal moves
	if depth == 0 {
		score, err := Eval(pos)
		return score.Rel(pos.ToMove), nil, err
	}

	results = recommended
	defer func() { results.SortFor(pos.ToMove) }()

	for _, r := range recommended {
		if r.err != nil {
			// The move is already known not to avoid a game-ending state; no need to search it further
			results = results.Update(r)
			continue
		}

		score, cont, err := s.negamax(ctx, Make(pos, r.move), r.cont, w.Neg(), depth-1)
		score *= -1

		results = results.Update(Result{move: r.move, score: score.Abs(pos.ToMove), depth: depth - 1, cont: cont, err: err})

		if depth >= 3 && ctx.Err() != nil {
			break
		}

		var constrained, ok bool
		w, constrained, ok = w.Constrain(score)
		if constrained {
			// improved lower bound
		}
		if !ok && s.allowCutoff {
			break
		}
	}
	if err, ok := results[0].err.(checkmateError); ok {
		return w.alpha, results, err.Prev()
	}
	return w.alpha, results, results[0].err
}

// IsPseudoLegal returns whether a Move is pseudo-legal in a Position.
// A move is pseudo-legal if the square to be moved from contains the specified piece
// and the piece is capable of moving to the target square if doing so would not put the king in check.
func IsPseudoLegal(pos Position, move Move) bool {
	for _, m := range PieceMoves[move.Piece](pos) {
		if m == move {
			return true
		}
	}
	return false
}

// IsLegal returns whether a Position is legal.
// A position is illegal if the king of the side that just moved is in check.
func IsLegal(pos Position) bool {
	return !IsAttacked(pos, pos.KingSquare[pos.Opp()], pos.ToMove)
}

// IsCheck returns whether the king of the side to move is in check.
func IsCheck(pos Position) bool {
	return IsAttacked(pos, pos.KingSquare[pos.ToMove], pos.Opp())
}

// IsTerminal returns whether or not a Position is checkmate or stalemate.
// A position is checkmate or stalemate if the side to move has no legal moves.
func IsTerminal(pos Position) bool { return !anyLegal(pos, Candidates(pos)) }

// anyLegal returns whether any of the given Moves are legal in the given Position.
func anyLegal(pos Position, moves []Move) bool {
	for _, m := range moves {
		if IsLegal(Make(pos, m)) {
			return true
		}
	}
	return false
}

func checkDone(pos Position, recommended Results) (RelScore, Results, error) {
	if len(recommended) == 0 {
		moves := Candidates(pos) // pseudo-legal
		recommended = make(Results, 0, len(moves))
		for _, m := range moves {
			if !IsLegal(Make(pos, m)) {
				continue
			}
			recommended = append(recommended, Result{move: m})
		}
	}
	if len(recommended) == 0 {
		// no legal moves
		if IsCheck(pos) {
			return -evalInf, nil, errCheckmate
		}
		return 0, nil, errStalemate
	}
	if pos.HalfMove == 100 {
		// fifty-move rule
		return 0, nil, errFiftyMove
	}
	return 0, recommended, nil
}

// A Result holds a searched Move along with its evaluated Score,
// the search depth, and the continuation Results of the search.
type Result struct {
	move  Move
	score Score
	depth int
	cont  Results
	err   error
}

// String returns a string representation of r, including its principal variation.
func (r Result) String() string {
	s := fmt.Sprintf("%v (%v) %v", r.score, r.depth, r.PV())
	if r.err != nil {
		return s + " " + r.err.Error()
	}
	return s
}

// PV returns a string representation of r's principal variation.
func (r Result) PV() string {
	if r.depth == 0 || len(r.cont) == 0 {
		return LongAlgebraic(r.move)
	}
	return LongAlgebraic(r.move) + " " + r.cont[0].PV()
}

// A Results contains the results of a search.
// Functions returning a Results should sort it before returning.
type Results []Result

// SortFor sorts rs beginning with the best move for c.
// The sort is not guaranteed to be stable.
func (rs Results) SortFor(c Color) {
	// Sort first by mate condition, then by depth decreasing,
	// then by Score increasing/decreasing for White/Black,
	// and then by origin and destination Square increasing
	sort.Slice(rs, func(i, j int) bool {
		if less, ok := rs.mateSort(i, j); ok {
			return less
		}
		if rs[i].depth != rs[j].depth {
			return rs[i].depth > rs[j].depth
		}
		if rs[i].score != rs[j].score {
			return (rs[i].score > rs[j].score) == (c == White)
		}
		if rs[i].move.From != rs[j].move.From {
			return rs[i].move.From < rs[j].move.From
		}
		return rs[i].move.To < rs[j].move.To
	})
}

// mateSort reports how two Result elements should be sorted if at least one of them leads to checkmate.
// ok reports whether either Result contains an error of type checkmateError.
// If ok is true, less reports whether the Result with index i should sort before the Result with index j.
// If ok is false, the value of less is meaningless.
func (rs Results) mateSort(i, j int) (less bool, ok bool) {
	ich, iok := rs[i].err.(checkmateError)
	jch, jok := rs[j].err.(checkmateError)
	switch {
	case iok && jok:
		switch iwin, jwin := ich&1 == 0, jch&1 == 0; {
		case iwin && jwin:
			return ich < jch, true
		case !iwin && !jwin:
			return ich > jch, true
		default:
			return iwin, true
		}
	case iok:
		return ich&1 == 0, true
	case jok:
		return jch&1 != 0, true
	}
	return false, false
}

// Update looks for a Result in rs with the same Move as r and replaces it with r,
// or appends r if there is none, and returns the possibly updated slice.
func (rs Results) Update(r Result) Results {
	for i := range rs {
		if rs[i].move == r.move {
			rs[i] = r
			return rs
		}
	}
	return append(rs, r)
}

// String returns a string representation of all Result values in r.
func (rs Results) String() string {
	var s string
	for _, r := range rs {
		s += fmt.Sprintf("%v\n", r.String())
	}
	return s
}

// A Window represents the bounds of a position's evaluation.
type Window struct{ alpha, beta RelScore }

// Constrain updates the lower bound of w, if applicable, and returns the updated window,
// whether the lower bound was changed, and whether the returned Window remains valid.
func (w Window) Constrain(n RelScore) (c Window, constrained bool, ok bool) {
	if n <= w.alpha {
		return w, false, true
	}
	return Window{n, w.beta}, true, n <= w.beta
}

// Neg returns the additive inverse of w.
func (w Window) Neg() Window { return Window{-w.beta, -w.alpha} }
<commit_msg>document checkDone<commit_after>package main

import (
	"context"
	"errors"
	"fmt"
	"sort"
)

const (
	evalInf = 50000
)

var (
	errCheckmate checkmateError
	errStalemate = errors.New("stalemate")
	errFiftyMove = errors.New("fifty-move rule")
	// TODO: threefold repetition
)

// A checkmateError indicates the number of plies until a forced checkmate can be delivered.
// The zero value of type checkmateError indicates that the current position is checkmate.
type checkmateError int

func (n checkmateError) Error() string {
	if n&1 == 0 {
		return fmt.Sprintf("-#%d", n/2)
	}
	return fmt.Sprintf("#%d", (n+1)/2)
}

// Prev returns the checkmateError corresponding to the previous ply.
func (n checkmateError) Prev() checkmateError { return n + 1 }

// SearchPosition searches a Position to the specified depth via iterative deepening
// and returns the search results.
func SearchPosition(ctx context.Context, pos Position, depth int) Results {
	var results Results
	s := Search{
		allowCutoff: true,
		counters:    make([]int, depth+1),
	}
	for d := 1; d <= depth; d++ {
		_, results, _ = s.negamax(ctx, pos, results, Window{-evalInf, evalInf}, d)
		if ctx.Err() != nil {
			break
		}
	}
	return results
}

type Search struct {
	allowCutoff bool
	counters    []int
}

// negamax recursively searches a Position to the specified depth and returns the evaluation score
// relative to the side to move and the search results. It employs alpha-beta pruning outside of
// the specified Window. If recommended is zero length, negamax will generate and search all legal
// moves; if recommended moves are provided, they must all be legal, and only they will be searched.
func (s *Search) negamax(
	ctx context.Context,
	pos Position,
	recommended Results,
	w Window,
	depth int,
) (bestScore RelScore, results Results, err error) {
	s.counters[len(s.counters)-1-depth]++

	score, recommended, err := checkDone(pos, recommended)
	if err != nil {
		return score, nil, err
	}
	// Invariant: len(recommended) > 0 and recommended contains only legal moves
	if depth == 0 {
		score, err := Eval(pos)
		return score.Rel(pos.ToMove), nil, err
	}

	results = recommended
	defer func() { results.SortFor(pos.ToMove) }()

	for _, r := range recommended {
		if r.err != nil {
			// The move is already known not to avoid a game-ending state; no need to search it further
			results = results.Update(r)
			continue
		}

		score, cont, err := s.negamax(ctx, Make(pos, r.move), r.cont, w.Neg(), depth-1)
		score *= -1

		results = results.Update(Result{move: r.move, score: score.Abs(pos.ToMove), depth: depth - 1, cont: cont, err: err})

		if depth >= 3 && ctx.Err() != nil {
			break
		}

		var constrained, ok bool
		w, constrained, ok = w.Constrain(score)
		if constrained {
			// improved lower bound
		}
		if !ok && s.allowCutoff {
			break
		}
	}
	if err, ok := results[0].err.(checkmateError); ok {
		return w.alpha, results, err.Prev()
	}
	return w.alpha, results, results[0].err
}

// checkDone reports whether pos represents a game-ending position.
// If so, the RelScore and error value indicate the type of ending.
// If not, the RelScore is 0 and the error is nil.
// In either case, checkDone returns a Results of all legal moves in pos.
// If recommended is not provided, checkDone generates the legal moves first.
func checkDone(pos Position, recommended Results) (RelScore, Results, error) {
	if len(recommended) == 0 {
		moves := Candidates(pos) // pseudo-legal
		recommended = make(Results, 0, len(moves))
		for _, m := range moves {
			if !IsLegal(Make(pos, m)) {
				continue
			}
			recommended = append(recommended, Result{move: m})
		}
	}
	if len(recommended) == 0 {
		// no legal moves
		if IsCheck(pos) {
			return -evalInf, nil, errCheckmate
		}
		return 0, nil, errStalemate
	}
	if pos.HalfMove == 100 {
		// fifty-move rule
		return 0, nil, errFiftyMove
	}
	return 0, recommended, nil
}

// IsPseudoLegal returns whether a Move is pseudo-legal in a Position.
// A move is pseudo-legal if the square to be moved from contains the specified piece
// and the piece is capable of moving to the target square if doing so would not put the king in check.
func IsPseudoLegal(pos Position, move Move) bool {
	for _, m := range PieceMoves[move.Piece](pos) {
		if m == move {
			return true
		}
	}
	return false
}

// IsLegal returns whether a Position is legal.
// A position is illegal if the king of the side that just moved is in check.
func IsLegal(pos Position) bool {
	return !IsAttacked(pos, pos.KingSquare[pos.Opp()], pos.ToMove)
}

// IsCheck returns whether the king of the side to move is in check.
func IsCheck(pos Position) bool {
	return IsAttacked(pos, pos.KingSquare[pos.ToMove], pos.Opp())
}

// IsTerminal returns whether or not a Position is checkmate or stalemate.
// A position is checkmate or stalemate if the side to move has no legal moves.
func IsTerminal(pos Position) bool { return !anyLegal(pos, Candidates(pos)) }

// anyLegal returns whether any of the given Moves are legal in the given Position.
func anyLegal(pos Position, moves []Move) bool {
	for _, m := range moves {
		if IsLegal(Make(pos, m)) {
			return true
		}
	}
	return false
}

// A Result holds a searched Move along with its evaluated Score,
// the search depth, and the continuation Results of the search.
type Result struct {
	move  Move
	score Score
	depth int
	cont  Results
	err   error
}

// String returns a string representation of r, including its principal variation.
func (r Result) String() string {
	s := fmt.Sprintf("%v (%v) %v", r.score, r.depth, r.PV())
	if r.err != nil {
		return s + " " + r.err.Error()
	}
	return s
}

// PV returns a string representation of r's principal variation.
func (r Result) PV() string {
	if r.depth == 0 || len(r.cont) == 0 {
		return LongAlgebraic(r.move)
	}
	return LongAlgebraic(r.move) + " " + r.cont[0].PV()
}

// A Results contains the results of a search.
// Functions returning a Results should sort it before returning.
type Results []Result

// SortFor sorts rs beginning with the best move for c.
// The sort is not guaranteed to be stable.
func (rs Results) SortFor(c Color) {
	// Sort first by mate condition, then by depth decreasing,
	// then by Score increasing/decreasing for White/Black,
	// and then by origin and destination Square increasing
	sort.Slice(rs, func(i, j int) bool {
		if less, ok := rs.mateSort(i, j); ok {
			return less
		}
		if rs[i].depth != rs[j].depth {
			return rs[i].depth > rs[j].depth
		}
		if rs[i].score != rs[j].score {
			return (rs[i].score > rs[j].score) == (c == White)
		}
		if rs[i].move.From != rs[j].move.From {
			return rs[i].move.From < rs[j].move.From
		}
		return rs[i].move.To < rs[j].move.To
	})
}

// mateSort reports how two Result elements should be sorted if at least one of them leads to checkmate.
// ok reports whether either Result contains an error of type checkmateError.
// If ok is true, less reports whether the Result with index i should sort before the Result with index j.
// If ok is false, the value of less is meaningless.
func (rs Results) mateSort(i, j int) (less bool, ok bool) {
	ich, iok := rs[i].err.(checkmateError)
	jch, jok := rs[j].err.(checkmateError)
	switch {
	case iok && jok:
		switch iwin, jwin := ich&1 == 0, jch&1 == 0; {
		case iwin && jwin:
			return ich < jch, true
		case !iwin && !jwin:
			return ich > jch, true
		default:
			return iwin, true
		}
	case iok:
		return ich&1 == 0, true
	case jok:
		return jch&1 != 0, true
	}
	return false, false
}

// Update looks for a Result in rs with the same Move as r and replaces it with r,
// or appends r if there is none, and returns the possibly updated slice.
func (rs Results) Update(r Result) Results {
	for i := range rs {
		if rs[i].move == r.move {
			rs[i] = r
			return rs
		}
	}
	return append(rs, r)
}

// String returns a string representation of all Result values in r.
func (rs Results) String() string {
	var s string
	for _, r := range rs {
		s += fmt.Sprintf("%v\n", r.String())
	}
	return s
}

// A Window represents the bounds of a position's evaluation.
type Window struct{ alpha, beta RelScore }

// Constrain updates the lower bound of w, if applicable, and returns the updated window,
// whether the lower bound was changed, and whether the returned Window remains valid.
func (w Window) Constrain(n RelScore) (c Window, constrained bool, ok bool) {
	if n <= w.alpha {
		return w, false, true
	}
	return Window{n, w.beta}, true, n <= w.beta
}

// Neg returns the additive inverse of w.
func (w Window) Neg() Window { return Window{-w.beta, -w.alpha} }
<|endoftext|>


<commit_before>package algoliasearch

import (
	"bytes"
	_ "crypto/sha512" // Fix certificates
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"
)

const (
	version = "2.6.0"
)

// Define the constants used to specify the type of request.
const (
	search = 1 << iota
	write
	read
)

// Seed the RNG used to shuffle the hosts slice (see `defaultHosts` function).
func init() {
	rand.Seed(int64(time.Now().Nanosecond()))
}

// Transport is responsible for the connection and the retry strategy to
// Algolia servers.
type Transport struct {
	activeHost        string
	activeSince       time.Time
	apiKey            string
	appId             string
	dialTimeout       time.Duration
	headers           map[string]string
	hosts             []string
	httpClient        *http.Client
	keepAliveDuration time.Duration
}

// NewTransport instantiates a new Transport with the default Algolia hosts to
// connect to.
func NewTransport(appId, apiKey string) *Transport {
	return &Transport{
		activeHost:        "",
		apiKey:            apiKey,
		appId:             appId,
		dialTimeout:       1 * time.Second,
		headers:           defaultHeaders(appId, apiKey),
		hosts:             defaultHosts(appId),
		httpClient:        defaultHttpClient(),
		keepAliveDuration: 5 * 60 * time.Second,
	}
}

// NewTransport instantiates a new Transport with the specificed hosts as main
// servers to connect to.
func NewTransportWithHosts(appId, apiKey string, hosts []string) *Transport {
	return &Transport{
		activeHost:        "",
		apiKey:            apiKey,
		appId:             appId,
		dialTimeout:       1 * time.Second,
		headers:           defaultHeaders(appId, apiKey),
		hosts:             hosts,
		httpClient:        defaultHttpClient(),
		keepAliveDuration: 5 * 60 * time.Second,
	}
}

// defaultHeaders is used to set the default HTTP headers to use with each
// requests.
func defaultHeaders(appId, apiKey string) map[string]string {
	return map[string]string{
		"Connection":               "keep-alive",
		"User-Agent":               "Algolia for Go (" + version + ")",
		"X-Algolia-API-Key":        apiKey,
		"X-Algolia-Application-Id": appId,
	}
}

// defaultHosts returns the list of the default Algolia hosts to use. The
// entries are shuffled.
func defaultHosts(appId string) []string {
	hosts := []string{
		appId + "-1.algolianet.com",
		appId + "-2.algolianet.com",
		appId + "-3.algolianet.com",
	}

	shuffled := make([]string, len(hosts))
	for i, v := range rand.Perm(len(hosts)) {
		shuffled[i] = hosts[v]
	}

	return shuffled
}

// defaultHttpClient returns the `*http.Client` which will perform all the
// requests. All the timeout settings are explicitely defined here.
func defaultHttpClient() *http.Client {
	return &http.Client{
		Timeout:   time.Second * 30,
		Transport: defaultTransport(1 * time.Second),
	}
}

// defaultTransport returns the `*http.Transport` which starts and maintain the
// connection with the server. The `dialTimeout` is used to specify the timeout
// beyond which the connection is considered as failed (used to control DNS
// lookup timeouts).
func defaultTransport(dialTimeout time.Duration) *http.Transport {
	return &http.Transport{
		Dial: (&net.Dialer{
			KeepAlive: 180 * time.Second,
			Timeout:   dialTimeout,
		}).Dial,
		DisableKeepAlives:   false,
		MaxIdleConnsPerHost: 2,
		TLSHandshakeTimeout: 2 * time.Second,
	}
}

// addHeaders add the key/value pairs from `headers` to the header list of the
// `req` request.
func addHeaders(req *http.Request, headers map[string]string) {
	for k, v := range headers {
		req.Header.Add(k, v)
	}
}

// setExtraHeader lets the user (through the exported `Client.SetExtraHeader`)
// add custom headers to the requests.
func (t *Transport) setExtraHeader(key, value string) {
	t.headers[key] = value
}

// setTimeout lets the user (through the exported `Client.SetTimeout`) replace
// the default values of `TLSHandshakeTimeout` (via `connectTimeout`) and
// `ResponseHeaderTimeout` (via `readTimeout`).
func (t *Transport) setTimeout(connectTimeout, readTimeout time.Duration) {
	switch transport := t.httpClient.Transport.(type) {
	case *http.Transport:
		transport.TLSHandshakeTimeout = connectTimeout
		transport.ResponseHeaderTimeout = readTimeout
	default:
		fmt.Fprintln(os.Stderr, "Timeouts not set for nonstandard underlying Transport")
	}
}

// request is the method used by the `Client` to perform the request against
// the Algolia servers (or to the list of specified hosts).
func (t *Transport) request(method, path string, body interface{}, typeCall int) ([]byte, error) {
	var res []byte
	var err error

	for _, host := range t.hostsToTry(typeCall) {
		res, err = t.tryRequest(method, host, path, body)
		if err == nil {
			t.resetDialTimeout()
			t.activeSince = time.Now()
			t.activeHost = host
			return res, nil
		}
		t.increaseDialTimeout()
	}

	t.activeHost = ""
	return nil, err
}

// hostsToTry returns the list of hosts to try ordered by priority according to
// the type of request (write vs. read/search) and if a previous host was
// marked as active.
func (t *Transport) hostsToTry(typeCall int) []string {
	// If it's a `write` request, we first try `APPID.algolia.net` and then the
	// `APPID-{1,2,3}.algolianet.com` endpoints (shuffled).
	if typeCall == write {
		return append([]string{t.appId + ".algolia.net"}, t.hosts...)
	}

	var hosts []string

	// If a previous host was active since less than k.keepAliveDuration
	// seconds and was not the `APPID-dsn.algolia.net` endpoint, then it is
	// used first in the list of endpoints to try.
	if t.activeHost != "" &&
		!strings.Contains(t.activeHost, "-dsn.algolia.net") &&
		time.Now().Sub(t.activeSince) <= t.keepAliveDuration {
		hosts = []string{t.activeHost}
	}

	// Whether or not an active host is first in the list of hosts to try, the
	// following entries in the list is the DSN endpoint followed by the
	// `APPID-{1,2,3}.algolianet.com` endpoints (shuffled).
	hosts = append(hosts, t.appId+"-dsn.algolia.net")
	hosts = append(hosts, t.hosts...)
	return hosts

}

// tryRequest is the underlying method which actually performs the request. It
// returns the response as a byte slice or a non-nil error if anything went
// wrong.
func (t *Transport) tryRequest(method, host, path string, body interface{}) ([]byte, error) {
	// Build the request
	req, err := t.buildRequest(method, host, path, body)
	if err != nil {
		return nil, err
	}

	// Perform the request
	res, err := t.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("Cannot perform request [%s] %s (%s): %s", method, path, host, err)
	}

	// Read response's body
	bodyRes, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, fmt.Errorf("Cannot read response body: %s", err)
	}
	res.Body.Close()

	// Return the body as an error if the status code is not 2XX
	code := res.StatusCode
	if code < 200 || 300 <= code {
		return nil, errors.New(string(bodyRes))
	}

	return bodyRes, nil
}

// buildRequest returns a valid `http.Request` with the headers and body (if
// any) correctly set. The return error is non-nil if the request is invalid or
// if the body, if non-nil, is not a valid JSON.
func (t *Transport) buildRequest(method, host, path string, body interface{}) (*http.Request, error) {
	var req *http.Request
	var err error
	urlStr := "https://" + host + path

	if body == nil {
		// As the body is nil, an empty body request is instantiated
		req, err = http.NewRequest(method, urlStr, nil)
		if err != nil {
			return nil, fmt.Errorf("Cannot instantiate request: [%s] %s", method, urlStr)
		}
	} else {
		// As the body is non-nil, the content is read
		data, err := json.Marshal(body)
		if err != nil {
			return nil, errors.New("Invalid JSON in the query")
		}
		reader := bytes.NewReader(data)

		// The request is then instantiated with the body content
		req, err = http.NewRequest(method, urlStr, reader)
		if err != nil {
			return nil, fmt.Errorf("Cannot instantiate request: [%s] %s", method, urlStr)
		}

		// Add content specific headers
		req.Header.Add("Content-Length", strconv.Itoa(len(string(data))))
		req.Header.Add("Content-Type", "application/json; charset=utf-8")
	}

	// Add default and Algolia specific headers
	addHeaders(req, t.headers)

	if strings.Contains(path, "/*/") {
		req.URL = &url.URL{
			Scheme: "https",
			Host:   host,
			Opaque: "//" + host + path, //Remove url encoding
		}
	}

	return req, nil
}

// resetDialTimeout increases the `Timeout` value of the underlying dialer by 1
// second.
func (t *Transport) increaseDialTimeout() {
	t.dialTimeout = t.dialTimeout + time.Second
	t.httpClient.Transport = defaultTransport(t.dialTimeout)
}

// resetDialTimeout resets the `Timeout` value of the underlying dialer to 1
// second.
func (t *Transport) resetDialTimeout() {
	t.dialTimeout = 1 * time.Second
	t.httpClient.Transport = defaultTransport(t.dialTimeout)
}
<commit_msg>refactor: Make status code checking more readable<commit_after>package algoliasearch

import (
	"bytes"
	_ "crypto/sha512" // Fix certificates
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"
)

const (
	version = "2.6.0"
)

// Define the constants used to specify the type of request.
const (
	search = 1 << iota
	write
	read
)

// Seed the RNG used to shuffle the hosts slice (see `defaultHosts` function).
func init() {
	rand.Seed(int64(time.Now().Nanosecond()))
}

// Transport is responsible for the connection and the retry strategy to
// Algolia servers.
type Transport struct {
	activeHost        string
	activeSince       time.Time
	apiKey            string
	appId             string
	dialTimeout       time.Duration
	headers           map[string]string
	hosts             []string
	httpClient        *http.Client
	keepAliveDuration time.Duration
}

// NewTransport instantiates a new Transport with the default Algolia hosts to
// connect to.
func NewTransport(appId, apiKey string) *Transport {
	return &Transport{
		activeHost:        "",
		apiKey:            apiKey,
		appId:             appId,
		dialTimeout:       1 * time.Second,
		headers:           defaultHeaders(appId, apiKey),
		hosts:             defaultHosts(appId),
		httpClient:        defaultHttpClient(),
		keepAliveDuration: 5 * 60 * time.Second,
	}
}

// NewTransport instantiates a new Transport with the specificed hosts as main
// servers to connect to.
func NewTransportWithHosts(appId, apiKey string, hosts []string) *Transport {
	return &Transport{
		activeHost:        "",
		apiKey:            apiKey,
		appId:             appId,
		dialTimeout:       1 * time.Second,
		headers:           defaultHeaders(appId, apiKey),
		hosts:             hosts,
		httpClient:        defaultHttpClient(),
		keepAliveDuration: 5 * 60 * time.Second,
	}
}

// defaultHeaders is used to set the default HTTP headers to use with each
// requests.
func defaultHeaders(appId, apiKey string) map[string]string {
	return map[string]string{
		"Connection":               "keep-alive",
		"User-Agent":               "Algolia for Go (" + version + ")",
		"X-Algolia-API-Key":        apiKey,
		"X-Algolia-Application-Id": appId,
	}
}

// defaultHosts returns the list of the default Algolia hosts to use. The
// entries are shuffled.
func defaultHosts(appId string) []string {
	hosts := []string{
		appId + "-1.algolianet.com",
		appId + "-2.algolianet.com",
		appId + "-3.algolianet.com",
	}

	shuffled := make([]string, len(hosts))
	for i, v := range rand.Perm(len(hosts)) {
		shuffled[i] = hosts[v]
	}

	return shuffled
}

// defaultHttpClient returns the `*http.Client` which will perform all the
// requests. All the timeout settings are explicitely defined here.
func defaultHttpClient() *http.Client {
	return &http.Client{
		Timeout:   time.Second * 30,
		Transport: defaultTransport(1 * time.Second),
	}
}

// defaultTransport returns the `*http.Transport` which starts and maintain the
// connection with the server. The `dialTimeout` is used to specify the timeout
// beyond which the connection is considered as failed (used to control DNS
// lookup timeouts).
func defaultTransport(dialTimeout time.Duration) *http.Transport {
	return &http.Transport{
		Dial: (&net.Dialer{
			KeepAlive: 180 * time.Second,
			Timeout:   dialTimeout,
		}).Dial,
		DisableKeepAlives:   false,
		MaxIdleConnsPerHost: 2,
		TLSHandshakeTimeout: 2 * time.Second,
	}
}

// addHeaders add the key/value pairs from `headers` to the header list of the
// `req` request.
func addHeaders(req *http.Request, headers map[string]string) {
	for k, v := range headers {
		req.Header.Add(k, v)
	}
}

// setExtraHeader lets the user (through the exported `Client.SetExtraHeader`)
// add custom headers to the requests.
func (t *Transport) setExtraHeader(key, value string) {
	t.headers[key] = value
}

// setTimeout lets the user (through the exported `Client.SetTimeout`) replace
// the default values of `TLSHandshakeTimeout` (via `connectTimeout`) and
// `ResponseHeaderTimeout` (via `readTimeout`).
func (t *Transport) setTimeout(connectTimeout, readTimeout time.Duration) {
	switch transport := t.httpClient.Transport.(type) {
	case *http.Transport:
		transport.TLSHandshakeTimeout = connectTimeout
		transport.ResponseHeaderTimeout = readTimeout
	default:
		fmt.Fprintln(os.Stderr, "Timeouts not set for nonstandard underlying Transport")
	}
}

// request is the method used by the `Client` to perform the request against
// the Algolia servers (or to the list of specified hosts).
func (t *Transport) request(method, path string, body interface{}, typeCall int) ([]byte, error) {
	var res []byte
	var err error

	for _, host := range t.hostsToTry(typeCall) {
		res, err = t.tryRequest(method, host, path, body)
		if err == nil {
			t.resetDialTimeout()
			t.activeSince = time.Now()
			t.activeHost = host
			return res, nil
		}
		t.increaseDialTimeout()
	}

	t.activeHost = ""
	return nil, err
}

// hostsToTry returns the list of hosts to try ordered by priority according to
// the type of request (write vs. read/search) and if a previous host was
// marked as active.
func (t *Transport) hostsToTry(typeCall int) []string {
	// If it's a `write` request, we first try `APPID.algolia.net` and then the
	// `APPID-{1,2,3}.algolianet.com` endpoints (shuffled).
	if typeCall == write {
		return append([]string{t.appId + ".algolia.net"}, t.hosts...)
	}

	var hosts []string

	// If a previous host was active since less than k.keepAliveDuration
	// seconds and was not the `APPID-dsn.algolia.net` endpoint, then it is
	// used first in the list of endpoints to try.
	if t.activeHost != "" &&
		!strings.Contains(t.activeHost, "-dsn.algolia.net") &&
		time.Now().Sub(t.activeSince) <= t.keepAliveDuration {
		hosts = []string{t.activeHost}
	}

	// Whether or not an active host is first in the list of hosts to try, the
	// following entries in the list is the DSN endpoint followed by the
	// `APPID-{1,2,3}.algolianet.com` endpoints (shuffled).
	hosts = append(hosts, t.appId+"-dsn.algolia.net")
	hosts = append(hosts, t.hosts...)
	return hosts

}

// tryRequest is the underlying method which actually performs the request. It
// returns the response as a byte slice or a non-nil error if anything went
// wrong.
func (t *Transport) tryRequest(method, host, path string, body interface{}) ([]byte, error) {
	// Build the request
	req, err := t.buildRequest(method, host, path, body)
	if err != nil {
		return nil, err
	}

	// Perform the request
	res, err := t.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("Cannot perform request [%s] %s (%s): %s", method, path, host, err)
	}

	// Read response's body
	bodyRes, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, fmt.Errorf("Cannot read response body: %s", err)
	}
	res.Body.Close()

	// Return the body as an error if the status code is not 2XX
	code := res.StatusCode
	if !(200 <= code && code < 300) {
		return nil, errors.New(string(bodyRes))
	}

	return bodyRes, nil
}

// buildRequest returns a valid `http.Request` with the headers and body (if
// any) correctly set. The return error is non-nil if the request is invalid or
// if the body, if non-nil, is not a valid JSON.
func (t *Transport) buildRequest(method, host, path string, body interface{}) (*http.Request, error) {
	var req *http.Request
	var err error
	urlStr := "https://" + host + path

	if body == nil {
		// As the body is nil, an empty body request is instantiated
		req, err = http.NewRequest(method, urlStr, nil)
		if err != nil {
			return nil, fmt.Errorf("Cannot instantiate request: [%s] %s", method, urlStr)
		}
	} else {
		// As the body is non-nil, the content is read
		data, err := json.Marshal(body)
		if err != nil {
			return nil, errors.New("Invalid JSON in the query")
		}
		reader := bytes.NewReader(data)

		// The request is then instantiated with the body content
		req, err = http.NewRequest(method, urlStr, reader)
		if err != nil {
			return nil, fmt.Errorf("Cannot instantiate request: [%s] %s", method, urlStr)
		}

		// Add content specific headers
		req.Header.Add("Content-Length", strconv.Itoa(len(string(data))))
		req.Header.Add("Content-Type", "application/json; charset=utf-8")
	}

	// Add default and Algolia specific headers
	addHeaders(req, t.headers)

	if strings.Contains(path, "/*/") {
		req.URL = &url.URL{
			Scheme: "https",
			Host:   host,
			Opaque: "//" + host + path, //Remove url encoding
		}
	}

	return req, nil
}

// resetDialTimeout increases the `Timeout` value of the underlying dialer by 1
// second.
func (t *Transport) increaseDialTimeout() {
	t.dialTimeout = t.dialTimeout + time.Second
	t.httpClient.Transport = defaultTransport(t.dialTimeout)
}

// resetDialTimeout resets the `Timeout` value of the underlying dialer to 1
// second.
func (t *Transport) resetDialTimeout() {
	t.dialTimeout = 1 * time.Second
	t.httpClient.Transport = defaultTransport(t.dialTimeout)
}
<|endoftext|>


<commit_before>package master

import (
	"errors"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/PaddlePaddle/recordio"
)

// Service is the master server service.
type Service struct {
	chunksPerTask int
	timeoutDur    time.Duration
	timeoutMax    int
	ready         chan struct{}

	mu         sync.Mutex
	initDone   bool
	taskQueues taskQueues
}

// Recover recovers service state from etcd.
func Recover() (*Service, error) {
	// TODO(helin): recover from snapshot state from etcd.
	return nil, nil
}

func partition(chunks []Chunk, chunksPerTask int) []taskEntry {
	id := 0
	if chunksPerTask <= 0 {
		chunksPerTask = 1
	}

	var result []taskEntry
	var cur taskEntry
	for i, c := range chunks {
		if i%chunksPerTask == 0 && len(cur.Task.Chunks) > 0 {
			cur.Task.ID = id
			id++
			result = append(result, cur)
			cur.Task.Chunks = nil
		}

		cur.Task.Chunks = append(cur.Task.Chunks, c)
	}

	if len(cur.Task.Chunks) > 0 {
		cur.Task.ID = id
		result = append(result, cur)
	}

	return result
}

// NewService creates a new service.
func NewService(chunksPerTask int, timeoutDur time.Duration, timeoutMax int) *Service {
	s := &Service{}
	s.chunksPerTask = chunksPerTask
	s.timeoutDur = timeoutDur
	s.timeoutMax = timeoutMax
	s.taskQueues = taskQueues{}
	s.taskQueues.Pending = make(map[int]taskEntry)
	s.ready = make(chan struct{})
	return s
}

// Chunk is a chunk of data consisted of several data instances.
type Chunk struct {
	Path  string
	Index recordio.Index // chunk index
}

// Task is the basic unit of data instances assigned to trainers.
type Task struct {
	ID     int
	Chunks []Chunk
}

type taskEntry struct {
	Epoch      int
	NumTimeout int
	Task       Task
}

type taskQueues struct {
	Todo    []taskEntry
	Pending map[int]taskEntry // map from task ID to task entry
	Done    []taskEntry
	Failed  []Task
}

// *must* be called with s.mu being held.
func (s *Service) snapshot() error {
	// TODO(helin): snapshot state on etcd.
	return nil
}

func getChunks(globPaths []string) ([]Chunk, error) {
	var chunks []Chunk
	var paths []string

	for _, s := range globPaths {
		match, err := filepath.Glob(s)
		if err != nil {
			return nil, err
		}
		paths = append(paths, match...)
	}

	if len(paths) == 0 {
		return nil, errors.New("no valid datset specified")
	}

	for _, path := range paths {
		f, err := os.Open(path)
		if err != nil {
			return nil, err
		}

		index, err := recordio.LoadIndex(f)
		if err != nil {
			return nil, err
		}
		err = f.Close()
		if err != nil {
			return nil, err
		}

		count := index.NumChunks()
		for i := 0; i < count; i++ {
			chunk := Chunk{
				Path:  path,
				Index: *index.ChunkIndex(i),
			}
			chunks = append(chunks, chunk)
		}
	}

	return chunks, nil
}

// SetDataset sets dataset to dispatch for the master server.
//
// SetDataset can be call multiple times. But only the first call will
// be honored.
func (s *Service) SetDataset(globPaths []string, dummy *int) error {
	if len(globPaths) == 0 {
		return errors.New("no dataset specified")
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	if s.initDone {
		// Already initialized. All trainer will call
		// SetDataset, but we only handle the first one. Treat
		// other calls as successful but do nothing.
		return nil
	}

	chunks, err := getChunks(globPaths)
	if err != nil {
		return err
	}

	s.taskQueues.Todo = partition(chunks, s.chunksPerTask)

	err = s.snapshot()
	if err != nil {
		return err
	}

	close(s.ready)
	s.initDone = true
	return nil
}

// GetTask gets a new task from the service.
func (s *Service) GetTask(dummy int, task *Task) error {
	select {
	case <-s.ready:
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if len(s.taskQueues.Todo) == 0 {
		if len(s.taskQueues.Done) == 0 {
			if len(s.taskQueues.Pending) == 0 {
				return errors.New("all task failed")
			}

			// TODO(helin): client need to retry in this
			// error case. Gotcha: RPC client can't
			// compare returned error with predefined
			// errors like io.EOF. Because interface don't
			// have same dynamic value when in different
			// process.
			return errors.New("no more available task")
		}
		s.taskQueues.Todo = s.taskQueues.Done
		s.taskQueues.Todo = nil
	}

	t := s.taskQueues.Todo[0]
	t.Epoch++
	s.taskQueues.Todo = s.taskQueues.Todo[1:]
	s.taskQueues.Pending[t.Task.ID] = t
	err := s.snapshot()
	if err != nil {
		return err
	}

	*task = t.Task

	time.AfterFunc(s.timeoutDur, func(taskID int, epoch int) func() {
		return func() {
			s.mu.Lock()
			defer s.mu.Unlock()

			t, ok := s.taskQueues.Pending[taskID]
			if !ok {
				return
			}

			if t.Epoch != epoch {
				// new epoch, task launched after the
				// schedule of this timeout check.
				return
			}

			defer func() {
				err := s.snapshot()
				if err != nil {
					log.Println(err)
				}
			}()

			delete(s.taskQueues.Pending, t.Task.ID)

			t.NumTimeout++
			if t.NumTimeout > s.timeoutMax {
				s.taskQueues.Failed = append(s.taskQueues.Failed, t.Task)
				return
			}

			s.taskQueues.Todo = append(s.taskQueues.Todo, t)
		}
	}(t.Task.ID, t.Epoch))
	return nil
}

// TaskFinished tell the service that a task is finished.
func (s *Service) TaskFinished(taskID int, dummy *int) error {
	select {
	case <-s.ready:
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	t, ok := s.taskQueues.Pending[taskID]
	if !ok {
		return errors.New("pending task not found")
	}

	// task finished, reset timeout
	t.NumTimeout = 0
	s.taskQueues.Done = append(s.taskQueues.Done, t)
	delete(s.taskQueues.Pending, taskID)

	if len(s.taskQueues.Pending) == 0 {
		s.taskQueues.Todo = append(s.taskQueues.Todo, s.taskQueues.Done...)
		s.taskQueues.Done = nil
	}

	return s.snapshot()
}
<commit_msg>rename getChunks to readChunks<commit_after>package master

import (
	"errors"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/PaddlePaddle/recordio"
)

// Service is the master server service.
type Service struct {
	chunksPerTask int
	timeoutDur    time.Duration
	timeoutMax    int
	ready         chan struct{}

	mu         sync.Mutex
	initDone   bool
	taskQueues taskQueues
}

// Recover recovers service state from etcd.
func Recover() (*Service, error) {
	// TODO(helin): recover from snapshot state from etcd.
	return nil, nil
}

func partition(chunks []Chunk, chunksPerTask int) []taskEntry {
	id := 0
	if chunksPerTask <= 0 {
		chunksPerTask = 1
	}

	var result []taskEntry
	var cur taskEntry
	for i, c := range chunks {
		if i%chunksPerTask == 0 && len(cur.Task.Chunks) > 0 {
			cur.Task.ID = id
			id++
			result = append(result, cur)
			cur.Task.Chunks = nil
		}

		cur.Task.Chunks = append(cur.Task.Chunks, c)
	}

	if len(cur.Task.Chunks) > 0 {
		cur.Task.ID = id
		result = append(result, cur)
	}

	return result
}

// NewService creates a new service.
func NewService(chunksPerTask int, timeoutDur time.Duration, timeoutMax int) *Service {
	s := &Service{}
	s.chunksPerTask = chunksPerTask
	s.timeoutDur = timeoutDur
	s.timeoutMax = timeoutMax
	s.taskQueues = taskQueues{}
	s.taskQueues.Pending = make(map[int]taskEntry)
	s.ready = make(chan struct{})
	return s
}

// Chunk is a chunk of data consisted of several data instances.
type Chunk struct {
	Path  string
	Index recordio.Index // chunk index
}

// Task is the basic unit of data instances assigned to trainers.
type Task struct {
	ID     int
	Chunks []Chunk
}

type taskEntry struct {
	Epoch      int
	NumTimeout int
	Task       Task
}

type taskQueues struct {
	Todo    []taskEntry
	Pending map[int]taskEntry // map from task ID to task entry
	Done    []taskEntry
	Failed  []Task
}

// *must* be called with s.mu being held.
func (s *Service) snapshot() error {
	// TODO(helin): snapshot state on etcd.
	return nil
}

func readChunks(globPaths []string) ([]Chunk, error) {
	var chunks []Chunk
	var paths []string

	for _, s := range globPaths {
		match, err := filepath.Glob(s)
		if err != nil {
			return nil, err
		}
		paths = append(paths, match...)
	}

	if len(paths) == 0 {
		return nil, errors.New("no valid datset specified")
	}

	for _, path := range paths {
		f, err := os.Open(path)
		if err != nil {
			return nil, err
		}

		index, err := recordio.LoadIndex(f)
		if err != nil {
			return nil, err
		}
		err = f.Close()
		if err != nil {
			return nil, err
		}

		count := index.NumChunks()
		for i := 0; i < count; i++ {
			chunk := Chunk{
				Path:  path,
				Index: *index.ChunkIndex(i),
			}
			chunks = append(chunks, chunk)
		}
	}

	return chunks, nil
}

// SetDataset sets dataset to dispatch for the master server.
//
// SetDataset can be call multiple times. But only the first call will
// be honored.
func (s *Service) SetDataset(globPaths []string, dummy *int) error {
	if len(globPaths) == 0 {
		return errors.New("no dataset specified")
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	if s.initDone {
		// Already initialized. All trainer will call
		// SetDataset, but we only handle the first one. Treat
		// other calls as successful but do nothing.
		return nil
	}

	chunks, err := readChunks(globPaths)
	if err != nil {
		return err
	}

	s.taskQueues.Todo = partition(chunks, s.chunksPerTask)

	err = s.snapshot()
	if err != nil {
		return err
	}

	close(s.ready)
	s.initDone = true
	return nil
}

// GetTask gets a new task from the service.
func (s *Service) GetTask(dummy int, task *Task) error {
	select {
	case <-s.ready:
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if len(s.taskQueues.Todo) == 0 {
		if len(s.taskQueues.Done) == 0 {
			if len(s.taskQueues.Pending) == 0 {
				return errors.New("all task failed")
			}

			// TODO(helin): client need to retry in this
			// error case. Gotcha: RPC client can't
			// compare returned error with predefined
			// errors like io.EOF. Because interface don't
			// have same dynamic value when in different
			// process.
			return errors.New("no more available task")
		}
		s.taskQueues.Todo = s.taskQueues.Done
		s.taskQueues.Todo = nil
	}

	t := s.taskQueues.Todo[0]
	t.Epoch++
	s.taskQueues.Todo = s.taskQueues.Todo[1:]
	s.taskQueues.Pending[t.Task.ID] = t
	err := s.snapshot()
	if err != nil {
		return err
	}

	*task = t.Task

	time.AfterFunc(s.timeoutDur, func(taskID int, epoch int) func() {
		return func() {
			s.mu.Lock()
			defer s.mu.Unlock()

			t, ok := s.taskQueues.Pending[taskID]
			if !ok {
				return
			}

			if t.Epoch != epoch {
				// new epoch, task launched after the
				// schedule of this timeout check.
				return
			}

			defer func() {
				err := s.snapshot()
				if err != nil {
					log.Println(err)
				}
			}()

			delete(s.taskQueues.Pending, t.Task.ID)

			t.NumTimeout++
			if t.NumTimeout > s.timeoutMax {
				s.taskQueues.Failed = append(s.taskQueues.Failed, t.Task)
				return
			}

			s.taskQueues.Todo = append(s.taskQueues.Todo, t)
		}
	}(t.Task.ID, t.Epoch))
	return nil
}

// TaskFinished tell the service that a task is finished.
func (s *Service) TaskFinished(taskID int, dummy *int) error {
	select {
	case <-s.ready:
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	t, ok := s.taskQueues.Pending[taskID]
	if !ok {
		return errors.New("pending task not found")
	}

	// task finished, reset timeout
	t.NumTimeout = 0
	s.taskQueues.Done = append(s.taskQueues.Done, t)
	delete(s.taskQueues.Pending, taskID)

	if len(s.taskQueues.Pending) == 0 {
		s.taskQueues.Todo = append(s.taskQueues.Todo, s.taskQueues.Done...)
		s.taskQueues.Done = nil
	}

	return s.snapshot()
}
<|endoftext|>


<commit_before>package storage

import (
	"code.google.com/p/weed-fs/go/util"
	"encoding/hex"
	"errors"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"path"
	"strconv"
	"strings"
	"time"
)

const (
	NeedleHeaderSize   = 16 //should never change this
	NeedlePaddingSize  = 8
	NeedleChecksumSize = 4
)

type Needle struct {
	Cookie uint32 `comment:"random number to mitigate brute force lookups"`
	Id     uint64 `comment:"needle id"`
	Size   uint32 `comment:"sum of DataSize,Data,NameSize,Name,MimeSize,Mime"`

	DataSize     uint32 `comment:"Data size"` //version2
	Data         []byte `comment:"The actual file data"`
	Flags        byte   `comment:"boolean flags"` //version2
	NameSize     uint8  //version2
	Name         []byte `comment:"maximum 256 characters"` //version2
	MimeSize     uint8  //version2
	Mime         []byte `comment:"maximum 256 characters"` //version2
	LastModified uint64 //only store LastModifiedBytesLength bytes, which is 5 bytes to disk

	Checksum CRC    `comment:"CRC32 to check integrity"`
	Padding  []byte `comment:"Aligned to 8 bytes"`
}

func NewNeedle(r *http.Request) (n *Needle, e error) {

	n = new(Needle)
	form, fe := r.MultipartReader()
	if fe != nil {
		log.Println("MultipartReader [ERROR]", fe)
		e = fe
		return
	}
	part, fe := form.NextPart()
	if fe != nil {
		log.Println("Reading Multi part [ERROR]", fe)
		e = fe
		return
	}
	fname := part.FileName()
	if fname != "" {
		fname = path.Base(part.FileName())
	} else {
		e = errors.New("No file found!")
		return
	}
	data, ioe := ioutil.ReadAll(part)
	if ioe != nil {
		e = ioe
		log.Println("Reading Content [ERROR]", ioe)
		return
	}
	dotIndex := strings.LastIndex(fname, ".")
	ext, mtype := "", ""
	if dotIndex > 0 {
		ext = fname[dotIndex:]
		mtype = mime.TypeByExtension(ext)
	}
	contentType := part.Header.Get("Content-Type")
	if contentType != "" && mtype != contentType && len(contentType) < 256 {
		n.Mime = []byte(contentType)
		n.SetHasMime()
		mtype = contentType
	}
	if part.Header.Get("Content-Encoding") == "gzip" {
		n.SetGzipped()
	} else if IsGzippable(ext, mtype) {
		if data, e = GzipData(data); e != nil {
			return
		}
		n.SetGzipped()
	}
	if ext == ".gz" {
		n.SetGzipped()
	}
	if len(fname) < 256 {
		if strings.HasSuffix(fname, ".gz") {
			n.Name = []byte(fname[:len(fname)-3])
		} else {
			n.Name = []byte(fname)
		}
		n.SetHasName()
	}
	var parseError error
	if n.LastModified, parseError = strconv.ParseUint(r.FormValue("ts"), 10, 64); parseError != nil {
		n.LastModified = uint64(time.Now().Unix())
	}
	n.SetHasLastModifiedDate()

	n.Data = data
	n.Checksum = NewCRC(data)

	commaSep := strings.LastIndex(r.URL.Path, ",")
	dotSep := strings.LastIndex(r.URL.Path, ".")
	fid := r.URL.Path[commaSep+1:]
	if dotSep > 0 {
		fid = r.URL.Path[commaSep+1 : dotSep]
	}

	n.ParsePath(fid)

	return
}
func (n *Needle) ParsePath(fid string) {
	length := len(fid)
	if length <= 8 {
		if length > 0 {
			log.Println("Invalid fid", fid, "length", length)
		}
		return
	}
	delta := ""
	deltaIndex := strings.LastIndex(fid, "_")
	if deltaIndex > 0 {
		fid, delta = fid[0:deltaIndex], fid[deltaIndex+1:]
	}
	n.Id, n.Cookie = ParseKeyHash(fid)
	if delta != "" {
		d, e := strconv.ParseUint(delta, 10, 64)
		if e == nil {
			n.Id += d
		}
	}
}

func ParseKeyHash(key_hash_string string) (uint64, uint32) {
	key_hash_bytes, khe := hex.DecodeString(key_hash_string)
	key_hash_len := len(key_hash_bytes)
	if khe != nil || key_hash_len <= 4 {
		log.Println("Invalid key_hash", key_hash_string, "length:", key_hash_len, "error", khe)
		return 0, 0
	}
	key := util.BytesToUint64(key_hash_bytes[0 : key_hash_len-4])
	hash := util.BytesToUint32(key_hash_bytes[key_hash_len-4 : key_hash_len])
	return key, hash
}
<commit_msg>refactor "content upload" out of needle creation<commit_after>package storage

import (
	"code.google.com/p/weed-fs/go/util"
	"encoding/hex"
	"errors"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"path"
	"strconv"
	"strings"
	"time"
)

const (
	NeedleHeaderSize   = 16 //should never change this
	NeedlePaddingSize  = 8
	NeedleChecksumSize = 4
)

type Needle struct {
	Cookie uint32 `comment:"random number to mitigate brute force lookups"`
	Id     uint64 `comment:"needle id"`
	Size   uint32 `comment:"sum of DataSize,Data,NameSize,Name,MimeSize,Mime"`

	DataSize     uint32 `comment:"Data size"` //version2
	Data         []byte `comment:"The actual file data"`
	Flags        byte   `comment:"boolean flags"` //version2
	NameSize     uint8  //version2
	Name         []byte `comment:"maximum 256 characters"` //version2
	MimeSize     uint8  //version2
	Mime         []byte `comment:"maximum 256 characters"` //version2
	LastModified uint64 //only store LastModifiedBytesLength bytes, which is 5 bytes to disk

	Checksum CRC    `comment:"CRC32 to check integrity"`
	Padding  []byte `comment:"Aligned to 8 bytes"`
}

func ParseUpload(r *http.Request) (fileName string, data []byte, mimeType string, isGzipped bool, modifiedTime uint64, e error) {
	form, fe := r.MultipartReader()
	if fe != nil {
		log.Println("MultipartReader [ERROR]", fe)
		e = fe
		return
	}
	part, fe := form.NextPart()
	if fe != nil {
		log.Println("Reading Multi part [ERROR]", fe)
		e = fe
		return
	}
	fileName = part.FileName()
	if fileName != "" {
		fileName = path.Base(fileName)
	} else {
		e = errors.New("No file found!")
		return
	}
	data, e = ioutil.ReadAll(part)
	if e != nil {
		log.Println("Reading Content [ERROR]", e)
		return
	}
	dotIndex := strings.LastIndex(fileName, ".")
	ext, mtype := "", ""
	if dotIndex > 0 {
		ext = strings.ToLower(fileName[dotIndex:])
		mtype = mime.TypeByExtension(ext)
	}
	contentType := part.Header.Get("Content-Type")
	if contentType != "" && mtype != contentType {
		mimeType = contentType //only return mime type if not deductable
		mtype = contentType
	}
	if part.Header.Get("Content-Encoding") == "gzip" {
		isGzipped = true
	} else if IsGzippable(ext, mtype) {
		if data, e = GzipData(data); e != nil {
			return
		}
		isGzipped = true
	}
	if ext == ".gz" {
		isGzipped = true
	}
	if strings.HasSuffix(fileName, ".gz") {
		fileName = fileName[:len(fileName)-3]
	}
	modifiedTime, _ = strconv.ParseUint(r.FormValue("ts"), 10, 64)
  return
}
func NewNeedle(r *http.Request) (n *Needle, e error) {
  fname, mimeType, isGzipped := "", "", false
	n = new(Needle)
	fname, n.Data, mimeType, isGzipped, n.LastModified, e = ParseUpload(r)
	if e != nil {
	  return
	}
  if len(fname) < 256 {
    n.Name = []byte(fname)
    n.SetHasName()
  }
	if len(mimeType) < 256 {
		n.Mime = []byte(mimeType)
		n.SetHasMime()
	}
	if isGzipped {
		n.SetGzipped()
	}
	if n.LastModified == 0 {
		n.LastModified = uint64(time.Now().Unix())
    n.SetHasLastModifiedDate()
	}

	n.Checksum = NewCRC(n.Data)

	commaSep := strings.LastIndex(r.URL.Path, ",")
	dotSep := strings.LastIndex(r.URL.Path, ".")
	fid := r.URL.Path[commaSep+1:]
	if dotSep > 0 {
		fid = r.URL.Path[commaSep+1 : dotSep]
	}

	n.ParsePath(fid)

	return
}
func (n *Needle) ParsePath(fid string) {
	length := len(fid)
	if length <= 8 {
		if length > 0 {
			log.Println("Invalid fid", fid, "length", length)
		}
		return
	}
	delta := ""
	deltaIndex := strings.LastIndex(fid, "_")
	if deltaIndex > 0 {
		fid, delta = fid[0:deltaIndex], fid[deltaIndex+1:]
	}
	n.Id, n.Cookie = ParseKeyHash(fid)
	if delta != "" {
		d, e := strconv.ParseUint(delta, 10, 64)
		if e == nil {
			n.Id += d
		}
	}
}

func ParseKeyHash(key_hash_string string) (uint64, uint32) {
	key_hash_bytes, khe := hex.DecodeString(key_hash_string)
	key_hash_len := len(key_hash_bytes)
	if khe != nil || key_hash_len <= 4 {
		log.Println("Invalid key_hash", key_hash_string, "length:", key_hash_len, "error", khe)
		return 0, 0
	}
	key := util.BytesToUint64(key_hash_bytes[0 : key_hash_len-4])
	hash := util.BytesToUint32(key_hash_bytes[key_hash_len-4 : key_hash_len])
	return key, hash
}
<|endoftext|>