<commit_before>// Copyright 2022 Google LLC 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 options

import (
	"log"
	"strings"
)

const bareBaseFlagsWarning = `WARNING!
-----------------------------------------------------------------
Both --base-import-paths and --bare were set.

--base-import-paths will take precedence and ignore --bare flag.

In a future release this will be an error.
-----------------------------------------------------------------
`

const localFlagsWarning = `WARNING!
-----------------------------------------------------------------
The --local flag is set and KO_DOCKER_REPO is set to ko.local

You can choose either one to build a local image.

The --local flag might be deprecated in the future.
-----------------------------------------------------------------
`

func Validate(po *PublishOptions, bo *BuildOptions) error {
	if po.Bare && po.BaseImportPaths {
		log.Print(bareBaseFlagsWarning)
		// TODO: return error when we decided to make this an error, for now it is a warning
	}

	if po.Local && strings.Contains(po.DockerRepo, "ko.local") {
		log.Print(localFlagsWarning)
	}

	return nil
}
<commit_msg>check if have all and other platforms set in the --platform flag (#630)<commit_after>// Copyright 2022 Google LLC 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 options

import (
	"errors"
	"log"
	"strings"
)

const bareBaseFlagsWarning = `WARNING!
-----------------------------------------------------------------
Both --base-import-paths and --bare were set.

--base-import-paths will take precedence and ignore --bare flag.

In a future release this will be an error.
-----------------------------------------------------------------
`

const localFlagsWarning = `WARNING!
-----------------------------------------------------------------
The --local flag is set and KO_DOCKER_REPO is set to ko.local

You can choose either one to build a local image.

The --local flag might be deprecated in the future.
-----------------------------------------------------------------
`

func Validate(po *PublishOptions, bo *BuildOptions) error {
	if po.Bare && po.BaseImportPaths {
		log.Print(bareBaseFlagsWarning)
		// TODO: return error when we decided to make this an error, for now it is a warning
	}

	if po.Local && strings.Contains(po.DockerRepo, "ko.local") {
		log.Print(localFlagsWarning)
	}

	if len(bo.Platforms) > 1 {
		for _, platform := range bo.Platforms {
			if platform == "all" {
				return errors.New("all or specific platforms should be used")
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative Authors

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

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

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

package handler

import (
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"
)

func TestTimeoutWriterAllowsForAdditionalWrites(t *testing.T) {
	recorder := httptest.NewRecorder()
	handler := &timeoutWriter{w: recorder}
	handler.WriteHeader(http.StatusOK)
	handler.timeoutAndWriteError("error")
	if _, err := io.WriteString(handler, "test"); err != nil {
		t.Fatalf("handler.Write() = %v, want no error", err)
	}

	if got, want := recorder.Code, http.StatusOK; got != want {
		t.Errorf("recorder.Status = %d, want %d", got, want)
	}
	if got, want := recorder.Body.String(), "test"; got != want {
		t.Errorf("recorder.Body = %s, want %s", got, want)
	}
}

func TestTimeoutWriterDoesntFlushAfterTimeout(t *testing.T) {
	recorder := httptest.NewRecorder()
	handler := &timeoutWriter{w: recorder}
	handler.timeoutAndWriteError("error")
	handler.Flush()

	if got, want := recorder.Flushed, false; got != want {
		t.Errorf("recorder.Flushed = %t, want %t", got, want)
	}
}

func TestTimeoutWriterFlushesAfterTimeout(t *testing.T) {
	recorder := httptest.NewRecorder()
	handler := &timeoutWriter{w: recorder}
	handler.Flush()

	if got, want := recorder.Flushed, true; got != want {
		t.Errorf("recorder.Flushed = %t, want %t", got, want)
	}
}

func TestTimeoutWriterErrorsWriteAfterTimeout(t *testing.T) {
	recorder := httptest.NewRecorder()
	handler := &timeoutWriter{w: recorder}
	handler.timeoutAndWriteError("error")
	if _, err := handler.Write([]byte("hello")); !errors.Is(err, http.ErrHandlerTimeout) {
		t.Errorf("ErrHandlerTimeout got %v, want: %s", err, http.ErrHandlerTimeout)
	}
}

func TestTimeToFirstByteTimeoutHandler(t *testing.T) {
	const (
		immediateTimeout = 0 * time.Millisecond
		longTimeout      = 1 * time.Minute // Super long, not supposed to hit this.
	)

	tests := []struct {
		name           string
		timeout        time.Duration
		handler        func(mux *sync.Mutex, writeErrors chan error) http.Handler
		timeoutMessage string
		wantStatus     int
		wantBody       string
		wantWriteError bool
		wantPanic      bool
	}{{
		name:    "all good",
		timeout: longTimeout,
		handler: func(*sync.Mutex, chan error) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte("hi"))
			})
		},
		wantStatus: http.StatusOK,
		wantBody:   "hi",
	}, {
		name:    "custom timeout message",
		timeout: immediateTimeout,
		handler: func(mux *sync.Mutex, writeErrors chan error) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				mux.Lock()
				defer mux.Unlock()
				_, werr := w.Write([]byte("hi"))
				writeErrors <- werr
			})
		},
		timeoutMessage: "request timeout",
		wantStatus:     http.StatusGatewayTimeout,
		wantBody:       "request timeout",
		wantWriteError: true,
	}, {
		name:    "propagate panic",
		timeout: longTimeout,
		handler: func(*sync.Mutex, chan error) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				panic(http.ErrAbortHandler)
			})
		},
		wantStatus: http.StatusGatewayTimeout,
		wantBody:   "request timeout",
		wantPanic:  true,
	}, {
		name:    "timeout before panic",
		timeout: immediateTimeout,
		handler: func(*sync.Mutex, chan error) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				time.Sleep(1 * time.Second)
				panic(http.ErrAbortHandler)
			})
		},
		timeoutMessage: "request timeout",
		wantStatus:     http.StatusGatewayTimeout,
		wantBody:       "request timeout",
		wantPanic:      false,
	}}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, "/", nil)

			var reqMux sync.Mutex
			writeErrors := make(chan error, 1)
			rr := httptest.NewRecorder()
			handler := NewTimeToFirstByteTimeoutHandler(test.handler(&reqMux, writeErrors), test.timeoutMessage, test.timeout)

			defer func() {
				if test.wantPanic {
					if recovered := recover(); recovered != http.ErrAbortHandler { //nolint // False positive for errors.Is check
						t.Errorf("Recover = %v, want: %v", recovered, http.ErrAbortHandler)
					}
				}
			}()

			reqMux.Lock() // Will cause an inner 'Lock' to block. ServeHTTP will exit early if the call times out.
			handler.ServeHTTP(rr, req)
			reqMux.Unlock() // Allows the inner 'Lock' to go through to complete potential writes.

			if status := rr.Code; status != test.wantStatus {
				t.Errorf("Handler returned wrong status code: got %v want %v", status, test.wantStatus)
			}

			if rr.Body.String() != test.wantBody {
				t.Errorf("Handler returned unexpected body: got %q want %q", rr.Body.String(), test.wantBody)
			}

			if test.wantWriteError {
				if err := <-writeErrors; !errors.Is(err, http.ErrHandlerTimeout) {
					t.Error("Expected a timeout error, got", err)
				}
			}
		})
	}
}
<commit_msg>Add a benchmark for the TimeToFirstByteTimeoutHandler (#11766)<commit_after>/*
Copyright 2020 The Knative Authors

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

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

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

package handler

import (
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"
)

func TestTimeoutWriterAllowsForAdditionalWrites(t *testing.T) {
	recorder := httptest.NewRecorder()
	handler := &timeoutWriter{w: recorder}
	handler.WriteHeader(http.StatusOK)
	handler.timeoutAndWriteError("error")
	if _, err := io.WriteString(handler, "test"); err != nil {
		t.Fatalf("handler.Write() = %v, want no error", err)
	}

	if got, want := recorder.Code, http.StatusOK; got != want {
		t.Errorf("recorder.Status = %d, want %d", got, want)
	}
	if got, want := recorder.Body.String(), "test"; got != want {
		t.Errorf("recorder.Body = %s, want %s", got, want)
	}
}

func TestTimeoutWriterDoesntFlushAfterTimeout(t *testing.T) {
	recorder := httptest.NewRecorder()
	handler := &timeoutWriter{w: recorder}
	handler.timeoutAndWriteError("error")
	handler.Flush()

	if got, want := recorder.Flushed, false; got != want {
		t.Errorf("recorder.Flushed = %t, want %t", got, want)
	}
}

func TestTimeoutWriterFlushesAfterTimeout(t *testing.T) {
	recorder := httptest.NewRecorder()
	handler := &timeoutWriter{w: recorder}
	handler.Flush()

	if got, want := recorder.Flushed, true; got != want {
		t.Errorf("recorder.Flushed = %t, want %t", got, want)
	}
}

func TestTimeoutWriterErrorsWriteAfterTimeout(t *testing.T) {
	recorder := httptest.NewRecorder()
	handler := &timeoutWriter{w: recorder}
	handler.timeoutAndWriteError("error")
	if _, err := handler.Write([]byte("hello")); !errors.Is(err, http.ErrHandlerTimeout) {
		t.Errorf("ErrHandlerTimeout got %v, want: %s", err, http.ErrHandlerTimeout)
	}
}

func TestTimeToFirstByteTimeoutHandler(t *testing.T) {
	const (
		immediateTimeout = 0 * time.Millisecond
		longTimeout      = 1 * time.Minute // Super long, not supposed to hit this.
	)

	tests := []struct {
		name           string
		timeout        time.Duration
		handler        func(mux *sync.Mutex, writeErrors chan error) http.Handler
		timeoutMessage string
		wantStatus     int
		wantBody       string
		wantWriteError bool
		wantPanic      bool
	}{{
		name:    "all good",
		timeout: longTimeout,
		handler: func(*sync.Mutex, chan error) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte("hi"))
			})
		},
		wantStatus: http.StatusOK,
		wantBody:   "hi",
	}, {
		name:    "custom timeout message",
		timeout: immediateTimeout,
		handler: func(mux *sync.Mutex, writeErrors chan error) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				mux.Lock()
				defer mux.Unlock()
				_, werr := w.Write([]byte("hi"))
				writeErrors <- werr
			})
		},
		timeoutMessage: "request timeout",
		wantStatus:     http.StatusGatewayTimeout,
		wantBody:       "request timeout",
		wantWriteError: true,
	}, {
		name:    "propagate panic",
		timeout: longTimeout,
		handler: func(*sync.Mutex, chan error) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				panic(http.ErrAbortHandler)
			})
		},
		wantStatus: http.StatusGatewayTimeout,
		wantBody:   "request timeout",
		wantPanic:  true,
	}, {
		name:    "timeout before panic",
		timeout: immediateTimeout,
		handler: func(*sync.Mutex, chan error) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				time.Sleep(1 * time.Second)
				panic(http.ErrAbortHandler)
			})
		},
		timeoutMessage: "request timeout",
		wantStatus:     http.StatusGatewayTimeout,
		wantBody:       "request timeout",
		wantPanic:      false,
	}}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, "/", nil)

			var reqMux sync.Mutex
			writeErrors := make(chan error, 1)
			rr := httptest.NewRecorder()
			handler := NewTimeToFirstByteTimeoutHandler(test.handler(&reqMux, writeErrors), test.timeoutMessage, test.timeout)

			defer func() {
				if test.wantPanic {
					if recovered := recover(); recovered != http.ErrAbortHandler { //nolint // False positive for errors.Is check
						t.Errorf("Recover = %v, want: %v", recovered, http.ErrAbortHandler)
					}
				}
			}()

			reqMux.Lock() // Will cause an inner 'Lock' to block. ServeHTTP will exit early if the call times out.
			handler.ServeHTTP(rr, req)
			reqMux.Unlock() // Allows the inner 'Lock' to go through to complete potential writes.

			if status := rr.Code; status != test.wantStatus {
				t.Errorf("Handler returned wrong status code: got %v want %v", status, test.wantStatus)
			}

			if rr.Body.String() != test.wantBody {
				t.Errorf("Handler returned unexpected body: got %q want %q", rr.Body.String(), test.wantBody)
			}

			if test.wantWriteError {
				if err := <-writeErrors; !errors.Is(err, http.ErrHandlerTimeout) {
					t.Error("Expected a timeout error, got", err)
				}
			}
		})
	}
}

func BenchmarkTimeToFirstByteTimeoutHandler(b *testing.B) {
	writes := [][]byte{[]byte("this"), []byte("is"), []byte("a"), []byte("test")}
	baseHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		for _, write := range writes {
			w.Write(write)
		}
	})
	handler := NewTimeToFirstByteTimeoutHandler(baseHandler, "test", 10*time.Minute)
	req := httptest.NewRequest(http.MethodPost, "http://example.com", nil)

	b.Run("sequential", func(b *testing.B) {
		resp := httptest.NewRecorder()
		for j := 0; j < b.N; j++ {
			handler.ServeHTTP(resp, req)
		}
	})

	b.Run("parallel", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			resp := httptest.NewRecorder()
			for pb.Next() {
				handler.ServeHTTP(resp, req)
			}
		})
	})
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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 attestation

import (
	"bytes"
	"encoding/base64"
	"fmt"

	"github.com/pkg/errors"
	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/armor"
	"golang.org/x/crypto/openpgp/packet"
)

// PgpKey struct converts the base64 encoded PEM keys into openpgp private and
// public keys. Kubernetes Secrets are stored as base64 encoded PEM keys.
type PgpKey struct {
	privateKey *packet.PrivateKey
	publicKey  *packet.PublicKey
}

func NewPgpKey(privateKeyEnc string, publicKeyEnc string) (*PgpKey, error) {
	var publicKey *packet.PublicKey
	var privateKey *packet.PrivateKey
	var err error

	if privateKeyEnc != "" {
		privateKey, err = parsePrivateKey(string(privateKeyEnc))
		if err != nil {
			return nil, errors.Wrap(err, "parsing private key")
		}
	}
	if publicKeyEnc != "" {
		publicKey, err = parsePublicKey(string(publicKeyEnc))
		if err != nil {
			return nil, errors.Wrap(err, "parsing public key")
		}
	}
	return &PgpKey{
		privateKey: privateKey,
		publicKey:  publicKey,
	}, nil
}

func (key *PgpKey) PublicKey() *packet.PublicKey {
	return key.publicKey
}

func (key *PgpKey) PrivateKey() *packet.PrivateKey {
	return key.privateKey
}

func parsePublicKey(publicKey string) (*packet.PublicKey, error) {
	pkt, err := parseKey(publicKey, openpgp.PublicKeyType)
	if err != nil {
		return nil, err
	}
	key, ok := pkt.(*packet.PublicKey)
	if !ok {
		return nil, fmt.Errorf("Not a public key")
	}
	return key, nil
}

func parsePrivateKey(privateKey string) (*packet.PrivateKey, error) {
	pkt, err := parseKey(privateKey, openpgp.PrivateKeyType)
	if err != nil {
		return nil, err
	}
	key, ok := pkt.(*packet.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("Not a private Key")
	}
	return key, nil
}

func parseKey(key string, keytype string) (packet.Packet, error) {
	pemKeyString, err := base64.StdEncoding.DecodeString(key)
	if err != nil {
		return nil, err
	}
	r := bytes.NewReader(pemKeyString)
	block, err := armor.Decode(r)
	if err != nil {
		return nil, err
	}
	if block.Type != keytype {
		return nil, err
	}
	reader := packet.NewReader(block.Body)
	return reader.Next()
}
<commit_msg>another nit<commit_after>/*
Copyright 2018 Google LLC

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 attestation

import (
	"bytes"
	"encoding/base64"
	"fmt"

	"github.com/pkg/errors"
	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/armor"
	"golang.org/x/crypto/openpgp/packet"
)

// PgpKey struct converts the base64 encoded PEM keys into openpgp private and
// public keys. Kubernetes Secrets are stored as base64 encoded PEM keys.
type PgpKey struct {
	privateKey *packet.PrivateKey
	publicKey  *packet.PublicKey
}

func NewPgpKey(privateKeyEnc string, publicKeyEnc string) (*PgpKey, error) {
	var publicKey *packet.PublicKey
	var privateKey *packet.PrivateKey
	var err error

	if privateKeyEnc != "" {
		privateKey, err = parsePrivateKey(string(privateKeyEnc))
		if err != nil {
			return nil, errors.Wrap(err, "parsing private key")
		}
	}
	if publicKeyEnc != "" {
		publicKey, err = parsePublicKey(string(publicKeyEnc))
		if err != nil {
			return nil, errors.Wrap(err, "parsing public key")
		}
	}
	return &PgpKey{
		privateKey: privateKey,
		publicKey:  publicKey,
	}, nil
}

func (key *PgpKey) PublicKey() *packet.PublicKey {
	return key.publicKey
}

func (key *PgpKey) PrivateKey() *packet.PrivateKey {
	return key.privateKey
}

func parsePublicKey(publicKey string) (*packet.PublicKey, error) {
	pkt, err := parseKey(publicKey, openpgp.PublicKeyType)
	if err != nil {
		return nil, err
	}
	key, ok := pkt.(*packet.PublicKey)
	if !ok {
		return nil, fmt.Errorf("Not a public key")
	}
	return key, nil
}

func parsePrivateKey(privateKey string) (*packet.PrivateKey, error) {
	pkt, err := parseKey(privateKey, openpgp.PrivateKeyType)
	if err != nil {
		return nil, err
	}
	key, ok := pkt.(*packet.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("Not a private Key")
	}
	return key, nil
}

func parseKey(key string, keytype string) (packet.Packet, error) {
	s, err := base64.StdEncoding.DecodeString(key)
	if err != nil {
		return nil, err
	}
	r := bytes.NewReader(s)
	block, err := armor.Decode(r)
	if err != nil {
		return nil, err
	}
	if block.Type != keytype {
		return nil, err
	}
	reader := packet.NewReader(block.Body)
	return reader.Next()
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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 violation

import (
	"fmt"

	"github.com/golang/glog"
	"github.com/grafeas/kritis/pkg/kritis/constants"
	"github.com/grafeas/kritis/pkg/kritis/pods"
	"github.com/grafeas/kritis/pkg/kritis/policy"
	"k8s.io/api/core/v1"
)

type Strategy interface {
	HandleViolation(image string, pod *v1.Pod, violations []policy.Violation) error
	HandleAttestation(image string, pod *v1.Pod, isAttested bool) error
}

type LoggingStrategy struct {
}

func (l *LoggingStrategy) HandleViolation(image string, pod *v1.Pod, violations []policy.Violation) error {
	glog.Info("HandleViolation via LoggingStrategy")
	if len(violations) == 0 {
		return nil
	}
	glog.Warningf("Found violations in image %s", image)
	for _, v := range violations {
		glog.Warning(v.Reason())
	}
	return nil
}

func (l *LoggingStrategy) HandleAttestation(image string, pod *v1.Pod, isAttested bool) error {
	glog.Info("Handling attestation via LoggingStrategy")
	if isAttested {
		glog.Infof("Image %s has one or more valid attestation(s)", image)
	} else {
		glog.Infof("No Valid Attestations Found for image %s. Proceeding with next checks", image)
	}
	return nil
}

// AnnotationStrategy adds "InvalidImageSecPolicy" as a label and an annotation, with detailed info in the annotations
type AnnotationStrategy struct {
}

func (a *AnnotationStrategy) HandleViolation(image string, pod *v1.Pod, violations []policy.Violation) error {
	// First, remove "kritis.grafeas.io/invalidImageSecPolicy" label/annotation in case it doesn't apply anymore
	if err := pods.DeleteLabelsAndAnnotations(*pod, []string{constants.InvalidImageSecPolicy}, []string{constants.InvalidImageSecPolicy}); err != nil {
		return err
	}
	if len(violations) == 0 {
		return nil
	}
	// Now, construct labels and annotations
	labelValue := constants.InvalidImageSecPolicyLabelValue
	annotationValue := fmt.Sprintf("found %d CVEs", len(violations))
	for _, v := range violations {
		if v.Type() == policy.UnqualifiedImageViolation {
			annotationValue += fmt.Sprintf(", %s", v.Reason())
			break
		}
	}
	labels := map[string]string{constants.InvalidImageSecPolicy: labelValue}
	annotations := map[string]string{constants.InvalidImageSecPolicy: annotationValue}
	glog.Info(fmt.Sprintf("Adding label %s and annotation %s", labelValue, annotationValue))
	return pods.AddLabelsAndAnnotations(*pod, labels, annotations)
}

func (a *AnnotationStrategy) HandleAttestation(image string, pod *v1.Pod, isAttested bool) error {
	// First, remove "kritis.grafeas.io/attestation" label/annotation in case it doesn't apply anymore
	if err := pods.DeleteLabelsAndAnnotations(*pod, []string{constants.ImageAttestation}, []string{constants.ImageAttestation}); err != nil {
		return err
	}
	lValue := constants.NoAttestationsLabelValue
	aValue := constants.NoAttestationsAnnotation
	if isAttested {
		lValue = constants.PreviouslyAttestedLabelValue
		aValue = constants.PreviouslyAttestedAnnotation
	}
	labels := map[string]string{constants.ImageAttestation: lValue}
	annotations := map[string]string{constants.ImageAttestation: aValue}
	glog.Info(fmt.Sprintf("Adding label %s and annotation %s", lValue, aValue))
	return pods.AddLabelsAndAnnotations(*pod, labels, annotations)
}

// For unit testing.
type MemoryStrategy struct {
	Violations   map[string]bool
	Attestations map[string]bool
}

func (ms *MemoryStrategy) HandleViolation(image string, p *v1.Pod, v []policy.Violation) error {
	ms.Violations[image] = true
	return nil
}

func (ms *MemoryStrategy) HandleAttestation(image string, pod *v1.Pod, isAttested bool) error {
	ms.Attestations[image] = isAttested
	return nil
}
<commit_msg>minor formatting fix in log msg<commit_after>/*
Copyright 2018 Google LLC

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 violation

import (
	"fmt"

	"github.com/golang/glog"
	"github.com/grafeas/kritis/pkg/kritis/constants"
	"github.com/grafeas/kritis/pkg/kritis/pods"
	"github.com/grafeas/kritis/pkg/kritis/policy"
	"k8s.io/api/core/v1"
)

type Strategy interface {
	HandleViolation(image string, pod *v1.Pod, violations []policy.Violation) error
	HandleAttestation(image string, pod *v1.Pod, isAttested bool) error
}

type LoggingStrategy struct {
}

func (l *LoggingStrategy) HandleViolation(image string, pod *v1.Pod, violations []policy.Violation) error {
	glog.Info("HandleViolation via LoggingStrategy")
	if len(violations) == 0 {
		return nil
	}
	glog.Warningf("Found violations in image %s", image)
	for _, v := range violations {
		glog.Warning(v.Reason())
	}
	return nil
}

func (l *LoggingStrategy) HandleAttestation(image string, pod *v1.Pod, isAttested bool) error {
	glog.Info("Handling attestation via LoggingStrategy")
	if isAttested {
		glog.Infof("Image %s has one or more valid attestation(s)", image)
	} else {
		glog.Infof("No valid attestations found for image %s. Proceeding with next checks", image)
	}
	return nil
}

// AnnotationStrategy adds "InvalidImageSecPolicy" as a label and an annotation, with detailed info in the annotations
type AnnotationStrategy struct {
}

func (a *AnnotationStrategy) HandleViolation(image string, pod *v1.Pod, violations []policy.Violation) error {
	// First, remove "kritis.grafeas.io/invalidImageSecPolicy" label/annotation in case it doesn't apply anymore
	if err := pods.DeleteLabelsAndAnnotations(*pod, []string{constants.InvalidImageSecPolicy}, []string{constants.InvalidImageSecPolicy}); err != nil {
		return err
	}
	if len(violations) == 0 {
		return nil
	}
	// Now, construct labels and annotations
	labelValue := constants.InvalidImageSecPolicyLabelValue
	annotationValue := fmt.Sprintf("found %d CVEs", len(violations))
	for _, v := range violations {
		if v.Type() == policy.UnqualifiedImageViolation {
			annotationValue += fmt.Sprintf(", %s", v.Reason())
			break
		}
	}
	labels := map[string]string{constants.InvalidImageSecPolicy: labelValue}
	annotations := map[string]string{constants.InvalidImageSecPolicy: annotationValue}
	glog.Info(fmt.Sprintf("Adding label %s and annotation %s", labelValue, annotationValue))
	return pods.AddLabelsAndAnnotations(*pod, labels, annotations)
}

func (a *AnnotationStrategy) HandleAttestation(image string, pod *v1.Pod, isAttested bool) error {
	// First, remove "kritis.grafeas.io/attestation" label/annotation in case it doesn't apply anymore
	if err := pods.DeleteLabelsAndAnnotations(*pod, []string{constants.ImageAttestation}, []string{constants.ImageAttestation}); err != nil {
		return err
	}
	lValue := constants.NoAttestationsLabelValue
	aValue := constants.NoAttestationsAnnotation
	if isAttested {
		lValue = constants.PreviouslyAttestedLabelValue
		aValue = constants.PreviouslyAttestedAnnotation
	}
	labels := map[string]string{constants.ImageAttestation: lValue}
	annotations := map[string]string{constants.ImageAttestation: aValue}
	glog.Info(fmt.Sprintf("Adding label %s and annotation %s", lValue, aValue))
	return pods.AddLabelsAndAnnotations(*pod, labels, annotations)
}

// For unit testing.
type MemoryStrategy struct {
	Violations   map[string]bool
	Attestations map[string]bool
}

func (ms *MemoryStrategy) HandleViolation(image string, p *v1.Pod, v []policy.Violation) error {
	ms.Violations[image] = true
	return nil
}

func (ms *MemoryStrategy) HandleAttestation(image string, pod *v1.Pod, isAttested bool) error {
	ms.Attestations[image] = isAttested
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2015 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 pod

import (
	"context"
	"fmt"

	v1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/klog/v2"
	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
	kubetypes "k8s.io/kubernetes/pkg/kubelet/types"
)

// MirrorClient knows how to create/delete a mirror pod in the API server.
type MirrorClient interface {
	// CreateMirrorPod creates a mirror pod in the API server for the given
	// pod or returns an error.  The mirror pod will have the same annotations
	// as the given pod as well as an extra annotation containing the hash of
	// the static pod.
	CreateMirrorPod(pod *v1.Pod) error
	// DeleteMirrorPod deletes the mirror pod with the given full name from
	// the API server or returns an error.
	DeleteMirrorPod(podFullName string, uid *types.UID) (bool, error)
}

// nodeGetter is a subset a NodeLister, simplified for testing.
type nodeGetter interface {
	// Get retrieves the Node for a given name.
	Get(name string) (*v1.Node, error)
}

// basicMirrorClient is a functional MirrorClient.  Mirror pods are stored in
// the kubelet directly because they need to be in sync with the internal
// pods.
type basicMirrorClient struct {
	apiserverClient clientset.Interface
	nodeGetter      nodeGetter
	nodeName        string
}

// NewBasicMirrorClient returns a new MirrorClient.
func NewBasicMirrorClient(apiserverClient clientset.Interface, nodeName string, nodeGetter nodeGetter) MirrorClient {
	return &basicMirrorClient{
		apiserverClient: apiserverClient,
		nodeName:        nodeName,
		nodeGetter:      nodeGetter,
	}
}

func (mc *basicMirrorClient) CreateMirrorPod(pod *v1.Pod) error {
	if mc.apiserverClient == nil {
		return nil
	}
	// Make a copy of the pod.
	copyPod := *pod
	copyPod.Annotations = make(map[string]string)

	for k, v := range pod.Annotations {
		copyPod.Annotations[k] = v
	}
	hash := getPodHash(pod)
	copyPod.Annotations[kubetypes.ConfigMirrorAnnotationKey] = hash

	// With the MirrorPodNodeRestriction feature, mirror pods are required to have an owner reference
	// to the owning node.
	// See http://git.k8s.io/enhancements/keps/sig-auth/20190916-noderestriction-pods.md
	nodeUID, err := mc.getNodeUID()
	if err != nil {
		return fmt.Errorf("failed to get node UID: %v", err)
	}
	controller := true
	copyPod.OwnerReferences = []metav1.OwnerReference{{
		APIVersion: v1.SchemeGroupVersion.String(),
		Kind:       "Node",
		Name:       mc.nodeName,
		UID:        nodeUID,
		Controller: &controller,
	}}

	apiPod, err := mc.apiserverClient.CoreV1().Pods(copyPod.Namespace).Create(context.TODO(), &copyPod, metav1.CreateOptions{})
	if err != nil && apierrors.IsAlreadyExists(err) {
		// Check if the existing pod is the same as the pod we want to create.
		if h, ok := apiPod.Annotations[kubetypes.ConfigMirrorAnnotationKey]; ok && h == hash {
			return nil
		}
	}
	return err
}

// DeleteMirrorPod deletes a mirror pod.
// It takes the full name of the pod and optionally a UID.  If the UID
// is non-nil, the pod is deleted only if its UID matches the supplied UID.
// It returns whether the pod was actually deleted, and any error returned
// while parsing the name of the pod.
// Non-existence of the pod or UID mismatch is not treated as an error; the
// routine simply returns false in that case.
func (mc *basicMirrorClient) DeleteMirrorPod(podFullName string, uid *types.UID) (bool, error) {
	if mc.apiserverClient == nil {
		return false, nil
	}
	name, namespace, err := kubecontainer.ParsePodFullName(podFullName)
	if err != nil {
		klog.Errorf("Failed to parse a pod full name %q", podFullName)
		return false, err
	}
	klog.V(2).Infof("Deleting a mirror pod %q (uid %#v)", podFullName, uid)
	var GracePeriodSeconds int64
	if err := mc.apiserverClient.CoreV1().Pods(namespace).Delete(context.TODO(), name, metav1.DeleteOptions{GracePeriodSeconds: &GracePeriodSeconds, Preconditions: &metav1.Preconditions{UID: uid}}); err != nil {
		// Unfortunately, there's no generic error for failing a precondition
		if !(apierrors.IsNotFound(err) || apierrors.IsConflict(err)) {
			// We should return the error here, but historically this routine does
			// not return an error unless it can't parse the pod name
			klog.Errorf("Failed deleting a mirror pod %q: %v", podFullName, err)
		}
		return false, nil
	}
	return true, nil
}

func (mc *basicMirrorClient) getNodeUID() (types.UID, error) {
	node, err := mc.nodeGetter.Get(mc.nodeName)
	if err != nil {
		return "", err
	}
	if node.UID == "" {
		return "", fmt.Errorf("UID unset for node %s", mc.nodeName)
	}
	return node.UID, nil
}

// IsStaticPod returns true if the passed Pod is static.
func IsStaticPod(pod *v1.Pod) bool {
	source, err := kubetypes.GetPodSource(pod)
	return err == nil && source != kubetypes.ApiserverSource
}

func getHashFromMirrorPod(pod *v1.Pod) (string, bool) {
	hash, ok := pod.Annotations[kubetypes.ConfigMirrorAnnotationKey]
	return hash, ok
}

func getPodHash(pod *v1.Pod) string {
	// The annotation exists for all static pods.
	return pod.Annotations[kubetypes.ConfigHashAnnotationKey]
}
<commit_msg>Migrate `pkg/kubelet/pod` to structured logging<commit_after>/*
Copyright 2015 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 pod

import (
	"context"
	"fmt"

	v1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/klog/v2"
	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
	kubetypes "k8s.io/kubernetes/pkg/kubelet/types"
)

// MirrorClient knows how to create/delete a mirror pod in the API server.
type MirrorClient interface {
	// CreateMirrorPod creates a mirror pod in the API server for the given
	// pod or returns an error.  The mirror pod will have the same annotations
	// as the given pod as well as an extra annotation containing the hash of
	// the static pod.
	CreateMirrorPod(pod *v1.Pod) error
	// DeleteMirrorPod deletes the mirror pod with the given full name from
	// the API server or returns an error.
	DeleteMirrorPod(podFullName string, uid *types.UID) (bool, error)
}

// nodeGetter is a subset a NodeLister, simplified for testing.
type nodeGetter interface {
	// Get retrieves the Node for a given name.
	Get(name string) (*v1.Node, error)
}

// basicMirrorClient is a functional MirrorClient.  Mirror pods are stored in
// the kubelet directly because they need to be in sync with the internal
// pods.
type basicMirrorClient struct {
	apiserverClient clientset.Interface
	nodeGetter      nodeGetter
	nodeName        string
}

// NewBasicMirrorClient returns a new MirrorClient.
func NewBasicMirrorClient(apiserverClient clientset.Interface, nodeName string, nodeGetter nodeGetter) MirrorClient {
	return &basicMirrorClient{
		apiserverClient: apiserverClient,
		nodeName:        nodeName,
		nodeGetter:      nodeGetter,
	}
}

func (mc *basicMirrorClient) CreateMirrorPod(pod *v1.Pod) error {
	if mc.apiserverClient == nil {
		return nil
	}
	// Make a copy of the pod.
	copyPod := *pod
	copyPod.Annotations = make(map[string]string)

	for k, v := range pod.Annotations {
		copyPod.Annotations[k] = v
	}
	hash := getPodHash(pod)
	copyPod.Annotations[kubetypes.ConfigMirrorAnnotationKey] = hash

	// With the MirrorPodNodeRestriction feature, mirror pods are required to have an owner reference
	// to the owning node.
	// See http://git.k8s.io/enhancements/keps/sig-auth/20190916-noderestriction-pods.md
	nodeUID, err := mc.getNodeUID()
	if err != nil {
		return fmt.Errorf("failed to get node UID: %v", err)
	}
	controller := true
	copyPod.OwnerReferences = []metav1.OwnerReference{{
		APIVersion: v1.SchemeGroupVersion.String(),
		Kind:       "Node",
		Name:       mc.nodeName,
		UID:        nodeUID,
		Controller: &controller,
	}}

	apiPod, err := mc.apiserverClient.CoreV1().Pods(copyPod.Namespace).Create(context.TODO(), &copyPod, metav1.CreateOptions{})
	if err != nil && apierrors.IsAlreadyExists(err) {
		// Check if the existing pod is the same as the pod we want to create.
		if h, ok := apiPod.Annotations[kubetypes.ConfigMirrorAnnotationKey]; ok && h == hash {
			return nil
		}
	}
	return err
}

// DeleteMirrorPod deletes a mirror pod.
// It takes the full name of the pod and optionally a UID.  If the UID
// is non-nil, the pod is deleted only if its UID matches the supplied UID.
// It returns whether the pod was actually deleted, and any error returned
// while parsing the name of the pod.
// Non-existence of the pod or UID mismatch is not treated as an error; the
// routine simply returns false in that case.
func (mc *basicMirrorClient) DeleteMirrorPod(podFullName string, uid *types.UID) (bool, error) {
	if mc.apiserverClient == nil {
		return false, nil
	}
	name, namespace, err := kubecontainer.ParsePodFullName(podFullName)
	if err != nil {
		klog.ErrorS(err, "Failed to parse a pod full name", "podFullName", podFullName)
		return false, err
	}
	klog.V(2).InfoS("Deleting a mirror pod", "pod", klog.KRef(namespace, name), "podUID", uid)
	var GracePeriodSeconds int64
	if err := mc.apiserverClient.CoreV1().Pods(namespace).Delete(context.TODO(), name, metav1.DeleteOptions{GracePeriodSeconds: &GracePeriodSeconds, Preconditions: &metav1.Preconditions{UID: uid}}); err != nil {
		// Unfortunately, there's no generic error for failing a precondition
		if !(apierrors.IsNotFound(err) || apierrors.IsConflict(err)) {
			// We should return the error here, but historically this routine does
			// not return an error unless it can't parse the pod name
			klog.ErrorS(err, "Failed deleting a mirror pod", "pod", klog.KRef(namespace, name))
		}
		return false, nil
	}
	return true, nil
}

func (mc *basicMirrorClient) getNodeUID() (types.UID, error) {
	node, err := mc.nodeGetter.Get(mc.nodeName)
	if err != nil {
		return "", err
	}
	if node.UID == "" {
		return "", fmt.Errorf("UID unset for node %s", mc.nodeName)
	}
	return node.UID, nil
}

// IsStaticPod returns true if the passed Pod is static.
func IsStaticPod(pod *v1.Pod) bool {
	source, err := kubetypes.GetPodSource(pod)
	return err == nil && source != kubetypes.ApiserverSource
}

func getHashFromMirrorPod(pod *v1.Pod) (string, bool) {
	hash, ok := pod.Annotations[kubetypes.ConfigMirrorAnnotationKey]
	return hash, ok
}

func getPodHash(pod *v1.Pod) string {
	// The annotation exists for all static pods.
	return pod.Annotations[kubetypes.ConfigHashAnnotationKey]
}
<|endoftext|>


<commit_before>package kubernetes

import (
	"regexp"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/grafana/tanka/pkg/kubernetes/manifest"
	"github.com/grafana/tanka/pkg/spec/v1alpha1"
)

func TestExtract(t *testing.T) {
	tests := []struct {
		name string
		data testData
		err  error
	}{
		{
			name: "regular",
			data: testDataRegular(),
		},
		{
			name: "flat",
			data: testDataFlat(),
		},
		{
			name: "primitive",
			data: testDataPrimitive(),
			err:  ErrorPrimitiveReached{path: ".service", key: "note", primitive: "invalid because apiVersion and kind are missing"},
		},
		{
			name: "deep",
			data: testDataDeep(),
		},
		{
			name: "array",
			data: testDataArray(),
		},
		{
			name: "nil",
			data: func() testData {
				d := testDataRegular()
				d.Deep.(map[string]interface{})["disabledObject"] = nil
				return d
			}(),
			err: nil, // we expect no error, just the result of testDataRegular
		},
	}

	for _, c := range tests {
		t.Run(c.name, func(t *testing.T) {
			extracted, err := extract(c.data.Deep)

			require.Equal(t, c.err, err)
			assert.EqualValues(t, c.data.Flat, extracted)
		})
	}
}

func mkobj(kind string, name string, ns string) map[string]interface{} {
	ret := map[string]interface{}{
		"kind":       kind,
		"apiVersion": "apiversion",
		"metadata": map[string]interface{}{
			"name": name,
		},
	}
	if ns != "" {
		ret["metadata"].(map[string]interface{})["namespace"] = ns
	}

	return ret
}

func TestReconcileSorting(t *testing.T) {
	tests := []struct {
		raw     map[string]interface{}
		targets []*regexp.Regexp
		state   manifest.List
		err     error
	}{
		{
			// sorting by kinds in the `kindOrder` list
			raw: map[string]interface{}{
				"a": mkobj("Service", "service", "default"),
				"b": mkobj("Deployment", "deployment", "default"),
				"c": mkobj("CustomResourceDefinition", "crd", ""),
			},
			state: manifest.List{
				mkobj("CustomResourceDefinition", "crd", ""),
				mkobj("Service", "service", "default"),
				mkobj("Deployment", "deployment", "default"),
			},
		},
		{
			// alphabtical sorting by kinds outside `kindOrder` list
			raw: map[string]interface{}{
				"a": mkobj("B", "b", "default"),
				"b": mkobj("C", "c", "default"),
				"c": mkobj("A", "a", "default"),
			},
			state: manifest.List{
				mkobj("A", "a", "default"),
				mkobj("B", "b", "default"),
				mkobj("C", "c", "default"),
			},
		},
		{
			// sorting by the namespace if kinds match
			raw: map[string]interface{}{
				"a": mkobj("Service", "service", "default2"),
				"b": mkobj("Service", "service", "default"),
				"c": mkobj("Service", "service", "default1"),
			},
			state: manifest.List{
				mkobj("Service", "service", "default"),
				mkobj("Service", "service", "default1"),
				mkobj("Service", "service", "default2"),
			},
		},
		{
			// sorting by the names if both kinds and namespaces match
			raw: map[string]interface{}{
				"a": mkobj("Service", "service2", "default"),
				"b": mkobj("Service", "service", "default"),
				"c": mkobj("Service", "service1", "default"),
			},
			state: manifest.List{
				mkobj("Service", "service", "default"),
				mkobj("Service", "service1", "default"),
				mkobj("Service", "service2", "default"),
			},
		},
		{
			// sorting by the names if both kinds match and there are no namespaces
			raw: map[string]interface{}{
				"a": mkobj("CustomResourceDefinition", "crd2", ""),
				"b": mkobj("CustomResourceDefinition", "crd", ""),
				"c": mkobj("CustomResourceDefinition", "crd1", ""),
			},
			state: manifest.List{
				mkobj("CustomResourceDefinition", "crd", ""),
				mkobj("CustomResourceDefinition", "crd1", ""),
				mkobj("CustomResourceDefinition", "crd2", ""),
			},
		},
		{
			raw: map[string]interface{}{
				"a": mkobj("Deployment", "b", "a"),
				"b": mkobj("ConfigMap", "a", "a"),
				"c": mkobj("Issuer", "a", "a"),
				"d": mkobj("Service", "b", "a"),
				"e": mkobj("Service", "a", "a"),
				"f": mkobj("Deployment", "a", "a"),
				"g": mkobj("ConfigMap", "a", "b"),
				"h": mkobj("Issuer", "b", "a"),
				"i": mkobj("Service", "b", "b"),
				"j": mkobj("Deployment", "a", "b"),
				"k": mkobj("ConfigMap", "b", "a"),
				"l": mkobj("Issuer", "a", "b"),
			},
			state: manifest.List{
				mkobj("ConfigMap", "a", "a"),
				mkobj("ConfigMap", "b", "a"),
				mkobj("ConfigMap", "a", "b"),
				mkobj("Service", "a", "a"),
				mkobj("Service", "b", "a"),
				mkobj("Service", "b", "b"),
				mkobj("Deployment", "a", "a"),
				mkobj("Deployment", "b", "a"),
				mkobj("Deployment", "a", "b"),
				mkobj("Issuer", "a", "a"),
				mkobj("Issuer", "b", "a"),
				mkobj("Issuer", "a", "b"),
			},
		},
	}

	for _, test := range tests {
		res, err := Reconcile(test.raw, v1alpha1.New().Spec, test.targets)

		require.NoError(t, err)
		require.Equal(t, test.state, res)
	}
}
<commit_msg>fix(test): use full config instead of spec (#266)<commit_after>package kubernetes

import (
	"regexp"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/grafana/tanka/pkg/kubernetes/manifest"
	"github.com/grafana/tanka/pkg/spec/v1alpha1"
)

func TestExtract(t *testing.T) {
	tests := []struct {
		name string
		data testData
		err  error
	}{
		{
			name: "regular",
			data: testDataRegular(),
		},
		{
			name: "flat",
			data: testDataFlat(),
		},
		{
			name: "primitive",
			data: testDataPrimitive(),
			err:  ErrorPrimitiveReached{path: ".service", key: "note", primitive: "invalid because apiVersion and kind are missing"},
		},
		{
			name: "deep",
			data: testDataDeep(),
		},
		{
			name: "array",
			data: testDataArray(),
		},
		{
			name: "nil",
			data: func() testData {
				d := testDataRegular()
				d.Deep.(map[string]interface{})["disabledObject"] = nil
				return d
			}(),
			err: nil, // we expect no error, just the result of testDataRegular
		},
	}

	for _, c := range tests {
		t.Run(c.name, func(t *testing.T) {
			extracted, err := extract(c.data.Deep)

			require.Equal(t, c.err, err)
			assert.EqualValues(t, c.data.Flat, extracted)
		})
	}
}

func mkobj(kind string, name string, ns string) map[string]interface{} {
	ret := map[string]interface{}{
		"kind":       kind,
		"apiVersion": "apiversion",
		"metadata": map[string]interface{}{
			"name": name,
		},
	}
	if ns != "" {
		ret["metadata"].(map[string]interface{})["namespace"] = ns
	}

	return ret
}

func TestReconcileSorting(t *testing.T) {
	tests := []struct {
		raw     map[string]interface{}
		targets []*regexp.Regexp
		state   manifest.List
		err     error
	}{
		{
			// sorting by kinds in the `kindOrder` list
			raw: map[string]interface{}{
				"a": mkobj("Service", "service", "default"),
				"b": mkobj("Deployment", "deployment", "default"),
				"c": mkobj("CustomResourceDefinition", "crd", ""),
			},
			state: manifest.List{
				mkobj("CustomResourceDefinition", "crd", ""),
				mkobj("Service", "service", "default"),
				mkobj("Deployment", "deployment", "default"),
			},
		},
		{
			// alphabtical sorting by kinds outside `kindOrder` list
			raw: map[string]interface{}{
				"a": mkobj("B", "b", "default"),
				"b": mkobj("C", "c", "default"),
				"c": mkobj("A", "a", "default"),
			},
			state: manifest.List{
				mkobj("A", "a", "default"),
				mkobj("B", "b", "default"),
				mkobj("C", "c", "default"),
			},
		},
		{
			// sorting by the namespace if kinds match
			raw: map[string]interface{}{
				"a": mkobj("Service", "service", "default2"),
				"b": mkobj("Service", "service", "default"),
				"c": mkobj("Service", "service", "default1"),
			},
			state: manifest.List{
				mkobj("Service", "service", "default"),
				mkobj("Service", "service", "default1"),
				mkobj("Service", "service", "default2"),
			},
		},
		{
			// sorting by the names if both kinds and namespaces match
			raw: map[string]interface{}{
				"a": mkobj("Service", "service2", "default"),
				"b": mkobj("Service", "service", "default"),
				"c": mkobj("Service", "service1", "default"),
			},
			state: manifest.List{
				mkobj("Service", "service", "default"),
				mkobj("Service", "service1", "default"),
				mkobj("Service", "service2", "default"),
			},
		},
		{
			// sorting by the names if both kinds match and there are no namespaces
			raw: map[string]interface{}{
				"a": mkobj("CustomResourceDefinition", "crd2", ""),
				"b": mkobj("CustomResourceDefinition", "crd", ""),
				"c": mkobj("CustomResourceDefinition", "crd1", ""),
			},
			state: manifest.List{
				mkobj("CustomResourceDefinition", "crd", ""),
				mkobj("CustomResourceDefinition", "crd1", ""),
				mkobj("CustomResourceDefinition", "crd2", ""),
			},
		},
		{
			raw: map[string]interface{}{
				"a": mkobj("Deployment", "b", "a"),
				"b": mkobj("ConfigMap", "a", "a"),
				"c": mkobj("Issuer", "a", "a"),
				"d": mkobj("Service", "b", "a"),
				"e": mkobj("Service", "a", "a"),
				"f": mkobj("Deployment", "a", "a"),
				"g": mkobj("ConfigMap", "a", "b"),
				"h": mkobj("Issuer", "b", "a"),
				"i": mkobj("Service", "b", "b"),
				"j": mkobj("Deployment", "a", "b"),
				"k": mkobj("ConfigMap", "b", "a"),
				"l": mkobj("Issuer", "a", "b"),
			},
			state: manifest.List{
				mkobj("ConfigMap", "a", "a"),
				mkobj("ConfigMap", "b", "a"),
				mkobj("ConfigMap", "a", "b"),
				mkobj("Service", "a", "a"),
				mkobj("Service", "b", "a"),
				mkobj("Service", "b", "b"),
				mkobj("Deployment", "a", "a"),
				mkobj("Deployment", "b", "a"),
				mkobj("Deployment", "a", "b"),
				mkobj("Issuer", "a", "a"),
				mkobj("Issuer", "b", "a"),
				mkobj("Issuer", "a", "b"),
			},
		},
	}

	for _, test := range tests {
		res, err := Reconcile(test.raw, *v1alpha1.New(), test.targets)

		require.NoError(t, err)
		require.Equal(t, test.state, res)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The CDI Authors.

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

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

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

package controller

import (
	"context"
	"fmt"

	"github.com/go-logr/logr"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	routev1 "github.com/openshift/api/route/v1"

	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	"sigs.k8s.io/controller-runtime/pkg/source"

	cdiv1alpha1 "kubevirt.io/containerized-data-importer/pkg/apis/core/v1alpha1"
)

const (
	uploadProxyServiceName = "cdi-uploadproxy"
	uploadProxyRouteName   = uploadProxyServiceName
	uploadProxyCASecret    = "cdi-upload-proxy-ca-key"
)

func (r *ReconcileCDI) watchRoutes(c controller.Controller) error {
	eventHandler := &handler.EnqueueRequestForOwner{
		IsController: true,
		OwnerType:    &cdiv1alpha1.CDI{},
	}

	err := c.Watch(&source.Kind{Type: &routev1.Route{}}, eventHandler)
	if err != nil {
		if errors.IsNotFound(err) || meta.IsNoMatchError(err) {
			log.Info("Not watching Routes")
			return nil
		}

		return err
	}

	return nil
}

func (r *ReconcileCDI) ensureUploadProxyRouteExists(logger logr.Logger, cr *cdiv1alpha1.CDI) error {
	key := client.ObjectKey{Namespace: r.namespace, Name: uploadProxyRouteName}

	err := r.client.Get(context.TODO(), key, &routev1.Route{})
	if err == nil {
		// route already exists, so do nothing (user can mutate this)
		return nil
	}

	if meta.IsNoMatchError(err) {
		// not in openshift
		logger.V(3).Info("No match error for Route, must not be in openshift")
		return nil
	}

	if !errors.IsNotFound(err) {
		return err
	}

	secret := &corev1.Secret{}
	key = client.ObjectKey{Namespace: r.namespace, Name: uploadProxyCASecret}

	if err = r.client.Get(context.TODO(), key, secret); err != nil {
		return err
	}

	cert, exists := secret.Data["tls.crt"]
	if !exists {
		return fmt.Errorf("Unexpected secret format, 'tls.crt' key missing")
	}

	route := &routev1.Route{
		ObjectMeta: metav1.ObjectMeta{
			Name:      uploadProxyRouteName,
			Namespace: r.namespace,
		},
		Spec: routev1.RouteSpec{
			To: routev1.RouteTargetReference{
				Kind: "Service",
				Name: uploadProxyServiceName,
			},
			TLS: &routev1.TLSConfig{
				Termination:              routev1.TLSTerminationReencrypt,
				DestinationCACertificate: string(cert),
			},
		},
	}

	if err = controllerutil.SetControllerReference(cr, route, r.scheme); err != nil {
		return err
	}

	return r.client.Create(context.TODO(), route)
}
<commit_msg>increase uploadproxy route timeout (#854)<commit_after>/*
Copyright 2018 The CDI Authors.

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

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

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

package controller

import (
	"context"
	"fmt"

	"github.com/go-logr/logr"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	routev1 "github.com/openshift/api/route/v1"

	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	"sigs.k8s.io/controller-runtime/pkg/source"

	cdiv1alpha1 "kubevirt.io/containerized-data-importer/pkg/apis/core/v1alpha1"
)

const (
	uploadProxyServiceName = "cdi-uploadproxy"
	uploadProxyRouteName   = uploadProxyServiceName
	uploadProxyCASecret    = "cdi-upload-proxy-ca-key"
)

func (r *ReconcileCDI) watchRoutes(c controller.Controller) error {
	eventHandler := &handler.EnqueueRequestForOwner{
		IsController: true,
		OwnerType:    &cdiv1alpha1.CDI{},
	}

	err := c.Watch(&source.Kind{Type: &routev1.Route{}}, eventHandler)
	if err != nil {
		if errors.IsNotFound(err) || meta.IsNoMatchError(err) {
			log.Info("Not watching Routes")
			return nil
		}

		return err
	}

	return nil
}

func (r *ReconcileCDI) ensureUploadProxyRouteExists(logger logr.Logger, cr *cdiv1alpha1.CDI) error {
	key := client.ObjectKey{Namespace: r.namespace, Name: uploadProxyRouteName}

	err := r.client.Get(context.TODO(), key, &routev1.Route{})
	if err == nil {
		// route already exists, so do nothing (user can mutate this)
		return nil
	}

	if meta.IsNoMatchError(err) {
		// not in openshift
		logger.V(3).Info("No match error for Route, must not be in openshift")
		return nil
	}

	if !errors.IsNotFound(err) {
		return err
	}

	secret := &corev1.Secret{}
	key = client.ObjectKey{Namespace: r.namespace, Name: uploadProxyCASecret}

	if err = r.client.Get(context.TODO(), key, secret); err != nil {
		return err
	}

	cert, exists := secret.Data["tls.crt"]
	if !exists {
		return fmt.Errorf("Unexpected secret format, 'tls.crt' key missing")
	}

	route := &routev1.Route{
		ObjectMeta: metav1.ObjectMeta{
			Name:      uploadProxyRouteName,
			Namespace: r.namespace,
			Annotations: map[string]string{
				"haproxy.router.openshift.io/timeout": "3m",
			},
		},
		Spec: routev1.RouteSpec{
			To: routev1.RouteTargetReference{
				Kind: "Service",
				Name: uploadProxyServiceName,
			},
			TLS: &routev1.TLSConfig{
				Termination:              routev1.TLSTerminationReencrypt,
				DestinationCACertificate: string(cert),
			},
		},
	}

	if err = controllerutil.SetControllerReference(cr, route, r.scheme); err != nil {
		return err
	}

	return r.client.Create(context.TODO(), route)
}
<|endoftext|>


<commit_before>package scheduler

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/go-redis/redis"
	"github.com/sirupsen/logrus"
)

// TriggersKey is the the key of the sorted set in redis used for triggers
// waiting to be activated
const TriggersKey = "triggers"

// SchedKey is the the key of the sorted set in redis used for triggers
// currently being executed
const SchedKey = "scheduling"

// pollInterval is the time interval between 2 redis polling
const pollInterval = 1 * time.Second

// eventLoopSize is the number of goroutines handling @events and triggering
// jobs.
const eventLoopSize = 50

// luaPoll returns the lua script used for polling triggers in redis.
// If a trigger is in the scheduling key for more than 10 seconds, it is
// an error and we can try again to schedule it.
const luaPoll = `
local w = KEYS[1] - 10
local s = redis.call("ZRANGEBYSCORE", "` + SchedKey + `", 0, w, "WITHSCORES", "LIMIT", 0, 1)
if #s > 0 then
  redis.call("ZADD", "` + SchedKey + `", KEYS[1], s[1])
  return s
end
local t = redis.call("ZRANGEBYSCORE", "` + TriggersKey + `", 0, KEYS[1], "WITHSCORES", "LIMIT", 0, 1)
if #t > 0 then
  redis.call("ZREM", "` + TriggersKey + `", t[1])
  redis.call("ZADD", "` + SchedKey + `", t[2], t[1])
end
return t`

// RedisScheduler is a centralized scheduler of many triggers. It starts all of
// them and schedules jobs accordingly.
type RedisScheduler struct {
	broker  jobs.Broker
	client  redis.UniversalClient
	closed  chan struct{}
	stopped chan struct{}
	log     *logrus.Entry
}

// NewRedisScheduler creates a new scheduler that use redis to synchronize with
// other cozy-stack processes to schedule jobs.
func NewRedisScheduler(client redis.UniversalClient) *RedisScheduler {
	return &RedisScheduler{
		client:  client,
		log:     logger.WithNamespace("scheduler-redis"),
		stopped: make(chan struct{}),
	}
}

func redisKey(infos *TriggerInfos) string {
	return infos.Domain + "/" + infos.TID
}

func eventsKey(domain string) string {
	return "events-" + domain
}

// Start a goroutine that will fetch triggers in redis to schedule their jobs
func (s *RedisScheduler) Start(b jobs.Broker) error {
	s.broker = b
	s.closed = make(chan struct{})
	s.startEventDispatcher()
	go s.pollLoop()
	return nil
}

func (s *RedisScheduler) pollLoop() {
	ticker := time.NewTicker(pollInterval)
	for {
		select {
		case <-s.closed:
			ticker.Stop()
			s.stopped <- struct{}{}
			return
		case <-ticker.C:
			now := time.Now().UTC().Unix()
			if err := s.Poll(now); err != nil {
				s.log.Warnf("[scheduler] Failed to poll redis: %s", err)
			}
		}
	}
}

func (s *RedisScheduler) startEventDispatcher() {
	eventsCh := make(chan *realtime.Event, 100)
	go func() {
		c := realtime.GetHub().SubscribeLocalAll()
		defer func() {
			c.Close()
			close(eventsCh)
		}()
		for {
			select {
			case <-s.closed:
				return
			case event := <-c.Channel:
				eventsCh <- event
			}
		}
	}()
	for i := 0; i < eventLoopSize; i++ {
		go s.eventLoop(eventsCh)
	}
}

func (s *RedisScheduler) eventLoop(eventsCh <-chan *realtime.Event) {
	for event := range eventsCh {
		key := eventsKey(event.Domain)
		m, err := s.client.HGetAll(key).Result()
		if err != nil {
			s.log.Errorf("[scheduler] Could not fetch redis set %s: %s",
				key, err.Error())
			continue
		}
		for triggerID, args := range m {
			rule, err := permissions.UnmarshalRuleString(args)
			if err != nil {
				s.log.Warnf("[scheduler] Coud not unmarshal rule %s: %s",
					key, err.Error())
				continue
			}
			if !eventMatchPermission(event, &rule) {
				continue
			}
			t, err := s.Get(event.Domain, triggerID)
			if err != nil {
				s.log.Warnf("[scheduler] Could not fetch @event trigger %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
			et := t.(*EventTrigger)
			if et.infos.Debounce != "" {
				var d time.Duration
				if d, err = time.ParseDuration(et.infos.Debounce); err == nil {
					timestamp := time.Now().Add(d)
					s.client.ZAddNX(TriggersKey, redis.Z{
						Score:  float64(timestamp.UTC().Unix()),
						Member: redisKey(t.Infos()),
					})
					continue
				}
				s.log.Warnf("[scheduler] Trigger %s %s has an invalid debounce: %s",
					et.infos.Domain, et.infos.TID, et.infos.Debounce)
			}
			jobRequest, err := et.Infos().JobRequestWithEvent(event)
			if err != nil {
				s.log.Warnf("[scheduler] Could not encode realtime event %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
			_, err = s.broker.PushJob(jobRequest)
			if err != nil {
				s.log.Warnf("[scheduler] Could not push job trigger by event %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
		}
	}
}

// Shutdown the scheduling of triggers
func (s *RedisScheduler) Shutdown(ctx context.Context) error {
	if s.closed == nil {
		return nil
	}
	fmt.Print("  shutting down redis scheduler...")
	close(s.closed)
	select {
	case <-ctx.Done():
		fmt.Println("failed: ", ctx.Err())
		return ctx.Err()
	case <-s.stopped:
		fmt.Println("ok.")
	}
	return nil
}

// Poll redis to see if there are some triggers ready
func (s *RedisScheduler) Poll(now int64) error {
	keys := []string{strconv.FormatInt(now, 10)}
	for {
		res, err := s.client.Eval(luaPoll, keys).Result()
		if err != nil || res == nil {
			return err
		}
		results, ok := res.([]interface{})
		if !ok {
			return errors.New("Unexpected response from redis")
		}
		if len(results) < 2 {
			return nil
		}
		parts := strings.SplitN(results[0].(string), "/", 2)
		if len(parts) != 2 {
			s.client.ZRem(SchedKey, results[0])
			return fmt.Errorf("Invalid key %s", res)
		}
		t, err := s.Get(parts[0], parts[1])
		if err != nil {
			if err == ErrNotFoundTrigger {
				s.client.ZRem(SchedKey, results[0])
			}
			return err
		}
		switch t := t.(type) {
		case *EventTrigger: // Debounced
			job := t.Infos().JobRequest()
			job.Debounced = true
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
		case *AtTrigger:
			job := t.Infos().JobRequest()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			if err = s.deleteTrigger(t); err != nil {
				return err
			}
		case *CronTrigger:
			job := t.Infos().JobRequest()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			score, err := strconv.ParseInt(results[1].(string), 10, 64)
			var prev time.Time
			if err != nil {
				prev = time.Now()
			} else {
				prev = time.Unix(score, 0)
			}
			if err := s.addToRedis(t, prev); err != nil {
				return err
			}
		default:
			return errors.New("Not implemented yet")
		}
	}
}

// Add a trigger to the system, by persisting it and using redis for scheduling
// its jobs
func (s *RedisScheduler) Add(t Trigger) error {
	infos := t.Infos()
	db := couchdb.SimpleDatabasePrefix(infos.Domain)
	if err := couchdb.CreateDoc(db, infos); err != nil {
		return err
	}
	return s.addToRedis(t, time.Now())
}

func (s *RedisScheduler) addToRedis(t Trigger, prev time.Time) error {
	var timestamp time.Time
	switch t := t.(type) {
	case *EventTrigger:
		hKey := eventsKey(t.Infos().Domain)
		pipe := s.client.Pipeline()
		pipe.HSet(hKey, t.ID(), t.Infos().Arguments)
		pipe.ZRem(SchedKey, redisKey(t.Infos()))
		_, err := pipe.Exec()
		return err
	case *AtTrigger:
		timestamp = t.at
	case *CronTrigger:
		timestamp = t.NextExecution(prev)
		now := time.Now()
		if timestamp.Before(now) {
			timestamp = t.NextExecution(now)
		}
	default:
		return errors.New("Not implemented yet")
	}
	pipe := s.client.Pipeline()
	pipe.ZAdd(TriggersKey, redis.Z{
		Score:  float64(timestamp.UTC().Unix()),
		Member: redisKey(t.Infos()),
	})
	pipe.ZRem(SchedKey, redisKey(t.Infos()))
	_, err := pipe.Exec()
	return err
}

// Get returns the trigger with the specified ID.
func (s *RedisScheduler) Get(domain, id string) (Trigger, error) {
	var infos TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	if err := couchdb.GetDoc(db, consts.Triggers, id, &infos); err != nil {
		if couchdb.IsNotFoundError(err) {
			return nil, ErrNotFoundTrigger
		}
		return nil, err
	}
	return NewTrigger(&infos)
}

// Delete removes the trigger with the specified ID. The trigger is unscheduled
// and remove from the storage.
func (s *RedisScheduler) Delete(domain, id string) error {
	t, err := s.Get(domain, id)
	if err != nil {
		return err
	}
	return s.deleteTrigger(t)
}

func (s *RedisScheduler) deleteTrigger(t Trigger) error {
	db := couchdb.SimpleDatabasePrefix(t.Infos().Domain)
	if err := couchdb.DeleteDoc(db, t.Infos()); err != nil {
		return err
	}
	switch t.(type) {
	case *EventTrigger:
		return s.client.HDel(eventsKey(t.Infos().Domain), t.ID()).Err()
	case *AtTrigger, *CronTrigger:
		pipe := s.client.Pipeline()
		pipe.ZRem(TriggersKey, redisKey(t.Infos()))
		pipe.ZRem(SchedKey, redisKey(t.Infos()))
		_, err := pipe.Exec()
		return err
	}
	return nil
}

// GetAll returns all the triggers for a domain, from couch.
func (s *RedisScheduler) GetAll(domain string) ([]Trigger, error) {
	var infos []*TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	err := couchdb.ForeachDocs(db, consts.Triggers, func(data []byte) error {
		var t *TriggerInfos
		if err := json.Unmarshal(data, &t); err != nil {
			return err
		}
		infos = append(infos, t)
		return nil
	})
	if err != nil {
		if couchdb.IsNoDatabaseError(err) {
			return nil, nil
		}
		return nil, err
	}
	v := make([]Trigger, 0, len(infos))
	for _, info := range infos {
		t, err := NewTrigger(info)
		if err != nil {
			return nil, err
		}
		v = append(v, t)
	}
	return v, nil
}

// RebuildRedis puts all the triggers in redis (idempotent)
func (s *RedisScheduler) RebuildRedis(domain string) error {
	err := s.client.Del(TriggersKey, SchedKey).Err()
	if err != nil {
		return err
	}
	triggers, err := s.GetAll(domain)
	if err != nil {
		return err
	}
	for _, t := range triggers {
		if err = s.addToRedis(t, time.Now()); err != nil {
			return err
		}
	}
	return nil
}

var _ Scheduler = &RedisScheduler{}
<commit_msg>Revert "Fix debounced event triggers"<commit_after>package scheduler

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/go-redis/redis"
	"github.com/sirupsen/logrus"
)

// TriggersKey is the the key of the sorted set in redis used for triggers
// waiting to be activated
const TriggersKey = "triggers"

// SchedKey is the the key of the sorted set in redis used for triggers
// currently being executed
const SchedKey = "scheduling"

// pollInterval is the time interval between 2 redis polling
const pollInterval = 1 * time.Second

// eventLoopSize is the number of goroutines handling @events and triggering
// jobs.
const eventLoopSize = 50

// luaPoll returns the lua script used for polling triggers in redis.
// If a trigger is in the scheduling key for more than 10 seconds, it is
// an error and we can try again to schedule it.
const luaPoll = `
local w = KEYS[1] - 10
local s = redis.call("ZRANGEBYSCORE", "` + SchedKey + `", 0, w, "WITHSCORES", "LIMIT", 0, 1)
if #s > 0 then
  redis.call("ZADD", "` + SchedKey + `", KEYS[1], s[1])
  return s
end
local t = redis.call("ZRANGEBYSCORE", "` + TriggersKey + `", 0, KEYS[1], "WITHSCORES", "LIMIT", 0, 1)
if #t > 0 then
  redis.call("ZREM", "` + TriggersKey + `", t[1])
  redis.call("ZADD", "` + SchedKey + `", t[2], t[1])
end
return t`

// RedisScheduler is a centralized scheduler of many triggers. It starts all of
// them and schedules jobs accordingly.
type RedisScheduler struct {
	broker  jobs.Broker
	client  redis.UniversalClient
	closed  chan struct{}
	stopped chan struct{}
	log     *logrus.Entry
}

// NewRedisScheduler creates a new scheduler that use redis to synchronize with
// other cozy-stack processes to schedule jobs.
func NewRedisScheduler(client redis.UniversalClient) *RedisScheduler {
	return &RedisScheduler{
		client:  client,
		log:     logger.WithNamespace("scheduler-redis"),
		stopped: make(chan struct{}),
	}
}

func redisKey(infos *TriggerInfos) string {
	return infos.Domain + "/" + infos.TID
}

func eventsKey(domain string) string {
	return "events-" + domain
}

// Start a goroutine that will fetch triggers in redis to schedule their jobs
func (s *RedisScheduler) Start(b jobs.Broker) error {
	s.broker = b
	s.closed = make(chan struct{})
	s.startEventDispatcher()
	go s.pollLoop()
	return nil
}

func (s *RedisScheduler) pollLoop() {
	ticker := time.NewTicker(pollInterval)
	for {
		select {
		case <-s.closed:
			ticker.Stop()
			s.stopped <- struct{}{}
			return
		case <-ticker.C:
			now := time.Now().UTC().Unix()
			if err := s.Poll(now); err != nil {
				s.log.Warnf("[scheduler] Failed to poll redis: %s", err)
			}
		}
	}
}

func (s *RedisScheduler) startEventDispatcher() {
	eventsCh := make(chan *realtime.Event, 100)
	go func() {
		c := realtime.GetHub().SubscribeLocalAll()
		defer func() {
			c.Close()
			close(eventsCh)
		}()
		for {
			select {
			case <-s.closed:
				return
			case event := <-c.Channel:
				eventsCh <- event
			}
		}
	}()
	for i := 0; i < eventLoopSize; i++ {
		go s.eventLoop(eventsCh)
	}
}

func (s *RedisScheduler) eventLoop(eventsCh <-chan *realtime.Event) {
	for event := range eventsCh {
		key := eventsKey(event.Domain)
		m, err := s.client.HGetAll(key).Result()
		if err != nil {
			s.log.Errorf("[scheduler] Could not fetch redis set %s: %s",
				key, err.Error())
			continue
		}
		for triggerID, args := range m {
			rule, err := permissions.UnmarshalRuleString(args)
			if err != nil {
				s.log.Warnf("[scheduler] Coud not unmarshal rule %s: %s",
					key, err.Error())
				continue
			}
			if !eventMatchPermission(event, &rule) {
				continue
			}
			t, err := s.Get(event.Domain, triggerID)
			if err != nil {
				s.log.Warnf("[scheduler] Could not fetch @event trigger %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
			et := t.(*EventTrigger)
			if et.infos.Debounce != "" {
				var d time.Duration
				if d, err = time.ParseDuration(et.infos.Debounce); err == nil {
					timestamp := time.Now().Add(d)
					s.client.ZAddNX(TriggersKey, redis.Z{
						Score:  float64(timestamp.UTC().Unix()),
						Member: redisKey(t.Infos()),
					})
					continue
				}
				s.log.Warnf("[scheduler] Trigger %s %s has an invalid debounce: %s",
					et.infos.Domain, et.infos.TID, et.infos.Debounce)
			}
			jobRequest, err := et.Infos().JobRequestWithEvent(event)
			if err != nil {
				s.log.Warnf("[scheduler] Could not encode realtime event %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
			_, err = s.broker.PushJob(jobRequest)
			if err != nil {
				s.log.Warnf("[scheduler] Could not push job trigger by event %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
		}
	}
}

// Shutdown the scheduling of triggers
func (s *RedisScheduler) Shutdown(ctx context.Context) error {
	if s.closed == nil {
		return nil
	}
	fmt.Print("  shutting down redis scheduler...")
	close(s.closed)
	select {
	case <-ctx.Done():
		fmt.Println("failed: ", ctx.Err())
		return ctx.Err()
	case <-s.stopped:
		fmt.Println("ok.")
	}
	return nil
}

// Poll redis to see if there are some triggers ready
func (s *RedisScheduler) Poll(now int64) error {
	keys := []string{strconv.FormatInt(now, 10)}
	for {
		res, err := s.client.Eval(luaPoll, keys).Result()
		if err != nil || res == nil {
			return err
		}
		results, ok := res.([]interface{})
		if !ok {
			return errors.New("Unexpected response from redis")
		}
		if len(results) < 2 {
			return nil
		}
		parts := strings.SplitN(results[0].(string), "/", 2)
		if len(parts) != 2 {
			s.client.ZRem(SchedKey, results[0])
			return fmt.Errorf("Invalid key %s", res)
		}
		t, err := s.Get(parts[0], parts[1])
		if err != nil {
			if err == ErrNotFoundTrigger {
				s.client.ZRem(SchedKey, results[0])
			}
			return err
		}
		switch t := t.(type) {
		case *EventTrigger: // Debounced
			job := t.Infos().JobRequest()
			job.Debounced = true
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
		case *AtTrigger:
			job := t.Infos().JobRequest()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			if err = s.deleteTrigger(t); err != nil {
				return err
			}
		case *CronTrigger:
			job := t.Infos().JobRequest()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			score, err := strconv.ParseInt(results[1].(string), 10, 64)
			var prev time.Time
			if err != nil {
				prev = time.Now()
			} else {
				prev = time.Unix(score, 0)
			}
			if err := s.addToRedis(t, prev); err != nil {
				return err
			}
		default:
			return errors.New("Not implemented yet")
		}
	}
}

// Add a trigger to the system, by persisting it and using redis for scheduling
// its jobs
func (s *RedisScheduler) Add(t Trigger) error {
	infos := t.Infos()
	db := couchdb.SimpleDatabasePrefix(infos.Domain)
	if err := couchdb.CreateDoc(db, infos); err != nil {
		return err
	}
	return s.addToRedis(t, time.Now())
}

func (s *RedisScheduler) addToRedis(t Trigger, prev time.Time) error {
	var timestamp time.Time
	switch t := t.(type) {
	case *EventTrigger:
		hKey := eventsKey(t.Infos().Domain)
		return s.client.HSet(hKey, t.ID(), t.Infos().Arguments).Err()
	case *AtTrigger:
		timestamp = t.at
	case *CronTrigger:
		timestamp = t.NextExecution(prev)
		now := time.Now()
		if timestamp.Before(now) {
			timestamp = t.NextExecution(now)
		}
	default:
		return errors.New("Not implemented yet")
	}
	pipe := s.client.Pipeline()
	pipe.ZAdd(TriggersKey, redis.Z{
		Score:  float64(timestamp.UTC().Unix()),
		Member: redisKey(t.Infos()),
	}).Err()
	pipe.ZRem(SchedKey, redisKey(t.Infos()))
	_, err := pipe.Exec()
	return err
}

// Get returns the trigger with the specified ID.
func (s *RedisScheduler) Get(domain, id string) (Trigger, error) {
	var infos TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	if err := couchdb.GetDoc(db, consts.Triggers, id, &infos); err != nil {
		if couchdb.IsNotFoundError(err) {
			return nil, ErrNotFoundTrigger
		}
		return nil, err
	}
	return NewTrigger(&infos)
}

// Delete removes the trigger with the specified ID. The trigger is unscheduled
// and remove from the storage.
func (s *RedisScheduler) Delete(domain, id string) error {
	t, err := s.Get(domain, id)
	if err != nil {
		return err
	}
	return s.deleteTrigger(t)
}

func (s *RedisScheduler) deleteTrigger(t Trigger) error {
	db := couchdb.SimpleDatabasePrefix(t.Infos().Domain)
	if err := couchdb.DeleteDoc(db, t.Infos()); err != nil {
		return err
	}
	switch t.(type) {
	case *EventTrigger:
		return s.client.HDel(eventsKey(t.Infos().Domain), t.ID()).Err()
	case *AtTrigger, *CronTrigger:
		pipe := s.client.Pipeline()
		pipe.ZRem(TriggersKey, redisKey(t.Infos()))
		pipe.ZRem(SchedKey, redisKey(t.Infos()))
		_, err := pipe.Exec()
		return err
	}
	return nil
}

// GetAll returns all the triggers for a domain, from couch.
func (s *RedisScheduler) GetAll(domain string) ([]Trigger, error) {
	var infos []*TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	err := couchdb.ForeachDocs(db, consts.Triggers, func(data []byte) error {
		var t *TriggerInfos
		if err := json.Unmarshal(data, &t); err != nil {
			return err
		}
		infos = append(infos, t)
		return nil
	})
	if err != nil {
		if couchdb.IsNoDatabaseError(err) {
			return nil, nil
		}
		return nil, err
	}
	v := make([]Trigger, 0, len(infos))
	for _, info := range infos {
		t, err := NewTrigger(info)
		if err != nil {
			return nil, err
		}
		v = append(v, t)
	}
	return v, nil
}

// RebuildRedis puts all the triggers in redis (idempotent)
func (s *RedisScheduler) RebuildRedis(domain string) error {
	err := s.client.Del(TriggersKey, SchedKey).Err()
	if err != nil {
		return err
	}
	triggers, err := s.GetAll(domain)
	if err != nil {
		return err
	}
	for _, t := range triggers {
		if err = s.addToRedis(t, time.Now()); err != nil {
			return err
		}
	}
	return nil
}

var _ Scheduler = &RedisScheduler{}
<|endoftext|>


<commit_before>/*
Copyright 2018 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 deploy

import (
	"bytes"
	"context"
	"io"
	"io/ioutil"
	"os/exec"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/v1alpha2"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/pkg/errors"
)

type KustomizeDeployer struct {
	*v1alpha2.KustomizeDeploy

	kubeContext string
}

func NewKustomizeDeployer(cfg *v1alpha2.KustomizeDeploy, kubeContext string) *KustomizeDeployer {
	return &KustomizeDeployer{
		KustomizeDeploy: cfg,
		kubeContext:     kubeContext,
	}
}

func (k *KustomizeDeployer) Labels() map[string]string {
	return map[string]string{
		constants.Labels.Deployer: "kustomize",
	}
}

func (k *KustomizeDeployer) Deploy(ctx context.Context, out io.Writer, builds []build.Artifact) ([]Artifact, error) {
	manifests, err := buildManifests(k.KustomizePath)
	if err != nil {
		return nil, errors.Wrap(err, "kustomize")
	}
	manifestList, err := newManifestList(manifests)
	if err != nil {
		return nil, errors.Wrap(err, "getting manifest list")
	}
	manifestList, err = manifestList.replaceImages(builds)
	if err != nil {
		return nil, errors.Wrap(err, "replacing images")
	}
	if err := kubectl(manifestList.reader(), out, k.kubeContext, "apply", "-f", "-"); err != nil {
		return nil, errors.Wrap(err, "running kubectl")
	}
	return parseManifestsForDeploys(manifestList)
}

func newManifestList(r io.Reader) (manifestList, error) {
	var manifests manifestList
	buf, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, errors.Wrap(err, "reading manifests")
	}

	parts := bytes.Split(buf, []byte("\n---"))
	for _, part := range parts {
		manifests = append(manifests, part)
	}

	return manifests, nil
}

func (k *KustomizeDeployer) Cleanup(ctx context.Context, out io.Writer) error {
	manifests, err := buildManifests(constants.DefaultKustomizationPath)
	if err != nil {
		return errors.Wrap(err, "kustomize")
	}
	if err := kubectl(manifests, out, k.kubeContext, "delete", "-f", "-"); err != nil {
		return errors.Wrap(err, "kubectl delete")
	}
	return nil
}

func (k *KustomizeDeployer) Dependencies() ([]string, error) {
	// TODO(r2d4): parse kustomization yaml and add base and patches as dependencies
	return []string{constants.DefaultKustomizationPath}, nil
}

func buildManifests(kustomization string) (io.Reader, error) {
	cmd := exec.Command("kustomize", "build", kustomization)
	out, err := util.DefaultExecCommand.RunCmdOut(cmd)
	if err != nil {
		return nil, errors.Wrap(err, "running kustomize build")
	}
	return bytes.NewReader(out), nil
}
<commit_msg>kustomize: use custom path in deploy deps<commit_after>/*
Copyright 2018 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 deploy

import (
	"bytes"
	"context"
	"io"
	"io/ioutil"
	"os/exec"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/v1alpha2"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/pkg/errors"
)

type KustomizeDeployer struct {
	*v1alpha2.KustomizeDeploy

	kubeContext string
}

func NewKustomizeDeployer(cfg *v1alpha2.KustomizeDeploy, kubeContext string) *KustomizeDeployer {
	return &KustomizeDeployer{
		KustomizeDeploy: cfg,
		kubeContext:     kubeContext,
	}
}

func (k *KustomizeDeployer) Labels() map[string]string {
	return map[string]string{
		constants.Labels.Deployer: "kustomize",
	}
}

func (k *KustomizeDeployer) Deploy(ctx context.Context, out io.Writer, builds []build.Artifact) ([]Artifact, error) {
	manifests, err := buildManifests(k.KustomizePath)
	if err != nil {
		return nil, errors.Wrap(err, "kustomize")
	}
	manifestList, err := newManifestList(manifests)
	if err != nil {
		return nil, errors.Wrap(err, "getting manifest list")
	}
	manifestList, err = manifestList.replaceImages(builds)
	if err != nil {
		return nil, errors.Wrap(err, "replacing images")
	}
	if err := kubectl(manifestList.reader(), out, k.kubeContext, "apply", "-f", "-"); err != nil {
		return nil, errors.Wrap(err, "running kubectl")
	}
	return parseManifestsForDeploys(manifestList)
}

func newManifestList(r io.Reader) (manifestList, error) {
	var manifests manifestList
	buf, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, errors.Wrap(err, "reading manifests")
	}

	parts := bytes.Split(buf, []byte("\n---"))
	for _, part := range parts {
		manifests = append(manifests, part)
	}

	return manifests, nil
}

func (k *KustomizeDeployer) Cleanup(ctx context.Context, out io.Writer) error {
	manifests, err := buildManifests(constants.DefaultKustomizationPath)
	if err != nil {
		return errors.Wrap(err, "kustomize")
	}
	if err := kubectl(manifests, out, k.kubeContext, "delete", "-f", "-"); err != nil {
		return errors.Wrap(err, "kubectl delete")
	}
	return nil
}

func (k *KustomizeDeployer) Dependencies() ([]string, error) {
	// TODO(r2d4): parse kustomization yaml and add base and patches as dependencies
	return []string{k.KustomizePath}, nil
}

func buildManifests(kustomization string) (io.Reader, error) {
	cmd := exec.Command("kustomize", "build", kustomization)
	out, err := util.DefaultExecCommand.RunCmdOut(cmd)
	if err != nil {
		return nil, errors.Wrap(err, "running kustomize build")
	}
	return bytes.NewReader(out), nil
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/csv"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
)

const VERSION = `0.1.1`

var debug = debugT(false)

type debugT bool

func (d debugT) Printf(format string, args ...interface{}) {
	if d {
		log.Printf(format, args...)
	}
}

// Run is main routine.
func (trdsql TRDSQL) Run(args []string) int {
	var (
		version bool
		odriver string
		odsn    string
		inSep   string
		ihead   bool
		iskip   int
		query   string
		driver  string
		dsn     string
		odebug  bool
	)
	flags := flag.NewFlagSet("trdsql", flag.ContinueOnError)
	driver = "sqlite3"
	dsn = ""
	cfgfile := configOpen()
	cfg, _ := loadConfig(cfgfile)
	flags.Usage = func() {
		fmt.Fprintf(os.Stderr, `
Usage: %s [OPTIONS] [SQL(SELECT...)]

Options:
`, os.Args[0])
		flags.PrintDefaults()
	}

	flags.StringVar(&cfg.Db, "db", cfg.Db, "Specify db name of the setting.")
	flags.StringVar(&odriver, "driver", "", "database driver.  [ "+strings.Join(sql.Drivers(), " | ")+" ]")
	flags.StringVar(&odsn, "dsn", "", "database connection option.")
	flags.StringVar(&inSep, "id", ",", "Field delimiter for input.")
	flags.StringVar(&trdsql.outSep, "od", ",", "Field delimiter for output.")
	flags.BoolVar(&ihead, "ih", false, "The first line is interpreted as column names.")
	flags.BoolVar(&trdsql.outHeader, "oh", false, "Output column name as header.")
	flags.IntVar(&iskip, "is", 0, "Skip header row.")
	flags.StringVar(&query, "q", "", "Read query from the provided filename.")
	flags.BoolVar(&version, "version", false, "display version information.")
	flags.BoolVar(&odebug, "debug", false, "debug print.")
	flags.Parse(args[1:])
	if version {
		fmt.Println(VERSION)
		return (0)
	}
	var sqlstr string
	if query != "" {
		bq, err := ioutil.ReadFile(query)
		if err != nil {
			log.Println("ERROR: ", err)
			return (1)
		}
		sqlstr = string(bq)
	} else {
		sqlstr = strings.Join(flags.Args(), " ")
	}
	if len(sqlstr) == 0 {
		flags.Usage()
		return (2)
	}
	if odebug {
		debug = true
	}
	if strings.HasSuffix(sqlstr, ";") {
		sqlstr = sqlstr[:len(sqlstr)-1]
	}

	if cfg.Db != "" {
		if cfg.Database[cfg.Db].Driver == "" {
			debug.Printf("ERROR: db[%s] does not found", cfg.Db)
		} else {
			driver = cfg.Database[cfg.Db].Driver
			dsn = cfg.Database[cfg.Db].Dsn
		}
	}
	if odriver != "" {
		driver = odriver
	}
	if odsn != "" {
		dsn = odsn
	}

	readerComma, err := getSeparator(inSep)
	if err != nil {
		log.Println(err)
	}
	debug.Printf("driver: %s, dsn: %s", driver, dsn)
	db, err := Connect(driver, dsn)
	if err != nil {
		log.Println("ERROR: ", err)
		return 1
	}
	defer db.Disconnect()

	tablenames := sqlparse(sqlstr)
	if len(tablenames) == 0 {
		// withou FROM clause. ex. SELECT 1+1;
		debug.Printf("table not found\n")
	}
	var reader *csv.Reader
	var header []string
	for _, tablename := range tablenames {
		reader, err = csvOpen(tablename, iskip)
		if err != nil {
			// no file
			continue
		}
		rtable := db.escapetable(tablename)
		sqlstr = rewrite(sqlstr, tablename, rtable)
		reader.Comma = readerComma
		header, err = headerRead(reader)
		if err != nil {
			log.Println(err)
			return 1
		}
		db.Create(rtable, header, ihead)
		err = db.ImportPrepare(rtable, header, ihead)
		if err != nil {
			log.Println(err)
			return 1
		}
		db.Import(reader, header, ihead)
	}

	return trdsql.write(db, sqlstr)
}

func (trdsql TRDSQL) write(db *DDB, sqlstr string) int {
	var err error
	writer := csv.NewWriter(trdsql.outStream)
	writer.Comma, err = getSeparator(trdsql.outSep)
	if err != nil {
		log.Println(err)
	}
	err = db.Output(writer, sqlstr, trdsql.outHeader)
	if err != nil {
		log.Println(err)
		return 1
	}
	return 0
}
<commit_msg>version 0.1.2<commit_after>package main

import (
	"database/sql"
	"encoding/csv"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
)

const VERSION = `0.1.2`

var debug = debugT(false)

type debugT bool

func (d debugT) Printf(format string, args ...interface{}) {
	if d {
		log.Printf(format, args...)
	}
}

// Run is main routine.
func (trdsql TRDSQL) Run(args []string) int {
	var (
		version bool
		odriver string
		odsn    string
		inSep   string
		ihead   bool
		iskip   int
		query   string
		driver  string
		dsn     string
		odebug  bool
	)
	flags := flag.NewFlagSet("trdsql", flag.ContinueOnError)
	driver = "sqlite3"
	dsn = ""
	cfgfile := configOpen()
	cfg, _ := loadConfig(cfgfile)
	flags.Usage = func() {
		fmt.Fprintf(os.Stderr, `
Usage: %s [OPTIONS] [SQL(SELECT...)]

Options:
`, os.Args[0])
		flags.PrintDefaults()
	}

	flags.StringVar(&cfg.Db, "db", cfg.Db, "Specify db name of the setting.")
	flags.StringVar(&odriver, "driver", "", "database driver.  [ "+strings.Join(sql.Drivers(), " | ")+" ]")
	flags.StringVar(&odsn, "dsn", "", "database connection option.")
	flags.StringVar(&inSep, "id", ",", "Field delimiter for input.")
	flags.StringVar(&trdsql.outSep, "od", ",", "Field delimiter for output.")
	flags.BoolVar(&ihead, "ih", false, "The first line is interpreted as column names.")
	flags.BoolVar(&trdsql.outHeader, "oh", false, "Output column name as header.")
	flags.IntVar(&iskip, "is", 0, "Skip header row.")
	flags.StringVar(&query, "q", "", "Read query from the provided filename.")
	flags.BoolVar(&version, "version", false, "display version information.")
	flags.BoolVar(&odebug, "debug", false, "debug print.")
	flags.Parse(args[1:])
	if version {
		fmt.Println(VERSION)
		return (0)
	}
	var sqlstr string
	if query != "" {
		bq, err := ioutil.ReadFile(query)
		if err != nil {
			log.Println("ERROR: ", err)
			return (1)
		}
		sqlstr = string(bq)
	} else {
		sqlstr = strings.Join(flags.Args(), " ")
	}
	if len(sqlstr) == 0 {
		flags.Usage()
		return (2)
	}
	if odebug {
		debug = true
	}
	if strings.HasSuffix(sqlstr, ";") {
		sqlstr = sqlstr[:len(sqlstr)-1]
	}

	if cfg.Db != "" {
		if cfg.Database[cfg.Db].Driver == "" {
			debug.Printf("ERROR: db[%s] does not found", cfg.Db)
		} else {
			driver = cfg.Database[cfg.Db].Driver
			dsn = cfg.Database[cfg.Db].Dsn
		}
	}
	if odriver != "" {
		driver = odriver
	}
	if odsn != "" {
		dsn = odsn
	}

	readerComma, err := getSeparator(inSep)
	if err != nil {
		log.Println(err)
	}
	debug.Printf("driver: %s, dsn: %s", driver, dsn)
	db, err := Connect(driver, dsn)
	if err != nil {
		log.Println("ERROR: ", err)
		return 1
	}
	defer db.Disconnect()

	tablenames := sqlparse(sqlstr)
	if len(tablenames) == 0 {
		// withou FROM clause. ex. SELECT 1+1;
		debug.Printf("table not found\n")
	}
	var reader *csv.Reader
	var header []string
	for _, tablename := range tablenames {
		reader, err = csvOpen(tablename, iskip)
		if err != nil {
			// no file
			continue
		}
		rtable := db.escapetable(tablename)
		sqlstr = rewrite(sqlstr, tablename, rtable)
		reader.Comma = readerComma
		header, err = headerRead(reader)
		if err != nil {
			log.Println(err)
			return 1
		}
		db.Create(rtable, header, ihead)
		err = db.ImportPrepare(rtable, header, ihead)
		if err != nil {
			log.Println(err)
			return 1
		}
		db.Import(reader, header, ihead)
	}

	return trdsql.write(db, sqlstr)
}

func (trdsql TRDSQL) write(db *DDB, sqlstr string) int {
	var err error
	writer := csv.NewWriter(trdsql.outStream)
	writer.Comma, err = getSeparator(trdsql.outSep)
	if err != nil {
		log.Println(err)
	}
	err = db.Output(writer, sqlstr, trdsql.outHeader)
	if err != nil {
		log.Println(err)
		return 1
	}
	return 0
}
<|endoftext|>