<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"
	"reflect"
	"time"

	"github.com/couchbaselabs/bleve/analysis"
	"github.com/couchbaselabs/bleve/document"
)

var tRUE = true

var fALSE = false

var DEFAULT_ID_FIELD = "_id"
var DEFAULT_TYPE_FIELD = "_type"
var DEFAULT_TYPE = "_default"
var DEFAULT_FIELD = "_all"

type IndexMapping struct {
	TypeMapping     map[string]*DocumentMapping `json:"types"`
	DefaultMapping  *DocumentMapping            `json:"default_mapping"`
	IdField         *string                     `json:"id_field"`
	TypeField       *string                     `json:"type_field"`
	DefaultType     *string                     `json:"default_type"`
	DefaultAnalyzer *string                     `json:"default_analyzer"`
	DefaultField    *string                     `json:"default_field"`
}

func (im *IndexMapping) GoString() string {
	return fmt.Sprintf("&bleve.IndexMapping{TypeMapping:%#v, TypeField:%s, DefaultType:%s}", im.TypeMapping, *im.TypeField, *im.DefaultType)
}

func NewIndexMapping() *IndexMapping {
	return &IndexMapping{
		TypeMapping:    make(map[string]*DocumentMapping),
		DefaultMapping: NewDocumentMapping(),
		IdField:        &DEFAULT_ID_FIELD,
		TypeField:      &DEFAULT_TYPE_FIELD,
		DefaultType:    &DEFAULT_TYPE,
		DefaultField:   &DEFAULT_FIELD,
	}
}

func (im *IndexMapping) AddDocumentMapping(doctype string, dm *DocumentMapping) *IndexMapping {
	im.TypeMapping[doctype] = dm
	return im
}

func (im *IndexMapping) SetTypeField(typeField string) *IndexMapping {
	im.TypeField = &typeField
	return im
}

func (im *IndexMapping) SetDefaultAnalyzer(analyzer string) *IndexMapping {
	im.DefaultAnalyzer = &analyzer
	return im
}

func (im *IndexMapping) SetDefaultField(field string) *IndexMapping {
	im.DefaultField = &field
	return im
}

func (im *IndexMapping) MappingForType(docType string) *DocumentMapping {
	docMapping := im.TypeMapping[docType]
	if docMapping == nil {
		docMapping = im.DefaultMapping
	}
	return docMapping
}

func (im *IndexMapping) UnmarshalJSON(data []byte) error {
	var tmp struct {
		TypeMapping     map[string]*DocumentMapping `json:"types"`
		DefaultMapping  *DocumentMapping            `json:"default_mapping"`
		IdField         *string                     `json:"id_field"`
		TypeField       *string                     `json:"type_field"`
		DefaultType     *string                     `json:"default_type"`
		DefaultAnalyzer *string                     `json:"default_analyzer"`
		DefaultField    *string                     `json:"default_field"`
	}
	err := json.Unmarshal(data, &tmp)
	if err != nil {
		return err
	}

	im.IdField = &DEFAULT_ID_FIELD
	if tmp.IdField != nil {
		im.IdField = tmp.IdField
	}

	im.TypeField = &DEFAULT_TYPE_FIELD
	if tmp.TypeField != nil {
		im.TypeField = tmp.TypeField
	}

	im.DefaultType = &DEFAULT_TYPE
	if tmp.DefaultType != nil {
		im.DefaultType = tmp.DefaultType
	}

	im.DefaultMapping = NewDocumentMapping()
	if tmp.DefaultMapping != nil {
		im.DefaultMapping = tmp.DefaultMapping
	}

	if tmp.DefaultAnalyzer != nil {
		im.DefaultAnalyzer = tmp.DefaultAnalyzer
	}

	if tmp.DefaultField != nil {
		im.DefaultField = tmp.DefaultField
	}

	im.TypeMapping = make(map[string]*DocumentMapping, len(tmp.TypeMapping))
	for typeName, typeDocMapping := range tmp.TypeMapping {
		im.TypeMapping[typeName] = typeDocMapping
	}
	return nil
}

func (im *IndexMapping) determineType(data interface{}) (string, bool) {
	// first see if the object implements Identifier
	classifier, ok := data.(Classifier)
	if ok {
		return classifier.Type(), true
	}

	// now see if we can find type using the mapping
	if im.TypeField != nil {
		typ, ok := mustString(lookupPropertyPath(data, *im.TypeField))
		if ok {
			return typ, true
		}
	}

	// fall back to default type if there was one
	if im.DefaultType != nil {
		return *im.DefaultType, true
	}

	return "", false
}

func (im *IndexMapping) MapDocument(doc *document.Document, data interface{}) error {
	docType, ok := im.determineType(data)
	if !ok {
		return ERROR_NO_TYPE
	}
	docMapping := im.MappingForType(docType)
	walkContext := newWalkContext(doc, docMapping)
	im.walkDocument(data, []string{}, walkContext)

	// see if the _all field was disabled
	allMapping := docMapping.DocumentMappingForPath("_all")
	if allMapping == nil || (allMapping.Enabled != nil && *allMapping.Enabled != false) {
		field := document.NewCompositeFieldWithIndexingOptions("_all", true, []string{}, walkContext.excludedFromAll, document.INDEX_FIELD|document.INCLUDE_TERM_VECTORS)
		doc.AddField(field)
	}

	return nil
}

type walkContext struct {
	doc             *document.Document
	dm              *DocumentMapping
	excludedFromAll []string
}

func newWalkContext(doc *document.Document, dm *DocumentMapping) *walkContext {
	return &walkContext{
		doc:             doc,
		dm:              dm,
		excludedFromAll: []string{},
	}
}

func (im *IndexMapping) walkDocument(data interface{}, path []string, context *walkContext) {
	val := reflect.ValueOf(data)
	typ := val.Type()
	switch typ.Kind() {
	case reflect.Map:
		// FIXME can add support for other map keys in the future
		if typ.Key().Kind() == reflect.String {
			for _, key := range val.MapKeys() {
				fieldName := key.String()
				fieldVal := val.MapIndex(key).Interface()
				im.processProperty(fieldVal, append(path, fieldName), context)
			}
		}
	case reflect.Struct:
		for i := 0; i < val.NumField(); i++ {
			field := typ.Field(i)
			fieldName := field.Name

			// if the field has a JSON name, prefer that
			jsonTag := field.Tag.Get("json")
			jsonFieldName := parseJSONTagName(jsonTag)
			if jsonFieldName != "" {
				fieldName = jsonFieldName
			}

			if val.Field(i).CanInterface() {
				fieldVal := val.Field(i).Interface()
				im.processProperty(fieldVal, append(path, fieldName), context)
			}
		}
	case reflect.Slice, reflect.Array:
		for i := 0; i < val.Len(); i++ {
			if val.Index(i).CanInterface() {
				fieldVal := val.Index(i).Interface()
				im.processProperty(fieldVal, path, context)
			}
		}
	case reflect.Ptr:
		ptrElem := val.Elem()
		if ptrElem.CanInterface() {
			im.walkDocument(ptrElem.Interface(), path, context)
		}
	}
}

func (im *IndexMapping) processProperty(property interface{}, path []string, context *walkContext) {
	pathString := encodePath(path)
	// look to see if there is a mapping for this field
	subDocMapping := context.dm.DocumentMappingForPath(pathString)

	// check tos see if we even need to do further processing
	if subDocMapping != nil && subDocMapping.Enabled != nil && !*subDocMapping.Enabled {
		return
	}

	propertyValue := reflect.ValueOf(property)
	propertyType := propertyValue.Type()
	switch propertyType.Kind() {
	case reflect.String:
		propertyValueString := propertyValue.String()
		if subDocMapping != nil {
			// index by explicit mapping
			for _, fieldMapping := range subDocMapping.Fields {
				fieldName := getFieldName(pathString, path, fieldMapping)
				if *fieldMapping.Type == "text" {
					options := fieldMapping.Options()
					analyzer := Config.Analysis.Analyzers[*fieldMapping.Analyzer]
					if analyzer != nil {
						field := document.NewTextFieldCustom(fieldName, []byte(propertyValueString), options, analyzer)
						context.doc.AddField(field)

						if fieldMapping.IncludeInAll != nil && !*fieldMapping.IncludeInAll {
							context.excludedFromAll = append(context.excludedFromAll, fieldName)
						}
					}
				} else if *fieldMapping.Type == "datetime" {
					options := fieldMapping.Options()
					dateTimeFormat := *Config.DefaultDateTimeFormat
					if fieldMapping.DateFormat != nil {
						dateTimeFormat = *fieldMapping.DateFormat
					}
					dateTimeParser := Config.Analysis.DateTimeParsers[dateTimeFormat]
					parsedDateTime, err := dateTimeParser.ParseDateTime(propertyValueString)
					if err != nil {
						field := document.NewDateTimeFieldWithIndexingOptions(fieldName, parsedDateTime, options)
						context.doc.AddField(field)
					}
				}
			}
		} else {
			// automatic indexing behavior

			// first see if it can be parsed by the default date parser
			// FIXME add support for index mapping overriding defaults
			dateTimeParser := Config.Analysis.DateTimeParsers[*Config.DefaultDateTimeFormat]
			parsedDateTime, err := dateTimeParser.ParseDateTime(propertyValueString)
			if err != nil {
				// index as plain text
				options := document.STORE_FIELD | document.INDEX_FIELD | document.INCLUDE_TERM_VECTORS
				analyzer := im.defaultAnalyzer(context.dm, path)
				field := document.NewTextFieldCustom(pathString, []byte(propertyValueString), options, analyzer)
				context.doc.AddField(field)
			} else {
				// index as datetime
				field := document.NewDateTimeField(pathString, parsedDateTime)
				context.doc.AddField(field)
			}
		}
	case reflect.Float64:
		propertyValFloat := propertyValue.Float()
		if subDocMapping != nil {
			// index by explicit mapping
			for _, fieldMapping := range subDocMapping.Fields {
				fieldName := getFieldName(pathString, path, fieldMapping)
				if *fieldMapping.Type == "number" {
					options := fieldMapping.Options()
					field := document.NewNumericFieldWithIndexingOptions(fieldName, propertyValFloat, options)
					context.doc.AddField(field)
				}
			}
		} else {
			// automatic indexing behavior
			field := document.NewNumericField(pathString, propertyValFloat)
			context.doc.AddField(field)
		}
	case reflect.Struct:
		switch property := property.(type) {
		case time.Time:
			// don't descend into the time struct
			if subDocMapping != nil {
				// index by explicit mapping
				for _, fieldMapping := range subDocMapping.Fields {
					fieldName := getFieldName(pathString, path, fieldMapping)
					if *fieldMapping.Type == "datetime" {
						options := fieldMapping.Options()
						field := document.NewDateTimeFieldWithIndexingOptions(fieldName, property, options)
						context.doc.AddField(field)
					}
				}
			} else {
				// automatic indexing behavior
				field := document.NewDateTimeField(pathString, property)
				context.doc.AddField(field)
			}

		default:
			im.walkDocument(property, path, context)
		}
	default:
		im.walkDocument(property, path, context)
	}
}

func (im *IndexMapping) defaultAnalyzer(dm *DocumentMapping, path []string) *analysis.Analyzer {
	// first see if the document mapping has an analyzer
	rv := dm.defaultAnalyzer(path)
	if rv == nil {
		if im.DefaultAnalyzer != nil {
			rv = Config.Analysis.Analyzers[*im.DefaultAnalyzer]
		} else if Config.DefaultAnalyzer != nil {
			rv = Config.Analysis.Analyzers[*Config.DefaultAnalyzer]
		}
	}
	return rv
}

// attempts to find the best analyzer to use with only a field name
// will walk all the document types, look for field mappings at the
// provided path, if one exists and it has an explicit analyzer
// that is returned
// nil should be an acceptable return value meaning we don't know
func (im *IndexMapping) analyzerForPath(path string) *analysis.Analyzer {

	// first we look for explicit mapping on the field
	for _, docMapping := range im.TypeMapping {
		pathMapping := docMapping.DocumentMappingForPath(path)
		if pathMapping != nil {
			if len(pathMapping.Fields) > 0 {
				if pathMapping.Fields[0].Analyzer != nil {
					return Config.Analysis.Analyzers[*pathMapping.Fields[0].Analyzer]
				}
			}
		}
	}

	// next we will try default analyzers for the path
	for _, docMapping := range im.TypeMapping {
		rv := im.defaultAnalyzer(docMapping, decodePath(path))
		if rv != nil {
			return rv
		}
	}

	// finally just return the system-wide default analyzer
	return Config.Analysis.Analyzers[*Config.DefaultAnalyzer]
}

func (im *IndexMapping) datetimeParserForPath(path string) analysis.DateTimeParser {

	// first we look for explicit mapping on the field
	for _, docMapping := range im.TypeMapping {
		pathMapping := docMapping.DocumentMappingForPath(path)
		if pathMapping != nil {
			if len(pathMapping.Fields) > 0 {
				if pathMapping.Fields[0].Analyzer != nil {
					return Config.Analysis.DateTimeParsers[*pathMapping.Fields[0].DateFormat]
				}
			}
		}
	}

	// next we will try default analyzers for the path
	// FIXME introduce default date time parsers at mapping leves

	// finally just return the system-wide default analyzer
	return Config.Analysis.DateTimeParsers[*Config.DefaultDateTimeFormat]
}

func (im *IndexMapping) defaultField() string {
	if im.DefaultField != nil {
		return *im.DefaultField
	} else if Config.DefaultField != nil {
		return *Config.DefaultField
	}
	return ""
}

func getFieldName(pathString string, path []string, fieldMapping *FieldMapping) string {
	fieldName := pathString
	if fieldMapping.Name != nil && *fieldMapping.Name != "" {
		parentName := ""
		if len(path) > 1 {
			parentName = encodePath(path[:len(path)-1]) + PATH_SEPARATOR
		}
		fieldName = parentName + *fieldMapping.Name
	}
	return fieldName
}
<commit_msg>fix broken test<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"
	"reflect"
	"time"

	"github.com/couchbaselabs/bleve/analysis"
	"github.com/couchbaselabs/bleve/document"
)

var tRUE = true

var fALSE = false

var DEFAULT_ID_FIELD = "_id"
var DEFAULT_TYPE_FIELD = "_type"
var DEFAULT_TYPE = "_default"
var DEFAULT_FIELD = "_all"

type IndexMapping struct {
	TypeMapping     map[string]*DocumentMapping `json:"types"`
	DefaultMapping  *DocumentMapping            `json:"default_mapping"`
	IdField         *string                     `json:"id_field"`
	TypeField       *string                     `json:"type_field"`
	DefaultType     *string                     `json:"default_type"`
	DefaultAnalyzer *string                     `json:"default_analyzer"`
	DefaultField    *string                     `json:"default_field"`
}

func (im *IndexMapping) GoString() string {
	return fmt.Sprintf("&bleve.IndexMapping{TypeMapping:%#v, TypeField:%s, DefaultType:%s}", im.TypeMapping, *im.TypeField, *im.DefaultType)
}

func NewIndexMapping() *IndexMapping {
	return &IndexMapping{
		TypeMapping:    make(map[string]*DocumentMapping),
		DefaultMapping: NewDocumentMapping(),
		IdField:        &DEFAULT_ID_FIELD,
		TypeField:      &DEFAULT_TYPE_FIELD,
		DefaultType:    &DEFAULT_TYPE,
		DefaultField:   &DEFAULT_FIELD,
	}
}

func (im *IndexMapping) AddDocumentMapping(doctype string, dm *DocumentMapping) *IndexMapping {
	im.TypeMapping[doctype] = dm
	return im
}

func (im *IndexMapping) SetTypeField(typeField string) *IndexMapping {
	im.TypeField = &typeField
	return im
}

func (im *IndexMapping) SetDefaultAnalyzer(analyzer string) *IndexMapping {
	im.DefaultAnalyzer = &analyzer
	return im
}

func (im *IndexMapping) SetDefaultField(field string) *IndexMapping {
	im.DefaultField = &field
	return im
}

func (im *IndexMapping) MappingForType(docType string) *DocumentMapping {
	docMapping := im.TypeMapping[docType]
	if docMapping == nil {
		docMapping = im.DefaultMapping
	}
	return docMapping
}

func (im *IndexMapping) UnmarshalJSON(data []byte) error {
	var tmp struct {
		TypeMapping     map[string]*DocumentMapping `json:"types"`
		DefaultMapping  *DocumentMapping            `json:"default_mapping"`
		IdField         *string                     `json:"id_field"`
		TypeField       *string                     `json:"type_field"`
		DefaultType     *string                     `json:"default_type"`
		DefaultAnalyzer *string                     `json:"default_analyzer"`
		DefaultField    *string                     `json:"default_field"`
	}
	err := json.Unmarshal(data, &tmp)
	if err != nil {
		return err
	}

	im.IdField = &DEFAULT_ID_FIELD
	if tmp.IdField != nil {
		im.IdField = tmp.IdField
	}

	im.TypeField = &DEFAULT_TYPE_FIELD
	if tmp.TypeField != nil {
		im.TypeField = tmp.TypeField
	}

	im.DefaultType = &DEFAULT_TYPE
	if tmp.DefaultType != nil {
		im.DefaultType = tmp.DefaultType
	}

	im.DefaultMapping = NewDocumentMapping()
	if tmp.DefaultMapping != nil {
		im.DefaultMapping = tmp.DefaultMapping
	}

	if tmp.DefaultAnalyzer != nil {
		im.DefaultAnalyzer = tmp.DefaultAnalyzer
	}

	im.DefaultField = &DEFAULT_FIELD
	if tmp.DefaultField != nil {
		im.DefaultField = tmp.DefaultField
	}

	im.TypeMapping = make(map[string]*DocumentMapping, len(tmp.TypeMapping))
	for typeName, typeDocMapping := range tmp.TypeMapping {
		im.TypeMapping[typeName] = typeDocMapping
	}
	return nil
}

func (im *IndexMapping) determineType(data interface{}) (string, bool) {
	// first see if the object implements Identifier
	classifier, ok := data.(Classifier)
	if ok {
		return classifier.Type(), true
	}

	// now see if we can find type using the mapping
	if im.TypeField != nil {
		typ, ok := mustString(lookupPropertyPath(data, *im.TypeField))
		if ok {
			return typ, true
		}
	}

	// fall back to default type if there was one
	if im.DefaultType != nil {
		return *im.DefaultType, true
	}

	return "", false
}

func (im *IndexMapping) MapDocument(doc *document.Document, data interface{}) error {
	docType, ok := im.determineType(data)
	if !ok {
		return ERROR_NO_TYPE
	}
	docMapping := im.MappingForType(docType)
	walkContext := newWalkContext(doc, docMapping)
	im.walkDocument(data, []string{}, walkContext)

	// see if the _all field was disabled
	allMapping := docMapping.DocumentMappingForPath("_all")
	if allMapping == nil || (allMapping.Enabled != nil && *allMapping.Enabled != false) {
		field := document.NewCompositeFieldWithIndexingOptions("_all", true, []string{}, walkContext.excludedFromAll, document.INDEX_FIELD|document.INCLUDE_TERM_VECTORS)
		doc.AddField(field)
	}

	return nil
}

type walkContext struct {
	doc             *document.Document
	dm              *DocumentMapping
	excludedFromAll []string
}

func newWalkContext(doc *document.Document, dm *DocumentMapping) *walkContext {
	return &walkContext{
		doc:             doc,
		dm:              dm,
		excludedFromAll: []string{},
	}
}

func (im *IndexMapping) walkDocument(data interface{}, path []string, context *walkContext) {
	val := reflect.ValueOf(data)
	typ := val.Type()
	switch typ.Kind() {
	case reflect.Map:
		// FIXME can add support for other map keys in the future
		if typ.Key().Kind() == reflect.String {
			for _, key := range val.MapKeys() {
				fieldName := key.String()
				fieldVal := val.MapIndex(key).Interface()
				im.processProperty(fieldVal, append(path, fieldName), context)
			}
		}
	case reflect.Struct:
		for i := 0; i < val.NumField(); i++ {
			field := typ.Field(i)
			fieldName := field.Name

			// if the field has a JSON name, prefer that
			jsonTag := field.Tag.Get("json")
			jsonFieldName := parseJSONTagName(jsonTag)
			if jsonFieldName != "" {
				fieldName = jsonFieldName
			}

			if val.Field(i).CanInterface() {
				fieldVal := val.Field(i).Interface()
				im.processProperty(fieldVal, append(path, fieldName), context)
			}
		}
	case reflect.Slice, reflect.Array:
		for i := 0; i < val.Len(); i++ {
			if val.Index(i).CanInterface() {
				fieldVal := val.Index(i).Interface()
				im.processProperty(fieldVal, path, context)
			}
		}
	case reflect.Ptr:
		ptrElem := val.Elem()
		if ptrElem.CanInterface() {
			im.walkDocument(ptrElem.Interface(), path, context)
		}
	}
}

func (im *IndexMapping) processProperty(property interface{}, path []string, context *walkContext) {
	pathString := encodePath(path)
	// look to see if there is a mapping for this field
	subDocMapping := context.dm.DocumentMappingForPath(pathString)

	// check tos see if we even need to do further processing
	if subDocMapping != nil && subDocMapping.Enabled != nil && !*subDocMapping.Enabled {
		return
	}

	propertyValue := reflect.ValueOf(property)
	propertyType := propertyValue.Type()
	switch propertyType.Kind() {
	case reflect.String:
		propertyValueString := propertyValue.String()
		if subDocMapping != nil {
			// index by explicit mapping
			for _, fieldMapping := range subDocMapping.Fields {
				fieldName := getFieldName(pathString, path, fieldMapping)
				if *fieldMapping.Type == "text" {
					options := fieldMapping.Options()
					analyzer := Config.Analysis.Analyzers[*fieldMapping.Analyzer]
					if analyzer != nil {
						field := document.NewTextFieldCustom(fieldName, []byte(propertyValueString), options, analyzer)
						context.doc.AddField(field)

						if fieldMapping.IncludeInAll != nil && !*fieldMapping.IncludeInAll {
							context.excludedFromAll = append(context.excludedFromAll, fieldName)
						}
					}
				} else if *fieldMapping.Type == "datetime" {
					options := fieldMapping.Options()
					dateTimeFormat := *Config.DefaultDateTimeFormat
					if fieldMapping.DateFormat != nil {
						dateTimeFormat = *fieldMapping.DateFormat
					}
					dateTimeParser := Config.Analysis.DateTimeParsers[dateTimeFormat]
					parsedDateTime, err := dateTimeParser.ParseDateTime(propertyValueString)
					if err != nil {
						field := document.NewDateTimeFieldWithIndexingOptions(fieldName, parsedDateTime, options)
						context.doc.AddField(field)
					}
				}
			}
		} else {
			// automatic indexing behavior

			// first see if it can be parsed by the default date parser
			// FIXME add support for index mapping overriding defaults
			dateTimeParser := Config.Analysis.DateTimeParsers[*Config.DefaultDateTimeFormat]
			parsedDateTime, err := dateTimeParser.ParseDateTime(propertyValueString)
			if err != nil {
				// index as plain text
				options := document.STORE_FIELD | document.INDEX_FIELD | document.INCLUDE_TERM_VECTORS
				analyzer := im.defaultAnalyzer(context.dm, path)
				field := document.NewTextFieldCustom(pathString, []byte(propertyValueString), options, analyzer)
				context.doc.AddField(field)
			} else {
				// index as datetime
				field := document.NewDateTimeField(pathString, parsedDateTime)
				context.doc.AddField(field)
			}
		}
	case reflect.Float64:
		propertyValFloat := propertyValue.Float()
		if subDocMapping != nil {
			// index by explicit mapping
			for _, fieldMapping := range subDocMapping.Fields {
				fieldName := getFieldName(pathString, path, fieldMapping)
				if *fieldMapping.Type == "number" {
					options := fieldMapping.Options()
					field := document.NewNumericFieldWithIndexingOptions(fieldName, propertyValFloat, options)
					context.doc.AddField(field)
				}
			}
		} else {
			// automatic indexing behavior
			field := document.NewNumericField(pathString, propertyValFloat)
			context.doc.AddField(field)
		}
	case reflect.Struct:
		switch property := property.(type) {
		case time.Time:
			// don't descend into the time struct
			if subDocMapping != nil {
				// index by explicit mapping
				for _, fieldMapping := range subDocMapping.Fields {
					fieldName := getFieldName(pathString, path, fieldMapping)
					if *fieldMapping.Type == "datetime" {
						options := fieldMapping.Options()
						field := document.NewDateTimeFieldWithIndexingOptions(fieldName, property, options)
						context.doc.AddField(field)
					}
				}
			} else {
				// automatic indexing behavior
				field := document.NewDateTimeField(pathString, property)
				context.doc.AddField(field)
			}

		default:
			im.walkDocument(property, path, context)
		}
	default:
		im.walkDocument(property, path, context)
	}
}

func (im *IndexMapping) defaultAnalyzer(dm *DocumentMapping, path []string) *analysis.Analyzer {
	// first see if the document mapping has an analyzer
	rv := dm.defaultAnalyzer(path)
	if rv == nil {
		if im.DefaultAnalyzer != nil {
			rv = Config.Analysis.Analyzers[*im.DefaultAnalyzer]
		} else if Config.DefaultAnalyzer != nil {
			rv = Config.Analysis.Analyzers[*Config.DefaultAnalyzer]
		}
	}
	return rv
}

// attempts to find the best analyzer to use with only a field name
// will walk all the document types, look for field mappings at the
// provided path, if one exists and it has an explicit analyzer
// that is returned
// nil should be an acceptable return value meaning we don't know
func (im *IndexMapping) analyzerForPath(path string) *analysis.Analyzer {

	// first we look for explicit mapping on the field
	for _, docMapping := range im.TypeMapping {
		pathMapping := docMapping.DocumentMappingForPath(path)
		if pathMapping != nil {
			if len(pathMapping.Fields) > 0 {
				if pathMapping.Fields[0].Analyzer != nil {
					return Config.Analysis.Analyzers[*pathMapping.Fields[0].Analyzer]
				}
			}
		}
	}

	// next we will try default analyzers for the path
	for _, docMapping := range im.TypeMapping {
		rv := im.defaultAnalyzer(docMapping, decodePath(path))
		if rv != nil {
			return rv
		}
	}

	// finally just return the system-wide default analyzer
	return Config.Analysis.Analyzers[*Config.DefaultAnalyzer]
}

func (im *IndexMapping) datetimeParserForPath(path string) analysis.DateTimeParser {

	// first we look for explicit mapping on the field
	for _, docMapping := range im.TypeMapping {
		pathMapping := docMapping.DocumentMappingForPath(path)
		if pathMapping != nil {
			if len(pathMapping.Fields) > 0 {
				if pathMapping.Fields[0].Analyzer != nil {
					return Config.Analysis.DateTimeParsers[*pathMapping.Fields[0].DateFormat]
				}
			}
		}
	}

	// next we will try default analyzers for the path
	// FIXME introduce default date time parsers at mapping leves

	// finally just return the system-wide default analyzer
	return Config.Analysis.DateTimeParsers[*Config.DefaultDateTimeFormat]
}

func (im *IndexMapping) defaultField() string {
	if im.DefaultField != nil {
		return *im.DefaultField
	} else if Config.DefaultField != nil {
		return *Config.DefaultField
	}
	return ""
}

func getFieldName(pathString string, path []string, fieldMapping *FieldMapping) string {
	fieldName := pathString
	if fieldMapping.Name != nil && *fieldMapping.Name != "" {
		parentName := ""
		if len(path) > 1 {
			parentName = encodePath(path[:len(path)-1]) + PATH_SEPARATOR
		}
		fieldName = parentName + *fieldMapping.Name
	}
	return fieldName
}
<|endoftext|>


<commit_before>package updates

import (
	"fmt"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/cozy/cozy-stack/pkg/apps"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/registry"
	"github.com/sirupsen/logrus"
)

const numUpdaters = 4
const numUpdatersSingleInstance = 4

var log = logger.WithNamespace("updates")

func init() {
	jobs.AddWorker(&jobs.WorkerConfig{
		WorkerType:   "updates",
		Concurrency:  1,
		MaxExecCount: 1,
		Timeout:      1 * time.Hour,
		WorkerFunc:   Worker,
	})
}

type updateError struct {
	domain string
	slug   string
	step   string
	reason error
}

func (u *updateError) toFields() logrus.Fields {
	fields := make(logrus.Fields, 4)
	fields["step"] = u.step
	fields["reason"] = u.reason.Error()
	if u.domain != "" {
		fields["domain"] = u.domain
	}
	if u.slug != "" {
		fields["slug"] = u.slug
	}
	return fields
}

func updateErrorFromInstaller(inst *apps.Installer, step string, reason error) *updateError {
	return &updateError{
		domain: inst.Domain(),
		slug:   inst.Slug(),
		step:   step,
		reason: reason,
	}
}

// Options is the option handler for updates:
//   - Slugs: allow to filter the application's slugs to update, if empty, all
//     applications are updated
//   - Force: forces the update, even if the user has not activated the auto-
//     update
//   - ForceRegistry: translates the git:// sourced application into
//     registry://
type Options struct {
	Slugs         []string
	Domain        string
	AllDomains    bool
	Force         bool
	ForceRegistry bool
	OnlyRegistry  bool
}

// Worker is the worker method to launch the updates.
func Worker(ctx *jobs.WorkerContext) error {
	var opts Options
	if err := ctx.UnmarshalMessage(&opts); err != nil {
		return err
	}
	if opts.AllDomains {
		return UpdateAll(ctx, &opts)
	}
	if opts.Domain != "" {
		inst, err := instance.Get(opts.Domain)
		if err != nil {
			return err
		}
		return UpdateInstance(ctx, inst, &opts)
	}
	return nil
}

// UpdateAll starts the auto-updates process for all instances. The slugs
// parameters can be used optionnaly to filter (whitelist) the applications'
// slug to update.
func UpdateAll(ctx *jobs.WorkerContext, opts *Options) error {
	insc := make(chan *apps.Installer)
	errc := make(chan *updateError)

	var g sync.WaitGroup
	g.Add(numUpdaters)

	for i := 0; i < numUpdaters; i++ {
		go func() {
			for installer := range insc {
				_, err := installer.RunSync()
				if err != nil {
					errc <- updateErrorFromInstaller(installer, "RunSync", err)
				} else {
					errc <- nil
				}
			}
			g.Done()
		}()
	}

	go func() {
		// TODO: filter instances that are AutoUpdate only
		instance.ForeachInstances(func(inst *instance.Instance) error {
			if opts.Force || !inst.NoAutoUpdate {
				installerPush(inst, insc, errc, opts)
			}
			return nil
		})
		close(insc)
		g.Wait()
		close(errc)
	}()

	errors := 0
	totals := 0
	for err := range errc {
		if err != nil {
			ctx.Logger().WithFields(err.toFields()).Error()
			errors++
		}
		totals++
	}

	if errors > 0 {
		return fmt.Errorf("At least one error has happened during the updates: "+
			"%d errors for %d updates", errors, totals)
	}
	return nil
}

// UpdateInstance starts the auto-update process on the given instance. The
// slugs parameters can be used to filter (whitelist) the applications' slug
func UpdateInstance(ctx *jobs.WorkerContext, inst *instance.Instance, opts *Options) error {
	insc := make(chan *apps.Installer)
	errc := make(chan *updateError)

	var g sync.WaitGroup
	g.Add(numUpdatersSingleInstance)

	for i := 0; i < numUpdatersSingleInstance; i++ {
		go func() {
			for installer := range insc {
				_, err := installer.RunSync()
				if err != nil {
					errc <- updateErrorFromInstaller(installer, "RunSync", err)
				} else {
					errc <- nil
				}
			}
			g.Done()
		}()
	}

	go func() {
		installerPush(inst, insc, errc, opts)
		close(insc)
		g.Wait()
		close(errc)
	}()

	errors := 0
	totals := 0
	for err := range errc {
		if err != nil {
			ctx.Logger().WithFields(err.toFields()).Error()
			errors++
		}
		totals++
	}

	if errors > 0 {
		return fmt.Errorf("At least one error has happened during the updates: "+
			"%d errors for %d updates", errors, totals)
	}
	return nil
}

func installerPush(inst *instance.Instance, insc chan *apps.Installer, errc chan *updateError, opts *Options) {
	registries, err := inst.Registries()
	if err != nil {
		errc <- &updateError{step: "Registries", reason: err}
		return
	}

	var g sync.WaitGroup
	g.Add(2)

	go func() {
		defer g.Done()
		webapps, err := apps.ListWebapps(inst)
		if err != nil {
			errc <- &updateError{
				domain: inst.Domain,
				step:   "ListWebapps",
				reason: err,
			}
			return
		}
		for _, app := range webapps {
			if filterSlug(app.Slug(), opts.Slugs) {
				continue
			}
			if opts.OnlyRegistry && strings.HasPrefix(app.Source(), "registry://") {
				continue
			}
			installer, err := createInstaller(inst, registries, app, opts)
			if err != nil {
				errc <- &updateError{
					domain: inst.Domain,
					slug:   app.Slug(),
					step:   "CreateInstaller",
					reason: err,
				}
			} else {
				insc <- installer
			}
		}
	}()

	go func() {
		defer g.Done()
		konnectors, err := apps.ListKonnectors(inst)
		if err != nil {
			errc <- &updateError{
				domain: inst.Domain,
				step:   "ListKonnectors",
				reason: err,
			}
			return
		}
		for _, app := range konnectors {
			if filterSlug(app.Slug(), opts.Slugs) {
				continue
			}
			if opts.OnlyRegistry && strings.HasPrefix(app.Source(), "registry://") {
				continue
			}
			installer, err := createInstaller(inst, registries, app, opts)
			if err != nil {
				errc <- &updateError{
					domain: inst.Domain,
					slug:   app.Slug(),
					step:   "CreateInstaller",
					reason: err,
				}
			} else {
				insc <- installer
			}
		}
	}()

	g.Wait()
}

func filterSlug(slug string, slugs []string) bool {
	if len(slugs) == 0 {
		return false
	}
	for _, s := range slugs {
		if s == slug {
			return false
		}
	}
	return true
}

func createInstaller(inst *instance.Instance, registries []*url.URL, man apps.Manifest, opts *Options) (*apps.Installer, error) {
	var sourceURL string
	if opts.ForceRegistry {
		originalSourceURL, err := url.Parse(man.Source())
		if err == nil && (originalSourceURL.Scheme == "git" ||
			originalSourceURL.Scheme == "git+ssh" ||
			originalSourceURL.Scheme == "ssh+git") {
			var channel string
			if man.AppType() == apps.Webapp && strings.HasPrefix(originalSourceURL.Fragment, "build") {
				channel = "dev"
			} else {
				channel = "stable"
			}
			_, err := registry.GetLatestVersion(man.Slug(), channel, registries)
			if err == nil {
				sourceURL = fmt.Sprintf("registry://%s/%s", man.Slug(), channel)
			}
		}
	}
	return apps.NewInstaller(inst, inst.AppsCopier(man.AppType()),
		&apps.InstallerOptions{
			Operation:  apps.Update,
			Manifest:   man,
			Registries: registries,
			SourceURL:  sourceURL,
		},
	)
}
<commit_msg>Remove unused var<commit_after>package updates

import (
	"fmt"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/cozy/cozy-stack/pkg/apps"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/registry"
	"github.com/sirupsen/logrus"
)

const numUpdaters = 4
const numUpdatersSingleInstance = 4

func init() {
	jobs.AddWorker(&jobs.WorkerConfig{
		WorkerType:   "updates",
		Concurrency:  1,
		MaxExecCount: 1,
		Timeout:      1 * time.Hour,
		WorkerFunc:   Worker,
	})
}

type updateError struct {
	domain string
	slug   string
	step   string
	reason error
}

func (u *updateError) toFields() logrus.Fields {
	fields := make(logrus.Fields, 4)
	fields["step"] = u.step
	fields["reason"] = u.reason.Error()
	if u.domain != "" {
		fields["domain"] = u.domain
	}
	if u.slug != "" {
		fields["slug"] = u.slug
	}
	return fields
}

func updateErrorFromInstaller(inst *apps.Installer, step string, reason error) *updateError {
	return &updateError{
		domain: inst.Domain(),
		slug:   inst.Slug(),
		step:   step,
		reason: reason,
	}
}

// Options is the option handler for updates:
//   - Slugs: allow to filter the application's slugs to update, if empty, all
//     applications are updated
//   - Force: forces the update, even if the user has not activated the auto-
//     update
//   - ForceRegistry: translates the git:// sourced application into
//     registry://
type Options struct {
	Slugs         []string
	Domain        string
	AllDomains    bool
	Force         bool
	ForceRegistry bool
	OnlyRegistry  bool
}

// Worker is the worker method to launch the updates.
func Worker(ctx *jobs.WorkerContext) error {
	var opts Options
	if err := ctx.UnmarshalMessage(&opts); err != nil {
		return err
	}
	if opts.AllDomains {
		return UpdateAll(ctx, &opts)
	}
	if opts.Domain != "" {
		inst, err := instance.Get(opts.Domain)
		if err != nil {
			return err
		}
		return UpdateInstance(ctx, inst, &opts)
	}
	return nil
}

// UpdateAll starts the auto-updates process for all instances. The slugs
// parameters can be used optionnaly to filter (whitelist) the applications'
// slug to update.
func UpdateAll(ctx *jobs.WorkerContext, opts *Options) error {
	insc := make(chan *apps.Installer)
	errc := make(chan *updateError)

	var g sync.WaitGroup
	g.Add(numUpdaters)

	for i := 0; i < numUpdaters; i++ {
		go func() {
			for installer := range insc {
				_, err := installer.RunSync()
				if err != nil {
					errc <- updateErrorFromInstaller(installer, "RunSync", err)
				} else {
					errc <- nil
				}
			}
			g.Done()
		}()
	}

	go func() {
		// TODO: filter instances that are AutoUpdate only
		instance.ForeachInstances(func(inst *instance.Instance) error {
			if opts.Force || !inst.NoAutoUpdate {
				installerPush(inst, insc, errc, opts)
			}
			return nil
		})
		close(insc)
		g.Wait()
		close(errc)
	}()

	errors := 0
	totals := 0
	for err := range errc {
		if err != nil {
			ctx.Logger().WithFields(err.toFields()).Error()
			errors++
		}
		totals++
	}

	if errors > 0 {
		return fmt.Errorf("At least one error has happened during the updates: "+
			"%d errors for %d updates", errors, totals)
	}
	return nil
}

// UpdateInstance starts the auto-update process on the given instance. The
// slugs parameters can be used to filter (whitelist) the applications' slug
func UpdateInstance(ctx *jobs.WorkerContext, inst *instance.Instance, opts *Options) error {
	insc := make(chan *apps.Installer)
	errc := make(chan *updateError)

	var g sync.WaitGroup
	g.Add(numUpdatersSingleInstance)

	for i := 0; i < numUpdatersSingleInstance; i++ {
		go func() {
			for installer := range insc {
				_, err := installer.RunSync()
				if err != nil {
					errc <- updateErrorFromInstaller(installer, "RunSync", err)
				} else {
					errc <- nil
				}
			}
			g.Done()
		}()
	}

	go func() {
		installerPush(inst, insc, errc, opts)
		close(insc)
		g.Wait()
		close(errc)
	}()

	errors := 0
	totals := 0
	for err := range errc {
		if err != nil {
			ctx.Logger().WithFields(err.toFields()).Error()
			errors++
		}
		totals++
	}

	if errors > 0 {
		return fmt.Errorf("At least one error has happened during the updates: "+
			"%d errors for %d updates", errors, totals)
	}
	return nil
}

func installerPush(inst *instance.Instance, insc chan *apps.Installer, errc chan *updateError, opts *Options) {
	registries, err := inst.Registries()
	if err != nil {
		errc <- &updateError{step: "Registries", reason: err}
		return
	}

	var g sync.WaitGroup
	g.Add(2)

	go func() {
		defer g.Done()
		webapps, err := apps.ListWebapps(inst)
		if err != nil {
			errc <- &updateError{
				domain: inst.Domain,
				step:   "ListWebapps",
				reason: err,
			}
			return
		}
		for _, app := range webapps {
			if filterSlug(app.Slug(), opts.Slugs) {
				continue
			}
			if opts.OnlyRegistry && strings.HasPrefix(app.Source(), "registry://") {
				continue
			}
			installer, err := createInstaller(inst, registries, app, opts)
			if err != nil {
				errc <- &updateError{
					domain: inst.Domain,
					slug:   app.Slug(),
					step:   "CreateInstaller",
					reason: err,
				}
			} else {
				insc <- installer
			}
		}
	}()

	go func() {
		defer g.Done()
		konnectors, err := apps.ListKonnectors(inst)
		if err != nil {
			errc <- &updateError{
				domain: inst.Domain,
				step:   "ListKonnectors",
				reason: err,
			}
			return
		}
		for _, app := range konnectors {
			if filterSlug(app.Slug(), opts.Slugs) {
				continue
			}
			if opts.OnlyRegistry && strings.HasPrefix(app.Source(), "registry://") {
				continue
			}
			installer, err := createInstaller(inst, registries, app, opts)
			if err != nil {
				errc <- &updateError{
					domain: inst.Domain,
					slug:   app.Slug(),
					step:   "CreateInstaller",
					reason: err,
				}
			} else {
				insc <- installer
			}
		}
	}()

	g.Wait()
}

func filterSlug(slug string, slugs []string) bool {
	if len(slugs) == 0 {
		return false
	}
	for _, s := range slugs {
		if s == slug {
			return false
		}
	}
	return true
}

func createInstaller(inst *instance.Instance, registries []*url.URL, man apps.Manifest, opts *Options) (*apps.Installer, error) {
	var sourceURL string
	if opts.ForceRegistry {
		originalSourceURL, err := url.Parse(man.Source())
		if err == nil && (originalSourceURL.Scheme == "git" ||
			originalSourceURL.Scheme == "git+ssh" ||
			originalSourceURL.Scheme == "ssh+git") {
			var channel string
			if man.AppType() == apps.Webapp && strings.HasPrefix(originalSourceURL.Fragment, "build") {
				channel = "dev"
			} else {
				channel = "stable"
			}
			_, err := registry.GetLatestVersion(man.Slug(), channel, registries)
			if err == nil {
				sourceURL = fmt.Sprintf("registry://%s/%s", man.Slug(), channel)
			}
		}
	}
	return apps.NewInstaller(inst, inst.AppsCopier(man.AppType()),
		&apps.InstallerOptions{
			Operation:  apps.Update,
			Manifest:   man,
			Registries: registries,
			SourceURL:  sourceURL,
		},
	)
}
<|endoftext|>


<commit_before>package market

import (
	"fmt"
	"log"
	"time"
)

const (
	//	雅虎财经的历史分时数据没有超过90天的
	lastestDays          = 90
	companyGCCount       = 64
	retryTimes           = 50
	retryIntervalSeconds = 10
)

//	市场更新
type Market interface {
	//	名称
	Name() string
	//	时区
	Timezone() string
	//	获取上市公司列表
	Companies() ([]Company, error)

	//	抓取任务(每日)
	Crawl(market Market, company Company, day time.Time) error
}

var markets = []Market{}

//	添加市场
func Add(market Market) {
	markets = append(markets, market)

	log.Printf("市场[%s]已经加入监视列表", market.Name())
}

//	监视市场(所有操作的入口)
func Monitor() error {
	log.Print("启动监视")

	//	查询已经保存的Raw60
	err := loadSavedRaw60()
	if err != nil {
		return fmt.Errorf("查询已经保存的Raw60失败:%s", err.Error())
	}

	for _, m := range markets {

		//	启动每日定时任务
		go func(market Market) {
			//	所处时区距明日0点的间隔
			now := locationNow(market)
			du := locationYesterdayZero(market).Add(time.Hour * 48).Sub(now)

			log.Printf("[%s]\t定时任务已启动，将于%s后激活首次任务", market.Name(), du.String())
			time.AfterFunc(du, func() {
				//	立即运行一次
				go dailyTask(market)

				//	每天运行一次
				ticker := time.NewTicker(time.Hour * 24)
				for _ = range ticker.C {
					dailyTask(market)
				}
			})

		}(m)

		//	启动历史数据获取任务
		go func(market Market) {
			historyTask(market, locationYesterdayZero(market))
		}(m)
	}

	return nil
}

//	所处时区当前时间
func locationNow(market Market) time.Time {
	now := time.Now()

	//	获取市场所在时区的今天0点
	location, err := time.LoadLocation(market.Timezone())
	if err != nil {
		return now
	}

	return now.In(location)
}

//	昨天0点
func locationYesterdayZero(market Market) time.Time {
	now := locationNow(market)
	year, month, day := now.Add(-time.Hour * 24).Date()

	return time.Date(year, month, day, 0, 0, 0, 0, now.Location())
}

//	每日定时任务
func dailyTask(market Market) {

	//	昨天零点
	yesterday := locationYesterdayZero(market)
	log.Printf("[%s]\t%s数据获取任务已启动", market.Name(), yesterday.Format("20060102"))

	//	获取市场所有上市公司
	companies, err := getCompanies(market)
	if err != nil {
		log.Printf("[%s]\t获取上市公司失败: %s", market.Name(), err.Error())
		return
	}

	chanSend := make(chan int, companyGCCount)
	chanReceive := make(chan int)

	for _, c := range companies {
		//	并发抓取
		go func(company Company) {

			err := market.Crawl(market, company, yesterday)
			if err != nil {
				log.Printf("[%s]\t抓取上市公司[%s]数据失败: %s", market.Name(), company.Name, err.Error())
			}

			<-chanSend
			chanReceive <- 1
		}(c)

		chanSend <- 1
	}

	//	阻塞，直到抓取所有
	for _, _ = range companies {
		<-chanReceive
	}

	close(chanSend)
	close(chanReceive)

	log.Printf("[%s]\t%s数据获取任务已结束", market.Name(), yesterday.Format("20060102"))
}

//	历史数据获取任务
func historyTask(market Market, yesterday time.Time) {
	//	获取市场所有上市公司
	companies, err := getCompanies(market)
	if err != nil {
		log.Printf("[%s]\t获取上市公司失败: %s", market.Name(), err.Error())
		return
	}

	log.Printf("[%s]\t开始抓取%d家上市公司在%s之前的历史", market.Name(), len(companies), yesterday.Format("20060102"))

	chanSend := make(chan int, companyGCCount)
	chanReceive := make(chan int)

	for _, c := range companies {
		//	并发抓取
		go func(company Company) {

			for index := 0; index < lastestDays; index++ {
				day := yesterday.Add(-time.Hour * 24 * time.Duration(index))
				err := market.Crawl(market, company, day)
				if err != nil {
					log.Printf("[%s]\t抓取[%s]在%s的分时数据出错:%s", market.Name(), company.Code, day.Format("20060102"), err)
				}
			}

			<-chanSend
			chanReceive <- 1
		}(c)

		chanSend <- 1
	}

	//	阻塞，直到抓取所有
	for _, _ = range companies {
		<-chanReceive
	}

	close(chanSend)
	close(chanReceive)

	log.Printf("[%s]\t上市公司的历史分时数据已经抓取结束", market.Name())
}

//	抓取市场上市公司信息
func getCompanies(market Market) ([]Company, error) {

	cl := CompanyList{}
	//	尝试更新上市公司列表
	log.Printf("[%s]\t更新上市公司列表-开始", market.Name())
	companies, err := market.Companies()
	if err != nil {

		//	如果更新失败，则尝试从上次的存档文件中读取上市公司列表
		log.Printf("[%s]\t更新上市公司列表失败，尝试从存档读取:%v", market.Name(), err)
		err = cl.Load(market)
		if err != nil {
			return nil, fmt.Errorf("[%s]\t尝试从存档读取上市公司列表-失败:%v", market.Name(), err)
		}

		companies = cl
		log.Printf("[%s]\t尝试从存档读取上市公司列表-成功,共%d家上市公司", market.Name(), len(companies))

		return companies, nil
	}

	//	存档
	cl = CompanyList(companies)
	err = cl.Save(market)
	if err != nil {
		return nil, err
	}

	log.Printf("[%s]\t更新上市公司列表-成功,共%d家上市公司", market.Name(), len(companies))

	return companies, nil
}
<commit_msg>监听市场时，首先计算市场所处时区的差异<commit_after>package market

import (
	"fmt"
	"log"
	"time"
)

const (
	//	雅虎财经的历史分时数据没有超过90天的
	lastestDays          = 90
	companyGCCount       = 64
	retryTimes           = 50
	retryIntervalSeconds = 10
)

//	市场更新
type Market interface {
	//	名称
	Name() string
	//	时区
	Timezone() string
	//	获取上市公司列表
	Companies() ([]Company, error)

	//	抓取任务(每日)
	Crawl(market Market, company Company, day time.Time) error
}

var (
	markets                       = []Market{}
	marketOffset map[string]int64 = make(map[string]int64)
)

//	添加市场
func Add(market Market) {
	markets = append(markets, market)

	log.Printf("市场[%s]已经加入监视列表", market.Name())
}

//	监视市场(所有操作的入口)
func Monitor() error {
	log.Print("启动监视")

	for _, m := range markets {

		//	本地时间
		now := time.Now()
		_, offsetLocal := now.Zone()

		//	获取市场所在时区
		location, err := time.LoadLocation(m.Timezone())
		if err != nil {
			return err
		}

		//	市场所处时区当前时间
		marketNow := now.In(location)
		_, offsetMarket := marketNow.Zone()

		//	计算TimeZoneOffset
		marketOffset[m.Name()] = int64(offsetMarket - offsetLocal)

		//	启动每日定时任务
		go func(market Market) {
			//	所处时区距明日0点的间隔
			now := marketow(market)
			du := locationYesterdayZero(market).Add(time.Hour * 48).Sub(now)

			log.Printf("[%s]\t定时任务已启动，将于%s后激活首次任务", market.Name(), du.String())
			time.AfterFunc(du, func() {
				//	立即运行一次
				go dailyTask(market)

				//	每天运行一次
				ticker := time.NewTicker(time.Hour * 24)
				for _ = range ticker.C {
					dailyTask(market)
				}
			})

		}(m)

		//	启动历史数据获取任务
		go func(market Market) {
			historyTask(market, locationYesterdayZero(market))
		}(m)
	}

	return nil
}

//	市场所处时区当前时间
func marketow(market Market) time.Time {
	now := time.Now()

	//	获取市场所在时区
	location, err := time.LoadLocation(market.Timezone())
	if err != nil {
		return now
	}

	return now.In(location)
}

//	昨天0点
func locationYesterdayZero(market Market) time.Time {
	now := marketow(market)
	year, month, day := now.Add(-time.Hour * 24).Date()

	return time.Date(year, month, day, 0, 0, 0, 0, now.Location())
}

//	每日定时任务
func dailyTask(market Market) {

	//	昨天零点
	yesterday := locationYesterdayZero(market)
	log.Printf("[%s]\t%s数据获取任务已启动", market.Name(), yesterday.Format("20060102"))

	//	获取市场所有上市公司
	companies, err := getCompanies(market)
	if err != nil {
		log.Printf("[%s]\t获取上市公司失败: %s", market.Name(), err.Error())
		return
	}

	chanSend := make(chan int, companyGCCount)
	chanReceive := make(chan int)

	for _, c := range companies {
		//	并发抓取
		go func(company Company) {

			err := market.Crawl(market, company, yesterday)
			if err != nil {
				log.Printf("[%s]\t抓取上市公司[%s]数据失败: %s", market.Name(), company.Name, err.Error())
			}

			<-chanSend
			chanReceive <- 1
		}(c)

		chanSend <- 1
	}

	//	阻塞，直到抓取所有
	for _, _ = range companies {
		<-chanReceive
	}

	close(chanSend)
	close(chanReceive)

	log.Printf("[%s]\t%s数据获取任务已结束", market.Name(), yesterday.Format("20060102"))
}

//	历史数据获取任务
func historyTask(market Market, yesterday time.Time) {
	//	获取市场所有上市公司
	companies, err := getCompanies(market)
	if err != nil {
		log.Printf("[%s]\t获取上市公司失败: %s", market.Name(), err.Error())
		return
	}

	log.Printf("[%s]\t开始抓取%d家上市公司在%s之前的历史", market.Name(), len(companies), yesterday.Format("20060102"))

	chanSend := make(chan int, companyGCCount)
	chanReceive := make(chan int)

	for _, c := range companies {
		//	并发抓取
		go func(company Company) {

			for index := 0; index < lastestDays; index++ {
				day := yesterday.Add(-time.Hour * 24 * time.Duration(index))
				err := market.Crawl(market, company, day)
				if err != nil {
					log.Printf("[%s]\t抓取[%s]在%s的分时数据出错:%s", market.Name(), company.Code, day.Format("20060102"), err)
				}
			}

			<-chanSend
			chanReceive <- 1
		}(c)

		chanSend <- 1
	}

	//	阻塞，直到抓取所有
	for _, _ = range companies {
		<-chanReceive
	}

	close(chanSend)
	close(chanReceive)

	log.Printf("[%s]\t上市公司的历史分时数据已经抓取结束", market.Name())
}

//	抓取市场上市公司信息
func getCompanies(market Market) ([]Company, error) {

	cl := CompanyList{}
	//	尝试更新上市公司列表
	log.Printf("[%s]\t更新上市公司列表-开始", market.Name())
	companies, err := market.Companies()
	if err != nil {

		//	如果更新失败，则尝试从上次的存档文件中读取上市公司列表
		log.Printf("[%s]\t更新上市公司列表失败，尝试从存档读取:%v", market.Name(), err)
		err = cl.Load(market)
		if err != nil {
			return nil, fmt.Errorf("[%s]\t尝试从存档读取上市公司列表-失败:%v", market.Name(), err)
		}

		companies = cl
		log.Printf("[%s]\t尝试从存档读取上市公司列表-成功,共%d家上市公司", market.Name(), len(companies))

		return companies, nil
	}

	//	存档
	cl = CompanyList(companies)
	err = cl.Save(market)
	if err != nil {
		return nil, err
	}

	log.Printf("[%s]\t更新上市公司列表-成功,共%d家上市公司", market.Name(), len(companies))

	return companies, nil
}
<|endoftext|>


<commit_before>package feed

import (
	"bytes"
	"testing"
)

func TestError(t *testing.T) {
	str := `<?xml version="1.0" encoding="UTF-8"?><parent></no-parent>`
	r := bytes.NewBuffer([]byte(str))

	_, err := ParseXML(r)
	if err == nil {
		t.Errorf("Parser does not report error")
	}
}

func verifyElementNode(t *testing.T, node *Node, ns string, tag string, size int) {
	if node.Type != ELEMENT {
		t.Errorf("Invalid node type")
		return
	}

	if ns != node.NS {
		t.Errorf("Invalid namespace: %s -> %s", ns, node.NS)
	}
	if tag != node.Tag {
		t.Errorf("Invalid tag name: %s -> %s", tag, node.Tag)
	}
	if size != len(node.Children) {
		t.Errorf("Invalid number of children: %d -> %d",
			size, len(node.Children))
	}
}

func verifyTextNode(t *testing.T, node *Node, value string) {
	if node.Type != TEXT {
		t.Errorf("Invalid node type")
		return
	}

	if value != node.Value {
		t.Errorf("Value mismatch: (len=%d) %w -> (len=%d) %w",
			len(value), []byte(value), len(node.Value), []byte(node.Value))
	}
}

func verifySiblings(t *testing.T, nodes ...*Node) {
	for i := 0; i < len(nodes); i += 1 {
		if i == 0 && nodes[i].Prev != nil {
			t.Errorf("The first node's previous is not nil")
		}
		if i > 0 {
			if nodes[i-1].Next != nodes[i] {
				t.Errorf("Node #%d's next is not Node #%d", i-1, i)
			}
			if nodes[i].Prev != nodes[i-1] {
				t.Errorf("Node #%d's previous is not Node #%d", i, i-1)
			}
		}
		if i == len(nodes)-1 && nodes[i].Next != nil {
			t.Errorf("The last node's next is not nil")
		}
	}
}

func verifyAttribute(t *testing.T, node *Node, ns string, name string, value string) {
	if node.Type != ELEMENT {
		t.Errorf("Invalid node type")
		return
	}

	for _, attr := range node.Attributes {
		if ns == attr.NS && name == attr.Name {
			if value != attr.Value {
				t.Errorf("Attribute's value mismatch, ns=%s name=%s: %s -> %s",
					ns, name, value, attr.Value)
			}
			return
		}
	}

	t.Errorf("Attribute is not found, ns=%s name=%s", ns, name)
}

func TestBasic(t *testing.T) {
	str := `
<?xml version="1.0" encoding="UTF-8"?><parent><child-1>first child</child-1><child-2>second child</child-2></parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyElementNode(t, tree, "", "parent", 2)
	verifyElementNode(t, tree.Children[0], "", "child-1", 1)
	verifyElementNode(t, tree.Children[1], "", "child-2", 1)
	verifySiblings(t, tree.Children[0], tree.Children[1])
}

func TestTextNode(t *testing.T) {
	str := `<?xml version="1.0" encoding="UTF-8"?>
<parent>
  <child-1>first child</child-1>
  <child-2>second child</child-2>
</parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyElementNode(t, tree, "", "parent", 5)
	verifyTextNode(t, tree.Children[0], "\n  ")
	verifyElementNode(t, tree.Children[1], "", "child-1", 1)
	verifyTextNode(t, tree.Children[2], "\n  ")
	verifyElementNode(t, tree.Children[3], "", "child-2", 1)
	verifyTextNode(t, tree.Children[4], "\n")
	verifySiblings(t, tree.Children[0], tree.Children[1], tree.Children[2],
		tree.Children[3], tree.Children[4])
}

func TestNamespace(t *testing.T) {
	str := `<?xml version="1.0" encoding="UTF-8"?>
<parent xmlns="http://example.com/ns"
  xmlns:extra="http://extra">
  <child-1>uses default namespace</child-1>
  <extra:child-2>in the extra namespace</extra:child-2>
  <child-3 xmlns="http://change-default">
    <first>one</first>
    <second>two</second>
  </child-3>
</parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyElementNode(t, tree.Children[1], "http://example.com/ns", "child-1", 1)
	verifyElementNode(t, tree.Children[3], "http://extra", "child-2", 1)
	verifyElementNode(t, tree.Children[5], "http://change-default", "child-3", 5)
	verifyElementNode(t, tree.Children[5].Children[1],
		"http://change-default", "first", 1)
	verifyElementNode(t, tree.Children[5].Children[3],
		"http://change-default", "second", 1)
}

func TestNoDeclaration(t *testing.T) {
	str := `<parent><child-1>one</child-1><child-2>two</child-2></parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyElementNode(t, tree, "", "parent", 2)
	verifyElementNode(t, tree.Children[0], "", "child-1", 1)
	verifyElementNode(t, tree.Children[1], "", "child-2", 1)
}

func TestEmptyElement(t *testing.T) {
	str := `<parent><child-1></child-1><child-2/></parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyElementNode(t, tree, "", "parent", 2)
	verifyElementNode(t, tree.Children[0], "", "child-1", 0)
	verifyElementNode(t, tree.Children[1], "", "child-2", 0)
}

func TestAttribute(t *testing.T) {
	str := `<parent version="1.0"></parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyAttribute(t, tree, "", "version", "1.0")
}

func TestAttributeNamespace(t *testing.T) {
	str := `<parent version="2.0" xmlns:extra="http://extra"><child-1 name="one" extra:name="satu"/></parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyAttribute(t, tree, "", "version", "2.0")

	child := tree.Children[0]
	verifyAttribute(t, child, "", "name", "one")
	verifyAttribute(t, child, "http://extra", "name", "satu")
}
<commit_msg>Make sure that the element's default namespace is set<commit_after>package feed

import (
	"bytes"
	"testing"
)

func TestError(t *testing.T) {
	str := `<?xml version="1.0" encoding="UTF-8"?><parent></no-parent>`
	r := bytes.NewBuffer([]byte(str))

	_, err := ParseXML(r)
	if err == nil {
		t.Errorf("Parser does not report error")
	}
}

func verifyElementNode(t *testing.T, node *Node, ns string, tag string, size int) {
	if node.Type != ELEMENT {
		t.Errorf("Invalid node type")
		return
	}

	if ns != node.NS {
		t.Errorf("Invalid namespace: %s -> %s", ns, node.NS)
	}
	if tag != node.Tag {
		t.Errorf("Invalid tag name: %s -> %s", tag, node.Tag)
	}
	if size != len(node.Children) {
		t.Errorf("Invalid number of children: %d -> %d",
			size, len(node.Children))
	}
}

func verifyTextNode(t *testing.T, node *Node, value string) {
	if node.Type != TEXT {
		t.Errorf("Invalid node type")
		return
	}

	if value != node.Value {
		t.Errorf("Value mismatch: (len=%d) %w -> (len=%d) %w",
			len(value), []byte(value), len(node.Value), []byte(node.Value))
	}
}

func verifySiblings(t *testing.T, nodes ...*Node) {
	for i := 0; i < len(nodes); i += 1 {
		if i == 0 && nodes[i].Prev != nil {
			t.Errorf("The first node's previous is not nil")
		}
		if i > 0 {
			if nodes[i-1].Next != nodes[i] {
				t.Errorf("Node #%d's next is not Node #%d", i-1, i)
			}
			if nodes[i].Prev != nodes[i-1] {
				t.Errorf("Node #%d's previous is not Node #%d", i, i-1)
			}
		}
		if i == len(nodes)-1 && nodes[i].Next != nil {
			t.Errorf("The last node's next is not nil")
		}
	}
}

func verifyAttribute(t *testing.T, node *Node, ns string, name string, value string) {
	if node.Type != ELEMENT {
		t.Errorf("Invalid node type")
		return
	}

	for _, attr := range node.Attributes {
		if ns == attr.NS && name == attr.Name {
			if value != attr.Value {
				t.Errorf("Attribute's value mismatch, ns=%s name=%s: %s -> %s",
					ns, name, value, attr.Value)
			}
			return
		}
	}

	t.Errorf("Attribute is not found, ns=%s name=%s", ns, name)
}

func TestBasic(t *testing.T) {
	str := `
<?xml version="1.0" encoding="UTF-8"?><parent><child-1>first child</child-1><child-2>second child</child-2></parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyElementNode(t, tree, "", "parent", 2)
	verifyElementNode(t, tree.Children[0], "", "child-1", 1)
	verifyElementNode(t, tree.Children[1], "", "child-2", 1)
	verifySiblings(t, tree.Children[0], tree.Children[1])
}

func TestTextNode(t *testing.T) {
	str := `<?xml version="1.0" encoding="UTF-8"?>
<parent>
  <child-1>first child</child-1>
  <child-2>second child</child-2>
</parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyElementNode(t, tree, "", "parent", 5)
	verifyTextNode(t, tree.Children[0], "\n  ")
	verifyElementNode(t, tree.Children[1], "", "child-1", 1)
	verifyTextNode(t, tree.Children[2], "\n  ")
	verifyElementNode(t, tree.Children[3], "", "child-2", 1)
	verifyTextNode(t, tree.Children[4], "\n")
	verifySiblings(t, tree.Children[0], tree.Children[1], tree.Children[2],
		tree.Children[3], tree.Children[4])
}

func TestNamespace(t *testing.T) {
	str := `<?xml version="1.0" encoding="UTF-8"?>
<parent xmlns="http://example.com/ns"
  xmlns:extra="http://extra">
  <child-1>uses default namespace</child-1>
  <extra:child-2>in the extra namespace</extra:child-2>
  <child-3 xmlns="http://change-default">
    <first>one</first>
    <second>two</second>
  </child-3>
</parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyElementNode(t, tree, "http://example.com/ns", "parent", 7)
	verifyElementNode(t, tree.Children[1], "http://example.com/ns", "child-1", 1)
	verifyElementNode(t, tree.Children[3], "http://extra", "child-2", 1)
	verifyElementNode(t, tree.Children[5], "http://change-default", "child-3", 5)
	verifyElementNode(t, tree.Children[5].Children[1],
		"http://change-default", "first", 1)
	verifyElementNode(t, tree.Children[5].Children[3],
		"http://change-default", "second", 1)
}

func TestNoDeclaration(t *testing.T) {
	str := `<parent><child-1>one</child-1><child-2>two</child-2></parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyElementNode(t, tree, "", "parent", 2)
	verifyElementNode(t, tree.Children[0], "", "child-1", 1)
	verifyElementNode(t, tree.Children[1], "", "child-2", 1)
}

func TestEmptyElement(t *testing.T) {
	str := `<parent><child-1></child-1><child-2/></parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyElementNode(t, tree, "", "parent", 2)
	verifyElementNode(t, tree.Children[0], "", "child-1", 0)
	verifyElementNode(t, tree.Children[1], "", "child-2", 0)
}

func TestAttribute(t *testing.T) {
	str := `<parent version="1.0"></parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyAttribute(t, tree, "", "version", "1.0")
}

func TestAttributeNamespace(t *testing.T) {
	str := `<parent version="2.0" xmlns:extra="http://extra"><child-1 name="one" extra:name="satu"/></parent>`
	r := bytes.NewBuffer([]byte(str))

	tree, err := ParseXML(r)
	if err != nil {
		t.Errorf("Parser should not report error: %v", err)
		return
	}

	verifyAttribute(t, tree, "", "version", "2.0")

	child := tree.Children[0]
	verifyAttribute(t, child, "", "name", "one")
	verifyAttribute(t, child, "http://extra", "name", "satu")
}
<|endoftext|>


<commit_before>package mastodon

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
	"time"
)

func TestDoAPI(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Query().Get("max_id") == "999" {
			w.Header().Set("Link", `<:>; rel="next"`)
		} else {
			w.Header().Set("Link", `<http://example.com?max_id=234>; rel="next", <http://example.com?since_id=890>; rel="prev"`)
		}
		fmt.Fprintln(w, `[{"username": "foo"}, {"username": "bar"}]`)
	}))
	defer ts.Close()

	c := NewClient(&Config{Server: ts.URL})
	var accounts []Account
	err := c.doAPI(context.Background(), http.MethodGet, "/", nil, &accounts, &Pagination{
		MaxID: "999",
	})
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	pg := &Pagination{
		MaxID:   "123",
		SinceID: "789",
		Limit:   10,
	}
	err = c.doAPI(context.Background(), http.MethodGet, "/", url.Values{}, &accounts, pg)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if pg.MaxID != "234" {
		t.Fatalf("want %q but %q", "234", pg.MaxID)
	}
	if pg.SinceID != "890" {
		t.Fatalf("want %q but %q", "890", pg.SinceID)
	}
	if accounts[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", accounts[0].Username)
	}
	if accounts[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", accounts[1].Username)
	}

	pg = &Pagination{
		MaxID:   "123",
		SinceID: "789",
		Limit:   10,
	}
	err = c.doAPI(context.Background(), http.MethodGet, "/", nil, &accounts, pg)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if pg.MaxID != "234" {
		t.Fatalf("want %q but %q", "234", pg.MaxID)
	}
	if pg.SinceID != "890" {
		t.Fatalf("want %q but %q", "890", pg.SinceID)
	}
	if accounts[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", accounts[0].Username)
	}
	if accounts[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", accounts[1].Username)
	}

	// *Pagination is nil
	err = c.doAPI(context.Background(), http.MethodGet, "/", nil, &accounts, nil)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if accounts[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", accounts[0].Username)
	}
	if accounts[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", accounts[1].Username)
	}
}

func TestAuthenticate(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.FormValue("username") != "valid" || r.FormValue("password") != "user" {
			http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
			return
		}
		fmt.Fprintln(w, `{"access_token": "zoo"}`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	err := client.Authenticate(context.Background(), "invalid", "user")
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	client = NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	err = client.Authenticate(context.Background(), "valid", "user")
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
}

func TestAuthenticateWithCancel(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(3 * time.Second)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	err := client.Authenticate(ctx, "invalid", "user")
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	if want := "Post " + ts.URL + "/oauth/token: context canceled"; want != err.Error() {
		t.Fatalf("want %q but %q", want, err.Error())
	}
}

func TestPostStatus(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer zoo" {
			http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
			return
		}
		fmt.Fprintln(w, `{"access_token": "zoo"}`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	_, err := client.PostStatus(context.Background(), &Toot{
		Status: "foobar",
	})
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	client = NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	_, err = client.PostStatus(context.Background(), &Toot{
		Status: "foobar",
	})
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
}

func TestPostStatusWithCancel(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(3 * time.Second)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := client.PostStatus(ctx, &Toot{
		Status: "foobar",
	})
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	if want := "Post " + ts.URL + "/api/v1/statuses: context canceled"; want != err.Error() {
		t.Fatalf("want %q but %q", want, err.Error())
	}
}

func TestGetTimelineHome(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, `[{"content": "foo"}, {"content": "bar"}]`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	_, err := client.PostStatus(context.Background(), &Toot{
		Status: "foobar",
	})
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	client = NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	tl, err := client.GetTimelineHome(context.Background(), nil)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if len(tl) != 2 {
		t.Fatalf("result should be two: %d", len(tl))
	}
	if tl[0].Content != "foo" {
		t.Fatalf("want %q but %q", "foo", tl[0].Content)
	}
	if tl[1].Content != "bar" {
		t.Fatalf("want %q but %q", "bar", tl[1].Content)
	}
}

func TestGetTimelineHomeWithCancel(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(3 * time.Second)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := client.GetTimelineHome(ctx, nil)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	if want := "Get " + ts.URL + "/api/v1/timelines/home: context canceled"; want != err.Error() {
		t.Fatalf("want %q but %q", want, err.Error())
	}
}

func TestForTheCoverages(t *testing.T) {
	(*UpdateEvent)(nil).event()
	(*NotificationEvent)(nil).event()
	(*DeleteEvent)(nil).event()
	(*ErrorEvent)(nil).event()
	_ = (&ErrorEvent{io.EOF}).Error()
}

func TestNewPagination(t *testing.T) {
	_, err := newPagination("")
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	_, err = newPagination(`<:>; rel="next"`)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	_, err = newPagination(`<:>; rel="prev"`)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	pg, err := newPagination(`<http://example.com?max_id=123>; rel="next", <http://example.com?since_id=789>; rel="prev"`)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if pg.MaxID != "123" {
		t.Fatalf("want %q but %q", "123", pg.MaxID)
	}
	if pg.SinceID != "789" {
		t.Fatalf("want %q but %q", "789", pg.SinceID)
	}
}

func TestGetPaginationID(t *testing.T) {
	_, err := getPaginationID(":", "max_id")
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	_, err = getPaginationID("http://example.com?max_id=abc", "max_id")
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	id, err := getPaginationID("http://example.com?max_id=123", "max_id")
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if id != "123" {
		t.Fatalf("want %q but %q", "123", id)
	}
}

func TestPaginationSetValues(t *testing.T) {
	p := &Pagination{
		MaxID:   "123",
		SinceID: "456",
		MinID:   "789",
		Limit:   10,
	}
	before := url.Values{"key": {"value"}}
	after := p.setValues(before)
	if after.Get("key") != "value" {
		t.Fatalf("want %q but %q", "value", after.Get("key"))
	}
	if after.Get("max_id") != "123" {
		t.Fatalf("want %q but %q", "123", after.Get("max_id"))
	}
	if after.Get("since_id") != "456" {
		t.Fatalf("result should be empty string: %q", after.Get("since_id"))
	}
	if after.Get("min_id") != "789" {
		t.Fatalf("result should be empty string: %q", after.Get("since_id"))
	}
	if after.Get("limit") != "10" {
		t.Fatalf("want %q but %q", "10", after.Get("limit"))
	}

	p = &Pagination{
		MaxID:   "",
		SinceID: "789",
	}
	before = url.Values{}
	after = p.setValues(before)
	if after.Get("max_id") != "" {
		t.Fatalf("result should be empty string: %q", after.Get("max_id"))
	}
	if after.Get("since_id") != "789" {
		t.Fatalf("want %q but %q", "789", after.Get("since_id"))
	}
}
<commit_msg>Fix TestPaginationSetValues<commit_after>package mastodon

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
	"time"
)

func TestDoAPI(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Query().Get("max_id") == "999" {
			w.Header().Set("Link", `<:>; rel="next"`)
		} else {
			w.Header().Set("Link", `<http://example.com?max_id=234>; rel="next", <http://example.com?since_id=890>; rel="prev"`)
		}
		fmt.Fprintln(w, `[{"username": "foo"}, {"username": "bar"}]`)
	}))
	defer ts.Close()

	c := NewClient(&Config{Server: ts.URL})
	var accounts []Account
	err := c.doAPI(context.Background(), http.MethodGet, "/", nil, &accounts, &Pagination{
		MaxID: "999",
	})
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	pg := &Pagination{
		MaxID:   "123",
		SinceID: "789",
		Limit:   10,
	}
	err = c.doAPI(context.Background(), http.MethodGet, "/", url.Values{}, &accounts, pg)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if pg.MaxID != "234" {
		t.Fatalf("want %q but %q", "234", pg.MaxID)
	}
	if pg.SinceID != "890" {
		t.Fatalf("want %q but %q", "890", pg.SinceID)
	}
	if accounts[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", accounts[0].Username)
	}
	if accounts[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", accounts[1].Username)
	}

	pg = &Pagination{
		MaxID:   "123",
		SinceID: "789",
		Limit:   10,
	}
	err = c.doAPI(context.Background(), http.MethodGet, "/", nil, &accounts, pg)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if pg.MaxID != "234" {
		t.Fatalf("want %q but %q", "234", pg.MaxID)
	}
	if pg.SinceID != "890" {
		t.Fatalf("want %q but %q", "890", pg.SinceID)
	}
	if accounts[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", accounts[0].Username)
	}
	if accounts[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", accounts[1].Username)
	}

	// *Pagination is nil
	err = c.doAPI(context.Background(), http.MethodGet, "/", nil, &accounts, nil)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if accounts[0].Username != "foo" {
		t.Fatalf("want %q but %q", "foo", accounts[0].Username)
	}
	if accounts[1].Username != "bar" {
		t.Fatalf("want %q but %q", "bar", accounts[1].Username)
	}
}

func TestAuthenticate(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.FormValue("username") != "valid" || r.FormValue("password") != "user" {
			http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
			return
		}
		fmt.Fprintln(w, `{"access_token": "zoo"}`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	err := client.Authenticate(context.Background(), "invalid", "user")
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	client = NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	err = client.Authenticate(context.Background(), "valid", "user")
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
}

func TestAuthenticateWithCancel(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(3 * time.Second)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	err := client.Authenticate(ctx, "invalid", "user")
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	if want := "Post " + ts.URL + "/oauth/token: context canceled"; want != err.Error() {
		t.Fatalf("want %q but %q", want, err.Error())
	}
}

func TestPostStatus(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer zoo" {
			http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
			return
		}
		fmt.Fprintln(w, `{"access_token": "zoo"}`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	_, err := client.PostStatus(context.Background(), &Toot{
		Status: "foobar",
	})
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	client = NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	_, err = client.PostStatus(context.Background(), &Toot{
		Status: "foobar",
	})
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
}

func TestPostStatusWithCancel(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(3 * time.Second)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := client.PostStatus(ctx, &Toot{
		Status: "foobar",
	})
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	if want := "Post " + ts.URL + "/api/v1/statuses: context canceled"; want != err.Error() {
		t.Fatalf("want %q but %q", want, err.Error())
	}
}

func TestGetTimelineHome(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, `[{"content": "foo"}, {"content": "bar"}]`)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
	})
	_, err := client.PostStatus(context.Background(), &Toot{
		Status: "foobar",
	})
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	client = NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	tl, err := client.GetTimelineHome(context.Background(), nil)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if len(tl) != 2 {
		t.Fatalf("result should be two: %d", len(tl))
	}
	if tl[0].Content != "foo" {
		t.Fatalf("want %q but %q", "foo", tl[0].Content)
	}
	if tl[1].Content != "bar" {
		t.Fatalf("want %q but %q", "bar", tl[1].Content)
	}
}

func TestGetTimelineHomeWithCancel(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(3 * time.Second)
		return
	}))
	defer ts.Close()

	client := NewClient(&Config{
		Server:       ts.URL,
		ClientID:     "foo",
		ClientSecret: "bar",
		AccessToken:  "zoo",
	})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := client.GetTimelineHome(ctx, nil)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}
	if want := "Get " + ts.URL + "/api/v1/timelines/home: context canceled"; want != err.Error() {
		t.Fatalf("want %q but %q", want, err.Error())
	}
}

func TestForTheCoverages(t *testing.T) {
	(*UpdateEvent)(nil).event()
	(*NotificationEvent)(nil).event()
	(*DeleteEvent)(nil).event()
	(*ErrorEvent)(nil).event()
	_ = (&ErrorEvent{io.EOF}).Error()
}

func TestNewPagination(t *testing.T) {
	_, err := newPagination("")
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	_, err = newPagination(`<:>; rel="next"`)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	_, err = newPagination(`<:>; rel="prev"`)
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	pg, err := newPagination(`<http://example.com?max_id=123>; rel="next", <http://example.com?since_id=789>; rel="prev"`)
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if pg.MaxID != "123" {
		t.Fatalf("want %q but %q", "123", pg.MaxID)
	}
	if pg.SinceID != "789" {
		t.Fatalf("want %q but %q", "789", pg.SinceID)
	}
}

func TestGetPaginationID(t *testing.T) {
	_, err := getPaginationID(":", "max_id")
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	_, err = getPaginationID("http://example.com?max_id=abc", "max_id")
	if err == nil {
		t.Fatalf("should be fail: %v", err)
	}

	id, err := getPaginationID("http://example.com?max_id=123", "max_id")
	if err != nil {
		t.Fatalf("should not be fail: %v", err)
	}
	if id != "123" {
		t.Fatalf("want %q but %q", "123", id)
	}
}

func TestPaginationSetValues(t *testing.T) {
	p := &Pagination{
		MaxID:   "123",
		SinceID: "456",
		MinID:   "789",
		Limit:   10,
	}
	before := url.Values{"key": {"value"}}
	after := p.setValues(before)
	if after.Get("key") != "value" {
		t.Fatalf("want %q but %q", "value", after.Get("key"))
	}
	if after.Get("max_id") != "123" {
		t.Fatalf("want %q but %q", "123", after.Get("max_id"))
	}
	if after.Get("since_id") != "456" {
		t.Fatalf("want %q but %q", "456", after.Get("since_id"))
	}
	if after.Get("min_id") != "789" {
		t.Fatalf("want %q but %q", "789", after.Get("min_id"))
	}
	if after.Get("limit") != "10" {
		t.Fatalf("want %q but %q", "10", after.Get("limit"))
	}

	p = &Pagination{
		MaxID:   "",
		SinceID: "789",
	}
	before = url.Values{}
	after = p.setValues(before)
	if after.Get("max_id") != "" {
		t.Fatalf("result should be empty string: %q", after.Get("max_id"))
	}
	if after.Get("since_id") != "789" {
		t.Fatalf("want %q but %q", "789", after.Get("since_id"))
	}
	if after.Get("min_id") != "" {
		t.Fatalf("result should be empty string: %q", after.Get("min_id"))
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package controller

import (
	"context"

	"github.com/kelseyhightower/envconfig"
	"go.uber.org/zap"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/tools/cache"

	"knative.dev/eventing-kafka/pkg/apis/messaging/v1beta1"
	"knative.dev/eventing-kafka/pkg/channel/consolidated/status"
	kafkamessagingv1beta1 "knative.dev/eventing-kafka/pkg/client/informers/externalversions/messaging/v1beta1"
	kafkaChannelClient "knative.dev/eventing-kafka/pkg/client/injection/client"
	"knative.dev/eventing-kafka/pkg/client/injection/informers/messaging/v1beta1/kafkachannel"
	kafkaChannelReconciler "knative.dev/eventing-kafka/pkg/client/injection/reconciler/messaging/v1beta1/kafkachannel"
	commonconfig "knative.dev/eventing-kafka/pkg/common/config"
	eventingduckv1 "knative.dev/eventing/pkg/apis/duck/v1"
	eventingClient "knative.dev/eventing/pkg/client/injection/client"
	kubeclient "knative.dev/pkg/client/injection/kube/client"
	"knative.dev/pkg/client/injection/kube/informers/apps/v1/deployment"
	endpointsinformer "knative.dev/pkg/client/injection/kube/informers/core/v1/endpoints"
	podinformer "knative.dev/pkg/client/injection/kube/informers/core/v1/pod"
	"knative.dev/pkg/client/injection/kube/informers/core/v1/service"
	"knative.dev/pkg/client/injection/kube/informers/core/v1/serviceaccount"
	"knative.dev/pkg/client/injection/kube/informers/rbac/v1/rolebinding"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/controller"
	"knative.dev/pkg/logging"
	knativeReconciler "knative.dev/pkg/reconciler"
	"knative.dev/pkg/system"
)

const (
	channelLabelKey   = "messaging.knative.dev/channel"
	channelLabelValue = "kafka-channel"
	roleLabelKey      = "messaging.knative.dev/role"
	roleLabelValue    = "dispatcher"
)

// NewController initializes the controller and is called by the generated code.
// Registers event handlers to enqueue events.
func NewController(
	ctx context.Context,
	cmw configmap.Watcher,
) *controller.Impl {
	logger := logging.FromContext(ctx)
	kafkaChannelInformer := kafkachannel.Get(ctx)
	deploymentInformer := deployment.Get(ctx)
	endpointsInformer := endpointsinformer.Get(ctx)
	serviceAccountInformer := serviceaccount.Get(ctx)
	roleBindingInformer := rolebinding.Get(ctx)
	serviceInformer := service.Get(ctx)
	podInformer := podinformer.Get(ctx)

	r := &Reconciler{
		systemNamespace:      system.Namespace(),
		KubeClientSet:        kubeclient.Get(ctx),
		kafkaClientSet:       kafkaChannelClient.Get(ctx),
		EventingClientSet:    eventingClient.Get(ctx),
		kafkachannelLister:   kafkaChannelInformer.Lister(),
		kafkachannelInformer: kafkaChannelInformer.Informer(),
		deploymentLister:     deploymentInformer.Lister(),
		serviceLister:        serviceInformer.Lister(),
		endpointsLister:      endpointsInformer.Lister(),
		serviceAccountLister: serviceAccountInformer.Lister(),
		roleBindingLister:    roleBindingInformer.Lister(),
	}

	env := &envConfig{}
	if err := envconfig.Process("", env); err != nil {
		logger.Panicf("unable to process Kafka channel's required environment variables: %v", err)
	}

	r.dispatcherImage = env.Image
	r.dispatcherServiceAccount = env.DispatcherServiceAccount

	impl := kafkaChannelReconciler.NewImpl(ctx, r)

	statusProber := status.NewProber(
		logger.Named("status-manager"),
		NewProbeTargetLister(logger, endpointsInformer.Lister()),
		func(c v1beta1.KafkaChannel, s eventingduckv1.SubscriberSpec) {
			logger.Debugf("Ready callback triggered for channel: %s/%s subscription: %s", c.Namespace, c.Name, string(s.UID))
			impl.EnqueueKey(types.NamespacedName{Namespace: c.Namespace, Name: c.Name})
		},
	)
	r.statusManager = statusProber
	statusProber.Start(ctx.Done())
	// Get and Watch the Kakfa config map and dynamically update Kafka configuration.
	err := commonconfig.InitializeKafkaConfigMapWatcher(ctx, cmw, logger, r.updateKafkaConfig, system.Namespace())
	if err != nil {
		logger.Fatal("Failed To Initialize ConfigMap Watcher", zap.Error(err))
	}

	logger.Info("Setting up event handlers")
	kafkaChannelInformer.Informer().AddEventHandler(controller.HandleAll(impl.Enqueue))

	// Set up watches for dispatcher resources we care about, since any changes to these
	// resources will affect our Channels. So, set up a watch here, that will cause
	// a global Resync for all the channels to take stock of their health when these change.
	filterFn := controller.FilterWithName(dispatcherName)

	// Call GlobalResync on kafkachannels.
	grCh := func(obj interface{}) {
		impl.GlobalResync(kafkaChannelInformer.Informer())
	}

	deploymentInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: filterFn,
		Handler:    controller.HandleAll(grCh),
	})
	serviceInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: filterFn,
		Handler:    controller.HandleAll(grCh),
	})
	endpointsInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: filterFn,
		Handler:    controller.HandleAll(grCh),
	})
	serviceAccountInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: filterFn,
		Handler:    controller.HandleAll(grCh),
	})
	roleBindingInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: filterFn,
		Handler:    controller.HandleAll(grCh),
	})

	podInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: knativeReconciler.ChainFilterFuncs(
			knativeReconciler.LabelFilterFunc(channelLabelKey, channelLabelValue, false),
			knativeReconciler.LabelFilterFunc(roleLabelKey, roleLabelValue, false),
		),
		Handler: cache.ResourceEventHandlerFuncs{
			// Cancel probing when a Pod is deleted
			DeleteFunc: getPodInformerEventHandler(ctx, logger, statusProber, impl, kafkaChannelInformer, "Delete"),
			AddFunc:    getPodInformerEventHandler(ctx, logger, statusProber, impl, kafkaChannelInformer, "Add"),
		},
	})

	return impl
}

func getPodInformerEventHandler(ctx context.Context, logger *zap.SugaredLogger, statusProber *status.Prober, impl *controller.Impl, kafkaChannelInformer kafkamessagingv1beta1.KafkaChannelInformer, handlerType string) func(obj interface{}) {
	return func(obj interface{}) {
		pod, ok := obj.(*corev1.Pod)
		if ok && pod != nil {
			logger.Debugw("%s pods. Refreshing pod probing.", handlerType,
				zap.String("pod", pod.GetName()))
			statusProber.RefreshPodProbing(ctx)
			impl.GlobalResync(kafkaChannelInformer.Informer())
		}
	}
}
<commit_msg>Global resync Kafka channels when config-kafka is changed (#486)<commit_after>/*
Copyright 2020 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package controller

import (
	"context"

	"github.com/kelseyhightower/envconfig"
	"go.uber.org/zap"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/tools/cache"

	"knative.dev/eventing-kafka/pkg/apis/messaging/v1beta1"
	"knative.dev/eventing-kafka/pkg/channel/consolidated/status"
	kafkamessagingv1beta1 "knative.dev/eventing-kafka/pkg/client/informers/externalversions/messaging/v1beta1"
	kafkaChannelClient "knative.dev/eventing-kafka/pkg/client/injection/client"
	"knative.dev/eventing-kafka/pkg/client/injection/informers/messaging/v1beta1/kafkachannel"
	kafkaChannelReconciler "knative.dev/eventing-kafka/pkg/client/injection/reconciler/messaging/v1beta1/kafkachannel"
	commonconfig "knative.dev/eventing-kafka/pkg/common/config"
	eventingduckv1 "knative.dev/eventing/pkg/apis/duck/v1"
	eventingClient "knative.dev/eventing/pkg/client/injection/client"
	kubeclient "knative.dev/pkg/client/injection/kube/client"
	"knative.dev/pkg/client/injection/kube/informers/apps/v1/deployment"
	endpointsinformer "knative.dev/pkg/client/injection/kube/informers/core/v1/endpoints"
	podinformer "knative.dev/pkg/client/injection/kube/informers/core/v1/pod"
	"knative.dev/pkg/client/injection/kube/informers/core/v1/service"
	"knative.dev/pkg/client/injection/kube/informers/core/v1/serviceaccount"
	"knative.dev/pkg/client/injection/kube/informers/rbac/v1/rolebinding"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/controller"
	"knative.dev/pkg/logging"
	knativeReconciler "knative.dev/pkg/reconciler"
	"knative.dev/pkg/system"
)

const (
	channelLabelKey   = "messaging.knative.dev/channel"
	channelLabelValue = "kafka-channel"
	roleLabelKey      = "messaging.knative.dev/role"
	roleLabelValue    = "dispatcher"
)

// NewController initializes the controller and is called by the generated code.
// Registers event handlers to enqueue events.
func NewController(
	ctx context.Context,
	cmw configmap.Watcher,
) *controller.Impl {
	logger := logging.FromContext(ctx)
	kafkaChannelInformer := kafkachannel.Get(ctx)
	deploymentInformer := deployment.Get(ctx)
	endpointsInformer := endpointsinformer.Get(ctx)
	serviceAccountInformer := serviceaccount.Get(ctx)
	roleBindingInformer := rolebinding.Get(ctx)
	serviceInformer := service.Get(ctx)
	podInformer := podinformer.Get(ctx)

	r := &Reconciler{
		systemNamespace:      system.Namespace(),
		KubeClientSet:        kubeclient.Get(ctx),
		kafkaClientSet:       kafkaChannelClient.Get(ctx),
		EventingClientSet:    eventingClient.Get(ctx),
		kafkachannelLister:   kafkaChannelInformer.Lister(),
		kafkachannelInformer: kafkaChannelInformer.Informer(),
		deploymentLister:     deploymentInformer.Lister(),
		serviceLister:        serviceInformer.Lister(),
		endpointsLister:      endpointsInformer.Lister(),
		serviceAccountLister: serviceAccountInformer.Lister(),
		roleBindingLister:    roleBindingInformer.Lister(),
	}

	env := &envConfig{}
	if err := envconfig.Process("", env); err != nil {
		logger.Panicf("unable to process Kafka channel's required environment variables: %v", err)
	}

	r.dispatcherImage = env.Image
	r.dispatcherServiceAccount = env.DispatcherServiceAccount

	impl := kafkaChannelReconciler.NewImpl(ctx, r)

	statusProber := status.NewProber(
		logger.Named("status-manager"),
		NewProbeTargetLister(logger, endpointsInformer.Lister()),
		func(c v1beta1.KafkaChannel, s eventingduckv1.SubscriberSpec) {
			logger.Debugf("Ready callback triggered for channel: %s/%s subscription: %s", c.Namespace, c.Name, string(s.UID))
			impl.EnqueueKey(types.NamespacedName{Namespace: c.Namespace, Name: c.Name})
		},
	)
	r.statusManager = statusProber
	statusProber.Start(ctx.Done())

	// Call GlobalResync on kafkachannels.
	grCh := func(obj interface{}) {
		logger.Info("Changes detected, doing global resync")
		impl.GlobalResync(kafkaChannelInformer.Informer())
	}

	handleKafkaConfigMapChange := func(ctx context.Context, configMap *corev1.ConfigMap) {
		logger.Info("Configmap is updated or, it is being read for the first time")
		r.updateKafkaConfig(ctx, configMap)
		grCh(configMap)
	}

	// Get and Watch the Kakfa config map and dynamically update Kafka configuration.
	err := commonconfig.InitializeKafkaConfigMapWatcher(ctx, cmw, logger, handleKafkaConfigMapChange, system.Namespace())
	if err != nil {
		logger.Fatal("Failed To Initialize ConfigMap Watcher", zap.Error(err))
	}

	logger.Info("Setting up event handlers")
	kafkaChannelInformer.Informer().AddEventHandler(controller.HandleAll(impl.Enqueue))

	// Set up watches for dispatcher resources we care about, since any changes to these
	// resources will affect our Channels. So, set up a watch here, that will cause
	// a global Resync for all the channels to take stock of their health when these change.
	filterFn := controller.FilterWithName(dispatcherName)

	deploymentInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: filterFn,
		Handler:    controller.HandleAll(grCh),
	})
	serviceInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: filterFn,
		Handler:    controller.HandleAll(grCh),
	})
	endpointsInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: filterFn,
		Handler:    controller.HandleAll(grCh),
	})
	serviceAccountInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: filterFn,
		Handler:    controller.HandleAll(grCh),
	})
	roleBindingInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: filterFn,
		Handler:    controller.HandleAll(grCh),
	})

	podInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		FilterFunc: knativeReconciler.ChainFilterFuncs(
			knativeReconciler.LabelFilterFunc(channelLabelKey, channelLabelValue, false),
			knativeReconciler.LabelFilterFunc(roleLabelKey, roleLabelValue, false),
		),
		Handler: cache.ResourceEventHandlerFuncs{
			// Cancel probing when a Pod is deleted
			DeleteFunc: getPodInformerEventHandler(ctx, logger, statusProber, impl, kafkaChannelInformer, "Delete"),
			AddFunc:    getPodInformerEventHandler(ctx, logger, statusProber, impl, kafkaChannelInformer, "Add"),
		},
	})

	return impl
}

func getPodInformerEventHandler(ctx context.Context, logger *zap.SugaredLogger, statusProber *status.Prober, impl *controller.Impl, kafkaChannelInformer kafkamessagingv1beta1.KafkaChannelInformer, handlerType string) func(obj interface{}) {
	return func(obj interface{}) {
		pod, ok := obj.(*corev1.Pod)
		if ok && pod != nil {
			logger.Debugw("%s pods. Refreshing pod probing.", handlerType,
				zap.String("pod", pod.GetName()))
			statusProber.RefreshPodProbing(ctx)
			impl.GlobalResync(kafkaChannelInformer.Informer())
		}
	}
}
<|endoftext|>