<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package config

import (
	"fmt"

	"github.com/spf13/pflag"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// ContainerRuntimeOptions defines options for the container runtime.
type ContainerRuntimeOptions struct {
	// General Options.

	// ContainerRuntime is the container runtime to use.
	ContainerRuntime string
	// RuntimeCgroups that container runtime is expected to be isolated in.
	RuntimeCgroups string

	// Docker-specific options.

	// DockershimRootDirectory is the path to the dockershim root directory. Defaults to
	// /var/lib/dockershim if unset. Exposed for integration testing (e.g. in OpenShift).
	DockershimRootDirectory string
	// PodSandboxImage is the image whose network/ipc namespaces
	// containers in each pod will use.
	PodSandboxImage string
	// DockerEndpoint is the path to the docker endpoint to communicate with.
	DockerEndpoint string
	// If no pulling progress is made before the deadline imagePullProgressDeadline,
	// the image pulling will be cancelled. Defaults to 1m0s.
	// +optional
	ImagePullProgressDeadline metav1.Duration

	// Network plugin options.

	// networkPluginName is the name of the network plugin to be invoked for
	// various events in kubelet/pod lifecycle
	NetworkPluginName string
	// NetworkPluginMTU is the MTU to be passed to the network plugin,
	// and overrides the default MTU for cases where it cannot be automatically
	// computed (such as IPSEC).
	NetworkPluginMTU int32
	// CNIConfDir is the full path of the directory in which to search for
	// CNI config files
	CNIConfDir string
	// CNIBinDir is the full path of the directory in which to search for
	// CNI plugin binaries
	CNIBinDir string
	// CNICacheDir is the full path of the directory in which CNI should store
	// cache files
	CNICacheDir string

	// Image credential provider plugin options

	// ImageCredentialProviderConfigFile is the path to the credential provider plugin config file.
	// This config file is a specification for what credential providers are enabled and invokved
	// by the kubelet. The plugin config should contain information about what plugin binary
	// to execute and what container images the plugin should be called for.
	// +optional
	ImageCredentialProviderConfigFile string
	// ImageCredentialProviderBinDir is the path to the directory where credential provider plugin
	// binaries exist. The name of each plugin binary is expected to match the name of the plugin
	// specified in imageCredentialProviderConfigFile.
	// +optional
	ImageCredentialProviderBinDir string
}

// AddFlags adds flags to the container runtime, according to ContainerRuntimeOptions.
func (s *ContainerRuntimeOptions) AddFlags(fs *pflag.FlagSet) {
	dockerOnlyWarning := "This docker-specific flag only works when container-runtime is set to docker."

	// General settings.
	fs.StringVar(&s.ContainerRuntime, "container-runtime", s.ContainerRuntime, "The container runtime to use. Possible values: 'docker', 'remote'.")
	fs.StringVar(&s.RuntimeCgroups, "runtime-cgroups", s.RuntimeCgroups, "Optional absolute name of cgroups to create and run the runtime in.")

	// Docker-specific settings.
	fs.StringVar(&s.DockershimRootDirectory, "experimental-dockershim-root-directory", s.DockershimRootDirectory, "Path to the dockershim root directory.")
	fs.MarkHidden("experimental-dockershim-root-directory")
	fs.StringVar(&s.PodSandboxImage, "pod-infra-container-image", s.PodSandboxImage, fmt.Sprintf("Specified image will not be pruned by the image garbage collector. "+
		"When container-runtime is set to 'docker', all containers in each pod will use the network/ipc namespaces from this image. Other CRI implementations have their own configuration to set this image."))
	fs.StringVar(&s.DockerEndpoint, "docker-endpoint", s.DockerEndpoint, fmt.Sprintf("Use this for the docker endpoint to communicate with. %s", dockerOnlyWarning))
	fs.MarkDeprecated("docker-endpoint", "will be removed along with dockershim.")
	fs.DurationVar(&s.ImagePullProgressDeadline.Duration, "image-pull-progress-deadline", s.ImagePullProgressDeadline.Duration, fmt.Sprintf("If no pulling progress is made before this deadline, the image pulling will be cancelled. %s", dockerOnlyWarning))
	fs.MarkDeprecated("image-pull-progress-deadline", "will be removed along with dockershim.")

	// Network plugin settings for Docker.
	fs.StringVar(&s.NetworkPluginName, "network-plugin", s.NetworkPluginName, fmt.Sprintf("The name of the network plugin to be invoked for various events in kubelet/pod lifecycle. %s", dockerOnlyWarning))
	fs.MarkDeprecated("network-plugin", "will be removed along with dockershim.")
	fs.StringVar(&s.CNIConfDir, "cni-conf-dir", s.CNIConfDir, fmt.Sprintf("The full path of the directory in which to search for CNI config files. %s", dockerOnlyWarning))
	fs.MarkDeprecated("cni-conf-dir", "will be removed along with dockershim.")
	fs.StringVar(&s.CNIBinDir, "cni-bin-dir", s.CNIBinDir, fmt.Sprintf("A comma-separated list of full paths of directories in which to search for CNI plugin binaries. %s", dockerOnlyWarning))
	fs.MarkDeprecated("cni-bin-dir", "will be removed along with dockershim.")
	fs.StringVar(&s.CNICacheDir, "cni-cache-dir", s.CNICacheDir, fmt.Sprintf("The full path of the directory in which CNI should store cache files. %s", dockerOnlyWarning))
	fs.MarkDeprecated("cni-cache-dir", "will be removed along with dockershim.")
	fs.Int32Var(&s.NetworkPluginMTU, "network-plugin-mtu", s.NetworkPluginMTU, fmt.Sprintf("The MTU to be passed to the network plugin, to override the default. Set to 0 to use the default 1460 MTU. %s", dockerOnlyWarning))
	fs.MarkDeprecated("network-plugin-mtu", "will be removed along with dockershim.")

	// Image credential provider settings.
	fs.StringVar(&s.ImageCredentialProviderConfigFile, "image-credential-provider-config", s.ImageCredentialProviderConfigFile, "The path to the credential provider plugin config file.")
	fs.StringVar(&s.ImageCredentialProviderBinDir, "image-credential-provider-bin-dir", s.ImageCredentialProviderBinDir, "The path to the directory where credential provider plugin binaries are located.")
}
<commit_msg>fix a typo in the comment of ImageCredentialProviderConfigFile<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package config

import (
	"fmt"

	"github.com/spf13/pflag"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// ContainerRuntimeOptions defines options for the container runtime.
type ContainerRuntimeOptions struct {
	// General Options.

	// ContainerRuntime is the container runtime to use.
	ContainerRuntime string
	// RuntimeCgroups that container runtime is expected to be isolated in.
	RuntimeCgroups string

	// Docker-specific options.

	// DockershimRootDirectory is the path to the dockershim root directory. Defaults to
	// /var/lib/dockershim if unset. Exposed for integration testing (e.g. in OpenShift).
	DockershimRootDirectory string
	// PodSandboxImage is the image whose network/ipc namespaces
	// containers in each pod will use.
	PodSandboxImage string
	// DockerEndpoint is the path to the docker endpoint to communicate with.
	DockerEndpoint string
	// If no pulling progress is made before the deadline imagePullProgressDeadline,
	// the image pulling will be cancelled. Defaults to 1m0s.
	// +optional
	ImagePullProgressDeadline metav1.Duration

	// Network plugin options.

	// networkPluginName is the name of the network plugin to be invoked for
	// various events in kubelet/pod lifecycle
	NetworkPluginName string
	// NetworkPluginMTU is the MTU to be passed to the network plugin,
	// and overrides the default MTU for cases where it cannot be automatically
	// computed (such as IPSEC).
	NetworkPluginMTU int32
	// CNIConfDir is the full path of the directory in which to search for
	// CNI config files
	CNIConfDir string
	// CNIBinDir is the full path of the directory in which to search for
	// CNI plugin binaries
	CNIBinDir string
	// CNICacheDir is the full path of the directory in which CNI should store
	// cache files
	CNICacheDir string

	// Image credential provider plugin options

	// ImageCredentialProviderConfigFile is the path to the credential provider plugin config file.
	// This config file is a specification for what credential providers are enabled and invoked
	// by the kubelet. The plugin config should contain information about what plugin binary
	// to execute and what container images the plugin should be called for.
	// +optional
	ImageCredentialProviderConfigFile string
	// ImageCredentialProviderBinDir is the path to the directory where credential provider plugin
	// binaries exist. The name of each plugin binary is expected to match the name of the plugin
	// specified in imageCredentialProviderConfigFile.
	// +optional
	ImageCredentialProviderBinDir string
}

// AddFlags adds flags to the container runtime, according to ContainerRuntimeOptions.
func (s *ContainerRuntimeOptions) AddFlags(fs *pflag.FlagSet) {
	dockerOnlyWarning := "This docker-specific flag only works when container-runtime is set to docker."

	// General settings.
	fs.StringVar(&s.ContainerRuntime, "container-runtime", s.ContainerRuntime, "The container runtime to use. Possible values: 'docker', 'remote'.")
	fs.StringVar(&s.RuntimeCgroups, "runtime-cgroups", s.RuntimeCgroups, "Optional absolute name of cgroups to create and run the runtime in.")

	// Docker-specific settings.
	fs.StringVar(&s.DockershimRootDirectory, "experimental-dockershim-root-directory", s.DockershimRootDirectory, "Path to the dockershim root directory.")
	fs.MarkHidden("experimental-dockershim-root-directory")
	fs.StringVar(&s.PodSandboxImage, "pod-infra-container-image", s.PodSandboxImage, fmt.Sprintf("Specified image will not be pruned by the image garbage collector. "+
		"When container-runtime is set to 'docker', all containers in each pod will use the network/ipc namespaces from this image. Other CRI implementations have their own configuration to set this image."))
	fs.StringVar(&s.DockerEndpoint, "docker-endpoint", s.DockerEndpoint, fmt.Sprintf("Use this for the docker endpoint to communicate with. %s", dockerOnlyWarning))
	fs.MarkDeprecated("docker-endpoint", "will be removed along with dockershim.")
	fs.DurationVar(&s.ImagePullProgressDeadline.Duration, "image-pull-progress-deadline", s.ImagePullProgressDeadline.Duration, fmt.Sprintf("If no pulling progress is made before this deadline, the image pulling will be cancelled. %s", dockerOnlyWarning))
	fs.MarkDeprecated("image-pull-progress-deadline", "will be removed along with dockershim.")

	// Network plugin settings for Docker.
	fs.StringVar(&s.NetworkPluginName, "network-plugin", s.NetworkPluginName, fmt.Sprintf("The name of the network plugin to be invoked for various events in kubelet/pod lifecycle. %s", dockerOnlyWarning))
	fs.MarkDeprecated("network-plugin", "will be removed along with dockershim.")
	fs.StringVar(&s.CNIConfDir, "cni-conf-dir", s.CNIConfDir, fmt.Sprintf("The full path of the directory in which to search for CNI config files. %s", dockerOnlyWarning))
	fs.MarkDeprecated("cni-conf-dir", "will be removed along with dockershim.")
	fs.StringVar(&s.CNIBinDir, "cni-bin-dir", s.CNIBinDir, fmt.Sprintf("A comma-separated list of full paths of directories in which to search for CNI plugin binaries. %s", dockerOnlyWarning))
	fs.MarkDeprecated("cni-bin-dir", "will be removed along with dockershim.")
	fs.StringVar(&s.CNICacheDir, "cni-cache-dir", s.CNICacheDir, fmt.Sprintf("The full path of the directory in which CNI should store cache files. %s", dockerOnlyWarning))
	fs.MarkDeprecated("cni-cache-dir", "will be removed along with dockershim.")
	fs.Int32Var(&s.NetworkPluginMTU, "network-plugin-mtu", s.NetworkPluginMTU, fmt.Sprintf("The MTU to be passed to the network plugin, to override the default. Set to 0 to use the default 1460 MTU. %s", dockerOnlyWarning))
	fs.MarkDeprecated("network-plugin-mtu", "will be removed along with dockershim.")

	// Image credential provider settings.
	fs.StringVar(&s.ImageCredentialProviderConfigFile, "image-credential-provider-config", s.ImageCredentialProviderConfigFile, "The path to the credential provider plugin config file.")
	fs.StringVar(&s.ImageCredentialProviderBinDir, "image-credential-provider-bin-dir", s.ImageCredentialProviderBinDir, "The path to the directory where credential provider plugin binaries are located.")
}
<|endoftext|>


<commit_before>package model

import (
	"net/http"
	"testing"
	"time"

	"github.com/labstack/echo"
	sd "github.com/labstack/echo/engine/standard"
)

const (
	etcdAddr   = "http://192.168.70.13:2379"
	etcdPrefix = "/gateway2"
)

var (
	serverAddr    = "127.0.0.1:12345"
	apiURL        = "/api/test"
	checkDuration = 3
	checkTimeout  = 2
	clusterName   = "app"
	lbName        = "ROUNDROBIN"
	sleep         = false
)

var rt *RouteTable

func createRouteTable(t *testing.T) {
	store, err := NewEtcdStore([]string{etcdAddr}, etcdPrefix)

	if nil != err {
		t.Fatalf("create etcd store err.addr:<%s>", err)
	}

	store.Clean()

	rt = NewRouteTable(store)
	time.Sleep(time.Second * 1)
}

func createLocalServer() {
	e := echo.New()

	e.Get("/check", func() echo.HandlerFunc {
		return func(c echo.Context) error {
			if sleep {
				time.Sleep(time.Second * time.Duration(checkTimeout+1))
			}

			return c.String(http.StatusOK, "OK")
		}
	}())

	e.Run(sd.New(serverAddr))
}

func waitNotify() {
	time.Sleep(time.Second * 1)
}

func TestCreateRouteTable(t *testing.T) {
	createRouteTable(t)
}

func TestEtcdWatchNewServer(t *testing.T) {
	go createLocalServer()

	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout,
		MaxQPS:          1500,
		HalfToOpen:      10,
		HalfTrafficRate: 10,
		CloseCount:      100,
	}

	err := rt.store.SaveServer(server)

	if nil != err {
		t.Error("add server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 1 {
		t.Errorf("expect:<1>, acture:<%d>", len(rt.svrs))
		return
	}

	if rt.svrs[serverAddr].lock == nil {
		t.Error("server init error.")
		return
	}
}

func TestServerCheckOk(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration))

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check ok err.expect:<UP>, acture:<%v>", Down)
	}
}

func TestServerCheckTimeout(t *testing.T) {
	defer func() {
		sleep = false
	}()

	sleep = true
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Up {
		t.Errorf("status check timeout err.expect:<DOWN>, acture:<%v>", Up)
		return
	}
}

func TestServerCheckTimeoutRecovery(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check timeout recovery err.expect:<UP>, acture:<%v>", Up)
		return
	}
}

func TestEtcdWatchNewCluster(t *testing.T) {
	cluster := &Cluster{
		Name:    clusterName,
		Pattern: "/api/*",
		LbName:  lbName,
	}

	err := rt.store.SaveCluster(cluster)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.clusters))
}

func TestEtcdWatchNewBind(t *testing.T) {
	bind := &Bind{
		ClusterName: clusterName,
		ServerAddr:  serverAddr,
	}

	err := rt.store.SaveBind(bind)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.mapping) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>. %+v", len(rt.mapping), rt.mapping)
}

func TestEtcdWatchNewAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		URL:         "/api/node/test",
		ClusterName: clusterName,
	}

	err := rt.store.SaveAPI(&API{
		URL:   apiURL,
		Nodes: []*Node{n},
	})

	if nil != err {
		t.Error("add api err.")
		return
	}

	waitNotify()

	if len(rt.apis) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.apis))
}

func TestEtcdWatchUpdateServer(t *testing.T) {
	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout * 2,
		MaxQPS:          3000,
		HalfToOpen:      100,
		HalfTrafficRate: 30,
		CloseCount:      200,
	}

	err := rt.store.UpdateServer(server)

	if nil != err {
		t.Error("update server err.")
		return
	}

	waitNotify()

	svr := rt.svrs[serverAddr]

	if svr.MaxQPS != server.MaxQPS {
		t.Errorf("MaxQPS expect:<%d>, acture:<%d>. ", server.MaxQPS, svr.MaxQPS)
		return
	}

	if svr.HalfToOpen != server.HalfToOpen {
		t.Errorf("HalfToOpen expect:<%d>, acture:<%d>. ", server.HalfToOpen, svr.HalfToOpen)
		return
	}

	if svr.HalfTrafficRate != server.HalfTrafficRate {
		t.Errorf("HalfTrafficRate expect:<%d>, acture:<%d>. ", server.HalfTrafficRate, svr.HalfTrafficRate)
		return
	}

	if svr.CloseCount != server.CloseCount {
		t.Errorf("CloseCount expect:<%d>, acture:<%d>. ", server.CloseCount, svr.CloseCount)
		return
	}

	if svr.CheckTimeout == server.CheckTimeout {
		t.Errorf("CheckTimeout expect:<%d>, acture:<%d>. ", svr.CheckTimeout, server.CheckTimeout)
		return
	}
}

func TestEtcdWatchUpdateCluster(t *testing.T) {
	cluster := &Cluster{
		Name:    clusterName,
		Pattern: "/api/new/*",
		LbName:  lbName,
	}

	err := rt.store.UpdateCluster(cluster)

	if nil != err {
		t.Error("update cluster err.")
		return
	}

	waitNotify()

	existCluster := rt.clusters[clusterName]

	if existCluster.Pattern != cluster.Pattern {
		t.Errorf("Pattern expect:<%s>, acture:<%s>. ", cluster.Pattern, existCluster.Pattern)
		return
	}

	if existCluster.LbName != cluster.LbName {
		t.Errorf("LbName expect:<%s>, acture:<%s>. ", cluster.LbName, existCluster.LbName)
		return
	}
}

func TestEtcdWatchUpdateAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		URL:         "/api/node/test",
		ClusterName: clusterName,
	}

	n2 := &Node{
		AttrName:    "tes2t",
		URL:         "/api/node/test2",
		ClusterName: clusterName,
	}

	api := &API{
		URL:   apiURL,
		Nodes: []*Node{n, n2},
	}

	err := rt.store.UpdateAPI(api)

	if nil != err {
		t.Error("update api err.")
		return
	}

	waitNotify()

	existAPI, _ := rt.apis[api.URL]

	if len(existAPI.Nodes) != len(api.Nodes) {
		t.Errorf("Nodes expect:<%s>, acture:<%s>. ", len(existAPI.Nodes), len(api.Nodes))
		return
	}
}

func TestEtcdWatchDeleteCluster(t *testing.T) {
	err := rt.store.DeleteCluster(clusterName)

	if nil != err {
		t.Error("delete cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) != 0 {
		t.Errorf("clusters expect:<0>, acture:<%d>", len(rt.clusters))
		return
	}

	banded, _ := rt.mapping[serverAddr]

	if len(banded) != 0 {
		t.Errorf("banded expect:<0>, acture:<%d>", len(banded))
		return
	}
}

func TestEtcdWatchDeleteServer(t *testing.T) {
	err := rt.store.DeleteServer(serverAddr)

	if nil != err {
		t.Error("delete server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 0 {
		t.Errorf("svrs expect:<0>, acture:<%d>", len(rt.svrs))
		return
	}

	if len(rt.mapping) != 0 {
		t.Errorf("mapping expect:<0>, acture:<%d>", len(rt.mapping))
		return
	}
}

func TestEtcdWatchDeleteAPI(t *testing.T) {
	err := rt.store.DeleteAPI(apiURL)

	if nil != err {
		t.Error("delete api err.")
		return
	}

	waitNotify()

	if len(rt.apis) != 0 {
		t.Errorf("apis expect:<0>, acture:<%d>", len(rt.apis))
		return
	}
}

func TestEtcdWatchNewRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 100; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	waitNotify()

	if len(rt.routings) == 1 {
		delete(rt.routings, r.ID)
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.routings))
}

func TestEtcdWatchDeleteRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 3; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	time.Sleep(time.Second * 30)

	if len(rt.routings) == 0 {
		return
	}

	t.Errorf("expect:<0>, acture:<%d>", len(rt.routings))
}
<commit_msg>update test<commit_after>package model

import (
	"net/http"
	"testing"
	"time"

	"github.com/labstack/echo"
	sd "github.com/labstack/echo/engine/standard"
)

const (
	etcdAddr   = "http://192.168.70.13:2379"
	etcdPrefix = "/gateway2"
)

var (
	serverAddr    = "127.0.0.1:12345"
	apiURL        = "/api/test"
	apiMethod     = "GET"
	checkDuration = 3
	checkTimeout  = 2
	clusterName   = "app"
	lbName        = "ROUNDROBIN"
	sleep         = false
)

var rt *RouteTable

func createRouteTable(t *testing.T) {
	store, err := NewEtcdStore([]string{etcdAddr}, etcdPrefix)

	if nil != err {
		t.Fatalf("create etcd store err.addr:<%s>", err)
	}

	store.Clean()

	rt = NewRouteTable(store)
	time.Sleep(time.Second * 1)
}

func createLocalServer() {
	e := echo.New()

	e.Get("/check", func() echo.HandlerFunc {
		return func(c echo.Context) error {
			if sleep {
				time.Sleep(time.Second * time.Duration(checkTimeout+1))
			}

			return c.String(http.StatusOK, "OK")
		}
	}())

	e.Run(sd.New(serverAddr))
}

func waitNotify() {
	time.Sleep(time.Second * 1)
}

func TestCreateRouteTable(t *testing.T) {
	createRouteTable(t)
}

func TestEtcdWatchNewServer(t *testing.T) {
	go createLocalServer()

	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout,
		MaxQPS:          1500,
		HalfToOpen:      10,
		HalfTrafficRate: 10,
		CloseCount:      100,
	}

	err := rt.store.SaveServer(server)

	if nil != err {
		t.Error("add server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 1 {
		t.Errorf("expect:<1>, acture:<%d>", len(rt.svrs))
		return
	}

	if rt.svrs[serverAddr].lock == nil {
		t.Error("server init error.")
		return
	}
}

func TestServerCheckOk(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration))

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check ok err.expect:<UP>, acture:<%v>", Down)
	}
}

func TestServerCheckTimeout(t *testing.T) {
	defer func() {
		sleep = false
	}()

	sleep = true
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Up {
		t.Errorf("status check timeout err.expect:<DOWN>, acture:<%v>", Up)
		return
	}
}

func TestServerCheckTimeoutRecovery(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check timeout recovery err.expect:<UP>, acture:<%v>", Up)
		return
	}
}

func TestEtcdWatchNewCluster(t *testing.T) {
	cluster := &Cluster{
		Name:   clusterName,
		LbName: lbName,
	}

	err := rt.store.SaveCluster(cluster)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.clusters))
}

func TestEtcdWatchNewBind(t *testing.T) {
	bind := &Bind{
		ClusterName: clusterName,
		ServerAddr:  serverAddr,
	}

	err := rt.store.SaveBind(bind)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.mapping) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>. %+v", len(rt.mapping), rt.mapping)
}

func TestEtcdWatchNewAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		ClusterName: clusterName,
	}

	err := rt.store.SaveAPI(&API{
		URL:    apiURL,
		Method: apiMethod,
		Nodes:  []*Node{n},
	})

	if nil != err {
		t.Error("add api err.")
		return
	}

	waitNotify()

	if len(rt.apis) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.apis))
}

func TestEtcdWatchUpdateServer(t *testing.T) {
	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout * 2,
		MaxQPS:          3000,
		HalfToOpen:      100,
		HalfTrafficRate: 30,
		CloseCount:      200,
	}

	err := rt.store.UpdateServer(server)

	if nil != err {
		t.Error("update server err.")
		return
	}

	waitNotify()

	svr := rt.svrs[serverAddr]

	if svr.MaxQPS != server.MaxQPS {
		t.Errorf("MaxQPS expect:<%d>, acture:<%d>. ", server.MaxQPS, svr.MaxQPS)
		return
	}

	if svr.HalfToOpen != server.HalfToOpen {
		t.Errorf("HalfToOpen expect:<%d>, acture:<%d>. ", server.HalfToOpen, svr.HalfToOpen)
		return
	}

	if svr.HalfTrafficRate != server.HalfTrafficRate {
		t.Errorf("HalfTrafficRate expect:<%d>, acture:<%d>. ", server.HalfTrafficRate, svr.HalfTrafficRate)
		return
	}

	if svr.CloseCount != server.CloseCount {
		t.Errorf("CloseCount expect:<%d>, acture:<%d>. ", server.CloseCount, svr.CloseCount)
		return
	}

	if svr.CheckTimeout == server.CheckTimeout {
		t.Errorf("CheckTimeout expect:<%d>, acture:<%d>. ", svr.CheckTimeout, server.CheckTimeout)
		return
	}
}

func TestEtcdWatchUpdateCluster(t *testing.T) {
	cluster := &Cluster{
		Name:   clusterName,
		LbName: lbName,
	}

	err := rt.store.UpdateCluster(cluster)

	if nil != err {
		t.Error("update cluster err.")
		return
	}

	waitNotify()

	existCluster := rt.clusters[clusterName]

	if existCluster.LbName != cluster.LbName {
		t.Errorf("LbName expect:<%s>, acture:<%s>. ", cluster.LbName, existCluster.LbName)
		return
	}
}

func TestEtcdWatchUpdateAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		ClusterName: clusterName,
	}

	n2 := &Node{
		AttrName:    "tes2t",
		ClusterName: clusterName,
	}

	api := &API{
		URL:    apiURL,
		Method: apiMethod,
		Nodes:  []*Node{n, n2},
	}

	err := rt.store.UpdateAPI(api)

	if nil != err {
		t.Error("update api err.")
		return
	}

	waitNotify()

	existAPI, _ := rt.apis[getAPIKey(api.URL, api.Method)]

	if len(existAPI.Nodes) != len(api.Nodes) {
		t.Errorf("Nodes expect:<%s>, acture:<%s>. ", len(existAPI.Nodes), len(api.Nodes))
		return
	}
}

func TestEtcdWatchDeleteCluster(t *testing.T) {
	err := rt.store.DeleteCluster(clusterName)

	if nil != err {
		t.Error("delete cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) != 0 {
		t.Errorf("clusters expect:<0>, acture:<%d>", len(rt.clusters))
		return
	}

	banded, _ := rt.mapping[serverAddr]

	if len(banded) != 0 {
		t.Errorf("banded expect:<0>, acture:<%d>", len(banded))
		return
	}
}

func TestEtcdWatchDeleteServer(t *testing.T) {
	err := rt.store.DeleteServer(serverAddr)

	if nil != err {
		t.Error("delete server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 0 {
		t.Errorf("svrs expect:<0>, acture:<%d>", len(rt.svrs))
		return
	}

	if len(rt.mapping) != 0 {
		t.Errorf("mapping expect:<0>, acture:<%d>", len(rt.mapping))
		return
	}
}

func TestEtcdWatchDeleteAPI(t *testing.T) {
	err := rt.store.DeleteAPI(apiURL, apiMethod)

	if nil != err {
		t.Error("delete api err.")
		return
	}

	waitNotify()

	if len(rt.apis) != 0 {
		t.Errorf("apis expect:<0>, acture:<%d>", len(rt.apis))
		return
	}
}

func TestEtcdWatchNewRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 100; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	waitNotify()

	if len(rt.routings) == 1 {
		delete(rt.routings, r.ID)
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.routings))
}

func TestEtcdWatchDeleteRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 3; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	time.Sleep(time.Second * 30)

	if len(rt.routings) == 0 {
		return
	}

	t.Errorf("expect:<0>, acture:<%d>", len(rt.routings))
}
<|endoftext|>


<commit_before>package handler

import (
	"encoding/json"
	"net/http"

	"github.com/sirupsen/logrus"
	"github.com/skygeario/skygear-server/pkg/core/asset"
	"github.com/skygeario/skygear-server/pkg/core/auth"
	"github.com/skygeario/skygear-server/pkg/core/auth/authz"
	"github.com/skygeario/skygear-server/pkg/core/auth/authz/policy"
	"github.com/skygeario/skygear-server/pkg/core/db"
	"github.com/skygeario/skygear-server/pkg/core/handler"
	"github.com/skygeario/skygear-server/pkg/core/inject"
	"github.com/skygeario/skygear-server/pkg/core/model"
	"github.com/skygeario/skygear-server/pkg/core/server"
	recordGear "github.com/skygeario/skygear-server/pkg/record"
	"github.com/skygeario/skygear-server/pkg/record/dependency/record"
	"github.com/skygeario/skygear-server/pkg/server/skyerr"
)

func AttachQueryHandler(
	server *server.Server,
	recordDependency recordGear.DependencyMap,
) *server.Server {
	server.Handle("/query", &QueryHandlerFactory{
		recordDependency,
	}).Methods("POST")
	return server
}

type QueryHandlerFactory struct {
	Dependency recordGear.DependencyMap
}

func (f QueryHandlerFactory) NewHandler(request *http.Request) http.Handler {
	h := &QueryHandler{}
	inject.DefaultInject(h, f.Dependency, request)
	return handler.APIHandlerToHandler(h, h.TxContext)
}

func (f QueryHandlerFactory) ProvideAuthzPolicy() authz.Policy {
	return policy.AllOf(
		authz.PolicyFunc(policy.DenyNoAccessKey),
		authz.PolicyFunc(policy.RequireAuthenticated),
		authz.PolicyFunc(policy.DenyDisabledUser),
	)
}

type QueryRequestPayload struct {
	Query record.Query
}

func (p QueryRequestPayload) Validate() error {
	return nil
}

type QueryResult struct {
	Records interface{}            `json:"records"`
	Info    map[string]interface{} `json:"info,omitempty"`
}

/*
QueryHandler is dummy implementation on fetching Records
curl -X POST -H "Content-Type: application/json" \
  -d @- http://localhost:3000/query <<EOF
{
    "record_type": "note",
    "sort": [
        [{"$val": "noteOrder", "$type": "desc"}, "asc"]
    ]
}
EOF
*/
type QueryHandler struct {
	AuthContext auth.ContextGetter `dependency:"AuthContextGetter"`
	TxContext   db.TxContext       `dependency:"TxContext"`
	RecordStore record.Store       `dependency:"RecordStore"`
	Logger      *logrus.Entry      `dependency:"HandlerLogger"`
	AssetStore  asset.Store        `dependency:"AssetStore"`
}

func (h QueryHandler) WithTx() bool {
	return true
}

func (h QueryHandler) DecodeRequest(request *http.Request) (handler.RequestPayload, error) {
	payload := QueryRequestPayload{}
	parser := QueryParser{}
	authInfo := h.AuthContext.AuthInfo()
	if authInfo != nil {
		parser.UserID = authInfo.ID
	}
	data := map[string]interface{}{}
	json.NewDecoder(request.Body).Decode(&data)
	if err := parser.queryFromRaw(data, &payload.Query); err != nil {
		return nil, err
	}

	return payload, nil
}

func (h QueryHandler) Handle(req interface{}) (resp interface{}, err error) {
	payload := req.(QueryRequestPayload)

	accessControlOptions := &record.AccessControlOptions{
		ViewAsUser:          h.AuthContext.AuthInfo(),
		BypassAccessControl: h.AuthContext.AccessKeyType() == model.MasterAccessKey,
	}

	fieldACL := func() record.FieldACL {
		acl, err := h.RecordStore.GetRecordFieldAccess()
		if err != nil {
			panic(err)
		}
		return acl
	}()

	if !accessControlOptions.BypassAccessControl {
		visitor := &queryAccessVisitor{
			FieldACL:   fieldACL,
			RecordType: payload.Query.Type,
			AuthInfo:   accessControlOptions.ViewAsUser,
			ExpressionACLChecker: ExpressionACLChecker{
				FieldACL:    fieldACL,
				RecordType:  payload.Query.Type,
				AuthInfo:    h.AuthContext.AuthInfo(),
				RecordStore: h.RecordStore,
			},
		}
		payload.Query.Accept(visitor)
		if err = visitor.Error(); err != nil {
			return
		}
	}

	results, err := h.RecordStore.Query(&payload.Query, accessControlOptions)
	if err != nil {
		err = skyerr.MakeError(err)
		return
	}
	defer results.Close()

	records := []record.Record{}
	for results.Scan() {
		record := results.Record()
		records = append(records, record)
	}

	err = results.Err()
	if err != nil {
		err = skyerr.MakeError(err)
		return
	}

	// Scan does not query assets,
	// it only replaces them with assets then only have name,
	// so we replace them with some complete assets.
	MakeAssetsComplete(h.RecordStore, records)

	eagerRecords := h.doQueryEager(h.eagerIDs(records, payload.Query), accessControlOptions)

	recordResultFilter, err := NewRecordResultFilter(
		h.RecordStore,
		h.TxContext,
		h.AssetStore,
		h.AuthContext.AuthInfo(),
		h.AuthContext.AccessKeyType() == model.MasterAccessKey,
	)
	if err != nil {
		err = skyerr.MakeError(err)
		return
	}

	result := QueryResult{}
	resultFilter := QueryResultFilter{
		RecordStore:        h.RecordStore,
		Query:              payload.Query,
		EagerRecords:       eagerRecords,
		RecordResultFilter: recordResultFilter,
	}

	output := make([]interface{}, len(records))
	for i := range records {
		record := records[i]
		output[i] = resultFilter.JSONResult(&record)
	}

	result.Records = output

	resultInfo, err := QueryResultInfo(h.RecordStore, &payload.Query, accessControlOptions, results)
	if err != nil {
		err = skyerr.MakeError(err)
		return
	}

	if len(resultInfo) > 0 {
		result.Info = resultInfo
	}

	resp = result

	return
}

func (h QueryHandler) eagerIDs(records []record.Record, query record.Query) map[string][]record.ID {
	eagers := map[string][]record.ID{}
	for _, transientExpression := range query.ComputedKeys {
		if transientExpression.Type != record.KeyPath {
			continue
		}
		keyPath := transientExpression.Value.(string)
		eagers[keyPath] = make([]record.ID, len(records))
	}

	for i, record := range records {
		for keyPath := range eagers {
			ref := getReferenceWithKeyPath(h.RecordStore, &record, keyPath)
			if ref.IsEmpty() {
				continue
			}
			eagers[keyPath][i] = ref.ID
		}
	}
	return eagers
}

func (h QueryHandler) doQueryEager(eagersIDs map[string][]record.ID, accessControlOptions *record.AccessControlOptions) map[string]map[string]*record.Record {
	eagerRecords := map[string]map[string]*record.Record{}
	for keyPath, ids := range eagersIDs {
		h.Logger.Debugf("Getting value for keypath %v", keyPath)
		eagerScanner, err := h.RecordStore.GetByIDs(ids, accessControlOptions)
		if err != nil {
			h.Logger.Debugf("No Records found in the eager load key path: %s", keyPath)
			eagerRecords[keyPath] = map[string]*record.Record{}
			continue
		}
		for eagerScanner.Scan() {
			er := eagerScanner.Record()
			if eagerRecords[keyPath] == nil {
				eagerRecords[keyPath] = map[string]*record.Record{}
			}
			eagerRecords[keyPath][er.ID.Key] = &er
		}
		eagerScanner.Close()
	}

	return eagerRecords
}
<commit_msg>Remove requiring user login for record query<commit_after>package handler

import (
	"encoding/json"
	"net/http"

	"github.com/sirupsen/logrus"
	"github.com/skygeario/skygear-server/pkg/core/asset"
	"github.com/skygeario/skygear-server/pkg/core/auth"
	"github.com/skygeario/skygear-server/pkg/core/auth/authz"
	"github.com/skygeario/skygear-server/pkg/core/auth/authz/policy"
	"github.com/skygeario/skygear-server/pkg/core/db"
	"github.com/skygeario/skygear-server/pkg/core/handler"
	"github.com/skygeario/skygear-server/pkg/core/inject"
	"github.com/skygeario/skygear-server/pkg/core/model"
	"github.com/skygeario/skygear-server/pkg/core/server"
	recordGear "github.com/skygeario/skygear-server/pkg/record"
	"github.com/skygeario/skygear-server/pkg/record/dependency/record"
	"github.com/skygeario/skygear-server/pkg/server/skyerr"
)

func AttachQueryHandler(
	server *server.Server,
	recordDependency recordGear.DependencyMap,
) *server.Server {
	server.Handle("/query", &QueryHandlerFactory{
		recordDependency,
	}).Methods("POST")
	return server
}

type QueryHandlerFactory struct {
	Dependency recordGear.DependencyMap
}

func (f QueryHandlerFactory) NewHandler(request *http.Request) http.Handler {
	h := &QueryHandler{}
	inject.DefaultInject(h, f.Dependency, request)
	return handler.APIHandlerToHandler(h, h.TxContext)
}

func (f QueryHandlerFactory) ProvideAuthzPolicy() authz.Policy {
	return policy.AllOf(
		authz.PolicyFunc(policy.DenyNoAccessKey),
		authz.PolicyFunc(policy.DenyDisabledUser),
	)
}

type QueryRequestPayload struct {
	Query record.Query
}

func (p QueryRequestPayload) Validate() error {
	return nil
}

type QueryResult struct {
	Records interface{}            `json:"records"`
	Info    map[string]interface{} `json:"info,omitempty"`
}

/*
QueryHandler is dummy implementation on fetching Records
curl -X POST -H "Content-Type: application/json" \
  -d @- http://localhost:3000/query <<EOF
{
    "record_type": "note",
    "sort": [
        [{"$val": "noteOrder", "$type": "desc"}, "asc"]
    ]
}
EOF
*/
type QueryHandler struct {
	AuthContext auth.ContextGetter `dependency:"AuthContextGetter"`
	TxContext   db.TxContext       `dependency:"TxContext"`
	RecordStore record.Store       `dependency:"RecordStore"`
	Logger      *logrus.Entry      `dependency:"HandlerLogger"`
	AssetStore  asset.Store        `dependency:"AssetStore"`
}

func (h QueryHandler) WithTx() bool {
	return true
}

func (h QueryHandler) DecodeRequest(request *http.Request) (handler.RequestPayload, error) {
	payload := QueryRequestPayload{}
	parser := QueryParser{}
	authInfo := h.AuthContext.AuthInfo()
	if authInfo != nil {
		parser.UserID = authInfo.ID
	}
	data := map[string]interface{}{}
	json.NewDecoder(request.Body).Decode(&data)
	if err := parser.queryFromRaw(data, &payload.Query); err != nil {
		return nil, err
	}

	return payload, nil
}

func (h QueryHandler) Handle(req interface{}) (resp interface{}, err error) {
	payload := req.(QueryRequestPayload)

	accessControlOptions := &record.AccessControlOptions{
		ViewAsUser:          h.AuthContext.AuthInfo(),
		BypassAccessControl: h.AuthContext.AccessKeyType() == model.MasterAccessKey,
	}

	fieldACL := func() record.FieldACL {
		acl, err := h.RecordStore.GetRecordFieldAccess()
		if err != nil {
			panic(err)
		}
		return acl
	}()

	if !accessControlOptions.BypassAccessControl {
		visitor := &queryAccessVisitor{
			FieldACL:   fieldACL,
			RecordType: payload.Query.Type,
			AuthInfo:   accessControlOptions.ViewAsUser,
			ExpressionACLChecker: ExpressionACLChecker{
				FieldACL:    fieldACL,
				RecordType:  payload.Query.Type,
				AuthInfo:    h.AuthContext.AuthInfo(),
				RecordStore: h.RecordStore,
			},
		}
		payload.Query.Accept(visitor)
		if err = visitor.Error(); err != nil {
			return
		}
	}

	results, err := h.RecordStore.Query(&payload.Query, accessControlOptions)
	if err != nil {
		err = skyerr.MakeError(err)
		return
	}
	defer results.Close()

	records := []record.Record{}
	for results.Scan() {
		record := results.Record()
		records = append(records, record)
	}

	err = results.Err()
	if err != nil {
		err = skyerr.MakeError(err)
		return
	}

	// Scan does not query assets,
	// it only replaces them with assets then only have name,
	// so we replace them with some complete assets.
	MakeAssetsComplete(h.RecordStore, records)

	eagerRecords := h.doQueryEager(h.eagerIDs(records, payload.Query), accessControlOptions)

	recordResultFilter, err := NewRecordResultFilter(
		h.RecordStore,
		h.TxContext,
		h.AssetStore,
		h.AuthContext.AuthInfo(),
		h.AuthContext.AccessKeyType() == model.MasterAccessKey,
	)
	if err != nil {
		err = skyerr.MakeError(err)
		return
	}

	result := QueryResult{}
	resultFilter := QueryResultFilter{
		RecordStore:        h.RecordStore,
		Query:              payload.Query,
		EagerRecords:       eagerRecords,
		RecordResultFilter: recordResultFilter,
	}

	output := make([]interface{}, len(records))
	for i := range records {
		record := records[i]
		output[i] = resultFilter.JSONResult(&record)
	}

	result.Records = output

	resultInfo, err := QueryResultInfo(h.RecordStore, &payload.Query, accessControlOptions, results)
	if err != nil {
		err = skyerr.MakeError(err)
		return
	}

	if len(resultInfo) > 0 {
		result.Info = resultInfo
	}

	resp = result

	return
}

func (h QueryHandler) eagerIDs(records []record.Record, query record.Query) map[string][]record.ID {
	eagers := map[string][]record.ID{}
	for _, transientExpression := range query.ComputedKeys {
		if transientExpression.Type != record.KeyPath {
			continue
		}
		keyPath := transientExpression.Value.(string)
		eagers[keyPath] = make([]record.ID, len(records))
	}

	for i, record := range records {
		for keyPath := range eagers {
			ref := getReferenceWithKeyPath(h.RecordStore, &record, keyPath)
			if ref.IsEmpty() {
				continue
			}
			eagers[keyPath][i] = ref.ID
		}
	}
	return eagers
}

func (h QueryHandler) doQueryEager(eagersIDs map[string][]record.ID, accessControlOptions *record.AccessControlOptions) map[string]map[string]*record.Record {
	eagerRecords := map[string]map[string]*record.Record{}
	for keyPath, ids := range eagersIDs {
		h.Logger.Debugf("Getting value for keypath %v", keyPath)
		eagerScanner, err := h.RecordStore.GetByIDs(ids, accessControlOptions)
		if err != nil {
			h.Logger.Debugf("No Records found in the eager load key path: %s", keyPath)
			eagerRecords[keyPath] = map[string]*record.Record{}
			continue
		}
		for eagerScanner.Scan() {
			er := eagerScanner.Record()
			if eagerRecords[keyPath] == nil {
				eagerRecords[keyPath] = map[string]*record.Record{}
			}
			eagerRecords[keyPath][er.ID.Key] = &er
		}
		eagerScanner.Close()
	}

	return eagerRecords
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package wardle

import (
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apiserver/pkg/registry/generic"
	genericregistry "k8s.io/apiserver/pkg/registry/generic/registry"
	"k8s.io/sample-apiserver/pkg/apis/wardle"
)

// rest implements a RESTStorage for API services against etcd
type REST struct {
	*genericregistry.Store
}

// NewREST returns a RESTStorage object that will work against API services.
func NewREST(scheme *runtime.Scheme, optsGetter generic.RESTOptionsGetter) *REST {
	strategy := NewStrategy(scheme)

	store := &genericregistry.Store{
		Copier:      scheme,
		NewFunc:     func() runtime.Object { return &wardle.Flunder{} },
		NewListFunc: func() runtime.Object { return &wardle.FlunderList{} },
		ObjectNameFunc: func(obj runtime.Object) (string, error) {
			return obj.(*wardle.Flunder).Name, nil
		},
		PredicateFunc:     MatchFlunder,
		QualifiedResource: wardle.Resource("flunders"),

		CreateStrategy: strategy,
		UpdateStrategy: strategy,
		DeleteStrategy: strategy,
	}
	options := &generic.StoreOptions{RESTOptions: optsGetter, AttrFunc: GetAttrs}
	if err := store.CompleteWithOptions(options); err != nil {
		panic(err) // TODO: Propagate error up
	}
	return &REST{store}
}
<commit_msg>Default ObjectNameFunc for all REST Stores<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package wardle

import (
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apiserver/pkg/registry/generic"
	genericregistry "k8s.io/apiserver/pkg/registry/generic/registry"
	"k8s.io/sample-apiserver/pkg/apis/wardle"
)

// rest implements a RESTStorage for API services against etcd
type REST struct {
	*genericregistry.Store
}

// NewREST returns a RESTStorage object that will work against API services.
func NewREST(scheme *runtime.Scheme, optsGetter generic.RESTOptionsGetter) *REST {
	strategy := NewStrategy(scheme)

	store := &genericregistry.Store{
		Copier:            scheme,
		NewFunc:           func() runtime.Object { return &wardle.Flunder{} },
		NewListFunc:       func() runtime.Object { return &wardle.FlunderList{} },
		PredicateFunc:     MatchFlunder,
		QualifiedResource: wardle.Resource("flunders"),

		CreateStrategy: strategy,
		UpdateStrategy: strategy,
		DeleteStrategy: strategy,
	}
	options := &generic.StoreOptions{RESTOptions: optsGetter, AttrFunc: GetAttrs}
	if err := store.CompleteWithOptions(options); err != nil {
		panic(err) // TODO: Propagate error up
	}
	return &REST{store}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Skaffold Authors

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

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

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

package build

import (
	"bytes"
	"context"
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/GoogleContainerTools/skaffold/cmd/skaffold/app/cmd/config"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/color"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/docker"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/google/go-containerregistry/pkg/name"
	homedir "github.com/mitchellh/go-homedir"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	yaml "gopkg.in/yaml.v2"
)

// ImageDetails holds the Digest and ID of an image
type ImageDetails struct {
	Digest string `yaml:"digest,omitempty"`
	ID     string `yaml:"id,omitempty"`
}

// ArtifactCache is a map of [artifact dependencies hash : ImageDetails]
type ArtifactCache map[string]ImageDetails

// Cache holds any data necessary for accessing the cache
type Cache struct {
	artifactCache ArtifactCache
	client        docker.LocalDaemon
	cacheFile     string
	useCache      bool
}

var (
	// For testing
	hashForArtifact = getHashForArtifact
	localCluster    = config.GetLocalCluster
	remoteDigest    = docker.RemoteDigest
	noCache         = &Cache{}
)

// NewCache returns the current state of the cache
func NewCache(useCache bool, cacheFile string) *Cache {
	if !useCache {
		return noCache
	}
	cf, err := resolveCacheFile(cacheFile)
	if err != nil {
		logrus.Warnf("Error resolving cache file, not using skaffold cache: %v", err)
		return noCache
	}
	cache, err := retrieveArtifactCache(cf)
	if err != nil {
		logrus.Warnf("Error retrieving artifact cache, not using skaffold cache: %v", err)
		return noCache
	}
	client, err := docker.NewAPIClient()
	if err != nil {
		logrus.Warnf("Error retrieving local daemon client, not using skaffold cache: %v", err)
		return noCache
	}
	return &Cache{
		artifactCache: cache,
		cacheFile:     cf,
		useCache:      useCache,
		client:        client,
	}
}

// resolveCacheFile makes sure that either a passed in cache file or the default cache file exists
func resolveCacheFile(cacheFile string) (string, error) {
	if cacheFile != "" {
		return cacheFile, util.VerifyOrCreateFile(cacheFile)
	}
	home, err := homedir.Dir()
	if err != nil {
		return "", errors.Wrap(err, "retrieving home directory")
	}
	defaultFile := filepath.Join(home, constants.DefaultSkaffoldDir, constants.DefaultCacheFile)
	return defaultFile, util.VerifyOrCreateFile(defaultFile)
}

func retrieveArtifactCache(cacheFile string) (ArtifactCache, error) {
	cache := ArtifactCache{}
	contents, err := ioutil.ReadFile(cacheFile)
	if err != nil {
		return nil, err
	}
	if err := yaml.Unmarshal(contents, &cache); err != nil {
		return nil, err
	}
	return cache, nil
}

// RetrieveCachedArtifacts checks to see if artifacts are cached, and returns tags for cached images, otherwise a list of images to be built
func (c *Cache) RetrieveCachedArtifacts(ctx context.Context, out io.Writer, artifacts []*latest.Artifact) ([]*latest.Artifact, []Artifact) {
	if !c.useCache {
		return artifacts, nil
	}
	color.Default.Fprintln(out, "Checking cache...")
	var needToBuild []*latest.Artifact
	var built []Artifact
	for _, a := range artifacts {
		artifact, err := c.resolveCachedArtifact(ctx, out, a)
		if err != nil {
			logrus.Debugf("error retrieving cached artifact for %s: %v\n", a.ImageName, err)
			needToBuild = append(needToBuild, a)
			continue
		}
		if artifact == nil {
			needToBuild = append(needToBuild, a)
			continue
		}
		built = append(built, *artifact)
	}
	return needToBuild, built
}

func (c *Cache) resolveCachedArtifact(ctx context.Context, out io.Writer, a *latest.Artifact) (*Artifact, error) {
	details, err := c.retrieveCachedArtifactDetails(ctx, a)
	if err != nil {
		return nil, errors.Wrap(err, "getting cached artifact details")
	}
	if details.needsRebuild {
		return nil, nil
	}
	color.Green.Fprintf(out, "Found %s locally, retagging and pushing if necessary ...\n", a.ImageName)
	if details.needsRetag {
		if err := c.client.Tag(ctx, details.prebuiltImage, details.hashTag); err != nil {
			return nil, errors.Wrap(err, "retagging image")
		}
	}
	if details.needsPush {
		if _, err := c.client.Push(ctx, out, details.hashTag); err != nil {
			return nil, errors.Wrap(err, "pushing image")
		}
	}
	color.Green.Fprintf(out, "Resolved %s, skipping rebuild.\n", details.hashTag)
	return &Artifact{
		ImageName: a.ImageName,
		Tag:       details.hashTag,
	}, nil
}

type cachedArtifactDetails struct {
	needsRebuild  bool
	needsRetag    bool
	needsPush     bool
	prebuiltImage string
	hashTag       string
}

func (c *Cache) retrieveCachedArtifactDetails(ctx context.Context, a *latest.Artifact) (*cachedArtifactDetails, error) {
	hash, err := hashForArtifact(ctx, a)
	if err != nil {
		return nil, errors.Wrapf(err, "getting hash for artifact %s", a.ImageName)
	}
	a.WorkspaceHash = hash
	localCluster, _ := localCluster()
	imageDetails, cacheHit := c.artifactCache[hash]
	if !cacheHit {
		return &cachedArtifactDetails{
			needsRebuild: true,
		}, nil
	}
	hashTag := fmt.Sprintf("%s:%s", a.ImageName, hash)

	// Check if we are using a local cluster
	var existsRemotely bool
	if !localCluster {
		// Check if tagged image exists remotely with the same digest
		existsRemotely = imageExistsRemotely(hashTag, imageDetails.Digest)
	}

	// See if this image exists in the local daemon
	if c.client.ImageExists(ctx, hashTag) {
		return &cachedArtifactDetails{
			needsPush: !existsRemotely && !localCluster,
			hashTag:   hashTag,
		}, nil
	}
	// Check for a local image with the same digest as the image we want to build
	prebuiltImage, err := c.retrievePrebuiltImage(ctx, imageDetails)
	if err != nil {
		return nil, errors.Wrapf(err, "getting prebuilt image")
	}
	if prebuiltImage == "" {
		return nil, errors.New("no tagged prebuilt image")
	}

	return &cachedArtifactDetails{
		needsRetag:    true,
		needsPush:     !localCluster,
		prebuiltImage: prebuiltImage,
		hashTag:       hashTag,
	}, nil
}

func (c *Cache) retrievePrebuiltImage(ctx context.Context, details ImageDetails) (string, error) {
	// first, search for an image with the same image ID
	img, err := c.client.FindImageByID(ctx, details.ID)
	if err != nil {
		logrus.Debugf("error getting tagged image with id %s, checking digest: %v", details.ID, err)
	}
	if err == nil && img != "" {
		return img, nil
	}
	// else, search for an image with the same digest
	img, err = c.client.FindTaggedImageByDigest(ctx, details.Digest)
	if err != nil {
		return "", errors.Wrapf(err, "getting image from digest %s", details.Digest)
	}
	if img == "" {
		return "", errors.New("no prebuilt image")
	}
	return img, nil
}

func imageExistsRemotely(image, digest string) bool {
	if digest == "" {
		logrus.Debugf("Checking if %s exists remotely, but digest is empty", image)
		return false
	}
	d, err := remoteDigest(image)
	if err != nil {
		logrus.Debugf("Checking if %s exists remotely, can't get digest: %v", image, err)
		return false
	}
	return d == digest
}

// CacheArtifacts determines the hash for each artifact, stores it in the artifact cache, and saves the cache at the end
func (c *Cache) CacheArtifacts(ctx context.Context, artifacts []*latest.Artifact, buildArtifacts []Artifact) error {
	if !c.useCache {
		return nil
	}
	tags := map[string]string{}
	for _, t := range buildArtifacts {
		tags[t.ImageName] = t.Tag
	}
	for _, a := range artifacts {
		hash, err := hashForArtifact(ctx, a)
		if err != nil {
			continue
		}
		digest, err := c.retrieveImageDigest(ctx, tags[a.ImageName])
		if err != nil {
			logrus.Debugf("error getting id for %s: %v, will try to get image id (expected with a local cluster)", tags[a.ImageName], err)
		}
		if digest == "" {
			logrus.Debugf("couldn't get image digest for %s, will try to cache just image id (expected with a local cluster)", tags[a.ImageName])
		}
		id, err := c.client.ImageID(ctx, tags[a.ImageName])
		if err != nil {
			logrus.Debugf("couldn't get image id for %s", tags[a.ImageName])
		}
		if id == "" && digest == "" {
			logrus.Debugf("both image id and digest are empty for %s, skipping caching", tags[a.ImageName])
			continue
		}
		c.artifactCache[hash] = ImageDetails{
			Digest: digest,
			ID:     id,
		}
	}
	return c.save()
}

// Retag retags newly built images in the format [imageName:workspaceHash] and pushes them if using a remote cluster
func (c *Cache) Retag(ctx context.Context, out io.Writer, artifactsToBuild []*latest.Artifact, buildArtifacts []Artifact) {
	if !c.useCache || len(artifactsToBuild) == 0 {
		return
	}
	tags := map[string]string{}
	for _, t := range buildArtifacts {
		tags[t.ImageName] = t.Tag
	}
	local, _ := localCluster()
	color.Default.Fprintln(out, "Retagging cached images...")
	for _, artifact := range artifactsToBuild {
		hashTag := fmt.Sprintf("%s:%s", artifact.ImageName, artifact.WorkspaceHash)
		// Retag the image
		if err := c.client.Tag(ctx, tags[artifact.ImageName], hashTag); err != nil {
			logrus.Warnf("error retagging %s as %s, caching for this image may not work: %v", tags[artifact.ImageName], hashTag, err)
			continue
		}
		if local {
			continue
		}
		// Push the retagged image
		if _, err := c.client.Push(ctx, out, hashTag); err != nil {
			logrus.Warnf("error pushing %s, caching for this image may not work: %v", hashTag, err)
		}
	}
}

// Check local daemon for img digest
func (c *Cache) retrieveImageDigest(ctx context.Context, img string) (string, error) {
	repoDigest, err := c.client.RepoDigest(ctx, img)
	if err != nil {
		return docker.RemoteDigest(img)
	}
	ref, err := name.NewDigest(repoDigest, name.WeakValidation)
	return ref.DigestStr(), err
}

// Save saves the artifactCache to the cacheFile
func (c *Cache) save() error {
	data, err := yaml.Marshal(c.artifactCache)
	if err != nil {
		return errors.Wrap(err, "marshalling hashes")
	}
	return ioutil.WriteFile(c.cacheFile, data, 0755)
}

func getHashForArtifact(ctx context.Context, a *latest.Artifact) (string, error) {
	deps, err := DependenciesForArtifact(ctx, a)
	if err != nil {
		return "", errors.Wrapf(err, "getting dependencies for %s", a.ImageName)
	}
	hasher := cacheHasher()
	var hashes []string
	for _, d := range deps {
		h, err := hasher(d)
		if err != nil {
			return "", errors.Wrapf(err, "getting hash for %s", d)
		}
		hashes = append(hashes, h)
	}
	// get a key for the hashes
	c := bytes.NewBuffer([]byte{})
	enc := json.NewEncoder(c)
	enc.Encode(hashes)
	return util.SHA256(c)
}

// cacheHasher takes hashes the contents and name of a file
func cacheHasher() func(string) (string, error) {
	hasher := func(p string) (string, error) {
		h := md5.New()
		fi, err := os.Lstat(p)
		if err != nil {
			return "", err
		}
		h.Write([]byte(fi.Mode().String()))
		if fi.Mode().IsRegular() {
			f, err := os.Open(p)
			if err != nil {
				return "", err
			}
			defer f.Close()
			if _, err := io.Copy(h, f); err != nil {
				return "", err
			}
		}
		return hex.EncodeToString(h.Sum(nil)), nil
	}
	return hasher
}
<commit_msg>Improve logging and colors<commit_after>/*
Copyright 2019 The Skaffold Authors

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

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

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

package build

import (
	"bytes"
	"context"
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"github.com/GoogleContainerTools/skaffold/cmd/skaffold/app/cmd/config"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/color"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/docker"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/google/go-containerregistry/pkg/name"
	homedir "github.com/mitchellh/go-homedir"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	yaml "gopkg.in/yaml.v2"
)

// ImageDetails holds the Digest and ID of an image
type ImageDetails struct {
	Digest string `yaml:"digest,omitempty"`
	ID     string `yaml:"id,omitempty"`
}

// ArtifactCache is a map of [artifact dependencies hash : ImageDetails]
type ArtifactCache map[string]ImageDetails

// Cache holds any data necessary for accessing the cache
type Cache struct {
	artifactCache ArtifactCache
	client        docker.LocalDaemon
	cacheFile     string
	useCache      bool
}

var (
	// For testing
	hashForArtifact = getHashForArtifact
	localCluster    = config.GetLocalCluster
	remoteDigest    = docker.RemoteDigest
	noCache         = &Cache{}
)

// NewCache returns the current state of the cache
func NewCache(useCache bool, cacheFile string) *Cache {
	if !useCache {
		return noCache
	}
	cf, err := resolveCacheFile(cacheFile)
	if err != nil {
		logrus.Warnf("Error resolving cache file, not using skaffold cache: %v", err)
		return noCache
	}
	cache, err := retrieveArtifactCache(cf)
	if err != nil {
		logrus.Warnf("Error retrieving artifact cache, not using skaffold cache: %v", err)
		return noCache
	}
	client, err := docker.NewAPIClient()
	if err != nil {
		logrus.Warnf("Error retrieving local daemon client, not using skaffold cache: %v", err)
		return noCache
	}
	return &Cache{
		artifactCache: cache,
		cacheFile:     cf,
		useCache:      useCache,
		client:        client,
	}
}

// resolveCacheFile makes sure that either a passed in cache file or the default cache file exists
func resolveCacheFile(cacheFile string) (string, error) {
	if cacheFile != "" {
		return cacheFile, util.VerifyOrCreateFile(cacheFile)
	}
	home, err := homedir.Dir()
	if err != nil {
		return "", errors.Wrap(err, "retrieving home directory")
	}
	defaultFile := filepath.Join(home, constants.DefaultSkaffoldDir, constants.DefaultCacheFile)
	return defaultFile, util.VerifyOrCreateFile(defaultFile)
}

func retrieveArtifactCache(cacheFile string) (ArtifactCache, error) {
	cache := ArtifactCache{}
	contents, err := ioutil.ReadFile(cacheFile)
	if err != nil {
		return nil, err
	}
	if err := yaml.Unmarshal(contents, &cache); err != nil {
		return nil, err
	}
	return cache, nil
}

// RetrieveCachedArtifacts checks to see if artifacts are cached, and returns tags for cached images, otherwise a list of images to be built
func (c *Cache) RetrieveCachedArtifacts(ctx context.Context, out io.Writer, artifacts []*latest.Artifact) ([]*latest.Artifact, []Artifact) {
	if !c.useCache {
		return artifacts, nil
	}
	start := time.Now()
	color.Default.Fprintln(out, "Checking cache...")
	var needToBuild []*latest.Artifact
	var built []Artifact
	for _, a := range artifacts {
		artifact, err := c.resolveCachedArtifact(ctx, out, a)
		if err != nil {
			logrus.Debugf("error retrieving cached artifact for %s: %v\n", a.ImageName, err)
			needToBuild = append(needToBuild, a)
			continue
		}
		if artifact == nil {
			needToBuild = append(needToBuild, a)
			continue
		}
		built = append(built, *artifact)
	}
	color.Default.Fprintln(out, "Cache check complete in", time.Since(start))
	return needToBuild, built
}

func (c *Cache) resolveCachedArtifact(ctx context.Context, out io.Writer, a *latest.Artifact) (*Artifact, error) {
	details, err := c.retrieveCachedArtifactDetails(ctx, a)
	if err != nil {
		return nil, errors.Wrap(err, "getting cached artifact details")
	}
	if details.needsRebuild {
		return nil, nil
	}
	color.Default.Fprintf(out, "Found %s in cache, resolving...\n", a.ImageName)
	if details.needsRetag {
		color.Green.Fprintf(out, "Retagging image...\n")
		if err := c.client.Tag(ctx, details.prebuiltImage, details.hashTag); err != nil {
			return nil, errors.Wrap(err, "retagging image")
		}
	}
	if details.needsPush {
		color.Green.Fprintf(out, "Pushing %s...\n", a.ImageName)
		if _, err := c.client.Push(ctx, out, details.hashTag); err != nil {
			return nil, errors.Wrap(err, "pushing image")
		}
	}
	color.Default.Fprintf(out, "Resolved %s, skipping rebuild.\n", details.hashTag)
	return &Artifact{
		ImageName: a.ImageName,
		Tag:       details.hashTag,
	}, nil
}

type cachedArtifactDetails struct {
	needsRebuild  bool
	needsRetag    bool
	needsPush     bool
	prebuiltImage string
	hashTag       string
}

func (c *Cache) retrieveCachedArtifactDetails(ctx context.Context, a *latest.Artifact) (*cachedArtifactDetails, error) {
	hash, err := hashForArtifact(ctx, a)
	if err != nil {
		return nil, errors.Wrapf(err, "getting hash for artifact %s", a.ImageName)
	}
	a.WorkspaceHash = hash
	localCluster, _ := localCluster()
	imageDetails, cacheHit := c.artifactCache[hash]
	if !cacheHit {
		return &cachedArtifactDetails{
			needsRebuild: true,
		}, nil
	}
	hashTag := fmt.Sprintf("%s:%s", a.ImageName, hash)

	// Check if we are using a local cluster
	var existsRemotely bool
	if !localCluster {
		// Check if tagged image exists remotely with the same digest
		existsRemotely = imageExistsRemotely(hashTag, imageDetails.Digest)
	}

	// See if this image exists in the local daemon
	if c.client.ImageExists(ctx, hashTag) {
		return &cachedArtifactDetails{
			needsPush: !existsRemotely && !localCluster,
			hashTag:   hashTag,
		}, nil
	}
	// Check for a local image with the same digest as the image we want to build
	prebuiltImage, err := c.retrievePrebuiltImage(ctx, imageDetails)
	if err != nil {
		return nil, errors.Wrapf(err, "getting prebuilt image")
	}
	if prebuiltImage == "" {
		return nil, errors.New("no tagged prebuilt image")
	}

	return &cachedArtifactDetails{
		needsRetag:    true,
		needsPush:     !localCluster,
		prebuiltImage: prebuiltImage,
		hashTag:       hashTag,
	}, nil
}

func (c *Cache) retrievePrebuiltImage(ctx context.Context, details ImageDetails) (string, error) {
	// first, search for an image with the same image ID
	img, err := c.client.FindImageByID(ctx, details.ID)
	if err != nil {
		logrus.Debugf("error getting tagged image with id %s, checking digest: %v", details.ID, err)
	}
	if err == nil && img != "" {
		return img, nil
	}
	// else, search for an image with the same digest
	img, err = c.client.FindTaggedImageByDigest(ctx, details.Digest)
	if err != nil {
		return "", errors.Wrapf(err, "getting image from digest %s", details.Digest)
	}
	if img == "" {
		return "", errors.New("no prebuilt image")
	}
	return img, nil
}

func imageExistsRemotely(image, digest string) bool {
	if digest == "" {
		logrus.Debugf("Checking if %s exists remotely, but digest is empty", image)
		return false
	}
	d, err := remoteDigest(image)
	if err != nil {
		logrus.Debugf("Checking if %s exists remotely, can't get digest: %v", image, err)
		return false
	}
	return d == digest
}

// CacheArtifacts determines the hash for each artifact, stores it in the artifact cache, and saves the cache at the end
func (c *Cache) CacheArtifacts(ctx context.Context, artifacts []*latest.Artifact, buildArtifacts []Artifact) error {
	if !c.useCache {
		return nil
	}
	tags := map[string]string{}
	for _, t := range buildArtifacts {
		tags[t.ImageName] = t.Tag
	}
	for _, a := range artifacts {
		hash, err := hashForArtifact(ctx, a)
		if err != nil {
			continue
		}
		digest, err := c.retrieveImageDigest(ctx, tags[a.ImageName])
		if err != nil {
			logrus.Debugf("error getting id for %s: %v, will try to get image id (expected with a local cluster)", tags[a.ImageName], err)
		}
		if digest == "" {
			logrus.Debugf("couldn't get image digest for %s, will try to cache just image id (expected with a local cluster)", tags[a.ImageName])
		}
		id, err := c.client.ImageID(ctx, tags[a.ImageName])
		if err != nil {
			logrus.Debugf("couldn't get image id for %s", tags[a.ImageName])
		}
		if id == "" && digest == "" {
			logrus.Debugf("both image id and digest are empty for %s, skipping caching", tags[a.ImageName])
			continue
		}
		c.artifactCache[hash] = ImageDetails{
			Digest: digest,
			ID:     id,
		}
	}
	return c.save()
}

// Retag retags newly built images in the format [imageName:workspaceHash] and pushes them if using a remote cluster
func (c *Cache) Retag(ctx context.Context, out io.Writer, artifactsToBuild []*latest.Artifact, buildArtifacts []Artifact) {
	if !c.useCache || len(artifactsToBuild) == 0 {
		return
	}
	tags := map[string]string{}
	for _, t := range buildArtifacts {
		tags[t.ImageName] = t.Tag
	}
	local, _ := localCluster()
	color.Default.Fprintln(out, "Retagging cached images...")
	for _, artifact := range artifactsToBuild {
		hashTag := fmt.Sprintf("%s:%s", artifact.ImageName, artifact.WorkspaceHash)
		// Retag the image
		if err := c.client.Tag(ctx, tags[artifact.ImageName], hashTag); err != nil {
			logrus.Warnf("error retagging %s as %s, caching for this image may not work: %v", tags[artifact.ImageName], hashTag, err)
			continue
		}
		if local {
			continue
		}
		// Push the retagged image
		if _, err := c.client.Push(ctx, out, hashTag); err != nil {
			logrus.Warnf("error pushing %s, caching for this image may not work: %v", hashTag, err)
		}
	}
}

// Check local daemon for img digest
func (c *Cache) retrieveImageDigest(ctx context.Context, img string) (string, error) {
	repoDigest, err := c.client.RepoDigest(ctx, img)
	if err != nil {
		return docker.RemoteDigest(img)
	}
	ref, err := name.NewDigest(repoDigest, name.WeakValidation)
	return ref.DigestStr(), err
}

// Save saves the artifactCache to the cacheFile
func (c *Cache) save() error {
	data, err := yaml.Marshal(c.artifactCache)
	if err != nil {
		return errors.Wrap(err, "marshalling hashes")
	}
	return ioutil.WriteFile(c.cacheFile, data, 0755)
}

func getHashForArtifact(ctx context.Context, a *latest.Artifact) (string, error) {
	deps, err := DependenciesForArtifact(ctx, a)
	if err != nil {
		return "", errors.Wrapf(err, "getting dependencies for %s", a.ImageName)
	}
	hasher := cacheHasher()
	var hashes []string
	for _, d := range deps {
		h, err := hasher(d)
		if err != nil {
			return "", errors.Wrapf(err, "getting hash for %s", d)
		}
		hashes = append(hashes, h)
	}
	// get a key for the hashes
	c := bytes.NewBuffer([]byte{})
	enc := json.NewEncoder(c)
	enc.Encode(hashes)
	return util.SHA256(c)
}

// cacheHasher takes hashes the contents and name of a file
func cacheHasher() func(string) (string, error) {
	hasher := func(p string) (string, error) {
		h := md5.New()
		fi, err := os.Lstat(p)
		if err != nil {
			return "", err
		}
		h.Write([]byte(fi.Mode().String()))
		if fi.Mode().IsRegular() {
			f, err := os.Open(p)
			if err != nil {
				return "", err
			}
			defer f.Close()
			if _, err := io.Copy(h, f); err != nil {
				return "", err
			}
		}
		return hex.EncodeToString(h.Sum(nil)), nil
	}
	return hasher
}
<|endoftext|>


<commit_before>package mssql

import (
	"database/sql"
	"database/sql/driver"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	"strings"
	"sync"
	"time"
)

var partnersCache partners = partners{mu: sync.RWMutex{}, v: make(map[string]string)}

func init() {
	sql.Register("mssql", &MssqlDriver{})
}

type MssqlDriver struct {
	log *log.Logger
}

func (d *MssqlDriver) SetLogger(logger *log.Logger) {
	d.log = logger
}

func CheckBadConn(err error) error {
	if err == io.EOF {
		return driver.ErrBadConn
	}
	neterr, ok := err.(net.Error)
	if !ok || (!neterr.Timeout() && neterr.Temporary()) {
		return err
	}
	return driver.ErrBadConn
}

type MssqlConn struct {
	sess *tdsSession
}

func (c *MssqlConn) Commit() error {
	headers := []headerStruct{
		{hdrtype: dataStmHdrTransDescr,
			data: transDescrHdr{c.sess.tranid, 1}.pack()},
	}
	if err := sendCommitXact(c.sess.buf, headers, "", 0, 0, ""); err != nil {
		return err
	}

	tokchan := make(chan tokenStruct, 5)
	go processResponse(c.sess, tokchan)
	for tok := range tokchan {
		switch token := tok.(type) {
		case error:
			return token
		}
	}
	return nil
}

func (c *MssqlConn) Rollback() error {
	headers := []headerStruct{
		{hdrtype: dataStmHdrTransDescr,
			data: transDescrHdr{c.sess.tranid, 1}.pack()},
	}
	if err := sendRollbackXact(c.sess.buf, headers, "", 0, 0, ""); err != nil {
		return err
	}

	tokchan := make(chan tokenStruct, 5)
	go processResponse(c.sess, tokchan)
	for tok := range tokchan {
		switch token := tok.(type) {
		case error:
			return token
		}
	}
	return nil
}

func (c *MssqlConn) Begin() (driver.Tx, error) {
	headers := []headerStruct{
		{hdrtype: dataStmHdrTransDescr,
			data: transDescrHdr{0, 1}.pack()},
	}
	if err := sendBeginXact(c.sess.buf, headers, 0, ""); err != nil {
		return nil, CheckBadConn(err)
	}
	tokchan := make(chan tokenStruct, 5)
	go processResponse(c.sess, tokchan)
	for tok := range tokchan {
		switch token := tok.(type) {
		case error:
			if c.sess.tranid != 0 {
				return nil, token
			}
			return nil, CheckBadConn(token)
		}
	}
	// successful BEGINXACT request will return sess.tranid
	// for started transaction
	return c, nil
}

func parseConnectionString(dsn string) (res map[string]string) {
	res = map[string]string{}
	parts := strings.Split(dsn, ";")
	for _, part := range parts {
		if len(part) == 0 {
			continue
		}
		lst := strings.SplitN(part, "=", 2)
		name := strings.TrimSpace(strings.ToLower(lst[0]))
		if len(name) == 0 {
			continue
		}
		var value string = ""
		if len(lst) > 1 {
			value = strings.TrimSpace(lst[1])
		}
		res[name] = value
	}
	return res
}

func (d *MssqlDriver) Open(dsn string) (driver.Conn, error) {
	params := parseConnectionString(dsn)

	conn, err := open(dsn, params)
	if err != nil {
		return nil, err
	}

	conn.sess.log = (*Logger)(d.log)
	return conn, nil
}

func open(dsn string, params map[string]string) (*MssqlConn, error) {
	buf, err := connect(params)
	if err != nil {
		partner := partnersCache.Get(dsn)
		if partner == "" {
			partner = params["failoverpartner"]
			// remove the failoverpartner entry to prevent infinite recursion
			delete(params, "failoverpartner")
			if port, ok := params["failoverport"]; ok {
				params["port"] = port
			}
		}

		if partner != "" {
			params["server"] = partner
			return open(dsn, params)
		}

		return nil, err
	}

	if partner := buf.partner; partner != "" {
		// append an instance so the port will be ignored when this value is used;
		// tds does not provide the port number.
		if !strings.Contains(partner, `\`) {
			partner += `\.`
		}
		partnersCache.Set(dsn, partner)
	}

	return &MssqlConn{buf}, nil
}

func (c *MssqlConn) Close() error {
	return c.sess.buf.transport.Close()
}

type MssqlStmt struct {
	c          *MssqlConn
	query      string
	paramCount int
}

func (c *MssqlConn) Prepare(query string) (driver.Stmt, error) {
	q, paramCount := parseParams(query)
	return &MssqlStmt{c, q, paramCount}, nil
}

func (s *MssqlStmt) Close() error {
	return nil
}

func (s *MssqlStmt) NumInput() int {
	return s.paramCount
}

func (s *MssqlStmt) sendQuery(args []driver.Value) (err error) {
	headers := []headerStruct{
		{hdrtype: dataStmHdrTransDescr,
			data: transDescrHdr{s.c.sess.tranid, 1}.pack()},
	}
	if len(args) != s.paramCount {
		return errors.New(fmt.Sprintf("sql: expected %d parameters, got %d", s.paramCount, len(args)))
	}
	if s.c.sess.logFlags&logSQL != 0 {
		s.c.sess.log.Println(s.query)
	}
	if s.c.sess.logFlags&logParams != 0 && len(args) > 0 {
		for i := 0; i < len(args); i++ {
			s.c.sess.log.Printf("\t@p%d\t%v\n", i+1, args[i])
		}

	}
	if len(args) == 0 {
		if err = sendSqlBatch72(s.c.sess.buf, s.query, headers); err != nil {
			if s.c.sess.tranid != 0 {
				return err
			}
			return CheckBadConn(err)
		}
	} else {
		params := make([]Param, len(args)+2)
		decls := make([]string, len(args))
		params[0], err = s.makeParam(s.query)
		if err != nil {
			return
		}
		for i, val := range args {
			params[i+2], err = s.makeParam(val)
			if err != nil {
				return
			}
			name := fmt.Sprintf("@p%d", i+1)
			params[i+2].Name = name
			decls[i] = fmt.Sprintf("%s %s", name, makeDecl(params[i+2].ti))
		}
		params[1], err = s.makeParam(strings.Join(decls, ","))
		if err != nil {
			return
		}
		if err = sendRpc(s.c.sess.buf, headers, Sp_ExecuteSql, 0, params); err != nil {
			if s.c.sess.tranid != 0 {
				return err
			}
			return CheckBadConn(err)
		}
	}
	return
}

func (s *MssqlStmt) Query(args []driver.Value) (res driver.Rows, err error) {
	if err = s.sendQuery(args); err != nil {
		return
	}
	tokchan := make(chan tokenStruct, 5)
	go processResponse(s.c.sess, tokchan)
	// process metadata
	var cols []string
loop:
	for tok := range tokchan {
		switch token := tok.(type) {
		case doneStruct:
			break loop
		case []columnStruct:
			cols = make([]string, len(token))
			for i, col := range token {
				cols[i] = col.ColName
			}
			break loop
		case error:
			if s.c.sess.tranid != 0 {
				return nil, token
			}
			return nil, CheckBadConn(token)
		}
	}
	return &MssqlRows{sess: s.c.sess, tokchan: tokchan, cols: cols}, nil
}

func (s *MssqlStmt) Exec(args []driver.Value) (res driver.Result, err error) {
	if err = s.sendQuery(args); err != nil {
		return
	}
	tokchan := make(chan tokenStruct, 5)
	go processResponse(s.c.sess, tokchan)
	var rowCount int64
	for token := range tokchan {
		switch token := token.(type) {
		case doneInProcStruct:
			if token.Status&doneCount != 0 {
				rowCount = int64(token.RowCount)
			}
		case doneStruct:
			if token.Status&doneCount != 0 {
				rowCount = int64(token.RowCount)
			}
		case error:
			if s.c.sess.logFlags&logErrors != 0 {
				s.c.sess.log.Println("got error:", token)
			}
			if s.c.sess.tranid != 0 {
				return nil, token
			}
			return nil, CheckBadConn(token)
		}
	}
	return &MssqlResult{s.c, rowCount}, nil
}

type MssqlRows struct {
	sess    *tdsSession
	cols    []string
	tokchan chan tokenStruct
}

func (rc *MssqlRows) Close() error {
	for _ = range rc.tokchan {
	}
	rc.tokchan = nil
	return nil
}

func (rc *MssqlRows) Columns() (res []string) {
	return rc.cols
}

func (rc *MssqlRows) Next(dest []driver.Value) (err error) {
	for tok := range rc.tokchan {
		switch tokdata := tok.(type) {
		case []columnStruct:
			return streamErrorf("Unexpected token COLMETADATA")
		case []interface{}:
			for i := range dest {
				dest[i] = tokdata[i]
			}
			return nil
		case error:
			return tokdata
		}
	}
	return io.EOF
}

func (s *MssqlStmt) makeParam(val driver.Value) (res Param, err error) {
	if val == nil {
		res.ti.TypeId = typeNVarChar
		res.buffer = nil
		res.ti.Size = 2
		return
	}
	switch val := val.(type) {
	case int64:
		res.ti.TypeId = typeIntN
		res.buffer = make([]byte, 8)
		res.ti.Size = 8
		binary.LittleEndian.PutUint64(res.buffer, uint64(val))
	case float64:
		res.ti.TypeId = typeFltN
		res.ti.Size = 8
		res.buffer = make([]byte, 8)
		binary.LittleEndian.PutUint64(res.buffer, math.Float64bits(val))
	case []byte:
		res.ti.TypeId = typeBigVarBin
		res.ti.Size = len(val)
		res.buffer = val
	case string:
		res.ti.TypeId = typeNVarChar
		res.buffer = str2ucs2(val)
		res.ti.Size = len(res.buffer)
	case bool:
		res.ti.TypeId = typeBitN
		res.ti.Size = 1
		res.buffer = make([]byte, 1)
		if val {
			res.buffer[0] = 1
		}
	case time.Time:
		if s.c.sess.loginAck.TDSVersion >= verTDS73 {
			res.ti.TypeId = typeDateTimeOffsetN
			res.ti.Scale = 7
			res.ti.Size = 10
			buf := make([]byte, 10)
			res.buffer = buf
			days, ns := dateTime2(val)
			ns /= 100
			buf[0] = byte(ns)
			buf[1] = byte(ns >> 8)
			buf[2] = byte(ns >> 16)
			buf[3] = byte(ns >> 24)
			buf[4] = byte(ns >> 32)
			buf[5] = byte(days)
			buf[6] = byte(days >> 8)
			buf[7] = byte(days >> 16)
			_, offset := val.Zone()
			offset /= 60
			buf[8] = byte(offset)
			buf[9] = byte(offset >> 8)
		} else {
			res.ti.TypeId = typeDateTimeN
			res.ti.Size = 8
			res.buffer = make([]byte, 8)
			ref := time.Date(1900, 1, 1, 0, 0, 0, 0, time.UTC)
			dur := val.Sub(ref)
			days := dur / (24 * time.Hour)
			tm := (300 * (dur % (24 * time.Hour))) / time.Second
			binary.LittleEndian.PutUint32(res.buffer[0:4], uint32(days))
			binary.LittleEndian.PutUint32(res.buffer[4:8], uint32(tm))
		}
	default:
		err = fmt.Errorf("mssql: unknown type for %T", val)
		return
	}
	return
}

type MssqlResult struct {
	c            *MssqlConn
	rowsAffected int64
}

func (r *MssqlResult) RowsAffected() (int64, error) {
	return r.rowsAffected, nil
}

func (r *MssqlResult) LastInsertId() (int64, error) {
	s, err := r.c.Prepare("select cast(@@identity as bigint)")
	if err != nil {
		return 0, err
	}
	defer s.Close()
	rows, err := s.Query(nil)
	if err != nil {
		return 0, err
	}
	defer rows.Close()
	dest := make([]driver.Value, 1)
	err = rows.Next(dest)
	if err != nil {
		return 0, err
	}
	if dest[0] == nil {
		return -1, errors.New("There is no generated identity value")
	}
	lastInsertId := dest[0].(int64)
	return lastInsertId, nil
}

type partners struct {
	mu sync.RWMutex
	v  map[string]string
}

func (p partners) Set(key, value string) error {
	p.mu.Lock()
	defer p.mu.Unlock()
	if _, ok := p.v[key]; ok {
		return errors.New("key already exists")
	}

	p.v[key] = value
	return nil
}

func (p partners) Get(key string) (value string) {
	p.mu.RLock()
	value = p.v[key]
	p.mu.RUnlock()
	return
}
<commit_msg>renamed open to openConnection<commit_after>package mssql

import (
	"database/sql"
	"database/sql/driver"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	"strings"
	"sync"
	"time"
)

var partnersCache partners = partners{mu: sync.RWMutex{}, v: make(map[string]string)}

func init() {
	sql.Register("mssql", &MssqlDriver{})
}

type MssqlDriver struct {
	log *log.Logger
}

func (d *MssqlDriver) SetLogger(logger *log.Logger) {
	d.log = logger
}

func CheckBadConn(err error) error {
	if err == io.EOF {
		return driver.ErrBadConn
	}
	neterr, ok := err.(net.Error)
	if !ok || (!neterr.Timeout() && neterr.Temporary()) {
		return err
	}
	return driver.ErrBadConn
}

type MssqlConn struct {
	sess *tdsSession
}

func (c *MssqlConn) Commit() error {
	headers := []headerStruct{
		{hdrtype: dataStmHdrTransDescr,
			data: transDescrHdr{c.sess.tranid, 1}.pack()},
	}
	if err := sendCommitXact(c.sess.buf, headers, "", 0, 0, ""); err != nil {
		return err
	}

	tokchan := make(chan tokenStruct, 5)
	go processResponse(c.sess, tokchan)
	for tok := range tokchan {
		switch token := tok.(type) {
		case error:
			return token
		}
	}
	return nil
}

func (c *MssqlConn) Rollback() error {
	headers := []headerStruct{
		{hdrtype: dataStmHdrTransDescr,
			data: transDescrHdr{c.sess.tranid, 1}.pack()},
	}
	if err := sendRollbackXact(c.sess.buf, headers, "", 0, 0, ""); err != nil {
		return err
	}

	tokchan := make(chan tokenStruct, 5)
	go processResponse(c.sess, tokchan)
	for tok := range tokchan {
		switch token := tok.(type) {
		case error:
			return token
		}
	}
	return nil
}

func (c *MssqlConn) Begin() (driver.Tx, error) {
	headers := []headerStruct{
		{hdrtype: dataStmHdrTransDescr,
			data: transDescrHdr{0, 1}.pack()},
	}
	if err := sendBeginXact(c.sess.buf, headers, 0, ""); err != nil {
		return nil, CheckBadConn(err)
	}
	tokchan := make(chan tokenStruct, 5)
	go processResponse(c.sess, tokchan)
	for tok := range tokchan {
		switch token := tok.(type) {
		case error:
			if c.sess.tranid != 0 {
				return nil, token
			}
			return nil, CheckBadConn(token)
		}
	}
	// successful BEGINXACT request will return sess.tranid
	// for started transaction
	return c, nil
}

func parseConnectionString(dsn string) (res map[string]string) {
	res = map[string]string{}
	parts := strings.Split(dsn, ";")
	for _, part := range parts {
		if len(part) == 0 {
			continue
		}
		lst := strings.SplitN(part, "=", 2)
		name := strings.TrimSpace(strings.ToLower(lst[0]))
		if len(name) == 0 {
			continue
		}
		var value string = ""
		if len(lst) > 1 {
			value = strings.TrimSpace(lst[1])
		}
		res[name] = value
	}
	return res
}

func (d *MssqlDriver) Open(dsn string) (driver.Conn, error) {
	params := parseConnectionString(dsn)

	conn, err := openConnection(dsn, params)
	if err != nil {
		return nil, err
	}

	conn.sess.log = (*Logger)(d.log)
	return conn, nil
}

func openConnection(dsn string, params map[string]string) (*MssqlConn, error) {
	buf, err := connect(params)
	if err != nil {
		partner := partnersCache.Get(dsn)
		if partner == "" {
			partner = params["failoverpartner"]
			// remove the failoverpartner entry to prevent infinite recursion
			delete(params, "failoverpartner")
			if port, ok := params["failoverport"]; ok {
				params["port"] = port
			}
		}

		if partner != "" {
			params["server"] = partner
			return openConnection(dsn, params)
		}

		return nil, err
	}

	if partner := buf.partner; partner != "" {
		// append an instance so the port will be ignored when this value is used;
		// tds does not provide the port number.
		if !strings.Contains(partner, `\`) {
			partner += `\.`
		}
		partnersCache.Set(dsn, partner)
	}

	return &MssqlConn{buf}, nil
}

func (c *MssqlConn) Close() error {
	return c.sess.buf.transport.Close()
}

type MssqlStmt struct {
	c          *MssqlConn
	query      string
	paramCount int
}

func (c *MssqlConn) Prepare(query string) (driver.Stmt, error) {
	q, paramCount := parseParams(query)
	return &MssqlStmt{c, q, paramCount}, nil
}

func (s *MssqlStmt) Close() error {
	return nil
}

func (s *MssqlStmt) NumInput() int {
	return s.paramCount
}

func (s *MssqlStmt) sendQuery(args []driver.Value) (err error) {
	headers := []headerStruct{
		{hdrtype: dataStmHdrTransDescr,
			data: transDescrHdr{s.c.sess.tranid, 1}.pack()},
	}
	if len(args) != s.paramCount {
		return errors.New(fmt.Sprintf("sql: expected %d parameters, got %d", s.paramCount, len(args)))
	}
	if s.c.sess.logFlags&logSQL != 0 {
		s.c.sess.log.Println(s.query)
	}
	if s.c.sess.logFlags&logParams != 0 && len(args) > 0 {
		for i := 0; i < len(args); i++ {
			s.c.sess.log.Printf("\t@p%d\t%v\n", i+1, args[i])
		}

	}
	if len(args) == 0 {
		if err = sendSqlBatch72(s.c.sess.buf, s.query, headers); err != nil {
			if s.c.sess.tranid != 0 {
				return err
			}
			return CheckBadConn(err)
		}
	} else {
		params := make([]Param, len(args)+2)
		decls := make([]string, len(args))
		params[0], err = s.makeParam(s.query)
		if err != nil {
			return
		}
		for i, val := range args {
			params[i+2], err = s.makeParam(val)
			if err != nil {
				return
			}
			name := fmt.Sprintf("@p%d", i+1)
			params[i+2].Name = name
			decls[i] = fmt.Sprintf("%s %s", name, makeDecl(params[i+2].ti))
		}
		params[1], err = s.makeParam(strings.Join(decls, ","))
		if err != nil {
			return
		}
		if err = sendRpc(s.c.sess.buf, headers, Sp_ExecuteSql, 0, params); err != nil {
			if s.c.sess.tranid != 0 {
				return err
			}
			return CheckBadConn(err)
		}
	}
	return
}

func (s *MssqlStmt) Query(args []driver.Value) (res driver.Rows, err error) {
	if err = s.sendQuery(args); err != nil {
		return
	}
	tokchan := make(chan tokenStruct, 5)
	go processResponse(s.c.sess, tokchan)
	// process metadata
	var cols []string
loop:
	for tok := range tokchan {
		switch token := tok.(type) {
		case doneStruct:
			break loop
		case []columnStruct:
			cols = make([]string, len(token))
			for i, col := range token {
				cols[i] = col.ColName
			}
			break loop
		case error:
			if s.c.sess.tranid != 0 {
				return nil, token
			}
			return nil, CheckBadConn(token)
		}
	}
	return &MssqlRows{sess: s.c.sess, tokchan: tokchan, cols: cols}, nil
}

func (s *MssqlStmt) Exec(args []driver.Value) (res driver.Result, err error) {
	if err = s.sendQuery(args); err != nil {
		return
	}
	tokchan := make(chan tokenStruct, 5)
	go processResponse(s.c.sess, tokchan)
	var rowCount int64
	for token := range tokchan {
		switch token := token.(type) {
		case doneInProcStruct:
			if token.Status&doneCount != 0 {
				rowCount = int64(token.RowCount)
			}
		case doneStruct:
			if token.Status&doneCount != 0 {
				rowCount = int64(token.RowCount)
			}
		case error:
			if s.c.sess.logFlags&logErrors != 0 {
				s.c.sess.log.Println("got error:", token)
			}
			if s.c.sess.tranid != 0 {
				return nil, token
			}
			return nil, CheckBadConn(token)
		}
	}
	return &MssqlResult{s.c, rowCount}, nil
}

type MssqlRows struct {
	sess    *tdsSession
	cols    []string
	tokchan chan tokenStruct
}

func (rc *MssqlRows) Close() error {
	for _ = range rc.tokchan {
	}
	rc.tokchan = nil
	return nil
}

func (rc *MssqlRows) Columns() (res []string) {
	return rc.cols
}

func (rc *MssqlRows) Next(dest []driver.Value) (err error) {
	for tok := range rc.tokchan {
		switch tokdata := tok.(type) {
		case []columnStruct:
			return streamErrorf("Unexpected token COLMETADATA")
		case []interface{}:
			for i := range dest {
				dest[i] = tokdata[i]
			}
			return nil
		case error:
			return tokdata
		}
	}
	return io.EOF
}

func (s *MssqlStmt) makeParam(val driver.Value) (res Param, err error) {
	if val == nil {
		res.ti.TypeId = typeNVarChar
		res.buffer = nil
		res.ti.Size = 2
		return
	}
	switch val := val.(type) {
	case int64:
		res.ti.TypeId = typeIntN
		res.buffer = make([]byte, 8)
		res.ti.Size = 8
		binary.LittleEndian.PutUint64(res.buffer, uint64(val))
	case float64:
		res.ti.TypeId = typeFltN
		res.ti.Size = 8
		res.buffer = make([]byte, 8)
		binary.LittleEndian.PutUint64(res.buffer, math.Float64bits(val))
	case []byte:
		res.ti.TypeId = typeBigVarBin
		res.ti.Size = len(val)
		res.buffer = val
	case string:
		res.ti.TypeId = typeNVarChar
		res.buffer = str2ucs2(val)
		res.ti.Size = len(res.buffer)
	case bool:
		res.ti.TypeId = typeBitN
		res.ti.Size = 1
		res.buffer = make([]byte, 1)
		if val {
			res.buffer[0] = 1
		}
	case time.Time:
		if s.c.sess.loginAck.TDSVersion >= verTDS73 {
			res.ti.TypeId = typeDateTimeOffsetN
			res.ti.Scale = 7
			res.ti.Size = 10
			buf := make([]byte, 10)
			res.buffer = buf
			days, ns := dateTime2(val)
			ns /= 100
			buf[0] = byte(ns)
			buf[1] = byte(ns >> 8)
			buf[2] = byte(ns >> 16)
			buf[3] = byte(ns >> 24)
			buf[4] = byte(ns >> 32)
			buf[5] = byte(days)
			buf[6] = byte(days >> 8)
			buf[7] = byte(days >> 16)
			_, offset := val.Zone()
			offset /= 60
			buf[8] = byte(offset)
			buf[9] = byte(offset >> 8)
		} else {
			res.ti.TypeId = typeDateTimeN
			res.ti.Size = 8
			res.buffer = make([]byte, 8)
			ref := time.Date(1900, 1, 1, 0, 0, 0, 0, time.UTC)
			dur := val.Sub(ref)
			days := dur / (24 * time.Hour)
			tm := (300 * (dur % (24 * time.Hour))) / time.Second
			binary.LittleEndian.PutUint32(res.buffer[0:4], uint32(days))
			binary.LittleEndian.PutUint32(res.buffer[4:8], uint32(tm))
		}
	default:
		err = fmt.Errorf("mssql: unknown type for %T", val)
		return
	}
	return
}

type MssqlResult struct {
	c            *MssqlConn
	rowsAffected int64
}

func (r *MssqlResult) RowsAffected() (int64, error) {
	return r.rowsAffected, nil
}

func (r *MssqlResult) LastInsertId() (int64, error) {
	s, err := r.c.Prepare("select cast(@@identity as bigint)")
	if err != nil {
		return 0, err
	}
	defer s.Close()
	rows, err := s.Query(nil)
	if err != nil {
		return 0, err
	}
	defer rows.Close()
	dest := make([]driver.Value, 1)
	err = rows.Next(dest)
	if err != nil {
		return 0, err
	}
	if dest[0] == nil {
		return -1, errors.New("There is no generated identity value")
	}
	lastInsertId := dest[0].(int64)
	return lastInsertId, nil
}

type partners struct {
	mu sync.RWMutex
	v  map[string]string
}

func (p partners) Set(key, value string) error {
	p.mu.Lock()
	defer p.mu.Unlock()
	if _, ok := p.v[key]; ok {
		return errors.New("key already exists")
	}

	p.v[key] = value
	return nil
}

func (p partners) Get(key string) (value string) {
	p.mu.RLock()
	value = p.v[key]
	p.mu.RUnlock()
	return
}
<|endoftext|>