<commit_before>/*-
 * Copyright (c) 2016-2021, F5 Networks, 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 as3

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"time"

	log "github.com/F5Networks/k8s-bigip-ctlr/pkg/vlogger"
	routeclient "github.com/openshift/client-go/route/clientset/versioned/typed/route/v1"
)

const (
	timeoutNill   = 0 * time.Second
	timeoutSmall  = 3 * time.Second
	timeoutMedium = 30 * time.Second
	timeoutLarge  = 60 * time.Second
)

const (
	responseStatusOk                  = "statusOK"
	responseStatusCommon              = "statusCommonResponse"
	responseStatusNotFound            = "statusNotFound"
	responseStatusServiceUnavailable  = "statusServiceUnavailable"
	responseStatusUnprocessableEntity = "statusUnprocessableEntity"
	responseStatusDummy               = "dummy"
)

type PostManager struct {
	postChan   chan config
	httpClient *http.Client
	activeCfg  config
	PostParams
	Tenants map[string]bool
}

type PostParams struct {
	BIGIPUsername string
	BIGIPPassword string
	BIGIPURL      string
	TrustedCerts  string
	SSLInsecure   bool
	AS3PostDelay  int
	//Log the AS3 response body in Controller logs
	LogResponse   bool
	RouteClientV1 routeclient.RouteV1Interface
}

type config struct {
	data      string
	as3APIURL string
}

func NewPostManager(params PostParams) *PostManager {
	pm := &PostManager{
		postChan:   make(chan config, 1),
		PostParams: params,
		Tenants:    make(map[string]bool),
	}
	pm.setupBIGIPRESTClient()

	return pm
}

func (postMgr *PostManager) setupBIGIPRESTClient() {
	// Get the SystemCertPool, continue with an empty pool on error
	rootCAs, _ := x509.SystemCertPool()
	if rootCAs == nil {
		rootCAs = x509.NewCertPool()
	}
	// TODO: Make sure appMgr sets certificates in bigipInfo
	certs := []byte(postMgr.TrustedCerts)

	// Append our certs to the system pool
	if ok := rootCAs.AppendCertsFromPEM(certs); !ok {
		log.Debug("[AS3] No certs appended, using only system certs")
	}

	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: postMgr.SSLInsecure,
			RootCAs:            rootCAs,
		},
	}

	postMgr.httpClient = &http.Client{
		Transport: tr,
		Timeout:   timeoutLarge,
	}
}

func (postMgr *PostManager) getAS3APIURL(tenants []string) string {
	apiURL := postMgr.BIGIPURL + "/mgmt/shared/appsvcs/declare/" + strings.Join(tenants, ",")
	return apiURL
}

func (postMgr *PostManager) getAS3VersionURL() string {
	apiURL := postMgr.BIGIPURL + "/mgmt/shared/appsvcs/info"
	return apiURL
}

func getTimeDurationForErrorResponse(errRsp string) time.Duration {
	duration := timeoutNill
	switch errRsp {
	case responseStatusCommon:
		duration = timeoutMedium
	case responseStatusServiceUnavailable:
		duration = timeoutSmall
	}
	return duration
}

func (postMgr *PostManager) postConfigRequests(data string, url string) (bool, string) {
	cfg := config{
		data:      data,
		as3APIURL: url,
	}
	httpReqBody := bytes.NewBuffer([]byte(cfg.data))

	req, err := http.NewRequest("POST", cfg.as3APIURL, httpReqBody)
	if err != nil {
		log.Errorf("[AS3] Creating new HTTP request error: %v ", err)
		return false, responseStatusCommon
	}
	log.Debugf("[AS3] posting request to %v", cfg.as3APIURL)
	req.SetBasicAuth(postMgr.BIGIPUsername, postMgr.BIGIPPassword)

	httpResp, responseMap := postMgr.httpReq(req)
	if httpResp == nil || responseMap == nil {
		return false, responseStatusCommon
	}

	switch httpResp.StatusCode {
	case http.StatusOK, http.StatusCreated, http.StatusAccepted:
		return postMgr.handleResponseStatusOK(responseMap)
	case http.StatusServiceUnavailable:
		return postMgr.handleResponseStatusServiceUnavailable(responseMap)
	case http.StatusNotFound:
		return postMgr.handleResponseStatusNotFound(responseMap)
	case http.StatusUnprocessableEntity:
		return postMgr.handleStatusUnprocessableEntity(responseMap)
	default:
		return postMgr.handleResponseOthers(responseMap)
	}
}

func (postMgr *PostManager) GetBigipAS3Version() (string, string, string, error) {
	url := postMgr.getAS3VersionURL()
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Errorf("[AS3] Creating new HTTP request error: %v ", err)
		return "", "", "", err
	}

	log.Debugf("[AS3] posting GET BIGIP AS3 Version request on %v", url)
	req.SetBasicAuth(postMgr.BIGIPUsername, postMgr.BIGIPPassword)

	httpResp, responseMap := postMgr.httpReq(req)
	if httpResp == nil || responseMap == nil {
		return "", "", "", fmt.Errorf("Internal Error")
	}

	switch httpResp.StatusCode {
	case http.StatusOK:
		if responseMap["version"] != nil {
			as3VersionStr := responseMap["version"].(string)
			as3versionreleaseStr := responseMap["release"].(string)
			as3SchemaVersion := responseMap["schemaCurrent"].(string)
			return as3VersionStr, as3versionreleaseStr, as3SchemaVersion, nil
		}
	case http.StatusNotFound:
		responseMap["code"] = int(responseMap["code"].(float64))
		if responseMap["code"] == http.StatusNotFound {
			return "", "", "", fmt.Errorf("AS3 RPM is not installed on BIGIP,"+
				" Error response from BIGIP with status code %v", httpResp.StatusCode)
		}
		// In case of 503 status code : CIS will exit and auto restart of the
		// controller might fetch the BIGIP version once BIGIP is available.
	}
	return "", "", "", fmt.Errorf("Error response from BIGIP with status code %v", httpResp.StatusCode)
}

// GetBigipRegKey ...
func (postMgr *PostManager) GetBigipRegKey() (string, error) {
	url := postMgr.getBigipRegKeyURL()
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Errorf("Creating new HTTP request error: %v ", err)
		return "", err
	}

	log.Debugf("Posting GET BIGIP Reg Key request on %v", url)
	req.SetBasicAuth(postMgr.BIGIPUsername, postMgr.BIGIPPassword)

	httpResp, responseMap := postMgr.httpReq(req)
	if httpResp == nil || responseMap == nil {
		return "", fmt.Errorf("Internal Error")
	}

	switch httpResp.StatusCode {
	case http.StatusOK:
		if responseMap["registrationKey"] != nil {
			registrationKey := responseMap["registrationKey"].(string)
			return registrationKey, nil
		}
	case http.StatusNotFound:
		if int(responseMap["code"].(float64)) == http.StatusNotFound {
			return "", fmt.Errorf("AS3 RPM is not installed on BIGIP,"+
				" Error response from BIGIP with status code %v", httpResp.StatusCode)
		}
	}
	return "", fmt.Errorf("Error response from BIGIP with status code %v", httpResp.StatusCode)
}

func (postMgr *PostManager) httpReq(request *http.Request) (*http.Response, map[string]interface{}) {
	httpResp, err := postMgr.httpClient.Do(request)
	if err != nil {
		log.Errorf("[AS3] REST call error: %v ", err)
		return nil, nil
	}
	defer httpResp.Body.Close()

	body, err := ioutil.ReadAll(httpResp.Body)
	if err != nil {
		log.Errorf("[AS3] REST call response error: %v ", err)
		return nil, nil
	}
	var response map[string]interface{}
	err = json.Unmarshal(body, &response)
	if err != nil {
		log.Errorf("[AS3] Response body unmarshal failed: %v\n", err)
		if postMgr.LogResponse {
			log.Errorf("[AS3] Raw response from Big-IP: %v", string(body))
		}
		return nil, nil
	}
	return httpResp, response
}

func (postMgr *PostManager) handleResponseStatusOK(responseMap map[string]interface{}) (bool, string) {
	//traverse all response results
	results := (responseMap["results"]).([]interface{})
	for _, value := range results {
		v := value.(map[string]interface{})
		//log result with code, tenant and message
		log.Debugf("[AS3] Response from BIG-IP: code: %v --- tenant:%v --- message: %v", v["code"], v["tenant"], v["message"])
	}
	return true, responseStatusOk
}

func (postMgr *PostManager) handleResponseStatusServiceUnavailable(responseMap map[string]interface{}) (bool, string) {
	log.Errorf("[AS3] Big-IP Responded with error code: %v", responseMap["code"])
	log.Debugf("[AS3] Response from BIG-IP: BIG-IP is busy, waiting %v seconds and re-posting the declaration", timeoutSmall)
	//return postMgr.postOnEventOrTimeout(timeoutSmall, cfg)
	return false, responseStatusServiceUnavailable
}

func (postMgr *PostManager) handleResponseStatusNotFound(responseMap map[string]interface{}) (bool, string) {
	if err, ok := (responseMap["error"]).(map[string]interface{}); ok {
		log.Errorf("[AS3] Big-IP Responded with error code: %v", err["code"])
	} else {
		log.Errorf("[AS3] Big-IP Responded with error code: %v", http.StatusNotFound)
	}

	if postMgr.LogResponse {
		log.Errorf("[AS3] Raw response from Big-IP: %v ", responseMap)
	}
	return true, responseStatusNotFound
}

func (postMgr *PostManager) handleStatusUnprocessableEntity(responseMap map[string]interface{}) (bool, string) {
	if err, ok := (responseMap["error"]).(map[string]interface{}); ok {
		log.Errorf("[AS3] Big-IP Responded with error code: %v", err["code"])
	} else {
		log.Errorf("[AS3] Big-IP Responded with error code: %v", http.StatusUnprocessableEntity)
	}

	if postMgr.LogResponse {
		log.Errorf("[AS3] Raw response from Big-IP: %v ", responseMap)
	}
	return true, responseStatusUnprocessableEntity
}

func (postMgr *PostManager) handleResponseOthers(responseMap map[string]interface{}) (bool, string) {
	if results, ok := (responseMap["results"]).([]interface{}); ok {
		for _, value := range results {
			v := value.(map[string]interface{})
			//log result with code, tenant and message
			log.Errorf("[AS3] Response from BIG-IP: code: %v --- tenant:%v --- message: %v", v["code"], v["tenant"], v["message"])
		}
	} else if err, ok := (responseMap["error"]).(map[string]interface{}); ok {
		log.Errorf("[AS3] Big-IP Responded with error code: %v", err["code"])
	} else {
		log.Errorf("[AS3] Big-IP Responded with code: %v", responseMap["code"])
	}

	if postMgr.LogResponse {
		log.Errorf("[AS3] Raw response from Big-IP: %v ", responseMap)
	}
	//return postMgr.postOnEventOrTimeout(timeoutMedium, cfg)
	return false, responseStatusCommon
}

func (postMgr *PostManager) getBigipRegKeyURL() string {
	apiURL := postMgr.BIGIPURL + "/mgmt/tm/shared/licensing/registration"
	return apiURL

}
<commit_msg>Fix for CIS not retrying when AS3 response code is 422 (#2276)<commit_after>/*-
 * Copyright (c) 2016-2021, F5 Networks, 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 as3

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"time"

	log "github.com/F5Networks/k8s-bigip-ctlr/pkg/vlogger"
	routeclient "github.com/openshift/client-go/route/clientset/versioned/typed/route/v1"
)

const (
	timeoutNill   = 0 * time.Second
	timeoutSmall  = 3 * time.Second
	timeoutMedium = 30 * time.Second
	timeoutLarge  = 60 * time.Second
)

const (
	responseStatusOk                  = "statusOK"
	responseStatusCommon              = "statusCommonResponse"
	responseStatusNotFound            = "statusNotFound"
	responseStatusServiceUnavailable  = "statusServiceUnavailable"
	responseStatusUnprocessableEntity = "statusUnprocessableEntity"
	responseStatusDummy               = "dummy"
)

type PostManager struct {
	postChan   chan config
	httpClient *http.Client
	activeCfg  config
	PostParams
	Tenants map[string]bool
}

type PostParams struct {
	BIGIPUsername string
	BIGIPPassword string
	BIGIPURL      string
	TrustedCerts  string
	SSLInsecure   bool
	AS3PostDelay  int
	//Log the AS3 response body in Controller logs
	LogResponse   bool
	RouteClientV1 routeclient.RouteV1Interface
}

type config struct {
	data      string
	as3APIURL string
}

func NewPostManager(params PostParams) *PostManager {
	pm := &PostManager{
		postChan:   make(chan config, 1),
		PostParams: params,
		Tenants:    make(map[string]bool),
	}
	pm.setupBIGIPRESTClient()

	return pm
}

func (postMgr *PostManager) setupBIGIPRESTClient() {
	// Get the SystemCertPool, continue with an empty pool on error
	rootCAs, _ := x509.SystemCertPool()
	if rootCAs == nil {
		rootCAs = x509.NewCertPool()
	}
	// TODO: Make sure appMgr sets certificates in bigipInfo
	certs := []byte(postMgr.TrustedCerts)

	// Append our certs to the system pool
	if ok := rootCAs.AppendCertsFromPEM(certs); !ok {
		log.Debug("[AS3] No certs appended, using only system certs")
	}

	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: postMgr.SSLInsecure,
			RootCAs:            rootCAs,
		},
	}

	postMgr.httpClient = &http.Client{
		Transport: tr,
		Timeout:   timeoutLarge,
	}
}

func (postMgr *PostManager) getAS3APIURL(tenants []string) string {
	apiURL := postMgr.BIGIPURL + "/mgmt/shared/appsvcs/declare/" + strings.Join(tenants, ",")
	return apiURL
}

func (postMgr *PostManager) getAS3VersionURL() string {
	apiURL := postMgr.BIGIPURL + "/mgmt/shared/appsvcs/info"
	return apiURL
}

func getTimeDurationForErrorResponse(errRsp string) time.Duration {
	duration := timeoutNill
	switch errRsp {
	case responseStatusCommon, responseStatusUnprocessableEntity:
		duration = timeoutMedium
	case responseStatusServiceUnavailable:
		duration = timeoutSmall
	}
	return duration
}

func (postMgr *PostManager) postConfigRequests(data string, url string) (bool, string) {
	cfg := config{
		data:      data,
		as3APIURL: url,
	}
	httpReqBody := bytes.NewBuffer([]byte(cfg.data))

	req, err := http.NewRequest("POST", cfg.as3APIURL, httpReqBody)
	if err != nil {
		log.Errorf("[AS3] Creating new HTTP request error: %v ", err)
		return false, responseStatusCommon
	}
	log.Debugf("[AS3] posting request to %v", cfg.as3APIURL)
	req.SetBasicAuth(postMgr.BIGIPUsername, postMgr.BIGIPPassword)

	httpResp, responseMap := postMgr.httpReq(req)
	if httpResp == nil || responseMap == nil {
		return false, responseStatusCommon
	}

	switch httpResp.StatusCode {
	case http.StatusOK, http.StatusCreated, http.StatusAccepted:
		return postMgr.handleResponseStatusOK(responseMap)
	case http.StatusServiceUnavailable:
		return postMgr.handleResponseStatusServiceUnavailable(responseMap)
	case http.StatusNotFound:
		return postMgr.handleResponseStatusNotFound(responseMap)
	case http.StatusUnprocessableEntity:
		return postMgr.handleStatusUnprocessableEntity(responseMap)
	default:
		return postMgr.handleResponseOthers(responseMap)
	}
}

func (postMgr *PostManager) GetBigipAS3Version() (string, string, string, error) {
	url := postMgr.getAS3VersionURL()
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Errorf("[AS3] Creating new HTTP request error: %v ", err)
		return "", "", "", err
	}

	log.Debugf("[AS3] posting GET BIGIP AS3 Version request on %v", url)
	req.SetBasicAuth(postMgr.BIGIPUsername, postMgr.BIGIPPassword)

	httpResp, responseMap := postMgr.httpReq(req)
	if httpResp == nil || responseMap == nil {
		return "", "", "", fmt.Errorf("Internal Error")
	}

	switch httpResp.StatusCode {
	case http.StatusOK:
		if responseMap["version"] != nil {
			as3VersionStr := responseMap["version"].(string)
			as3versionreleaseStr := responseMap["release"].(string)
			as3SchemaVersion := responseMap["schemaCurrent"].(string)
			return as3VersionStr, as3versionreleaseStr, as3SchemaVersion, nil
		}
	case http.StatusNotFound:
		responseMap["code"] = int(responseMap["code"].(float64))
		if responseMap["code"] == http.StatusNotFound {
			return "", "", "", fmt.Errorf("AS3 RPM is not installed on BIGIP,"+
				" Error response from BIGIP with status code %v", httpResp.StatusCode)
		}
		// In case of 503 status code : CIS will exit and auto restart of the
		// controller might fetch the BIGIP version once BIGIP is available.
	}
	return "", "", "", fmt.Errorf("Error response from BIGIP with status code %v", httpResp.StatusCode)
}

// GetBigipRegKey ...
func (postMgr *PostManager) GetBigipRegKey() (string, error) {
	url := postMgr.getBigipRegKeyURL()
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Errorf("Creating new HTTP request error: %v ", err)
		return "", err
	}

	log.Debugf("Posting GET BIGIP Reg Key request on %v", url)
	req.SetBasicAuth(postMgr.BIGIPUsername, postMgr.BIGIPPassword)

	httpResp, responseMap := postMgr.httpReq(req)
	if httpResp == nil || responseMap == nil {
		return "", fmt.Errorf("Internal Error")
	}

	switch httpResp.StatusCode {
	case http.StatusOK:
		if responseMap["registrationKey"] != nil {
			registrationKey := responseMap["registrationKey"].(string)
			return registrationKey, nil
		}
	case http.StatusNotFound:
		if int(responseMap["code"].(float64)) == http.StatusNotFound {
			return "", fmt.Errorf("AS3 RPM is not installed on BIGIP,"+
				" Error response from BIGIP with status code %v", httpResp.StatusCode)
		}
	}
	return "", fmt.Errorf("Error response from BIGIP with status code %v", httpResp.StatusCode)
}

func (postMgr *PostManager) httpReq(request *http.Request) (*http.Response, map[string]interface{}) {
	httpResp, err := postMgr.httpClient.Do(request)
	if err != nil {
		log.Errorf("[AS3] REST call error: %v ", err)
		return nil, nil
	}
	defer httpResp.Body.Close()

	body, err := ioutil.ReadAll(httpResp.Body)
	if err != nil {
		log.Errorf("[AS3] REST call response error: %v ", err)
		return nil, nil
	}
	var response map[string]interface{}
	err = json.Unmarshal(body, &response)
	if err != nil {
		log.Errorf("[AS3] Response body unmarshal failed: %v\n", err)
		if postMgr.LogResponse {
			log.Errorf("[AS3] Raw response from Big-IP: %v", string(body))
		}
		return nil, nil
	}
	return httpResp, response
}

func (postMgr *PostManager) handleResponseStatusOK(responseMap map[string]interface{}) (bool, string) {
	//traverse all response results
	results := (responseMap["results"]).([]interface{})
	for _, value := range results {
		v := value.(map[string]interface{})
		//log result with code, tenant and message
		log.Debugf("[AS3] Response from BIG-IP: code: %v --- tenant:%v --- message: %v", v["code"], v["tenant"], v["message"])
	}
	return true, responseStatusOk
}

func (postMgr *PostManager) handleResponseStatusServiceUnavailable(responseMap map[string]interface{}) (bool, string) {
	log.Errorf("[AS3] Big-IP Responded with error code: %v", responseMap["code"])
	log.Debugf("[AS3] Response from BIG-IP: BIG-IP is busy, waiting %v seconds and re-posting the declaration", timeoutSmall)
	//return postMgr.postOnEventOrTimeout(timeoutSmall, cfg)
	return false, responseStatusServiceUnavailable
}

func (postMgr *PostManager) handleResponseStatusNotFound(responseMap map[string]interface{}) (bool, string) {
	if err, ok := (responseMap["error"]).(map[string]interface{}); ok {
		log.Errorf("[AS3] Big-IP Responded with error code: %v", err["code"])
	} else {
		log.Errorf("[AS3] Big-IP Responded with error code: %v", http.StatusNotFound)
	}

	if postMgr.LogResponse {
		log.Errorf("[AS3] Raw response from Big-IP: %v ", responseMap)
	}
	return true, responseStatusNotFound
}

func (postMgr *PostManager) handleStatusUnprocessableEntity(responseMap map[string]interface{}) (bool, string) {
	if err, ok := (responseMap["error"]).(map[string]interface{}); ok {
		log.Errorf("[AS3] Big-IP Responded with error code: %v", err["code"])
	} else {
		log.Errorf("[AS3] Big-IP Responded with error code: %v", http.StatusUnprocessableEntity)
	}

	if postMgr.LogResponse {
		log.Errorf("[AS3] Raw response from Big-IP: %v ", responseMap)
	}
	return false, responseStatusUnprocessableEntity
}

func (postMgr *PostManager) handleResponseOthers(responseMap map[string]interface{}) (bool, string) {
	if results, ok := (responseMap["results"]).([]interface{}); ok {
		for _, value := range results {
			v := value.(map[string]interface{})
			//log result with code, tenant and message
			log.Errorf("[AS3] Response from BIG-IP: code: %v --- tenant:%v --- message: %v", v["code"], v["tenant"], v["message"])
		}
	} else if err, ok := (responseMap["error"]).(map[string]interface{}); ok {
		log.Errorf("[AS3] Big-IP Responded with error code: %v", err["code"])
	} else {
		log.Errorf("[AS3] Big-IP Responded with code: %v", responseMap["code"])
	}

	if postMgr.LogResponse {
		log.Errorf("[AS3] Raw response from Big-IP: %v ", responseMap)
	}
	//return postMgr.postOnEventOrTimeout(timeoutMedium, cfg)
	return false, responseStatusCommon
}

func (postMgr *PostManager) getBigipRegKeyURL() string {
	apiURL := postMgr.BIGIPURL + "/mgmt/tm/shared/licensing/registration"
	return apiURL

}
<|endoftext|>


<commit_before>package proxy

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	gmux "github.com/gorilla/mux"
	"github.com/rancher/rancher/pkg/features"
	v3 "github.com/rancher/rancher/pkg/generated/controllers/management.cattle.io/v3"
	managementv3 "github.com/rancher/rancher/pkg/generated/norman/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/steve/pkg/proxy"
	authzv1 "k8s.io/api/authorization/v1"
	"k8s.io/apiserver/pkg/authentication/user"
	"k8s.io/apiserver/pkg/authorization/authorizer"
	"k8s.io/apiserver/pkg/authorization/authorizerfactory"
	"k8s.io/apiserver/pkg/endpoints/request"
	v1 "k8s.io/client-go/kubernetes/typed/authorization/v1"
	"k8s.io/client-go/rest"
)

type Handler struct {
	authorizer    authorizer.Authorizer
	dialerFactory ClusterDialerFactory
	clusters      v3.ClusterCache
}

type ClusterDialerFactory interface {
	ClusterDialer(clusterID string) func(ctx context.Context, network, address string) (net.Conn, error)
}

func RewriteLocalCluster(next http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		if strings.HasPrefix(req.URL.Path, "/k8s/clusters/local") {
			req.URL.Path = strings.TrimPrefix(req.URL.Path, "/k8s/clusters/local")
			if req.URL.Path == "" {
				req.URL.Path = "/"
			}
		}
		next.ServeHTTP(rw, req)
	})
}

func NewProxyMiddleware(sar v1.SubjectAccessReviewInterface,
	dialerFactory ClusterDialerFactory,
	clusters v3.ClusterCache,
	localSupport bool,
	localCluster http.Handler) (func(http.Handler) http.Handler, error) {
	cfg := authorizerfactory.DelegatingAuthorizerConfig{
		SubjectAccessReviewClient: sar,
		AllowCacheTTL:             time.Second * time.Duration(settings.AuthorizationCacheTTLSeconds.GetInt()),
		DenyCacheTTL:              time.Second * time.Duration(settings.AuthorizationDenyCacheTTLSeconds.GetInt()),
	}

	authorizer, err := cfg.New()
	if err != nil {
		return nil, err
	}

	proxyHandler := NewProxyHandler(authorizer, dialerFactory, clusters)

	mux := gmux.NewRouter()
	mux.UseEncodedPath()
	mux.Path("/v1/management.cattle.io.clusters/{clusterID}").Queries("link", "shell").HandlerFunc(routeToShellProxy(localSupport, localCluster, mux, proxyHandler))
	mux.Path("/v3/clusters/{clusterID}").Queries("shell", "true").HandlerFunc(routeToShellProxy(localSupport, localCluster, mux, proxyHandler))
	mux.Path("/{prefix:k8s/clusters/[^/]+}{suffix:/v1.*}").MatcherFunc(proxyHandler.MatchNonLegacy("/k8s/clusters/", true)).Handler(proxyHandler)
	mux.Path("/{prefix:k8s/clusters/[^/]+}{suffix:.*}").MatcherFunc(proxyHandler.MatchNonLegacy("/k8s/clusters/", false)).Handler(proxyHandler)

	return func(handler http.Handler) http.Handler {
		return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
			mux.NotFoundHandler = handler
			mux.ServeHTTP(rw, req)
		})
	}, nil
}

func routeToShellLink(rw http.ResponseWriter, req *http.Request, next http.Handler) {
	clusterID := gmux.Vars(req)["id"]
	if clusterID == "local" {
		req.URL.RawPath = "/v1/management.cattle.io.clusters/local"
	} else {
		req.URL.RawPath = fmt.Sprintf("/k8s/clusters/%s/v1/management.cattle.io.clusters/local", clusterID)
	}
	req.URL.Path = req.URL.RawPath
	req.URL.RawQuery = url.Values{
		"link": []string{"shell"},
	}.Encode()
	next.ServeHTTP(rw, req)
}

func routeToShellProxy(localSupport bool, localCluster http.Handler, mux *gmux.Router, proxyHandler *Handler) func(rw http.ResponseWriter, r *http.Request) {
	return func(rw http.ResponseWriter, r *http.Request) {
		vars := gmux.Vars(r)
		cluster := vars["clusterID"]
		if cluster == "local" {
			if localSupport {
				localCluster.ServeHTTP(rw, r)
			} else {
				mux.NotFoundHandler.ServeHTTP(rw, r)
			}
			return
		}
		vars["prefix"] = "k8s/clusters/" + cluster
		vars["suffix"] = "/v1/management.cattle.io.clusters/local"
		// Ensure shell link is set
		q := r.URL.Query()
		q.Set("link", "shell")
		r.URL.RawQuery = q.Encode()
		r.URL.Path = "/k8s/clusters/" + cluster + "/v1/management.cattle.io.clusters/local"
		proxyHandler.ServeHTTP(rw, r)
	}
}

func NewProxyHandler(authorizer authorizer.Authorizer,
	dialerFactory ClusterDialerFactory,
	clusters v3.ClusterCache) *Handler {
	return &Handler{
		authorizer:    authorizer,
		dialerFactory: dialerFactory,
		clusters:      clusters,
	}
}

func (h *Handler) MatchNonLegacy(prefix string, force bool) gmux.MatcherFunc {
	return func(req *http.Request, match *gmux.RouteMatch) bool {
		if !features.SteveProxy.Enabled() && !force {
			return false
		}

		clusterID := strings.TrimPrefix(req.URL.Path, prefix)
		clusterID = strings.SplitN(clusterID, "/", 2)[0]
		if match.Vars == nil {
			match.Vars = map[string]string{}
		}
		match.Vars["clusterID"] = clusterID

		cluster, err := h.clusters.Get(clusterID)
		if err != nil {
			return true
		}

		// No steve means we are upgrading a node that doesn't have the right proxy
		return cluster.Status.AgentFeatures[features.Steve.Name()]
	}
}

func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	user, ok := request.UserFrom(req.Context())
	if !ok {
		rw.WriteHeader(http.StatusUnauthorized)
		return
	}

	prefix := "/" + gmux.Vars(req)["prefix"]
	clusterID := gmux.Vars(req)["clusterID"]

	if !h.canAccess(req.Context(), user, clusterID) {
		rw.WriteHeader(http.StatusUnauthorized)
		return
	}

	handler, err := h.next(clusterID, prefix)
	if err != nil {
		rw.WriteHeader(http.StatusInternalServerError)
		rw.Write([]byte(err.Error()))
		return
	}

	handler.ServeHTTP(rw, req)
}

func (h *Handler) dialer(ctx context.Context, network, address string) (net.Conn, error) {
	host, _, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}
	dialer := h.dialerFactory.ClusterDialer(host)
	return dialer(ctx, network, "127.0.0.1:6080")
}

func (h *Handler) next(clusterID, prefix string) (http.Handler, error) {
	cfg := &rest.Config{
		// this is bogus, the dialer will change it to 127.0.0.1:6080, but the clusterID is used to lookup the tunnel
		// connect
		Host:      "http://" + clusterID,
		UserAgent: rest.DefaultKubernetesUserAgent() + " cluster " + clusterID,
		// Ensure this function pointer does not change per invocation so that we don't
		// blow out the cache.
		Dial: h.dialer,
	}

	next := proxy.ImpersonatingHandler(prefix, cfg)
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		req.Header.Set("X-API-URL-Prefix", prefix)
		next.ServeHTTP(rw, req)
	}), nil
}

func (h *Handler) canAccess(ctx context.Context, user user.Info, clusterID string) bool {
	extra := map[string]authzv1.ExtraValue{}
	for k, v := range user.GetExtra() {
		extra[k] = v
	}

	resp, _, err := h.authorizer.Authorize(ctx, authorizer.AttributesRecord{
		ResourceRequest: true,
		User:            user,
		Verb:            "get",
		APIGroup:        managementv3.GroupName,
		APIVersion:      managementv3.Version,
		Resource:        "clusters",
		Name:            clusterID,
	})

	return err == nil && resp == authorizer.DecisionAllow
}
<commit_msg>Fix routing to local shell<commit_after>package proxy

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	gmux "github.com/gorilla/mux"
	"github.com/rancher/rancher/pkg/features"
	v3 "github.com/rancher/rancher/pkg/generated/controllers/management.cattle.io/v3"
	managementv3 "github.com/rancher/rancher/pkg/generated/norman/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/steve/pkg/proxy"
	authzv1 "k8s.io/api/authorization/v1"
	"k8s.io/apiserver/pkg/authentication/user"
	"k8s.io/apiserver/pkg/authorization/authorizer"
	"k8s.io/apiserver/pkg/authorization/authorizerfactory"
	"k8s.io/apiserver/pkg/endpoints/request"
	v1 "k8s.io/client-go/kubernetes/typed/authorization/v1"
	"k8s.io/client-go/rest"
)

type Handler struct {
	authorizer    authorizer.Authorizer
	dialerFactory ClusterDialerFactory
	clusters      v3.ClusterCache
}

type ClusterDialerFactory interface {
	ClusterDialer(clusterID string) func(ctx context.Context, network, address string) (net.Conn, error)
}

func RewriteLocalCluster(next http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		if strings.HasPrefix(req.URL.Path, "/k8s/clusters/local") {
			req.URL.Path = strings.TrimPrefix(req.URL.Path, "/k8s/clusters/local")
			if req.URL.Path == "" {
				req.URL.Path = "/"
			}
		}
		next.ServeHTTP(rw, req)
	})
}

func NewProxyMiddleware(sar v1.SubjectAccessReviewInterface,
	dialerFactory ClusterDialerFactory,
	clusters v3.ClusterCache,
	localSupport bool,
	localCluster http.Handler) (func(http.Handler) http.Handler, error) {
	cfg := authorizerfactory.DelegatingAuthorizerConfig{
		SubjectAccessReviewClient: sar,
		AllowCacheTTL:             time.Second * time.Duration(settings.AuthorizationCacheTTLSeconds.GetInt()),
		DenyCacheTTL:              time.Second * time.Duration(settings.AuthorizationDenyCacheTTLSeconds.GetInt()),
	}

	authorizer, err := cfg.New()
	if err != nil {
		return nil, err
	}

	proxyHandler := NewProxyHandler(authorizer, dialerFactory, clusters)

	mux := gmux.NewRouter()
	mux.UseEncodedPath()
	mux.Path("/v1/management.cattle.io.clusters/{clusterID}").Queries("link", "shell").HandlerFunc(routeToShellProxy(localSupport, localCluster, mux, proxyHandler))
	mux.Path("/v3/clusters/{clusterID}").Queries("shell", "true").HandlerFunc(routeToShellProxy(localSupport, localCluster, mux, proxyHandler))
	mux.Path("/{prefix:k8s/clusters/[^/]+}{suffix:/v1.*}").MatcherFunc(proxyHandler.MatchNonLegacy("/k8s/clusters/", true)).Handler(proxyHandler)
	mux.Path("/{prefix:k8s/clusters/[^/]+}{suffix:.*}").MatcherFunc(proxyHandler.MatchNonLegacy("/k8s/clusters/", false)).Handler(proxyHandler)

	return func(handler http.Handler) http.Handler {
		return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
			mux.NotFoundHandler = handler
			mux.ServeHTTP(rw, req)
		})
	}, nil
}

func routeToShellLink(rw http.ResponseWriter, req *http.Request, next http.Handler) {
	clusterID := gmux.Vars(req)["id"]
	if clusterID == "local" {
		req.URL.RawPath = "/v1/management.cattle.io.clusters/local"
	} else {
		req.URL.RawPath = fmt.Sprintf("/k8s/clusters/%s/v1/management.cattle.io.clusters/local", clusterID)
	}
	req.URL.Path = req.URL.RawPath
	req.URL.RawQuery = url.Values{
		"link": []string{"shell"},
	}.Encode()
	next.ServeHTTP(rw, req)
}

func routeToShellProxy(localSupport bool, localCluster http.Handler, mux *gmux.Router, proxyHandler *Handler) func(rw http.ResponseWriter, r *http.Request) {
	return func(rw http.ResponseWriter, r *http.Request) {
		vars := gmux.Vars(r)
		cluster := vars["clusterID"]
		if cluster == "local" {
			if localSupport {
				q := r.URL.Query()
				q.Set("link", "shell")
				r.URL.RawQuery = q.Encode()
				r.URL.Path = "/v1/management.cattle.io.clusters/local"
				localCluster.ServeHTTP(rw, r)
			} else {
				mux.NotFoundHandler.ServeHTTP(rw, r)
			}
			return
		}
		vars["prefix"] = "k8s/clusters/" + cluster
		vars["suffix"] = "/v1/management.cattle.io.clusters/local"
		// Ensure shell link is set
		q := r.URL.Query()
		q.Set("link", "shell")
		r.URL.RawQuery = q.Encode()
		r.URL.Path = "/k8s/clusters/" + cluster + "/v1/management.cattle.io.clusters/local"
		proxyHandler.ServeHTTP(rw, r)
	}
}

func NewProxyHandler(authorizer authorizer.Authorizer,
	dialerFactory ClusterDialerFactory,
	clusters v3.ClusterCache) *Handler {
	return &Handler{
		authorizer:    authorizer,
		dialerFactory: dialerFactory,
		clusters:      clusters,
	}
}

func (h *Handler) MatchNonLegacy(prefix string, force bool) gmux.MatcherFunc {
	return func(req *http.Request, match *gmux.RouteMatch) bool {
		if !features.SteveProxy.Enabled() && !force {
			return false
		}

		clusterID := strings.TrimPrefix(req.URL.Path, prefix)
		clusterID = strings.SplitN(clusterID, "/", 2)[0]
		if match.Vars == nil {
			match.Vars = map[string]string{}
		}
		match.Vars["clusterID"] = clusterID

		cluster, err := h.clusters.Get(clusterID)
		if err != nil {
			return true
		}

		// No steve means we are upgrading a node that doesn't have the right proxy
		return cluster.Status.AgentFeatures[features.Steve.Name()]
	}
}

func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	user, ok := request.UserFrom(req.Context())
	if !ok {
		rw.WriteHeader(http.StatusUnauthorized)
		return
	}

	prefix := "/" + gmux.Vars(req)["prefix"]
	clusterID := gmux.Vars(req)["clusterID"]

	if !h.canAccess(req.Context(), user, clusterID) {
		rw.WriteHeader(http.StatusUnauthorized)
		return
	}

	handler, err := h.next(clusterID, prefix)
	if err != nil {
		rw.WriteHeader(http.StatusInternalServerError)
		rw.Write([]byte(err.Error()))
		return
	}

	handler.ServeHTTP(rw, req)
}

func (h *Handler) dialer(ctx context.Context, network, address string) (net.Conn, error) {
	host, _, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}
	dialer := h.dialerFactory.ClusterDialer(host)
	return dialer(ctx, network, "127.0.0.1:6080")
}

func (h *Handler) next(clusterID, prefix string) (http.Handler, error) {
	cfg := &rest.Config{
		// this is bogus, the dialer will change it to 127.0.0.1:6080, but the clusterID is used to lookup the tunnel
		// connect
		Host:      "http://" + clusterID,
		UserAgent: rest.DefaultKubernetesUserAgent() + " cluster " + clusterID,
		// Ensure this function pointer does not change per invocation so that we don't
		// blow out the cache.
		Dial: h.dialer,
	}

	next := proxy.ImpersonatingHandler(prefix, cfg)
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		req.Header.Set("X-API-URL-Prefix", prefix)
		next.ServeHTTP(rw, req)
	}), nil
}

func (h *Handler) canAccess(ctx context.Context, user user.Info, clusterID string) bool {
	extra := map[string]authzv1.ExtraValue{}
	for k, v := range user.GetExtra() {
		extra[k] = v
	}

	resp, _, err := h.authorizer.Authorize(ctx, authorizer.AttributesRecord{
		ResourceRequest: true,
		User:            user,
		Verb:            "get",
		APIGroup:        managementv3.GroupName,
		APIVersion:      managementv3.Version,
		Resource:        "clusters",
		Name:            clusterID,
	})

	return err == nil && resp == authorizer.DecisionAllow
}
<|endoftext|>


<commit_before>package kubectl

import (
	"strings"
	"text/tabwriter"

	"github.com/spf13/cobra"

	kclient "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	kctl "github.com/GoogleCloudPlatform/kubernetes/pkg/kubectl"
	kubecmd "github.com/GoogleCloudPlatform/kubernetes/pkg/kubectl/cmd"
	buildapi "github.com/openshift/origin/pkg/build/api"
	"github.com/openshift/origin/pkg/client"
)

func DescriberFor(kind string, c *client.Client, cmd *cobra.Command) (kctl.Describer, bool) {
	switch kind {
	case "Build":
		return &BuildDescriber{c}, true
	case "BuildConfig":
		return &BuildConfigDescriber{c, cmd}, true
	case "Deployment":
		return &DeploymentDescriber{c}, true
	case "DeploymentConfig":
		return &DeploymentDescriber{c}, true
	case "Image":
		return &ImageDescriber{c}, true
	case "ImageRepository":
		return &ImageRepositoryDescriber{c}, true
	case "Route":
		return &RouteDescriber{c}, true
	case "Project":
		return &ProjectDescriber{c}, true
	}
	return nil, false
}

// BuildDescriber generates information about a build
type BuildDescriber struct {
	client.Interface
}

func (d *BuildDescriber) DescribeParameters(p buildapi.BuildParameters, out *tabwriter.Writer) {
	formatString(out, "Strategy", p.Strategy.Type)
	formatString(out, "Source Type", p.Source.Type)
	if p.Source.Git != nil {
		formatString(out, "URL", p.Source.Git.URI)
		if len(p.Source.Git.Ref) > 0 {
			formatString(out, "Ref", p.Source.Git.Ref)
		}
	}
	formatString(out, "Output Image", p.Output.ImageTag)
	formatString(out, "Output Registry", p.Output.Registry)
}

func (d *BuildDescriber) Describe(namespace, name string) (string, error) {
	c := d.Builds(namespace)
	build, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, build.ObjectMeta)
		formatString(out, "Status", bold(build.Status))
		formatString(out, "Build Pod", build.PodName)
		d.DescribeParameters(build.Parameters, out)
		return nil
	})
}

// BuildConfigDescriber generates information about a buildConfig
type BuildConfigDescriber struct {
	client.Interface
	*cobra.Command
}

func (d *BuildConfigDescriber) Describe(namespace, name string) (string, error) {
	c := d.BuildConfigs(namespace)
	buildConfig, err := c.Get(name)
	if err != nil {
		return "", err
	}

	var kubeConfig *kclient.Config
	if d.Command != nil {
		kubeConfig = kubecmd.GetKubeConfig(d.Command)
	}

	webhooks := webhookURL(buildConfig, kubeConfig)
	buildDescriber := &BuildDescriber{}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, buildConfig.ObjectMeta)
		buildDescriber.DescribeParameters(buildConfig.Parameters, out)
		for whType, whURL := range webhooks {
			t := strings.Title(whType)
			formatString(out, "Webhook "+t, whURL)
		}
		return nil
	})
}

// DeploymentDescriber generates information about a deployment
type DeploymentDescriber struct {
	client.Interface
}

func (d *DeploymentDescriber) Describe(namespace, name string) (string, error) {
	c := d.Deployments(namespace)
	deployment, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, deployment.ObjectMeta)
		formatString(out, "Status", bold(deployment.Status))
		formatString(out, "Strategy", deployment.Strategy.Type)
		causes := []string{}
		if deployment.Details != nil {
			for _, c := range deployment.Details.Causes {
				causes = append(causes, string(c.Type))
			}
		}
		formatString(out, "Causes", strings.Join(causes, ","))
		return nil
	})
}

// DeploymentConfigDescriber generates information about a DeploymentConfig
type DeploymentConfigDescriber struct {
	client.Interface
}

func (d *DeploymentConfigDescriber) Describe(namespace, name string) (string, error) {
	c := d.DeploymentConfigs(namespace)
	deploymentConfig, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, deploymentConfig.ObjectMeta)
		formatString(out, "Latest Version", deploymentConfig.LatestVersion)
		triggers := []string{}
		for _, t := range deploymentConfig.Triggers {
			triggers = append(triggers, string(t.Type))
		}
		formatString(out, "Triggers", strings.Join(triggers, ","))
		return nil
	})
}

// ImageDescriber generates information about a Image
type ImageDescriber struct {
	client.Interface
}

func (d *ImageDescriber) Describe(namespace, name string) (string, error) {
	c := d.Images(namespace)
	image, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, image.ObjectMeta)
		formatString(out, "Docker Image", image.DockerImageReference)
		return nil
	})
}

// ImageRepositoryDescriber generates information about a ImageRepository
type ImageRepositoryDescriber struct {
	client.Interface
}

func (d *ImageRepositoryDescriber) Describe(namespace, name string) (string, error) {
	c := d.ImageRepositories(namespace)
	imageRepository, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, imageRepository.ObjectMeta)
		formatString(out, "Tags", formatLabels(imageRepository.Tags))
		formatString(out, "Registry", imageRepository.DockerImageRepository)
		return nil
	})
}

// RouteDescriber generates information about a Route
type RouteDescriber struct {
	client.Interface
}

func (d *RouteDescriber) Describe(namespace, name string) (string, error) {
	c := d.Routes(namespace)
	route, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, route.ObjectMeta)
		formatString(out, "Host", route.Host)
		formatString(out, "Path", route.Path)
		formatString(out, "Service", route.ServiceName)
		return nil
	})
}

// ProjectDescriber generates information about a Project
type ProjectDescriber struct {
	client.Interface
}

func (d *ProjectDescriber) Describe(namespace, name string) (string, error) {
	c := d.Projects(namespace)
	project, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, project.ObjectMeta)
		formatString(out, "Display Name", project.DisplayName)
		return nil
	})
}
<commit_msg>Bug 1170545 - Error about deployment not found<commit_after>package kubectl

import (
	"strings"
	"text/tabwriter"

	"github.com/spf13/cobra"

	kclient "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	kctl "github.com/GoogleCloudPlatform/kubernetes/pkg/kubectl"
	kubecmd "github.com/GoogleCloudPlatform/kubernetes/pkg/kubectl/cmd"
	buildapi "github.com/openshift/origin/pkg/build/api"
	"github.com/openshift/origin/pkg/client"
)

func DescriberFor(kind string, c *client.Client, cmd *cobra.Command) (kctl.Describer, bool) {
	switch kind {
	case "Build":
		return &BuildDescriber{c}, true
	case "BuildConfig":
		return &BuildConfigDescriber{c, cmd}, true
	case "Deployment":
		return &DeploymentDescriber{c}, true
	case "DeploymentConfig":
		return &DeploymentConfigDescriber{c}, true
	case "Image":
		return &ImageDescriber{c}, true
	case "ImageRepository":
		return &ImageRepositoryDescriber{c}, true
	case "Route":
		return &RouteDescriber{c}, true
	case "Project":
		return &ProjectDescriber{c}, true
	}
	return nil, false
}

// BuildDescriber generates information about a build
type BuildDescriber struct {
	client.Interface
}

func (d *BuildDescriber) DescribeParameters(p buildapi.BuildParameters, out *tabwriter.Writer) {
	formatString(out, "Strategy", p.Strategy.Type)
	formatString(out, "Source Type", p.Source.Type)
	if p.Source.Git != nil {
		formatString(out, "URL", p.Source.Git.URI)
		if len(p.Source.Git.Ref) > 0 {
			formatString(out, "Ref", p.Source.Git.Ref)
		}
	}
	formatString(out, "Output Image", p.Output.ImageTag)
	formatString(out, "Output Registry", p.Output.Registry)
}

func (d *BuildDescriber) Describe(namespace, name string) (string, error) {
	c := d.Builds(namespace)
	build, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, build.ObjectMeta)
		formatString(out, "Status", bold(build.Status))
		formatString(out, "Build Pod", build.PodName)
		d.DescribeParameters(build.Parameters, out)
		return nil
	})
}

// BuildConfigDescriber generates information about a buildConfig
type BuildConfigDescriber struct {
	client.Interface
	*cobra.Command
}

func (d *BuildConfigDescriber) Describe(namespace, name string) (string, error) {
	c := d.BuildConfigs(namespace)
	buildConfig, err := c.Get(name)
	if err != nil {
		return "", err
	}

	var kubeConfig *kclient.Config
	if d.Command != nil {
		kubeConfig = kubecmd.GetKubeConfig(d.Command)
	}

	webhooks := webhookURL(buildConfig, kubeConfig)
	buildDescriber := &BuildDescriber{}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, buildConfig.ObjectMeta)
		buildDescriber.DescribeParameters(buildConfig.Parameters, out)
		for whType, whURL := range webhooks {
			t := strings.Title(whType)
			formatString(out, "Webhook "+t, whURL)
		}
		return nil
	})
}

// DeploymentDescriber generates information about a deployment
type DeploymentDescriber struct {
	client.Interface
}

func (d *DeploymentDescriber) Describe(namespace, name string) (string, error) {
	c := d.Deployments(namespace)
	deployment, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, deployment.ObjectMeta)
		formatString(out, "Status", bold(deployment.Status))
		formatString(out, "Strategy", deployment.Strategy.Type)
		causes := []string{}
		if deployment.Details != nil {
			for _, c := range deployment.Details.Causes {
				causes = append(causes, string(c.Type))
			}
		}
		formatString(out, "Causes", strings.Join(causes, ","))
		return nil
	})
}

// DeploymentConfigDescriber generates information about a DeploymentConfig
type DeploymentConfigDescriber struct {
	client.Interface
}

func (d *DeploymentConfigDescriber) Describe(namespace, name string) (string, error) {
	c := d.DeploymentConfigs(namespace)
	deploymentConfig, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, deploymentConfig.ObjectMeta)
		formatString(out, "Latest Version", deploymentConfig.LatestVersion)
		triggers := []string{}
		for _, t := range deploymentConfig.Triggers {
			triggers = append(triggers, string(t.Type))
		}
		formatString(out, "Triggers", strings.Join(triggers, ","))
		return nil
	})
}

// ImageDescriber generates information about a Image
type ImageDescriber struct {
	client.Interface
}

func (d *ImageDescriber) Describe(namespace, name string) (string, error) {
	c := d.Images(namespace)
	image, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, image.ObjectMeta)
		formatString(out, "Docker Image", image.DockerImageReference)
		return nil
	})
}

// ImageRepositoryDescriber generates information about a ImageRepository
type ImageRepositoryDescriber struct {
	client.Interface
}

func (d *ImageRepositoryDescriber) Describe(namespace, name string) (string, error) {
	c := d.ImageRepositories(namespace)
	imageRepository, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, imageRepository.ObjectMeta)
		formatString(out, "Tags", formatLabels(imageRepository.Tags))
		formatString(out, "Registry", imageRepository.DockerImageRepository)
		return nil
	})
}

// RouteDescriber generates information about a Route
type RouteDescriber struct {
	client.Interface
}

func (d *RouteDescriber) Describe(namespace, name string) (string, error) {
	c := d.Routes(namespace)
	route, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, route.ObjectMeta)
		formatString(out, "Host", route.Host)
		formatString(out, "Path", route.Path)
		formatString(out, "Service", route.ServiceName)
		return nil
	})
}

// ProjectDescriber generates information about a Project
type ProjectDescriber struct {
	client.Interface
}

func (d *ProjectDescriber) Describe(namespace, name string) (string, error) {
	c := d.Projects(namespace)
	project, err := c.Get(name)
	if err != nil {
		return "", err
	}

	return tabbedString(func(out *tabwriter.Writer) error {
		formatMeta(out, project.ObjectMeta)
		formatString(out, "Display Name", project.DisplayName)
		return nil
	})
}
<|endoftext|>


<commit_before>package edit

import "testing"

func TestBindingMap(t *testing.T) {
	// TODO
}
<commit_msg>pkg/edit: Improve coverage of binding_map.go.<commit_after>package edit

import (
	"testing"

	"src.elv.sh/pkg/eval"
	. "src.elv.sh/pkg/eval/evaltest"
)

func TestBindingMap(t *testing.T) {
	// TODO
	TestWithSetup(t, func(ev *eval.Evaler) {
		ev.AddBuiltin(eval.NsBuilder{}.AddGoFn("", "binding-map", makeBindingMap).Ns())
	},
		// Checking key and value when constructing
		That("binding-map [&[]={ }]").
			Throws(ErrorWithMessage("must be key or string")),
		That("binding-map [&foo={ }]").
			Throws(ErrorWithMessage("bad key: foo")),
		That("binding-map [&a=string]").
			Throws(ErrorWithMessage("value should be function")),

		// repr prints a binding-map like an ordinary map
		That("repr (binding-map [&])").Prints("[&]\n"),
		// Keys are always sorted
		That("repr (binding-map [&a=$nop~ &b=$nop~ &c=$nop~])").
			Prints("[&a=<builtin nop> &b=<builtin nop> &c=<builtin nop>]\n"),

		// Indexing
		That("eq $nop~ (binding-map [&a=$nop~])[a]").Puts(true),
		// Checking key
		That("put (binding-map [&a=$nop~])[foo]").
			Throws(ErrorWithMessage("bad key: foo")),

		// Assoc
		That("count (assoc (binding-map [&a=$nop~]) b $nop~)").Puts(2),
		// Checking key
		That("(assoc (binding-map [&a=$nop~]) foo $nop~)").
			Throws(ErrorWithMessage("bad key: foo")),
		// Checking value
		That("(assoc (binding-map [&a=$nop~]) b foo)").
			Throws(ErrorWithMessage("value should be function")),

		// Dissoc
		That("count (dissoc (binding-map [&a=$nop~]) a)").Puts(0),
		// Allows bad key - no op
		That("count (dissoc (binding-map [&a=$nop~]) foo)").Puts(1),
	)
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package prober

import (
	"math/rand"
	"time"

	"github.com/golang/glog"
	"k8s.io/kubernetes/pkg/api"
	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
	"k8s.io/kubernetes/pkg/kubelet/prober/results"
	"k8s.io/kubernetes/pkg/kubelet/util/format"
	"k8s.io/kubernetes/pkg/util/runtime"
)

// worker handles the periodic probing of its assigned container. Each worker has a go-routine
// associated with it which runs the probe loop until the container permanently terminates, or the
// stop channel is closed. The worker uses the probe Manager's statusManager to get up-to-date
// container IDs.
type worker struct {
	// Channel for stopping the probe.
	stopCh chan struct{}

	// The pod containing this probe (read-only)
	pod *api.Pod

	// The container to probe (read-only)
	container api.Container

	// Describes the probe configuration (read-only)
	spec *api.Probe

	// The type of the worker.
	probeType probeType

	// The probe value during the initial delay.
	initialValue results.Result

	// Where to store this workers results.
	resultsManager results.Manager
	probeManager   *manager

	// The last known container ID for this worker.
	containerID kubecontainer.ContainerID
	// The last probe result for this worker.
	lastResult results.Result
	// How many times in a row the probe has returned the same result.
	resultRun int
}

// Creates and starts a new probe worker.
func newWorker(
	m *manager,
	probeType probeType,
	pod *api.Pod,
	container api.Container) *worker {

	w := &worker{
		stopCh:       make(chan struct{}, 1), // Buffer so stop() can be non-blocking.
		pod:          pod,
		container:    container,
		probeType:    probeType,
		probeManager: m,
	}

	switch probeType {
	case readiness:
		w.spec = container.ReadinessProbe
		w.resultsManager = m.readinessManager
		w.initialValue = results.Failure
	case liveness:
		w.spec = container.LivenessProbe
		w.resultsManager = m.livenessManager
		w.initialValue = results.Success
	}

	return w
}

// run periodically probes the container.
func (w *worker) run() {
	probeTickerPeriod := time.Duration(w.spec.PeriodSeconds) * time.Second
	probeTicker := time.NewTicker(probeTickerPeriod)

	defer func() {
		// Clean up.
		probeTicker.Stop()
		if !w.containerID.IsEmpty() {
			w.resultsManager.Remove(w.containerID)
		}

		w.probeManager.removeWorker(w.pod.UID, w.container.Name, w.probeType)
	}()

	// If kubelet restarted the probes could be started in rapid succession.
	// Let the worker wait for a random portion of tickerPeriod before probing.
	time.Sleep(time.Duration(rand.Float64() * float64(probeTickerPeriod)))

probeLoop:
	for w.doProbe() {
		// Wait for next probe tick.
		select {
		case <-w.stopCh:
			break probeLoop
		case <-probeTicker.C:
			// continue
		}
	}
}

// stop stops the probe worker. The worker handles cleanup and removes itself from its manager.
// It is safe to call stop multiple times.
func (w *worker) stop() {
	select {
	case w.stopCh <- struct{}{}:
	default: // Non-blocking.
	}
}

// doProbe probes the container once and records the result.
// Returns whether the worker should continue.
func (w *worker) doProbe() (keepGoing bool) {
	defer runtime.HandleCrash(func(_ interface{}) { keepGoing = true })

	status, ok := w.probeManager.statusManager.GetPodStatus(w.pod.UID)
	if !ok {
		// Either the pod has not been created yet, or it was already deleted.
		glog.V(3).Infof("No status for pod: %v", format.Pod(w.pod))
		return true
	}

	// Worker should terminate if pod is terminated.
	if status.Phase == api.PodFailed || status.Phase == api.PodSucceeded {
		glog.V(3).Infof("Pod %v %v, exiting probe worker",
			format.Pod(w.pod), status.Phase)
		return false
	}

	c, ok := api.GetContainerStatus(status.ContainerStatuses, w.container.Name)
	if !ok {
		// Either the container has not been created yet, or it was deleted.
		glog.V(3).Infof("Non-existant container probed: %v - %v",
			format.Pod(w.pod), w.container.Name)
		return true // Wait for more information.
	}

	if w.containerID.String() != c.ContainerID {
		if !w.containerID.IsEmpty() {
			w.resultsManager.Remove(w.containerID)
		}
		w.containerID = kubecontainer.ParseContainerID(c.ContainerID)
		w.resultsManager.Set(w.containerID, w.initialValue, w.pod)
	}

	if c.State.Running == nil {
		glog.V(3).Infof("Non-running container probed: %v - %v",
			format.Pod(w.pod), w.container.Name)
		if !w.containerID.IsEmpty() {
			w.resultsManager.Set(w.containerID, results.Failure, w.pod)
		}
		// Abort if the container will not be restarted.
		return c.State.Terminated == nil ||
			w.pod.Spec.RestartPolicy != api.RestartPolicyNever
	}

	if int(time.Since(c.State.Running.StartedAt.Time).Seconds()) < w.spec.InitialDelaySeconds {
		return true
	}

	result, err := w.probeManager.prober.probe(w.probeType, w.pod, status, w.container, w.containerID)
	if err != nil {
		// Prober error, throw away the result.
		return true
	}

	if w.lastResult == result {
		w.resultRun++
	} else {
		w.lastResult = result
		w.resultRun = 1
	}

	if (result == results.Failure && w.resultRun < w.spec.FailureThreshold) ||
		(result == results.Success && w.resultRun < w.spec.SuccessThreshold) {
		// Success or failure is below threshold - leave the probe state unchanged.
		return true
	}

	w.resultsManager.Set(w.containerID, result, w.pod)

	return true
}
<commit_msg>Skip probe if container ID not yet set<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package prober

import (
	"math/rand"
	"time"

	"github.com/golang/glog"
	"k8s.io/kubernetes/pkg/api"
	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
	"k8s.io/kubernetes/pkg/kubelet/prober/results"
	"k8s.io/kubernetes/pkg/kubelet/util/format"
	"k8s.io/kubernetes/pkg/util/runtime"
)

// worker handles the periodic probing of its assigned container. Each worker has a go-routine
// associated with it which runs the probe loop until the container permanently terminates, or the
// stop channel is closed. The worker uses the probe Manager's statusManager to get up-to-date
// container IDs.
type worker struct {
	// Channel for stopping the probe.
	stopCh chan struct{}

	// The pod containing this probe (read-only)
	pod *api.Pod

	// The container to probe (read-only)
	container api.Container

	// Describes the probe configuration (read-only)
	spec *api.Probe

	// The type of the worker.
	probeType probeType

	// The probe value during the initial delay.
	initialValue results.Result

	// Where to store this workers results.
	resultsManager results.Manager
	probeManager   *manager

	// The last known container ID for this worker.
	containerID kubecontainer.ContainerID
	// The last probe result for this worker.
	lastResult results.Result
	// How many times in a row the probe has returned the same result.
	resultRun int
}

// Creates and starts a new probe worker.
func newWorker(
	m *manager,
	probeType probeType,
	pod *api.Pod,
	container api.Container) *worker {

	w := &worker{
		stopCh:       make(chan struct{}, 1), // Buffer so stop() can be non-blocking.
		pod:          pod,
		container:    container,
		probeType:    probeType,
		probeManager: m,
	}

	switch probeType {
	case readiness:
		w.spec = container.ReadinessProbe
		w.resultsManager = m.readinessManager
		w.initialValue = results.Failure
	case liveness:
		w.spec = container.LivenessProbe
		w.resultsManager = m.livenessManager
		w.initialValue = results.Success
	}

	return w
}

// run periodically probes the container.
func (w *worker) run() {
	probeTickerPeriod := time.Duration(w.spec.PeriodSeconds) * time.Second
	probeTicker := time.NewTicker(probeTickerPeriod)

	defer func() {
		// Clean up.
		probeTicker.Stop()
		if !w.containerID.IsEmpty() {
			w.resultsManager.Remove(w.containerID)
		}

		w.probeManager.removeWorker(w.pod.UID, w.container.Name, w.probeType)
	}()

	// If kubelet restarted the probes could be started in rapid succession.
	// Let the worker wait for a random portion of tickerPeriod before probing.
	time.Sleep(time.Duration(rand.Float64() * float64(probeTickerPeriod)))

probeLoop:
	for w.doProbe() {
		// Wait for next probe tick.
		select {
		case <-w.stopCh:
			break probeLoop
		case <-probeTicker.C:
			// continue
		}
	}
}

// stop stops the probe worker. The worker handles cleanup and removes itself from its manager.
// It is safe to call stop multiple times.
func (w *worker) stop() {
	select {
	case w.stopCh <- struct{}{}:
	default: // Non-blocking.
	}
}

// doProbe probes the container once and records the result.
// Returns whether the worker should continue.
func (w *worker) doProbe() (keepGoing bool) {
	defer runtime.HandleCrash(func(_ interface{}) { keepGoing = true })

	status, ok := w.probeManager.statusManager.GetPodStatus(w.pod.UID)
	if !ok {
		// Either the pod has not been created yet, or it was already deleted.
		glog.V(3).Infof("No status for pod: %v", format.Pod(w.pod))
		return true
	}

	// Worker should terminate if pod is terminated.
	if status.Phase == api.PodFailed || status.Phase == api.PodSucceeded {
		glog.V(3).Infof("Pod %v %v, exiting probe worker",
			format.Pod(w.pod), status.Phase)
		return false
	}

	c, ok := api.GetContainerStatus(status.ContainerStatuses, w.container.Name)
	if !ok || len(c.ContainerID) == 0 {
		// Either the container has not been created yet, or it was deleted.
		glog.V(3).Infof("Probe target container not found: %v - %v",
			format.Pod(w.pod), w.container.Name)
		return true // Wait for more information.
	}

	if w.containerID.String() != c.ContainerID {
		if !w.containerID.IsEmpty() {
			w.resultsManager.Remove(w.containerID)
		}
		w.containerID = kubecontainer.ParseContainerID(c.ContainerID)
		w.resultsManager.Set(w.containerID, w.initialValue, w.pod)
	}

	if c.State.Running == nil {
		glog.V(3).Infof("Non-running container probed: %v - %v",
			format.Pod(w.pod), w.container.Name)
		if !w.containerID.IsEmpty() {
			w.resultsManager.Set(w.containerID, results.Failure, w.pod)
		}
		// Abort if the container will not be restarted.
		return c.State.Terminated == nil ||
			w.pod.Spec.RestartPolicy != api.RestartPolicyNever
	}

	if int(time.Since(c.State.Running.StartedAt.Time).Seconds()) < w.spec.InitialDelaySeconds {
		return true
	}

	result, err := w.probeManager.prober.probe(w.probeType, w.pod, status, w.container, w.containerID)
	if err != nil {
		// Prober error, throw away the result.
		return true
	}

	if w.lastResult == result {
		w.resultRun++
	} else {
		w.lastResult = result
		w.resultRun = 1
	}

	if (result == results.Failure && w.resultRun < w.spec.FailureThreshold) ||
		(result == results.Success && w.resultRun < w.spec.SuccessThreshold) {
		// Success or failure is below threshold - leave the probe state unchanged.
		return true
	}

	w.resultsManager.Set(w.containerID, result, w.pod)

	return true
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Juniper Networks, Inc. All rights reserved.

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

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

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

package network

import (
	"fmt"
	"os/exec"
	"strconv"

	"github.com/docker/libcontainer/netlink"
	"github.com/milosgajdos83/tenus"
)

type NetnsManager interface {
	CreateInterface(dockerId, macAddress, ipAddress, gateway string) (string, error)
	DeleteInterface(dockerId string) error
}

type NetnsManagerImpl struct {
}

func NewNetnsManager() NetnsManager {
	m := new(NetnsManagerImpl)
	return m
}

func (m *NetnsManagerImpl) CreateInterface(dockerId, macAddress, ipAddress, gateway string) (string, error) {
	masterName := fmt.Sprintf("veth-%s", dockerId[0:10])
	veth, err := tenus.NewVethPairWithOptions(masterName, tenus.VethOptions{PeerName: "veth0"})
	if err != nil {
		return "", err
	}
	pid, err := tenus.DockerPidByName(dockerId, "/var/run/docker.sock")
	if err != nil {
		return "", err
	}
	veth.SetPeerLinkNsPid(pid)
	peer := veth.PeerNetInterface()
	netlink.NetworkSetMacAddress(peer, macAddress)
	cmd := exec.Command("nsenter", "-n", "-t", strconv.Itoa(pid), "ip", "addr", "add",
		fmt.Sprintf("%s/32", ipAddress), "dev", "veth0")
	err = cmd.Run()
	if err != nil {
		return "", err
	}
	cmd = exec.Command("nsenter", "-n", "-t", strconv.Itoa(pid), "ip", "route", "add",
		"default", "via", gateway)
	err = cmd.Run()

	veth.SetPeerLinkUp()
	return masterName, nil
}

func (m *NetnsManagerImpl) DeleteInterface(dockerId string) error {
	// masterName := fmt.Sprintf("veth-%s", dockerId[0:10])
	return nil
}
<commit_msg>Ensure that we set link up on the master instance side<commit_after>/*
Copyright 2015 Juniper Networks, Inc. All rights reserved.

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

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

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

package network

import (
	"fmt"
	"os/exec"
	"strconv"

	"github.com/docker/libcontainer/netlink"
	"github.com/milosgajdos83/tenus"
)

type NetnsManager interface {
	CreateInterface(dockerId, macAddress, ipAddress, gateway string) (string, error)
	DeleteInterface(dockerId string) error
}

type NetnsManagerImpl struct {
}

func NewNetnsManager() NetnsManager {
	m := new(NetnsManagerImpl)
	return m
}

func (m *NetnsManagerImpl) CreateInterface(dockerId, macAddress, ipAddress, gateway string) (string, error) {
	masterName := fmt.Sprintf("veth-%s", dockerId[0:10])
	veth, err := tenus.NewVethPairWithOptions(masterName, tenus.VethOptions{PeerName: "veth0"})
	if err != nil {
		return "", err
	}
	pid, err := tenus.DockerPidByName(dockerId, "/var/run/docker.sock")
	if err != nil {
		return "", err
	}
	veth.SetPeerLinkNsPid(pid)
	veth.SetLinkUp()
	peer := veth.PeerNetInterface()
	netlink.NetworkSetMacAddress(peer, macAddress)
	cmd := exec.Command("nsenter", "-n", "-t", strconv.Itoa(pid), "ip", "addr", "add",
		fmt.Sprintf("%s/32", ipAddress), "dev", "veth0")
	err = cmd.Run()
	if err != nil {
		return "", err
	}
	cmd = exec.Command("nsenter", "-n", "-t", strconv.Itoa(pid), "ip", "route", "add",
		"default", "via", gateway)
	err = cmd.Run()

	veth.SetPeerLinkUp()
	return masterName, nil
}

func (m *NetnsManagerImpl) DeleteInterface(dockerId string) error {
	// masterName := fmt.Sprintf("veth-%s", dockerId[0:10])
	return nil
}
<|endoftext|>


<commit_before>package mssql

import (
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/go-xorm/xorm"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/services/sqlstore/sqlutil"
	"github.com/grafana/grafana/pkg/tsdb"
	. "github.com/smartystreets/goconvey/convey"
)

// To run this test, remove the Skip from SkipConvey
// and set up a MSSQL db named grafana_tests and a user/password grafana/Password!
// and set the variable below to the IP address of the database
var serverIP string = "localhost"

func TestMSSQL(t *testing.T) {
	Convey("MSSQL", t, func() {
		x := InitMSSQLTestDB(t)

		endpoint := &MssqlQueryEndpoint{
			sqlEngine: &tsdb.DefaultSqlEngine{
				MacroEngine: NewMssqlMacroEngine(),
				XormEngine:  x,
			},
			log: log.New("tsdb.mssql"),
		}

		sess := x.NewSession()
		defer sess.Close()

		sql := `
			IF OBJECT_ID('dbo.[mssql_types]', 'U') IS NOT NULL
				DROP TABLE dbo.[mssql_types]

			CREATE TABLE [mssql_types] (
				c_bit bit,
				c_tinyint tinyint,
				c_smallint smallint,
				c_int int,
				c_bigint bigint,
				c_money money,
				c_smallmoney smallmoney,
				c_numeric numeric(10,5),

				c_real real,
				c_decimal decimal(10,2),
				c_float float,

				c_char char(10),
				c_varchar varchar(10),
				c_text text,

				c_nchar nchar(12),
				c_nvarchar nvarchar(12),
				c_ntext ntext,

				c_datetime datetime,
				c_datetime2 datetime2,
				c_smalldatetime smalldatetime,
				c_date date,
				c_time time,
				c_datetimeoffset datetimeoffset
			)

			IF OBJECT_ID('dbo.[metric]', 'U') IS NOT NULL
				DROP TABLE dbo.[metric]

			CREATE TABLE [metric] (
				time datetime,
				measurement nvarchar(100),
				value int
			)
		`

		_, err := sess.Exec(sql)
		So(err, ShouldBeNil)

		// type metric struct {
		// 	Time        time.Time
		// 	Measurement string
		// 	Value       int64
		// }

		// series := []*metric{}

		// from := time.Now().Truncate(60 * time.Minute).Add((-30 * time.Minute))
		// for _, t := range genTimeRangeByInterval(from, 10*time.Minute, 10*time.Second) {
		// 	series = append(series, &metric{
		// 		Time:        t,
		// 		Measurement: "test",
		// 		Value:       0,
		// 	})
		// }

		// for _, t := range genTimeRangeByInterval(from.Add(20*time.Minute), 10*time.Minute, 10*time.Second) {
		// 	series = append(series, &metric{
		// 		Time:        t,
		// 		Measurement: "test",
		// 		Value:       0,
		// 	})
		// }

		// rowsAffected, err := sess.InsertMulti(series)
		// So(err, ShouldBeNil)
		// So(rowsAffected, ShouldBeGreaterThan, 0)

		dt := time.Date(2018, 3, 14, 21, 20, 6, 527e6, time.UTC)
		dtFormat := "2006-01-02 15:04:05.999999999"
		d := dt.Format(dtFormat)
		dt2 := time.Date(2018, 3, 14, 21, 20, 6, 8896406e2, time.UTC)
		dt2Format := "2006-01-02 15:04:05.999999999 -07:00"
		d2 := dt2.Format(dt2Format)

		sql = fmt.Sprintf(`
			INSERT INTO [mssql_types]
			SELECT
        1, 5, 20020, 980300, 1420070400, '$20000.15', '£2.15', 12345.12,
        1.11, 2.22, 3.33,
				'char10', 'varchar10', 'text',
				N'☺nchar12☺', N'☺nvarchar12☺', N'☺text☺',
			  CAST('%s' AS DATETIME), CAST('%s' AS DATETIME2), CAST('%s' AS SMALLDATETIME), CAST('%s' AS DATE), CAST('%s' AS TIME), SWITCHOFFSET(CAST('%s' AS DATETIMEOFFSET), '-07:00')
    `, d, d2, d, d, d, d2)

		_, err = sess.Exec(sql)
		So(err, ShouldBeNil)

		Convey("Query with Table format should map MSSQL column types to Go types", func() {
			query := &tsdb.TsdbQuery{
				Queries: []*tsdb.Query{
					{
						Model: simplejson.NewFromAny(map[string]interface{}{
							"rawSql": "SELECT * FROM mssql_types",
							"format": "table",
						}),
						RefId: "A",
					},
				},
			}

			resp, err := endpoint.Query(nil, nil, query)
			queryResult := resp.Results["A"]
			So(err, ShouldBeNil)

			column := queryResult.Tables[0].Rows[0]

			So(column[0].(bool), ShouldEqual, true)
			So(column[1].(int64), ShouldEqual, 5)
			So(column[2].(int64), ShouldEqual, 20020)
			So(column[3].(int64), ShouldEqual, 980300)
			So(column[4].(int64), ShouldEqual, 1420070400)
			// So(column[5].(float64), ShouldEqual, 20000.15)
			// So(column[6].(float64), ShouldEqual, 2.15)
			//So(column[7].(float64), ShouldEqual, 12345.12)

			So(column[8].(float64), ShouldEqual, 1.1100000143051147) // MSSQL dose not have precision for "real" datatype
			// fix me: MSSQL driver puts the decimal inside an array of chars. and the test fails despite the values are correct.
			//So(column[9].([]uint8), ShouldEqual, []uint8{'2', '.', '2', '2'})
			So(column[10].(float64), ShouldEqual, 3.33)

			So(column[11].(string), ShouldEqual, "char10    ")
			So(column[12].(string), ShouldEqual, "varchar10")
			So(column[13].(string), ShouldEqual, "text")

			So(column[14].(string), ShouldEqual, "☺nchar12☺   ")
			So(column[15].(string), ShouldEqual, "☺nvarchar12☺")
			So(column[16].(string), ShouldEqual, "☺text☺")

			So(column[17].(time.Time), ShouldEqual, dt)
			So(column[18].(time.Time), ShouldEqual, dt2)
			So(column[19].(time.Time), ShouldEqual, dt.Truncate(time.Minute))
			So(column[20].(time.Time), ShouldEqual, dt.Truncate(24*time.Hour))
			So(column[21].(time.Time), ShouldEqual, time.Date(1, 1, 1, dt.Hour(), dt.Minute(), dt.Second(), dt.Nanosecond(), time.UTC))
			So(column[22].(time.Time), ShouldEqual, dt2.In(time.FixedZone("UTC", int(-7*time.Hour))))
		})

		Convey("stored procedure", func() {
			sql := `
				create procedure dbo.test_sp as
				begin
					select 1
				end
			`
			sess.Exec(sql)

			sql = `
				ALTER PROCEDURE dbo.test_sp
					@from int,
					@to 	int
				AS
				BEGIN
					select
						GETDATE() AS Time,
						1 as value,
						'metric' as metric
				END
			`
			_, err := sess.Exec(sql)
			So(err, ShouldBeNil)

			sql = `
				EXEC dbo.test_sp 1, 2
			`
			_, err = sess.Exec(sql)
			So(err, ShouldBeNil)
		})
	})
}

func InitMSSQLTestDB(t *testing.T) *xorm.Engine {
	x, err := xorm.NewEngine(sqlutil.TestDB_Mssql.DriverName, strings.Replace(sqlutil.TestDB_Mssql.ConnStr, "localhost", serverIP, 1))

	// x.ShowSQL()

	if err != nil {
		t.Fatalf("Failed to init mssql db %v", err)
	}

	sqlutil.CleanDB(x)

	return x
}

func genTimeRangeByInterval(from time.Time, duration time.Duration, interval time.Duration) []time.Time {
	durationSec := int64(duration.Seconds())
	intervalSec := int64(interval.Seconds())
	timeRange := []time.Time{}

	for i := int64(0); i <= durationSec; i += intervalSec {
		timeRange = append(timeRange, from)
		from = from.Add(time.Duration(int64(time.Second) * intervalSec))
	}

	return timeRange
}
<commit_msg>mssql: add timeGroup integration test<commit_after>package mssql

import (
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/go-xorm/xorm"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/services/sqlstore/sqlutil"
	"github.com/grafana/grafana/pkg/tsdb"
	. "github.com/smartystreets/goconvey/convey"
)

// To run this test, remove the Skip from SkipConvey
// and set up a MSSQL db named grafana_tests and a user/password grafana/Password!
// and set the variable below to the IP address of the database
var serverIP string = "localhost"

func TestMSSQL(t *testing.T) {
	Convey("MSSQL", t, func() {
		x := InitMSSQLTestDB(t)

		endpoint := &MssqlQueryEndpoint{
			sqlEngine: &tsdb.DefaultSqlEngine{
				MacroEngine: NewMssqlMacroEngine(),
				XormEngine:  x,
			},
			log: log.New("tsdb.mssql"),
		}

		sess := x.NewSession()
		defer sess.Close()

		Convey("Given a table with different native data types", func() {
			sql := `
				IF OBJECT_ID('dbo.[mssql_types]', 'U') IS NOT NULL
					DROP TABLE dbo.[mssql_types]

				CREATE TABLE [mssql_types] (
					c_bit bit,
					c_tinyint tinyint,
					c_smallint smallint,
					c_int int,
					c_bigint bigint,
					c_money money,
					c_smallmoney smallmoney,
					c_numeric numeric(10,5),

					c_real real,
					c_decimal decimal(10,2),
					c_float float,

					c_char char(10),
					c_varchar varchar(10),
					c_text text,

					c_nchar nchar(12),
					c_nvarchar nvarchar(12),
					c_ntext ntext,

					c_datetime datetime,
					c_datetime2 datetime2,
					c_smalldatetime smalldatetime,
					c_date date,
					c_time time,
					c_datetimeoffset datetimeoffset
				)
			`

			_, err := sess.Exec(sql)
			So(err, ShouldBeNil)

			dt := time.Date(2018, 3, 14, 21, 20, 6, 527e6, time.UTC)
			dtFormat := "2006-01-02 15:04:05.999999999"
			d := dt.Format(dtFormat)
			dt2 := time.Date(2018, 3, 14, 21, 20, 6, 8896406e2, time.UTC)
			dt2Format := "2006-01-02 15:04:05.999999999 -07:00"
			d2 := dt2.Format(dt2Format)

			sql = fmt.Sprintf(`
			INSERT INTO [mssql_types]
			SELECT
        1, 5, 20020, 980300, 1420070400, '$20000.15', '£2.15', 12345.12,
        1.11, 2.22, 3.33,
				'char10', 'varchar10', 'text',
				N'☺nchar12☺', N'☺nvarchar12☺', N'☺text☺',
			  CAST('%s' AS DATETIME), CAST('%s' AS DATETIME2), CAST('%s' AS SMALLDATETIME), CAST('%s' AS DATE), CAST('%s' AS TIME), SWITCHOFFSET(CAST('%s' AS DATETIMEOFFSET), '-07:00')
    `, d, d2, d, d, d, d2)

			_, err = sess.Exec(sql)
			So(err, ShouldBeNil)

			Convey("Query with Table format should map MSSQL column types to Go types", func() {
				query := &tsdb.TsdbQuery{
					Queries: []*tsdb.Query{
						{
							Model: simplejson.NewFromAny(map[string]interface{}{
								"rawSql": "SELECT * FROM mssql_types",
								"format": "table",
							}),
							RefId: "A",
						},
					},
				}

				resp, err := endpoint.Query(nil, nil, query)
				queryResult := resp.Results["A"]
				So(err, ShouldBeNil)

				column := queryResult.Tables[0].Rows[0]

				So(column[0].(bool), ShouldEqual, true)
				So(column[1].(int64), ShouldEqual, 5)
				So(column[2].(int64), ShouldEqual, 20020)
				So(column[3].(int64), ShouldEqual, 980300)
				So(column[4].(int64), ShouldEqual, 1420070400)
				// So(column[5].(float64), ShouldEqual, 20000.15)
				// So(column[6].(float64), ShouldEqual, 2.15)
				//So(column[7].(float64), ShouldEqual, 12345.12)

				So(column[8].(float64), ShouldEqual, 1.1100000143051147) // MSSQL dose not have precision for "real" datatype
				// fix me: MSSQL driver puts the decimal inside an array of chars. and the test fails despite the values are correct.
				//So(column[9].([]uint8), ShouldEqual, []uint8{'2', '.', '2', '2'})
				So(column[10].(float64), ShouldEqual, 3.33)

				So(column[11].(string), ShouldEqual, "char10    ")
				So(column[12].(string), ShouldEqual, "varchar10")
				So(column[13].(string), ShouldEqual, "text")

				So(column[14].(string), ShouldEqual, "☺nchar12☺   ")
				So(column[15].(string), ShouldEqual, "☺nvarchar12☺")
				So(column[16].(string), ShouldEqual, "☺text☺")

				So(column[17].(time.Time), ShouldEqual, dt)
				So(column[18].(time.Time), ShouldEqual, dt2)
				So(column[19].(time.Time), ShouldEqual, dt.Truncate(time.Minute))
				So(column[20].(time.Time), ShouldEqual, dt.Truncate(24*time.Hour))
				So(column[21].(time.Time), ShouldEqual, time.Date(1, 1, 1, dt.Hour(), dt.Minute(), dt.Second(), dt.Nanosecond(), time.UTC))
				So(column[22].(time.Time), ShouldEqual, dt2.In(time.FixedZone("UTC", int(-7*time.Hour))))
			})

			// Convey("stored procedure", func() {
			// 	sql := `
			// 	create procedure dbo.test_sp as
			// 	begin
			// 		select 1
			// 	end
			// `
			// 	sess.Exec(sql)

			// 	sql = `
			// 	ALTER PROCEDURE dbo.test_sp
			// 		@from int,
			// 		@to 	int
			// 	AS
			// 	BEGIN
			// 		select
			// 			GETDATE() AS Time,
			// 			1 as value,
			// 			'metric' as metric
			// 	END
			// `
			// 	_, err := sess.Exec(sql)
			// 	So(err, ShouldBeNil)

			// 	sql = `
			// 	EXEC dbo.test_sp 1, 2
			// `
			// 	_, err = sess.Exec(sql)
			// 	So(err, ShouldBeNil)
			// })
		})

		Convey("Given a table with metrics", func() {
			sql := `
				IF OBJECT_ID('dbo.[metric]', 'U') IS NOT NULL
					DROP TABLE dbo.[metric]

				CREATE TABLE [metric] (
					time datetime,
					measurement nvarchar(100),
					value int
				)
			`

			_, err := sess.Exec(sql)
			So(err, ShouldBeNil)

			type metric struct {
				Time        time.Time
				Measurement string
				Value       int64
			}

			series := []*metric{}

			fromStart := time.Date(2018, 3, 15, 13, 0, 0, 0, time.UTC)
			firstRange := genTimeRangeByInterval(fromStart, 10*time.Minute, 10*time.Second)
			secondRange := genTimeRangeByInterval(fromStart.Add(20*time.Minute), 10*time.Minute, 10*time.Second)

			for _, t := range firstRange {
				series = append(series, &metric{
					Time:        t,
					Measurement: "test",
					Value:       15,
				})
			}

			for _, t := range secondRange {
				series = append(series, &metric{
					Time:        t,
					Measurement: "test",
					Value:       20,
				})
			}

			dtFormat := "2006-01-02 15:04:05.999999999"
			for _, s := range series {
				sql = fmt.Sprintf(`
					INSERT INTO metric (time, measurement, value)
					VALUES(CAST('%s' AS DATETIME), '%s', %d)
				`, s.Time.Format(dtFormat), s.Measurement, s.Value)

				_, err = sess.Exec(sql)
				So(err, ShouldBeNil)
			}

			Convey("When doing a metric query using timeGroup", func() {
				query := &tsdb.TsdbQuery{
					Queries: []*tsdb.Query{
						{
							Model: simplejson.NewFromAny(map[string]interface{}{
								"rawSql": "SELECT $__timeGroup(time, '5m') AS time, measurement as metric, avg(value) as value FROM metric GROUP BY $__timeGroup(time, '5m'), measurement ORDER BY 1",
								"format": "time_series",
							}),
							RefId: "A",
						},
					},
				}

				resp, err := endpoint.Query(nil, nil, query)
				queryResult := resp.Results["A"]
				So(err, ShouldBeNil)
				So(queryResult.Error, ShouldBeNil)

				points := queryResult.Series[0].Points

				So(len(points), ShouldEqual, 4)
				actualValueFirst := points[0][0].Float64
				actualTimeFirst := time.Unix(int64(points[0][1].Float64)/1000, 0)
				So(actualValueFirst, ShouldEqual, 15)
				So(actualTimeFirst, ShouldEqual, fromStart)

				actualValueLast := points[3][0].Float64
				actualTimeLast := time.Unix(int64(points[3][1].Float64)/1000, 0)
				So(actualValueLast, ShouldEqual, 20)
				So(actualTimeLast, ShouldEqual, fromStart.Add(25*time.Minute))
			})
		})
	})
}

func InitMSSQLTestDB(t *testing.T) *xorm.Engine {
	x, err := xorm.NewEngine(sqlutil.TestDB_Mssql.DriverName, strings.Replace(sqlutil.TestDB_Mssql.ConnStr, "localhost", serverIP, 1))

	// x.ShowSQL()

	if err != nil {
		t.Fatalf("Failed to init mssql db %v", err)
	}

	sqlutil.CleanDB(x)

	return x
}

func genTimeRangeByInterval(from time.Time, duration time.Duration, interval time.Duration) []time.Time {
	durationSec := int64(duration.Seconds())
	intervalSec := int64(interval.Seconds())
	timeRange := []time.Time{}

	for i := int64(0); i < durationSec; i += intervalSec {
		timeRange = append(timeRange, from)
		from = from.Add(time.Duration(int64(time.Second) * intervalSec))
	}

	return timeRange
}
<|endoftext|>


<commit_before>package model

import (
	"errors"
	"log"
	"reflect"
	"strings"
	"time"

	"github.com/DanielRenne/GoCore/core/dbServices"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

const (
	ERROR_INVALID_ID_VALUE = "Invalid Id value for query."
)

type queryError func() error

type Range struct {
	Min interface{}
	Max interface{}
}

type Query struct {
	q          *mgo.Query
	m          bson.M
	limit      int
	skip       int
	sort       []string
	collection *mgo.Collection
	e          error
}

func Criteria(key string, value interface{}) map[string]interface{} {
	criteria := make(map[string]interface{}, 1)
	criteria[key] = value
	return criteria
}

func (self *Query) ById(val interface{}, x interface{}) error {

	objId, err := self.getIdHex(val)
	if err != nil {
		return err
	}

	q := self.collection.FindId(objId)

	err = q.One(x)

	if err != nil {
		callback := func() error {
			return q.One(x)
		}

		return self.handleQueryError(err, callback)
	}

	return err

}

func (self *Query) Filter(criteria map[string]interface{}) *Query {

	val, hasId := criteria["Id"]
	if hasId {

		objId, err := self.getIdHex(val)
		if err != nil {
			self.e = err
			return self
		}

		self.q = self.collection.FindId(objId)

		return self
	} else {

		if self.m == nil {
			self.m = make(bson.M)
		}

		for key, val := range criteria {
			if key != "" {
				self.m[key] = val
			}
		}

	}
	return self
}

func (self *Query) Exclude(criteria map[string]interface{}) *Query {

	if self.m == nil {
		self.m = make(bson.M)
	}

	self.inNot(criteria, "$nin")

	return self

}

func (self *Query) In(criteria map[string]interface{}) *Query {

	if self.m == nil {
		self.m = make(bson.M)
	}

	self.inNot(criteria, "$in")

	return self

}

func (self *Query) inNot(criteria map[string]interface{}, queryType string) {
	for key, value := range criteria {

		if key == "Id" {
			var ids []bson.ObjectId

			k := reflect.TypeOf(value).Kind()
			if k == reflect.Slice || k == reflect.Array {
				values := reflect.ValueOf(value)

				for i := 0; i < values.Len(); i++ {
					val := values.Index(i).Interface()
					objId, err := self.getIdHex(val)
					if err != nil {
						self.e = err
						continue
					}

					ids = append(ids, objId)
				}
			} else {
				objId, err := self.getIdHex(value)
				if err != nil {
					self.e = err
					continue
				}

				ids = append(ids, objId)
			}

			self.m["_id"] = bson.M{queryType: ids}

		} else {
			var valuesToQuery []interface{}

			k := reflect.TypeOf(value).Kind()
			if k == reflect.Slice || k == reflect.Array {
				values := reflect.ValueOf(value)
				for i := 0; i < values.Len(); i++ {
					val := values.Index(i).Interface()
					valuesToQuery = append(valuesToQuery, val)
				}

			} else {
				valuesToQuery = append(valuesToQuery, value)
			}

			self.m[key] = bson.M{queryType: valuesToQuery}
		}

	}
}

func (self *Query) Range(criteria map[string]Range) *Query {

	if self.m == nil {
		self.m = make(bson.M)
	}

	for key, val := range criteria {
		if key != "" {
			self.m[key] = bson.M{"$gte": val.Min, "$lte": val.Max}
		}
	}

	return self
}

func (self *Query) Sort(fields ...string) *Query {
	self.sort = fields
	return self
}

func (self *Query) Limit(val int) *Query {
	self.limit = val
	return self
}

func (self *Query) Skip(val int) *Query {
	self.skip = val
	return self
}

func (self *Query) All(x interface{}) error {

	if self.e != nil {
		return self.e
	}

	q := self.generateQuery()

	err := q.All(x)

	if err != nil {

		callback := func() error {
			return q.All(x)
		}

		return self.handleQueryError(err, callback)
	}
	return err
}

func (self *Query) One(x interface{}) error {

	if self.e != nil {
		return self.e
	}

	q := self.generateQuery()

	err := q.One(x)

	if err != nil {

		callback := func() error {
			return q.One(x)
		}

		return self.handleQueryError(err, callback)
	}
	return err
}

func (self *Query) Count(x interface{}) (int, error) {

	if self.e != nil {
		return 0, self.e
	}

	q := self.generateQuery()

	count, err := q.Count()

	if err != nil {

		callback := func() error {
			count, err = q.Count()
			return err
		}

		return count, self.handleQueryError(err, callback)
	}
	return count, err
}

func (self *Query) Distinct(key string, x interface{}) error {

	if self.e != nil {
		return self.e
	}

	q := self.generateQuery()

	err := q.Distinct(key, x)

	if err != nil {

		callback := func() error {
			return q.Distinct(key, x)
		}

		return self.handleQueryError(err, callback)
	}
	return err

}

func (self *Query) handleQueryError(err error, callback queryError) error {

	log.Println("Hanlding Query Error:  " + err.Error())

	if self.isDBConnectionError(err) {

		for i := 0; i < 2; i++ {

			log.Println("Attempting to Refresh Mongo Session")
			dbServices.MongoSession.Refresh()

			err = callback()
			if !self.isDBConnectionError(err) {
				return err
			}

			time.Sleep(200 * time.Millisecond)
		}
	}

	return err

}

func (self *Query) isDBConnectionError(err error) bool {
	if err == nil {
		return false
	}

	if strings.Contains(err.Error(), "Closed explicitly") || strings.Contains(err.Error(), "read: operation timed out") || strings.Contains(err.Error(), "read: connection reset by peer") {
		return true
	}
	return false
}

func (self *Query) generateQuery() *mgo.Query {

	q := self.collection.Find(bson.M{})

	if self.q != nil {
		q = self.q
	}

	if self.m != nil {
		q = self.collection.Find(self.m)
	}

	if self.limit > 0 {
		q = q.Limit(self.limit)
	}

	if self.skip > 0 {
		q = q.Skip(self.skip)
	}

	if len(self.sort) > 0 {
		q = q.Sort(self.sort...)
	}

	return q
}

func (self *Query) getIdHex(val interface{}) (bson.ObjectId, error) {

	myIdType := reflect.TypeOf(val)
	myIdInstance := reflect.ValueOf(val)
	myIdInstanceHex := myIdInstance.MethodByName("Hex")

	if myIdType.Name() != "ObjectId" && myIdType.Kind() == reflect.String && val.(string) != "" {
		return bson.ObjectIdHex(val.(string)), nil
	} else if myIdType.Name() == "ObjectId" && myIdType.Kind() == reflect.String {
		return bson.ObjectIdHex(myIdInstanceHex.Call([]reflect.Value{})[0].String()), nil
	} else {
		return bson.NewObjectId(), errors.New(ERROR_INVALID_ID_VALUE)
	}
}
<commit_msg>Removed Log<commit_after>package model

import (
	"errors"
	"log"
	"reflect"
	"strings"
	"time"

	"github.com/DanielRenne/GoCore/core/dbServices"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

const (
	ERROR_INVALID_ID_VALUE = "Invalid Id value for query."
)

type queryError func() error

type Range struct {
	Min interface{}
	Max interface{}
}

type Query struct {
	q          *mgo.Query
	m          bson.M
	limit      int
	skip       int
	sort       []string
	collection *mgo.Collection
	e          error
}

func Criteria(key string, value interface{}) map[string]interface{} {
	criteria := make(map[string]interface{}, 1)
	criteria[key] = value
	return criteria
}

func (self *Query) ById(val interface{}, x interface{}) error {

	objId, err := self.getIdHex(val)
	if err != nil {
		return err
	}

	q := self.collection.FindId(objId)

	err = q.One(x)

	if err != nil {
		callback := func() error {
			return q.One(x)
		}

		return self.handleQueryError(err, callback)
	}

	return err

}

func (self *Query) Filter(criteria map[string]interface{}) *Query {

	val, hasId := criteria["Id"]
	if hasId {

		objId, err := self.getIdHex(val)
		if err != nil {
			self.e = err
			return self
		}

		self.q = self.collection.FindId(objId)

		return self
	} else {

		if self.m == nil {
			self.m = make(bson.M)
		}

		for key, val := range criteria {
			if key != "" {
				self.m[key] = val
			}
		}

	}
	return self
}

func (self *Query) Exclude(criteria map[string]interface{}) *Query {

	if self.m == nil {
		self.m = make(bson.M)
	}

	self.inNot(criteria, "$nin")

	return self

}

func (self *Query) In(criteria map[string]interface{}) *Query {

	if self.m == nil {
		self.m = make(bson.M)
	}

	self.inNot(criteria, "$in")

	return self

}

func (self *Query) inNot(criteria map[string]interface{}, queryType string) {
	for key, value := range criteria {

		if key == "Id" {
			var ids []bson.ObjectId

			k := reflect.TypeOf(value).Kind()
			if k == reflect.Slice || k == reflect.Array {
				values := reflect.ValueOf(value)

				for i := 0; i < values.Len(); i++ {
					val := values.Index(i).Interface()
					objId, err := self.getIdHex(val)
					if err != nil {
						self.e = err
						continue
					}

					ids = append(ids, objId)
				}
			} else {
				objId, err := self.getIdHex(value)
				if err != nil {
					self.e = err
					continue
				}

				ids = append(ids, objId)
			}

			self.m["_id"] = bson.M{queryType: ids}

		} else {
			var valuesToQuery []interface{}

			k := reflect.TypeOf(value).Kind()
			if k == reflect.Slice || k == reflect.Array {
				values := reflect.ValueOf(value)
				for i := 0; i < values.Len(); i++ {
					val := values.Index(i).Interface()
					valuesToQuery = append(valuesToQuery, val)
				}

			} else {
				valuesToQuery = append(valuesToQuery, value)
			}

			self.m[key] = bson.M{queryType: valuesToQuery}
		}

	}
}

func (self *Query) Range(criteria map[string]Range) *Query {

	if self.m == nil {
		self.m = make(bson.M)
	}

	for key, val := range criteria {
		if key != "" {
			self.m[key] = bson.M{"$gte": val.Min, "$lte": val.Max}
		}
	}

	return self
}

func (self *Query) Sort(fields ...string) *Query {
	self.sort = fields
	return self
}

func (self *Query) Limit(val int) *Query {
	self.limit = val
	return self
}

func (self *Query) Skip(val int) *Query {
	self.skip = val
	return self
}

func (self *Query) All(x interface{}) error {

	if self.e != nil {
		return self.e
	}

	q := self.generateQuery()

	err := q.All(x)

	if err != nil {

		callback := func() error {
			return q.All(x)
		}

		return self.handleQueryError(err, callback)
	}
	return err
}

func (self *Query) One(x interface{}) error {

	if self.e != nil {
		return self.e
	}

	q := self.generateQuery()

	err := q.One(x)

	if err != nil {

		callback := func() error {
			return q.One(x)
		}

		return self.handleQueryError(err, callback)
	}
	return err
}

func (self *Query) Count(x interface{}) (int, error) {

	if self.e != nil {
		return 0, self.e
	}

	q := self.generateQuery()

	count, err := q.Count()

	if err != nil {

		callback := func() error {
			count, err = q.Count()
			return err
		}

		return count, self.handleQueryError(err, callback)
	}
	return count, err
}

func (self *Query) Distinct(key string, x interface{}) error {

	if self.e != nil {
		return self.e
	}

	q := self.generateQuery()

	err := q.Distinct(key, x)

	if err != nil {

		callback := func() error {
			return q.Distinct(key, x)
		}

		return self.handleQueryError(err, callback)
	}
	return err

}

func (self *Query) handleQueryError(err error, callback queryError) error {

	if self.isDBConnectionError(err) {

		for i := 0; i < 2; i++ {

			log.Println("Attempting to Refresh Mongo Session")
			dbServices.MongoSession.Refresh()

			err = callback()
			if !self.isDBConnectionError(err) {
				return err
			}

			time.Sleep(200 * time.Millisecond)
		}
	}

	return err

}

func (self *Query) isDBConnectionError(err error) bool {
	if err == nil {
		return false
	}

	if strings.Contains(err.Error(), "Closed explicitly") || strings.Contains(err.Error(), "read: operation timed out") || strings.Contains(err.Error(), "read: connection reset by peer") {
		return true
	}
	return false
}

func (self *Query) generateQuery() *mgo.Query {

	q := self.collection.Find(bson.M{})

	if self.q != nil {
		q = self.q
	}

	if self.m != nil {
		q = self.collection.Find(self.m)
	}

	if self.limit > 0 {
		q = q.Limit(self.limit)
	}

	if self.skip > 0 {
		q = q.Skip(self.skip)
	}

	if len(self.sort) > 0 {
		q = q.Sort(self.sort...)
	}

	return q
}

func (self *Query) getIdHex(val interface{}) (bson.ObjectId, error) {

	myIdType := reflect.TypeOf(val)
	myIdInstance := reflect.ValueOf(val)
	myIdInstanceHex := myIdInstance.MethodByName("Hex")

	if myIdType.Name() != "ObjectId" && myIdType.Kind() == reflect.String && val.(string) != "" {
		return bson.ObjectIdHex(val.(string)), nil
	} else if myIdType.Name() == "ObjectId" && myIdType.Kind() == reflect.String {
		return bson.ObjectIdHex(myIdInstanceHex.Call([]reflect.Value{})[0].String()), nil
	} else {
		return bson.NewObjectId(), errors.New(ERROR_INVALID_ID_VALUE)
	}
}
<|endoftext|>