<commit_before>/*
Copyright 2012 The Perkeep 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 sqlkv implements the sorted.KeyValue interface using an *sql.DB.
package sqlkv // import "perkeep.org/pkg/sorted/sqlkv"

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"strings"
	"sync"

	"go4.org/syncutil"
	"perkeep.org/internal/leak"
	"perkeep.org/pkg/sorted"
)

// KeyValue implements the sorted.KeyValue interface using an *sql.DB.
type KeyValue struct {
	DB *sql.DB

	// SetFunc is an optional func to use when REPLACE INTO does not exist
	SetFunc      func(*sql.DB, string, string) error
	BatchSetFunc func(*sql.Tx, string, string) error

	// PlaceHolderFunc optionally replaces ? placeholders
	// with the right ones for the rdbms in use.
	PlaceHolderFunc func(string) string

	// Gate optionally limits concurrent access.
	//
	// This originally existed just for SQLite, whose driver likes
	// to return "the database is locked"
	// (perkeep.org/issue/114), so this keeps some pressure
	// off. But we still trust SQLite to deal with concurrency in
	// most cases.
	//
	// It's also used to limit the number of MySQL connections.
	Gate *syncutil.Gate

	// TablePrefix optionally provides a prefix for SQL table
	// names. This is typically "dbname.", ending in a period.
	TablePrefix string

	queriesInitOnce sync.Once // guards initialization of both queries and replacer
	replacer        *strings.Replacer

	queriesMu sync.RWMutex
	queries   map[string]string
}

// sql returns the query, replacing placeholders using PlaceHolderFunc,
// and /*TPRE*/ with TablePrefix.
func (kv *KeyValue) sql(sqlStmt string) string {
	// string manipulation is done only once
	kv.queriesInitOnce.Do(func() {
		kv.queries = make(map[string]string, 8) // we have 8 queries in this file
		kv.replacer = strings.NewReplacer("/*TPRE*/", kv.TablePrefix)
	})
	kv.queriesMu.RLock()
	sqlQuery, ok := kv.queries[sqlStmt]
	kv.queriesMu.RUnlock()
	if ok {
		return sqlQuery
	}
	kv.queriesMu.Lock()
	// check again, now holding the lock
	if sqlQuery, ok = kv.queries[sqlStmt]; ok {
		kv.queriesMu.Unlock()
		return sqlQuery
	}
	sqlQuery = sqlStmt
	if f := kv.PlaceHolderFunc; f != nil {
		sqlQuery = f(sqlQuery)
	}
	sqlQuery = kv.replacer.Replace(sqlQuery)
	kv.queries[sqlStmt] = sqlQuery
	kv.queriesMu.Unlock()
	return sqlQuery
}

type batchTx struct {
	tx  *sql.Tx
	err error // sticky
	kv  *KeyValue
}

func (b *batchTx) Set(key, value string) {
	if b.err != nil {
		return
	}
	if err := sorted.CheckSizes(key, value); err != nil {
		log.Printf("Skipping storing (%q:%q): %v", key, value, err)
		return
	}
	if b.kv.BatchSetFunc != nil {
		b.err = b.kv.BatchSetFunc(b.tx, key, value)
		return
	}
	_, b.err = b.tx.Exec(b.kv.sql("REPLACE INTO /*TPRE*/rows (k, v) VALUES (?, ?)"), key, value)
}

func (b *batchTx) Delete(key string) {
	if b.err != nil {
		return
	}
	_, b.err = b.tx.Exec(b.kv.sql("DELETE FROM /*TPRE*/rows WHERE k=?"), key)
}

func (kv *KeyValue) BeginBatch() sorted.BatchMutation {
	if kv.Gate != nil {
		kv.Gate.Start()
	}
	tx, err := kv.DB.Begin()
	if err != nil {
		log.Printf("SQL BEGIN BATCH: %v", err)
	}
	return &batchTx{
		tx:  tx,
		err: err,
		kv:  kv,
	}
}

func (kv *KeyValue) CommitBatch(b sorted.BatchMutation) error {
	if kv.Gate != nil {
		defer kv.Gate.Done()
	}
	bt, ok := b.(*batchTx)
	if !ok {
		return fmt.Errorf("wrong BatchMutation type %T", b)
	}
	if bt.err != nil {
		if err := bt.tx.Rollback(); err != nil {
			log.Printf("Transaction rollback error: %v", err)
		}
		return bt.err
	}
	return bt.tx.Commit()
}

func (kv *KeyValue) Get(key string) (value string, err error) {
	if kv.Gate != nil {
		kv.Gate.Start()
		defer kv.Gate.Done()
	}
	err = kv.DB.QueryRow(kv.sql("SELECT v FROM /*TPRE*/rows WHERE k=?"), key).Scan(&value)
	if err == sql.ErrNoRows {
		err = sorted.ErrNotFound
	}
	return
}

func (kv *KeyValue) Set(key, value string) error {
	if err := sorted.CheckSizes(key, value); err != nil {
		log.Printf("Skipping storing (%q:%q): %v", key, value, err)
		return nil
	}
	if kv.Gate != nil {
		kv.Gate.Start()
		defer kv.Gate.Done()
	}
	if kv.SetFunc != nil {
		return kv.SetFunc(kv.DB, key, value)
	}
	_, err := kv.DB.Exec(kv.sql("REPLACE INTO /*TPRE*/rows (k, v) VALUES (?, ?)"), key, value)
	return err
}

func (kv *KeyValue) Delete(key string) error {
	if kv.Gate != nil {
		kv.Gate.Start()
		defer kv.Gate.Done()
	}
	_, err := kv.DB.Exec(kv.sql("DELETE FROM /*TPRE*/rows WHERE k=?"), key)
	return err
}

// TODO(mpl): implement Wipe for each of the SQLs, as it's done for MySQL, and
// remove this one below.

func (kv *KeyValue) Wipe() error {
	if kv.Gate != nil {
		kv.Gate.Start()
		defer kv.Gate.Done()
	}
	_, err := kv.DB.Exec(kv.sql("DELETE FROM /*TPRE*/rows"))
	return err
}

func (kv *KeyValue) Close() error { return kv.DB.Close() }

func (kv *KeyValue) Find(start, end string) sorted.Iterator {
	var releaseGate func() // nil if unused
	if kv.Gate != nil {
		kv.Gate.Start()
		releaseGate = kv.Gate.Done
	}
	var rows *sql.Rows
	var err error
	if end == "" {
		rows, err = kv.DB.Query(kv.sql("SELECT k, v FROM /*TPRE*/rows WHERE k >= ? ORDER BY k "), start)
	} else {
		rows, err = kv.DB.Query(kv.sql("SELECT k, v FROM /*TPRE*/rows WHERE k >= ? AND k < ? ORDER BY k "), start, end)
	}
	if err != nil {
		log.Printf("unexpected query error: %v", err)
		return &iter{err: err}
	}

	it := &iter{
		kv:          kv,
		rows:        rows,
		closeCheck:  leak.NewChecker(),
		releaseGate: releaseGate,
	}
	return it
}

// iter is a iterator over sorted key/value pairs in rows.
type iter struct {
	kv  *KeyValue
	err error // accumulated error, returned at Close

	closeCheck  *leak.Checker
	releaseGate func() // if non-nil, called on Close

	rows *sql.Rows // if non-nil, the rows we're reading from

	key        sql.RawBytes
	val        sql.RawBytes
	skey, sval *string // if non-nil, it's been stringified
}

var errClosed = errors.New("sqlkv: Iterator already closed")

func (t *iter) KeyBytes() []byte { return t.key }
func (t *iter) Key() string {
	if t.skey != nil {
		return *t.skey
	}
	str := string(t.key)
	t.skey = &str
	return str
}

func (t *iter) ValueBytes() []byte { return t.val }
func (t *iter) Value() string {
	if t.sval != nil {
		return *t.sval
	}
	str := string(t.val)
	t.sval = &str
	return str
}

func (t *iter) Close() error {
	t.closeCheck.Close()
	if t.rows != nil {
		t.rows.Close()
		t.rows = nil
	}
	if t.releaseGate != nil {
		t.releaseGate()
	}
	err := t.err
	t.err = errClosed
	return err
}

func (t *iter) Next() bool {
	if t.err != nil {
		return false
	}
	t.skey, t.sval = nil, nil
	if !t.rows.Next() {
		return false
	}
	t.err = t.rows.Scan(&t.key, &t.val)
	if t.err != nil {
		log.Printf("unexpected Scan error: %v", t.err)
		return false
	}
	return true
}
<commit_msg>pkg/sorted/sqlkv: make multiple close safe<commit_after>/*
Copyright 2012 The Perkeep 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 sqlkv implements the sorted.KeyValue interface using an *sql.DB.
package sqlkv // import "perkeep.org/pkg/sorted/sqlkv"

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"strings"
	"sync"

	"go4.org/syncutil"
	"perkeep.org/internal/leak"
	"perkeep.org/pkg/sorted"
)

// KeyValue implements the sorted.KeyValue interface using an *sql.DB.
type KeyValue struct {
	DB *sql.DB

	// SetFunc is an optional func to use when REPLACE INTO does not exist
	SetFunc      func(*sql.DB, string, string) error
	BatchSetFunc func(*sql.Tx, string, string) error

	// PlaceHolderFunc optionally replaces ? placeholders
	// with the right ones for the rdbms in use.
	PlaceHolderFunc func(string) string

	// Gate optionally limits concurrent access.
	//
	// This originally existed just for SQLite, whose driver likes
	// to return "the database is locked"
	// (perkeep.org/issue/114), so this keeps some pressure
	// off. But we still trust SQLite to deal with concurrency in
	// most cases.
	//
	// It's also used to limit the number of MySQL connections.
	Gate *syncutil.Gate

	// TablePrefix optionally provides a prefix for SQL table
	// names. This is typically "dbname.", ending in a period.
	TablePrefix string

	queriesInitOnce sync.Once // guards initialization of both queries and replacer
	replacer        *strings.Replacer

	queriesMu sync.RWMutex
	queries   map[string]string
}

// sql returns the query, replacing placeholders using PlaceHolderFunc,
// and /*TPRE*/ with TablePrefix.
func (kv *KeyValue) sql(sqlStmt string) string {
	// string manipulation is done only once
	kv.queriesInitOnce.Do(func() {
		kv.queries = make(map[string]string, 8) // we have 8 queries in this file
		kv.replacer = strings.NewReplacer("/*TPRE*/", kv.TablePrefix)
	})
	kv.queriesMu.RLock()
	sqlQuery, ok := kv.queries[sqlStmt]
	kv.queriesMu.RUnlock()
	if ok {
		return sqlQuery
	}
	kv.queriesMu.Lock()
	// check again, now holding the lock
	if sqlQuery, ok = kv.queries[sqlStmt]; ok {
		kv.queriesMu.Unlock()
		return sqlQuery
	}
	sqlQuery = sqlStmt
	if f := kv.PlaceHolderFunc; f != nil {
		sqlQuery = f(sqlQuery)
	}
	sqlQuery = kv.replacer.Replace(sqlQuery)
	kv.queries[sqlStmt] = sqlQuery
	kv.queriesMu.Unlock()
	return sqlQuery
}

type batchTx struct {
	tx  *sql.Tx
	err error // sticky
	kv  *KeyValue
}

func (b *batchTx) Set(key, value string) {
	if b.err != nil {
		return
	}
	if err := sorted.CheckSizes(key, value); err != nil {
		log.Printf("Skipping storing (%q:%q): %v", key, value, err)
		return
	}
	if b.kv.BatchSetFunc != nil {
		b.err = b.kv.BatchSetFunc(b.tx, key, value)
		return
	}
	_, b.err = b.tx.Exec(b.kv.sql("REPLACE INTO /*TPRE*/rows (k, v) VALUES (?, ?)"), key, value)
}

func (b *batchTx) Delete(key string) {
	if b.err != nil {
		return
	}
	_, b.err = b.tx.Exec(b.kv.sql("DELETE FROM /*TPRE*/rows WHERE k=?"), key)
}

func (kv *KeyValue) BeginBatch() sorted.BatchMutation {
	if kv.Gate != nil {
		kv.Gate.Start()
	}
	tx, err := kv.DB.Begin()
	if err != nil {
		log.Printf("SQL BEGIN BATCH: %v", err)
	}
	return &batchTx{
		tx:  tx,
		err: err,
		kv:  kv,
	}
}

func (kv *KeyValue) CommitBatch(b sorted.BatchMutation) error {
	if kv.Gate != nil {
		defer kv.Gate.Done()
	}
	bt, ok := b.(*batchTx)
	if !ok {
		return fmt.Errorf("wrong BatchMutation type %T", b)
	}
	if bt.err != nil {
		if err := bt.tx.Rollback(); err != nil {
			log.Printf("Transaction rollback error: %v", err)
		}
		return bt.err
	}
	return bt.tx.Commit()
}

func (kv *KeyValue) Get(key string) (value string, err error) {
	if kv.Gate != nil {
		kv.Gate.Start()
		defer kv.Gate.Done()
	}
	err = kv.DB.QueryRow(kv.sql("SELECT v FROM /*TPRE*/rows WHERE k=?"), key).Scan(&value)
	if err == sql.ErrNoRows {
		err = sorted.ErrNotFound
	}
	return
}

func (kv *KeyValue) Set(key, value string) error {
	if err := sorted.CheckSizes(key, value); err != nil {
		log.Printf("Skipping storing (%q:%q): %v", key, value, err)
		return nil
	}
	if kv.Gate != nil {
		kv.Gate.Start()
		defer kv.Gate.Done()
	}
	if kv.SetFunc != nil {
		return kv.SetFunc(kv.DB, key, value)
	}
	_, err := kv.DB.Exec(kv.sql("REPLACE INTO /*TPRE*/rows (k, v) VALUES (?, ?)"), key, value)
	return err
}

func (kv *KeyValue) Delete(key string) error {
	if kv.Gate != nil {
		kv.Gate.Start()
		defer kv.Gate.Done()
	}
	_, err := kv.DB.Exec(kv.sql("DELETE FROM /*TPRE*/rows WHERE k=?"), key)
	return err
}

// TODO(mpl): implement Wipe for each of the SQLs, as it's done for MySQL, and
// remove this one below.

func (kv *KeyValue) Wipe() error {
	if kv.Gate != nil {
		kv.Gate.Start()
		defer kv.Gate.Done()
	}
	_, err := kv.DB.Exec(kv.sql("DELETE FROM /*TPRE*/rows"))
	return err
}

func (kv *KeyValue) Close() error { return kv.DB.Close() }

func (kv *KeyValue) Find(start, end string) sorted.Iterator {
	var releaseGate func() // nil if unused
	if kv.Gate != nil {
		var once sync.Once
		kv.Gate.Start()
		releaseGate = func() {
			once.Do(kv.Gate.Done)
		}
	}
	var rows *sql.Rows
	var err error
	if end == "" {
		rows, err = kv.DB.Query(kv.sql("SELECT k, v FROM /*TPRE*/rows WHERE k >= ? ORDER BY k "), start)
	} else {
		rows, err = kv.DB.Query(kv.sql("SELECT k, v FROM /*TPRE*/rows WHERE k >= ? AND k < ? ORDER BY k "), start, end)
	}
	if err != nil {
		log.Printf("unexpected query error: %v", err)
		return &iter{err: err}
	}

	it := &iter{
		kv:          kv,
		rows:        rows,
		closeCheck:  leak.NewChecker(),
		releaseGate: releaseGate,
	}
	return it
}

// iter is a iterator over sorted key/value pairs in rows.
type iter struct {
	kv  *KeyValue
	err error // accumulated error, returned at Close

	closeCheck  *leak.Checker
	releaseGate func() // if non-nil, called on Close

	rows *sql.Rows // if non-nil, the rows we're reading from

	key        sql.RawBytes
	val        sql.RawBytes
	skey, sval *string // if non-nil, it's been stringified
}

var errClosed = errors.New("sqlkv: Iterator already closed")

func (t *iter) KeyBytes() []byte { return t.key }
func (t *iter) Key() string {
	if t.skey != nil {
		return *t.skey
	}
	str := string(t.key)
	t.skey = &str
	return str
}

func (t *iter) ValueBytes() []byte { return t.val }
func (t *iter) Value() string {
	if t.sval != nil {
		return *t.sval
	}
	str := string(t.val)
	t.sval = &str
	return str
}

func (t *iter) Close() error {
	t.closeCheck.Close()
	if t.rows != nil {
		t.rows.Close()
		t.rows = nil
	}
	if t.releaseGate != nil {
		t.releaseGate()
	}
	err := t.err
	t.err = errClosed
	return err
}

func (t *iter) Next() bool {
	if t.err != nil {
		return false
	}
	t.skey, t.sval = nil, nil
	if !t.rows.Next() {
		return false
	}
	t.err = t.rows.Scan(&t.key, &t.val)
	if t.err != nil {
		log.Printf("unexpected Scan error: %v", t.err)
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package storage

import (
	"strings"

	context "golang.org/x/net/context"
	"google.golang.org/grpc/status"

	opentracing "github.com/opentracing/opentracing-go"
	"gitlab.com/gitlab-org/gitaly/streamio"
	"google.golang.org/grpc/codes"
)

type sshService struct {
	storage Storage
}

// NOTE: #namingThings. And this does more than it should... for "simplicity"
func validateRepoGRERequest(ctx context.Context, s Storage, req *GRERequest, errIn error) (Repository, error) {
	if errIn != nil {
		return nil, status.Errorf(codes.Internal, errIn.Error())
	}
	if len(req.GetId()) == 0 {
		return nil, status.Errorf(codes.FailedPrecondition, "no repo id given")
	}
	if len(req.GetStdin()) != 0 {
		return nil, status.Errorf(codes.FailedPrecondition, "stdin given on first message")
	}
	id := strings.Replace(req.GetId(), "/", "", -1)
	repo, err := s.GetRepository(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.FailedPrecondition, "repo does not exist")
	}
	return repo, nil
}

func (s sshService) UploadPack(stream SSH_UploadPackServer) error {
	span, ctx := opentracing.StartSpanFromContext(stream.Context(), "storage.SSH.UploadPack")
	defer span.Finish()

	req, err := stream.Recv()
	repo, err := validateRepoGRERequest(ctx, s.storage, req, err)
	if err != nil {
		return err
	}

	span.SetTag("repo_hash", repo.GetID())

	ec, err := repo.UploadPack(ctx,
		streamio.NewReader(func() ([]byte, error) {
			request, err := stream.Recv()
			return request.GetStdin(), err
		}),
		streamio.NewWriter(func(p []byte) error {
			span.LogEvent(string(p))
			return stream.Send(&GREResponse{Stdout: p})
		}),
		streamio.NewWriter(func(p []byte) error {
			span.LogEvent(string(p))
			return stream.Send(&GREResponse{Stderr: p})
		}),
	)

	if err != nil {
		return status.Errorf(codes.Internal, err.Error())
	}
	return stream.Send(&GREResponse{ExitCode: &GREExitCode{ExitCode: ec}})
}

func (s sshService) ReceivePack(stream SSH_ReceivePackServer) error {
	span, ctx := opentracing.StartSpanFromContext(stream.Context(), "storage.SSH.ReceivePack")
	defer span.Finish()

	req, err := stream.Recv()
	repo, err := validateRepoGRERequest(ctx, s.storage, req, err)
	if err != nil {
		return err
	}

	span.SetTag("repo_hash", repo.GetID())

	ec, err := repo.ReceivePack(ctx,
		streamio.NewReader(func() ([]byte, error) {
			request, err := stream.Recv()
			span.LogEvent("stdin: " + string(request.GetStdin()))
			return request.GetStdin(), err
		}),
		streamio.NewWriter(func(p []byte) error {
			span.LogEvent("stdout: " + string(p))
			return stream.Send(&GREResponse{Stdout: p})
		}),
		streamio.NewWriter(func(p []byte) error {
			span.LogEvent("stderr: " + string(p))
			return stream.Send(&GREResponse{Stderr: p})
		}),
	)

	if err != nil {
		return status.Errorf(codes.Internal, err.Error())
	}
	return stream.Send(&GREResponse{ExitCode: &GREExitCode{ExitCode: ec}})
}
<commit_msg>remove debug-events..<commit_after>package storage

import (
	"strings"

	context "golang.org/x/net/context"
	"google.golang.org/grpc/status"

	opentracing "github.com/opentracing/opentracing-go"
	"gitlab.com/gitlab-org/gitaly/streamio"
	"google.golang.org/grpc/codes"
)

type sshService struct {
	storage Storage
}

// NOTE: #namingThings. And this does more than it should... for "simplicity"
func validateRepoGRERequest(ctx context.Context, s Storage, req *GRERequest, errIn error) (Repository, error) {
	if errIn != nil {
		return nil, status.Errorf(codes.Internal, errIn.Error())
	}
	if len(req.GetId()) == 0 {
		return nil, status.Errorf(codes.FailedPrecondition, "no repo id given")
	}
	if len(req.GetStdin()) != 0 {
		return nil, status.Errorf(codes.FailedPrecondition, "stdin given on first message")
	}
	id := strings.Replace(req.GetId(), "/", "", -1)
	repo, err := s.GetRepository(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.FailedPrecondition, "repo does not exist")
	}
	return repo, nil
}

func (s sshService) UploadPack(stream SSH_UploadPackServer) error {
	span, ctx := opentracing.StartSpanFromContext(stream.Context(), "storage.SSH.UploadPack")
	defer span.Finish()

	req, err := stream.Recv()
	repo, err := validateRepoGRERequest(ctx, s.storage, req, err)
	if err != nil {
		return err
	}

	span.SetTag("repo_hash", repo.GetID())

	ec, err := repo.UploadPack(ctx,
		streamio.NewReader(func() ([]byte, error) {
			request, err := stream.Recv()
			return request.GetStdin(), err
		}),
		streamio.NewWriter(func(p []byte) error {
			return stream.Send(&GREResponse{Stdout: p})
		}),
		streamio.NewWriter(func(p []byte) error {
			return stream.Send(&GREResponse{Stderr: p})
		}),
	)

	if err != nil {
		return status.Errorf(codes.Internal, err.Error())
	}
	return stream.Send(&GREResponse{ExitCode: &GREExitCode{ExitCode: ec}})
}

func (s sshService) ReceivePack(stream SSH_ReceivePackServer) error {
	span, ctx := opentracing.StartSpanFromContext(stream.Context(), "storage.SSH.ReceivePack")
	defer span.Finish()

	req, err := stream.Recv()
	repo, err := validateRepoGRERequest(ctx, s.storage, req, err)
	if err != nil {
		return err
	}

	span.SetTag("repo_hash", repo.GetID())

	ec, err := repo.ReceivePack(ctx,
		streamio.NewReader(func() ([]byte, error) {
			request, err := stream.Recv()
			return request.GetStdin(), err
		}),
		streamio.NewWriter(func(p []byte) error {
			return stream.Send(&GREResponse{Stdout: p})
		}),
		streamio.NewWriter(func(p []byte) error {
			return stream.Send(&GREResponse{Stderr: p})
		}),
	)

	if err != nil {
		return status.Errorf(codes.Internal, err.Error())
	}
	return stream.Send(&GREResponse{ExitCode: &GREExitCode{ExitCode: ec}})
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package containerdisk

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"

	kubev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"

	ephemeraldisk "kubevirt.io/kubevirt/pkg/ephemeral-disk"

	v1 "kubevirt.io/client-go/api/v1"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
	"kubevirt.io/kubevirt/pkg/util"
)

var containerDiskOwner = "qemu"

var podsBaseDir = util.KubeletPodsDir

var mountBaseDir = filepath.Join(util.VirtShareDir, "/container-disks")

type SocketPathGetter func(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error)

func GetLegacyVolumeMountDirOnHost(vmi *v1.VirtualMachineInstance) string {
	return filepath.Join(mountBaseDir, string(vmi.UID))
}

func GetVolumeMountDirOnGuest(vmi *v1.VirtualMachineInstance) string {
	return filepath.Join(mountBaseDir, string(vmi.UID))
}

func GetVolumeMountDirOnHost(vmi *v1.VirtualMachineInstance) (string, bool, error) {
	for podUID, _ := range vmi.Status.ActivePods {
		basepath := fmt.Sprintf("%s/%s/volumes/kubernetes.io~empty-dir/container-disks", podsBaseDir, string(podUID))
		exists, err := diskutils.FileExists(basepath)
		if err != nil {
			return "", false, err
		} else if exists {
			return basepath, true, nil
		}
	}
	return "", false, nil
}

func GetDiskTargetPathFromHostView(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error) {
	basepath, found, err := GetVolumeMountDirOnHost(vmi)
	if err != nil {
		return "", err
	} else if !found {
		return "", fmt.Errorf("container disk volume for vmi not found")
	}

	return fmt.Sprintf("%s/disk_%d.img", basepath, volumeIndex), nil
}

func GetDiskTargetPathFromLauncherView(volumeIndex int) string {
	return filepath.Join(mountBaseDir, fmt.Sprintf("disk_%d.img", volumeIndex))
}

func SetLocalDirectory(dir string) error {
	mountBaseDir = dir
	return os.MkdirAll(dir, 0755)
}

func SetKubeletPodsDirectory(dir string) {
	podsBaseDir = dir
}

// used for testing - we don't want to MkdirAll on a production host mount
func setPodsDirectory(dir string) error {
	podsBaseDir = dir
	return os.MkdirAll(dir, 0755)
}

// The unit test suite uses this function
func setLocalDataOwner(user string) {
	containerDiskOwner = user
}

// GetDiskTargetPartFromLauncherView returns (path to disk image, image type, and error)
func GetDiskTargetPartFromLauncherView(volumeIndex int) (string, error) {

	path := GetDiskTargetPathFromLauncherView(volumeIndex)
	exists, err := diskutils.FileExists(path)
	if err != nil {
		return "", err
	} else if exists {
		return path, nil
	}

	return "", fmt.Errorf("no supported file disk found for volume with index %d", volumeIndex)
}

// NewSocketPathGetter get the socket pat of a containerDisk. For testing a baseDir
// can be provided which can for instance point to /tmp.
func NewSocketPathGetter(baseDir string) SocketPathGetter {
	return func(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error) {
		for podUID, _ := range vmi.Status.ActivePods {
			basepath := fmt.Sprintf("%s/pods/%s/volumes/kubernetes.io~empty-dir/container-disks", baseDir, string(podUID))
			socketPath := filepath.Join(basepath, fmt.Sprintf("disk_%d.sock", volumeIndex))
			exists, _ := diskutils.FileExists(socketPath)
			if exists {
				return socketPath, nil
			}
		}
		return "", fmt.Errorf("container disk socket path not found for vmi")
	}
}

func GetImage(root string, imagePath string) (string, error) {
	fallbackPath := filepath.Join(root, DiskSourceFallbackPath)
	if imagePath != "" {
		imagePath = filepath.Join(root, imagePath)
		if _, err := os.Stat(imagePath); err != nil {
			if os.IsNotExist(err) {
				return "", fmt.Errorf("No image on path %s", imagePath)
			} else {
				return "", fmt.Errorf("Failed to check if an image exists at %s", imagePath)
			}
		}
	} else {
		files, err := ioutil.ReadDir(fallbackPath)
		if err != nil {
			return "", fmt.Errorf("Failed to check %s for disks: %v", fallbackPath, err)
		}
		if len(files) > 1 {
			return "", fmt.Errorf("More than one file found in folder %s, only one disk is allowed", DiskSourceFallbackPath)
		}
		imagePath = filepath.Join(fallbackPath, files[0].Name())
	}
	return imagePath, nil
}

// The controller uses this function to generate the container
// specs for hosting the container registry disks.
func GenerateContainers(vmi *v1.VirtualMachineInstance, podVolumeName string, binVolumeName string) []kubev1.Container {
	var containers []kubev1.Container

	// Make VirtualMachineInstance Image Wrapper Containers
	for index, volume := range vmi.Spec.Volumes {
		if volume.ContainerDisk != nil {

			volumeMountDir := GetVolumeMountDirOnGuest(vmi)
			diskContainerName := fmt.Sprintf("volume%s", volume.Name)
			diskContainerImage := volume.ContainerDisk.Image
			resources := kubev1.ResourceRequirements{}
			if vmi.IsCPUDedicated() || vmi.WantsToHaveQOSGuaranteed() {
				resources.Limits = make(kubev1.ResourceList)
				resources.Limits[kubev1.ResourceCPU] = resource.MustParse("10m")
				resources.Limits[kubev1.ResourceMemory] = resource.MustParse("40M")
				resources.Requests = make(kubev1.ResourceList)
				resources.Requests[kubev1.ResourceCPU] = resource.MustParse("10m")
				resources.Requests[kubev1.ResourceMemory] = resource.MustParse("40M")
			} else {
				resources.Limits = make(kubev1.ResourceList)
				resources.Limits[kubev1.ResourceCPU] = resource.MustParse("100m")
				resources.Limits[kubev1.ResourceMemory] = resource.MustParse("40M")
				resources.Requests = make(kubev1.ResourceList)
				resources.Requests[kubev1.ResourceCPU] = resource.MustParse("10m")
				resources.Requests[kubev1.ResourceMemory] = resource.MustParse("1M")
			}
			container := kubev1.Container{
				Name:            diskContainerName,
				Image:           diskContainerImage,
				ImagePullPolicy: volume.ContainerDisk.ImagePullPolicy,
				Command:         []string{"/usr/bin/container-disk"},
				Args:            []string{"--copy-path", volumeMountDir + "/disk_" + strconv.Itoa(index)},
				VolumeMounts: []kubev1.VolumeMount{
					{
						Name:      podVolumeName,
						MountPath: volumeMountDir,
					},
					{
						Name:      binVolumeName,
						MountPath: "/usr/bin",
					},
				},
				Resources: resources,
			}

			containers = append(containers, container)
		}
	}
	return containers
}

func CreateEphemeralImages(vmi *v1.VirtualMachineInstance) error {
	// The domain is setup to use the COW image instead of the base image. What we have
	// to do here is only create the image where the domain expects it (GetDiskTargetPartFromLauncherView)
	// for each disk that requires it.

	for i, volume := range vmi.Spec.Volumes {
		if volume.VolumeSource.ContainerDisk != nil {
			if backingFile, err := GetDiskTargetPartFromLauncherView(i); err != nil {
				return err
			} else if err := ephemeraldisk.CreateBackedImageForVolume(volume, backingFile); err != nil {
				return err
			}
		}
	}

	return nil
}
<commit_msg>ensure only we're mounting/unmounting the right pod's container disk during migration<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package containerdisk

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"

	kubev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"

	ephemeraldisk "kubevirt.io/kubevirt/pkg/ephemeral-disk"

	v1 "kubevirt.io/client-go/api/v1"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
	"kubevirt.io/kubevirt/pkg/util"
)

var containerDiskOwner = "qemu"

var podsBaseDir = util.KubeletPodsDir

var mountBaseDir = filepath.Join(util.VirtShareDir, "/container-disks")

type SocketPathGetter func(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error)

func GetLegacyVolumeMountDirOnHost(vmi *v1.VirtualMachineInstance) string {
	return filepath.Join(mountBaseDir, string(vmi.UID))
}

func GetVolumeMountDirOnGuest(vmi *v1.VirtualMachineInstance) string {
	return filepath.Join(mountBaseDir, string(vmi.UID))
}

func GetVolumeMountDirOnHost(vmi *v1.VirtualMachineInstance) (string, bool, error) {
	basepath := ""
	foundEntries := 0
	foundBasepath := ""
	for podUID, _ := range vmi.Status.ActivePods {
		basepath = fmt.Sprintf("%s/%s/volumes/kubernetes.io~empty-dir/container-disks", podsBaseDir, string(podUID))
		exists, err := diskutils.FileExists(basepath)
		if err != nil {
			return "", false, err
		} else if exists {
			foundEntries++
			foundBasepath = basepath
		}
	}

	if foundEntries == 1 {
		return foundBasepath, true, nil
	} else if foundEntries > 1 {
		// Don't mount until outdated pod environments are removed
		// otherwise we might stomp on a previous cleanup
		return "", false, fmt.Errorf("Found multiple pods active for vmi %s/%s. Waiting on outdated pod directories to be removed", vmi.Namespace, vmi.Name)
	}
	return "", false, nil
}

func GetDiskTargetPathFromHostView(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error) {
	basepath, found, err := GetVolumeMountDirOnHost(vmi)
	if err != nil {
		return "", err
	} else if !found {
		return "", fmt.Errorf("container disk volume for vmi not found")
	}

	return fmt.Sprintf("%s/disk_%d.img", basepath, volumeIndex), nil
}

func GetDiskTargetPathFromLauncherView(volumeIndex int) string {
	return filepath.Join(mountBaseDir, fmt.Sprintf("disk_%d.img", volumeIndex))
}

func SetLocalDirectory(dir string) error {
	mountBaseDir = dir
	return os.MkdirAll(dir, 0755)
}

func SetKubeletPodsDirectory(dir string) {
	podsBaseDir = dir
}

// used for testing - we don't want to MkdirAll on a production host mount
func setPodsDirectory(dir string) error {
	podsBaseDir = dir
	return os.MkdirAll(dir, 0755)
}

// The unit test suite uses this function
func setLocalDataOwner(user string) {
	containerDiskOwner = user
}

// GetDiskTargetPartFromLauncherView returns (path to disk image, image type, and error)
func GetDiskTargetPartFromLauncherView(volumeIndex int) (string, error) {

	path := GetDiskTargetPathFromLauncherView(volumeIndex)
	exists, err := diskutils.FileExists(path)
	if err != nil {
		return "", err
	} else if exists {
		return path, nil
	}

	return "", fmt.Errorf("no supported file disk found for volume with index %d", volumeIndex)
}

// NewSocketPathGetter get the socket pat of a containerDisk. For testing a baseDir
// can be provided which can for instance point to /tmp.
func NewSocketPathGetter(baseDir string) SocketPathGetter {
	return func(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error) {
		for podUID, _ := range vmi.Status.ActivePods {
			basepath := fmt.Sprintf("%s/pods/%s/volumes/kubernetes.io~empty-dir/container-disks", baseDir, string(podUID))
			socketPath := filepath.Join(basepath, fmt.Sprintf("disk_%d.sock", volumeIndex))
			exists, _ := diskutils.FileExists(socketPath)
			if exists {
				return socketPath, nil
			}
		}
		return "", fmt.Errorf("container disk socket path not found for vmi")
	}
}

func GetImage(root string, imagePath string) (string, error) {
	fallbackPath := filepath.Join(root, DiskSourceFallbackPath)
	if imagePath != "" {
		imagePath = filepath.Join(root, imagePath)
		if _, err := os.Stat(imagePath); err != nil {
			if os.IsNotExist(err) {
				return "", fmt.Errorf("No image on path %s", imagePath)
			} else {
				return "", fmt.Errorf("Failed to check if an image exists at %s", imagePath)
			}
		}
	} else {
		files, err := ioutil.ReadDir(fallbackPath)
		if err != nil {
			return "", fmt.Errorf("Failed to check %s for disks: %v", fallbackPath, err)
		}
		if len(files) > 1 {
			return "", fmt.Errorf("More than one file found in folder %s, only one disk is allowed", DiskSourceFallbackPath)
		}
		imagePath = filepath.Join(fallbackPath, files[0].Name())
	}
	return imagePath, nil
}

// The controller uses this function to generate the container
// specs for hosting the container registry disks.
func GenerateContainers(vmi *v1.VirtualMachineInstance, podVolumeName string, binVolumeName string) []kubev1.Container {
	var containers []kubev1.Container

	// Make VirtualMachineInstance Image Wrapper Containers
	for index, volume := range vmi.Spec.Volumes {
		if volume.ContainerDisk != nil {

			volumeMountDir := GetVolumeMountDirOnGuest(vmi)
			diskContainerName := fmt.Sprintf("volume%s", volume.Name)
			diskContainerImage := volume.ContainerDisk.Image
			resources := kubev1.ResourceRequirements{}
			if vmi.IsCPUDedicated() || vmi.WantsToHaveQOSGuaranteed() {
				resources.Limits = make(kubev1.ResourceList)
				resources.Limits[kubev1.ResourceCPU] = resource.MustParse("10m")
				resources.Limits[kubev1.ResourceMemory] = resource.MustParse("40M")
				resources.Requests = make(kubev1.ResourceList)
				resources.Requests[kubev1.ResourceCPU] = resource.MustParse("10m")
				resources.Requests[kubev1.ResourceMemory] = resource.MustParse("40M")
			} else {
				resources.Limits = make(kubev1.ResourceList)
				resources.Limits[kubev1.ResourceCPU] = resource.MustParse("100m")
				resources.Limits[kubev1.ResourceMemory] = resource.MustParse("40M")
				resources.Requests = make(kubev1.ResourceList)
				resources.Requests[kubev1.ResourceCPU] = resource.MustParse("10m")
				resources.Requests[kubev1.ResourceMemory] = resource.MustParse("1M")
			}
			container := kubev1.Container{
				Name:            diskContainerName,
				Image:           diskContainerImage,
				ImagePullPolicy: volume.ContainerDisk.ImagePullPolicy,
				Command:         []string{"/usr/bin/container-disk"},
				Args:            []string{"--copy-path", volumeMountDir + "/disk_" + strconv.Itoa(index)},
				VolumeMounts: []kubev1.VolumeMount{
					{
						Name:      podVolumeName,
						MountPath: volumeMountDir,
					},
					{
						Name:      binVolumeName,
						MountPath: "/usr/bin",
					},
				},
				Resources: resources,
			}

			containers = append(containers, container)
		}
	}
	return containers
}

func CreateEphemeralImages(vmi *v1.VirtualMachineInstance) error {
	// The domain is setup to use the COW image instead of the base image. What we have
	// to do here is only create the image where the domain expects it (GetDiskTargetPartFromLauncherView)
	// for each disk that requires it.

	for i, volume := range vmi.Spec.Volumes {
		if volume.VolumeSource.ContainerDisk != nil {
			if backingFile, err := GetDiskTargetPartFromLauncherView(i); err != nil {
				return err
			} else if err := ephemeraldisk.CreateBackedImageForVolume(volume, backingFile); err != nil {
				return err
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package explain

import (
	"bytes"
	"fmt"
	"io"

	"github.com/apprenda/kismatic-platform/pkg/ansible"
	"github.com/apprenda/kismatic-platform/pkg/util"
)

// AnsibleEventStreamExplainer explains the incoming ansible event stream
type AnsibleEventStreamExplainer struct {
	// Out is the destination where the explanations are written
	Out io.Writer
	// Verbose is used to control the output level
	Verbose bool
	// EventExplainer for processing ansible events
	EventExplainer AnsibleEventExplainer
}

// Explain the incoming ansible event stream
func (e *AnsibleEventStreamExplainer) Explain(events <-chan ansible.Event) error {
	for event := range events {
		exp := e.EventExplainer.ExplainEvent(event, e.Verbose)
		if exp != "" {
			fmt.Fprint(e.Out, exp)
		}
	}
	return nil
}

// AnsibleEventExplainer explains a single event
type AnsibleEventExplainer interface {
	ExplainEvent(e ansible.Event, verbose bool) string
}

// DefaultEventExplainer returns the default string explanation of a given event
type DefaultEventExplainer struct {
	// Keeping this state is necessary for supporting the current way of
	// printing output to the console... I am not a fan of this, but it'll
	// do for now...
	printPlayMessage bool
	printPlayStatus  bool
	lastPlay         string
	currentTask      string
}

// ExplainEvent returns an explanation for the given event
func (explainer *DefaultEventExplainer) ExplainEvent(e ansible.Event, verbose bool) string {
	buf := &bytes.Buffer{}
	switch event := e.(type) {
	case *ansible.PlayStartEvent:
		if verbose {
			// In verbose mode the status is printed as a whole line after all the tasks
			// Dont print message before first play
			if explainer.printPlayMessage {
				// No tasks were printed, add a new line: something is wrong
				if explainer.printPlayStatus {
					fmt.Fprintln(buf)
					util.PrintColor(buf, util.Red, "%s Finished with no tasks, are hosts reachable?\n", explainer.lastPlay)
				} else {
					util.PrintColor(buf, util.Green, "%s Finished\n", explainer.lastPlay)
				}
			}
		} else {
			// Do not print status on the first start event or when there is an ERROR
			if explainer.printPlayStatus {
				// In regular mode print the status
				util.PrintOkln(buf)
			}
		}
		// Print the play name
		util.PrettyPrint(buf, event.Name)
		// Set default state for the play
		explainer.lastPlay = event.Name
		explainer.printPlayStatus = true
		explainer.printPlayMessage = true
	case *ansible.RunnerFailedEvent:
		// An error
		// Print newline before first task status
		if explainer.printPlayStatus {
			fmt.Fprintln(buf)
			// Dont print play success status on error
			explainer.printPlayStatus = false
		}
		// Tasks only print at verbose level, on ERROR also print task name
		if !verbose {
			fmt.Fprintf(buf, "- Running task: %s\n", explainer.currentTask)
		}
		if event.IgnoreErrors {
			util.PrettyPrintErrorIgnored(buf, "  %s", event.Host)
		} else {
			util.PrettyPrintErr(buf, "  %s %s", event.Host, event.Result.Message)
		}
		if event.Result.Stdout != "" {
			util.PrintColor(buf, util.Red, "---- STDOUT ----\n%s\n", event.Result.Stdout)
		}
		if event.Result.Stderr != "" {
			util.PrintColor(buf, util.Red, "---- STDERR ----\n%s\n", event.Result.Stderr)
		}
		if event.Result.Stderr != "" || event.Result.Stdout != "" {
			util.PrintColor(buf, util.Red, "---------------\n")
		}
	case *ansible.RunnerUnreachableEvent:
		// Host is unreachable
		// Print newline before first task
		if explainer.printPlayStatus {
			fmt.Fprintln(buf)
			// Dont print play success status on error
			explainer.printPlayStatus = false
		}
		util.PrettyPrintUnreachable(buf, "  %s", event.Host)
	case *ansible.TaskStartEvent:
		if verbose {
			// Print newline before first task status
			if explainer.printPlayStatus {
				fmt.Fprintln(buf)
				// Dont print play success status on error
				explainer.printPlayStatus = false
			}
			fmt.Fprintf(buf, "- Running task: %s\n", event.Name)
		}
		// Set current task name
		explainer.currentTask = event.Name
	case *ansible.HandlerTaskStartEvent:
		if verbose {
			// Print newline before first task
			if explainer.printPlayStatus {
				fmt.Fprintln(buf)
				// Dont print play success status on error
				explainer.printPlayStatus = false
			}
			fmt.Fprintf(buf, "- Running task: %s\n", event.Name)
		}
		// Set current task name
		explainer.currentTask = event.Name
	case *ansible.PlaybookEndEvent:
		// Playbook ends, print the last play status
		if explainer.printPlayStatus {
			if verbose {
				util.PrintColor(buf, util.Green, "%s Finished\n", explainer.lastPlay)
			} else {
				util.PrintOkln(buf)
			}
		}
	case *ansible.RunnerSkippedEvent:
		if verbose {
			util.PrettyPrintSkipped(buf, "  %s", event.Host)
		}
	case *ansible.RunnerOKEvent:
		if verbose {
			util.PrettyPrintOk(buf, "  %s", event.Host)
		}
	case *ansible.RunnerItemOKEvent:
		if verbose {
			util.PrettyPrintOk(buf, "  %s", event.Host)
		}
	case *ansible.RunnerItemRetryEvent:
		return ""
	case *ansible.PlaybookStartEvent:
		return ""
	default:
		if verbose {
			util.PrintColor(buf, util.Orange, "Unhandled event: %T\n", event)
		}
	}
	return buf.String()
}
<commit_msg>KIS-42: Add comments<commit_after>package explain

import (
	"bytes"
	"fmt"
	"io"

	"github.com/apprenda/kismatic-platform/pkg/ansible"
	"github.com/apprenda/kismatic-platform/pkg/util"
)

// AnsibleEventStreamExplainer explains the incoming ansible event stream
type AnsibleEventStreamExplainer struct {
	// Out is the destination where the explanations are written
	Out io.Writer
	// Verbose is used to control the output level
	Verbose bool
	// EventExplainer for processing ansible events
	EventExplainer AnsibleEventExplainer
}

// Explain the incoming ansible event stream
func (e *AnsibleEventStreamExplainer) Explain(events <-chan ansible.Event) error {
	for event := range events {
		exp := e.EventExplainer.ExplainEvent(event, e.Verbose)
		if exp != "" {
			fmt.Fprint(e.Out, exp)
		}
	}
	return nil
}

// AnsibleEventExplainer explains a single event
type AnsibleEventExplainer interface {
	ExplainEvent(e ansible.Event, verbose bool) string
}

// DefaultEventExplainer returns the default string explanation of a given event
type DefaultEventExplainer struct {
	// Keeping this state is necessary for supporting the current way of
	// printing output to the console... I am not a fan of this, but it'll
	// do for now...
	printPlayMessage bool
	printPlayStatus  bool
	lastPlay         string
	currentTask      string
}

// ExplainEvent returns an explanation for the given event
func (explainer *DefaultEventExplainer) ExplainEvent(e ansible.Event, verbose bool) string {
	buf := &bytes.Buffer{}
	switch event := e.(type) {
	case *ansible.PlayStartEvent:
		// On a play start the previos play ends
		// Print a success status, but only when there were no errors
		if verbose {
			// In verbose mode the status is printed as a whole line after all the tasks
			// Dont print message before first play
			if explainer.printPlayMessage {
				// No tasks were printed, add a new line: something is wrong
				if explainer.printPlayStatus {
					fmt.Fprintln(buf)
					util.PrintColor(buf, util.Red, "%s Finished with no tasks, are hosts reachable?\n", explainer.lastPlay)
				} else {
					util.PrintColor(buf, util.Green, "%s Finished\n", explainer.lastPlay)
				}
			}
		} else {
			// Do not print status on the first start event or when there is an ERROR
			if explainer.printPlayStatus {
				// In regular mode print the status
				util.PrintOkln(buf)
			}
		}
		// Print the play name
		util.PrettyPrint(buf, event.Name)
		// Set default state for the play
		explainer.lastPlay = event.Name
		explainer.printPlayStatus = true
		explainer.printPlayMessage = true
	case *ansible.RunnerFailedEvent:
		// An error
		// Print newline before first task status
		if explainer.printPlayStatus {
			fmt.Fprintln(buf)
			// Dont print play success status on error
			explainer.printPlayStatus = false
		}
		// Tasks only print at verbose level, on ERROR also print task name
		if !verbose {
			fmt.Fprintf(buf, "- Running task: %s\n", explainer.currentTask)
		}
		if event.IgnoreErrors {
			util.PrettyPrintErrorIgnored(buf, "  %s", event.Host)
		} else {
			util.PrettyPrintErr(buf, "  %s %s", event.Host, event.Result.Message)
		}
		if event.Result.Stdout != "" {
			util.PrintColor(buf, util.Red, "---- STDOUT ----\n%s\n", event.Result.Stdout)
		}
		if event.Result.Stderr != "" {
			util.PrintColor(buf, util.Red, "---- STDERR ----\n%s\n", event.Result.Stderr)
		}
		if event.Result.Stderr != "" || event.Result.Stdout != "" {
			util.PrintColor(buf, util.Red, "---------------\n")
		}
	case *ansible.RunnerUnreachableEvent:
		// Host is unreachable
		// Print newline before first task
		if explainer.printPlayStatus {
			fmt.Fprintln(buf)
			// Dont print play success status on error
			explainer.printPlayStatus = false
		}
		util.PrettyPrintUnreachable(buf, "  %s", event.Host)
	case *ansible.TaskStartEvent:
		if verbose {
			// Print newline before first task status
			if explainer.printPlayStatus {
				fmt.Fprintln(buf)
				// Dont print play success status on error
				explainer.printPlayStatus = false
			}
			fmt.Fprintf(buf, "- Running task: %s\n", event.Name)
		}
		// Set current task name
		explainer.currentTask = event.Name
	case *ansible.HandlerTaskStartEvent:
		if verbose {
			// Print newline before first task
			if explainer.printPlayStatus {
				fmt.Fprintln(buf)
				// Dont print play success status on error
				explainer.printPlayStatus = false
			}
			fmt.Fprintf(buf, "- Running task: %s\n", event.Name)
		}
		// Set current task name
		explainer.currentTask = event.Name
	case *ansible.PlaybookEndEvent:
		// Playbook ends, print the last play status
		if explainer.printPlayStatus {
			if verbose {
				util.PrintColor(buf, util.Green, "%s Finished\n", explainer.lastPlay)
			} else {
				util.PrintOkln(buf)
			}
		}
	case *ansible.RunnerSkippedEvent:
		if verbose {
			util.PrettyPrintSkipped(buf, "  %s", event.Host)
		}
	case *ansible.RunnerOKEvent:
		if verbose {
			util.PrettyPrintOk(buf, "  %s", event.Host)
		}
	case *ansible.RunnerItemOKEvent:
		if verbose {
			util.PrettyPrintOk(buf, "  %s", event.Host)
		}
	case *ansible.RunnerItemRetryEvent:
		return ""
	case *ansible.PlaybookStartEvent:
		return ""
	default:
		if verbose {
			util.PrintColor(buf, util.Orange, "Unhandled event: %T\n", event)
		}
	}
	return buf.String()
}
<|endoftext|>


<commit_before>package gitlab

import (
	"fmt"
	"net/http"
	"testing"

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

func TestInstanceVariablesService_ListVariables(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/admin/ci/variables", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprintf(w, `
			[
				{
					"key": "TEST_VARIABLE_1",
					"variable_type": "env_var",
					"value": "TEST_1",
					"protected": false,
					"masked": false
				}
			]
		`)
	})

	want := []*InstanceVariable{{
		Key:          "TEST_VARIABLE_1",
		Value:        "TEST_1",
		VariableType: "env_var",
		Protected:    false,
		Masked:       false,
	}}

	ivs, resp, err := client.InstanceVariables.ListVariables(nil)
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.Equal(t, want, ivs)

	ivs, resp, err = client.InstanceVariables.ListVariables(nil, errorOption)
	require.EqualError(t, err, "RequestOptionFunc returns an error")
	require.Nil(t, resp)
	require.Nil(t, ivs)
}

func TestInstanceVariablesService_ListVariables_StatusNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/admin/ci/variables", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		w.WriteHeader(http.StatusNotFound)
	})

	ivs, resp, err := client.InstanceVariables.ListVariables(nil)
	require.Error(t, err)
	require.Nil(t, ivs)
	require.Equal(t, http.StatusNotFound, resp.StatusCode)
}

func TestInstanceVariablesService_GetVariable(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/admin/ci/variables/TEST_VARIABLE_1", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprintf(w, `
			{
				"key": "TEST_VARIABLE_1",
				"variable_type": "env_var",
				"value": "TEST_1",
				"protected": false,
				"masked": false
			}
		`)
	})

	want := &InstanceVariable{
		Key:          "TEST_VARIABLE_1",
		Value:        "TEST_1",
		VariableType: "env_var",
		Protected:    false,
		Masked:       false,
	}

	iv, resp, err := client.InstanceVariables.GetVariable("TEST_VARIABLE_1", nil)
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.Equal(t, want, iv)

	iv, resp, err = client.InstanceVariables.GetVariable("TEST_VARIABLE_1", nil, errorOption)
	require.EqualError(t, err, "RequestOptionFunc returns an error")
	require.Nil(t, resp)
	require.Nil(t, iv)

	iv, resp, err = client.InstanceVariables.GetVariable("TEST_VARIABLE_2", nil)
	require.Error(t, err)
	require.Nil(t, iv)
	require.Equal(t, http.StatusNotFound, resp.StatusCode)
}
<commit_msg>add tests for create, update, delete in instance_variables<commit_after>package gitlab

import (
	"fmt"
	"net/http"
	"testing"

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

func TestInstanceVariablesService_ListVariables(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/admin/ci/variables", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprintf(w, `
			[
				{
					"key": "TEST_VARIABLE_1",
					"variable_type": "env_var",
					"value": "TEST_1",
					"protected": false,
					"masked": false
				}
			]
		`)
	})

	want := []*InstanceVariable{{
		Key:          "TEST_VARIABLE_1",
		Value:        "TEST_1",
		VariableType: "env_var",
		Protected:    false,
		Masked:       false,
	}}

	ivs, resp, err := client.InstanceVariables.ListVariables(nil)
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.Equal(t, want, ivs)

	ivs, resp, err = client.InstanceVariables.ListVariables(nil, errorOption)
	require.EqualError(t, err, "RequestOptionFunc returns an error")
	require.Nil(t, resp)
	require.Nil(t, ivs)
}

func TestInstanceVariablesService_ListVariables_StatusNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/admin/ci/variables", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		w.WriteHeader(http.StatusNotFound)
	})

	ivs, resp, err := client.InstanceVariables.ListVariables(nil)
	require.Error(t, err)
	require.Nil(t, ivs)
	require.Equal(t, http.StatusNotFound, resp.StatusCode)
}

func TestInstanceVariablesService_GetVariable(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/admin/ci/variables/TEST_VARIABLE_1", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprintf(w, `
			{
				"key": "TEST_VARIABLE_1",
				"variable_type": "env_var",
				"value": "TEST_1",
				"protected": false,
				"masked": false
			}
		`)
	})

	want := &InstanceVariable{
		Key:          "TEST_VARIABLE_1",
		Value:        "TEST_1",
		VariableType: "env_var",
		Protected:    false,
		Masked:       false,
	}

	iv, resp, err := client.InstanceVariables.GetVariable("TEST_VARIABLE_1", nil)
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.Equal(t, want, iv)

	iv, resp, err = client.InstanceVariables.GetVariable("TEST_VARIABLE_1", nil, errorOption)
	require.EqualError(t, err, "RequestOptionFunc returns an error")
	require.Nil(t, resp)
	require.Nil(t, iv)

	iv, resp, err = client.InstanceVariables.GetVariable("TEST_VARIABLE_2", nil)
	require.Error(t, err)
	require.Nil(t, iv)
	require.Equal(t, http.StatusNotFound, resp.StatusCode)
}

func TestInstanceVariablesService_CreateVariable(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/admin/ci/variables", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprintf(w, `
			{
				"key": "NEW_VARIABLE",
				"value": "new value",
				"variable_type": "env_var",
				"protected": false,
				"masked": false
			}
		`)
	})

	want := &InstanceVariable{
		Key:          "NEW_VARIABLE",
		Value:        "new value",
		VariableType: "env_var",
		Protected:    false,
		Masked:       false,
	}

	iv, resp, err := client.InstanceVariables.CreateVariable(nil)
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.Equal(t, want, iv)

	iv, resp, err = client.InstanceVariables.CreateVariable(nil, errorOption)
	require.EqualError(t, err, "RequestOptionFunc returns an error")
	require.Nil(t, resp)
	require.Nil(t, iv)
}

func TestInstanceVariablesService_StatusInternalServerError(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/admin/ci/variables", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusInternalServerError)
	})

	iv, resp, err := client.InstanceVariables.CreateVariable(nil)
	require.Error(t, err)
	require.Nil(t, iv)
	require.Equal(t, http.StatusInternalServerError, resp.StatusCode)
}

func TestInstanceVariablesService_UpdateVariable(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/admin/ci/variables/NEW_VARIABLE", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPut)
		fmt.Fprintf(w, `
			{
				"key": "NEW_VARIABLE",
				"value": "updated value",
				"variable_type": "env_var",
				"protected": false,
				"masked": false
			}
		`)
	})

	want := &InstanceVariable{
		Key:          "NEW_VARIABLE",
		Value:        "updated value",
		VariableType: "env_var",
		Protected:    false,
		Masked:       false,
	}

	iv, resp, err := client.InstanceVariables.UpdateVariable("NEW_VARIABLE", nil)
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.Equal(t, want, iv)

	iv, resp, err = client.InstanceVariables.UpdateVariable("NEW_VARIABLE", nil, errorOption)
	require.EqualError(t, err, "RequestOptionFunc returns an error")
	require.Nil(t, resp)
	require.Nil(t, iv)

	iv, resp, err = client.InstanceVariables.UpdateVariable("NEW_VARIABLE_1", nil)
	require.Error(t, err)
	require.Nil(t, iv)
	require.Equal(t, http.StatusNotFound, resp.StatusCode)
}

func TestInstanceVariablesService_RemoveVariable(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/admin/ci/variables/NEW_VARIABLE", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
	})

	resp, err := client.InstanceVariables.RemoveVariable("NEW_VARIABLE", nil)
	require.NoError(t, err)
	require.NotNil(t, resp)

	resp, err = client.InstanceVariables.RemoveVariable("NEW_VARIABLE", nil, errorOption)
	require.EqualError(t, err, "RequestOptionFunc returns an error")
	require.Nil(t, resp)

	resp, err = client.InstanceVariables.RemoveVariable("NEW_VARIABLE_1", nil)
	require.Error(t, err)
	require.Equal(t, http.StatusNotFound, resp.StatusCode)
}
<|endoftext|>


<commit_before>package pinterest

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/url"
	"strings"
)

var (
	endPoint           = "http://api.pinterest.com/v1/urls/count.json?callback=&url="
	errInvalidResponse = errors.New("invalid JSON response")
)

type LinkStats struct {
	Normalized string
	Count      int
}

func GetLinkStats(u string) (*LinkStats, error) {
	if !strings.HasPrefix(u, "http://") && !strings.HasPrefix(u, "https://") {
		u = "http://" + u
	}
	resp, err := Client.Get(endPoint + url.QueryEscape(u))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	// Pinterest will return the results wrapped for a callback. Unfortunately, it
	// seems there's no way to remove that.
	d, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if len(d) < 2 {
		return nil, errInvalidResponse
	}
	d = d[1 : len(d)-1]
	var res map[string]interface{}
	if err := json.Unmarshal(d, &res); err != nil {
		return nil, err
	}
	normalized, ok := res["url"].(string)
	if !ok {
		return nil, errInvalidResponse
	}
	count, ok := res["count"].(float64)
	if !ok {
		return nil, errInvalidResponse
	}
	return &LinkStats{
		Normalized: normalized,
		Count:      int(count),
	}, nil
}

type result struct {
	url   string
	stats *LinkStats
	err   error
}

func GetLinksStats(urls []string) (map[string]*LinkStats, error) {
	count := len(urls)
	ch := make(chan *result, count)
	for _, v := range urls {
		go func(u string) {
			stats, err := GetLinkStats(u)
			ch <- &result{
				url:   u,
				stats: stats,
				err:   err,
			}
		}(v)
	}
	results := make(map[string]*LinkStats, count)
	for ii := 0; ii < len(urls); ii++ {
		res := <-ch
		if res.err != nil {
			return nil, res.err
		}
		results[res.url] = res.stats
	}
	close(ch)
	return results, nil
}
<commit_msg>Fix code for fetching pinterest stats<commit_after>package pinterest

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/url"
	"strings"
)

var (
	// callback can't be empty
	endPoint           = "http://api.pinterest.com/v1/urls/count.json?callback=cb&url="
	errInvalidResponse = errors.New("invalid JSON response")
)

type LinkStats struct {
	Normalized string
	Count      int
}

func GetLinkStats(u string) (*LinkStats, error) {
	if !strings.HasPrefix(u, "http://") && !strings.HasPrefix(u, "https://") {
		u = "http://" + u
	}
	resp, err := Client.Get(endPoint + url.QueryEscape(u))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	// Pinterest will return the results wrapped for a callback. Unfortunately, it
	// seems there's no way to remove that.
	d, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if len(d) < 2 {
		return nil, errInvalidResponse
	}
	d = d[3 : len(d)-1]
	var res map[string]interface{}
	if err := json.Unmarshal(d, &res); err != nil {
		return nil, err
	}
	normalized, ok := res["url"].(string)
	if !ok {
		return nil, errInvalidResponse
	}
	count, ok := res["count"].(float64)
	if !ok {
		return nil, errInvalidResponse
	}
	return &LinkStats{
		Normalized: normalized,
		Count:      int(count),
	}, nil
}

type result struct {
	url   string
	stats *LinkStats
	err   error
}

func GetLinksStats(urls []string) (map[string]*LinkStats, error) {
	count := len(urls)
	ch := make(chan *result, count)
	for _, v := range urls {
		go func(u string) {
			stats, err := GetLinkStats(u)
			ch <- &result{
				url:   u,
				stats: stats,
				err:   err,
			}
		}(v)
	}
	results := make(map[string]*LinkStats, count)
	for ii := 0; ii < len(urls); ii++ {
		res := <-ch
		if res.err != nil {
			return nil, res.err
		}
		results[res.url] = res.stats
	}
	close(ch)
	return results, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Rook Authors. All rights reserved.

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

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

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

// Package config provides default configurations which Rook will set in Ceph clusters.
package config

import (
	"github.com/rook/rook/pkg/operator/ceph/version"
)

// DefaultFlags returns the default configuration flags Rook will set on the command line for all
// calls to Ceph daemons and tools. Values specified here will not be able to be overridden using
// the mon's central KV store, and that is (and should be) by intent.
func DefaultFlags(fsid, mountedKeyringPath string, cephVersion version.CephVersion) []string {
	flags := []string{
		// fsid unnecessary but is a safety to make sure daemons can only connect to their cluster
		NewFlag("fsid", fsid),
		NewFlag("keyring", mountedKeyringPath),
		// For containers, we're expected to log everything to stderr
		NewFlag("log-to-stderr", "true"),
		NewFlag("err-to-stderr", "true"),
		NewFlag("mon-cluster-log-to-stderr", "true"),
		// differentiate debug text from audit text, and the space after 'debug' is critical
		NewFlag("log-stderr-prefix", "debug "),
		NewFlag("default-log-to-file", "false"),
		NewFlag("default-mon-cluster-log-to-file", "false"),
	}

	flags = append(flags, StoredMonHostEnvVarFlags()...)

	return flags
}

// makes it possible to be slightly less verbose to create a ConfigOverride here
func configOverride(who, option, value string) Option {
	return Option{Who: who, Option: option, Value: value}
}

// DefaultCentralizedConfigs returns the default configuration options Rook will set in Ceph's
// centralized config store.
func DefaultCentralizedConfigs(cephVersion version.CephVersion) []Option {
	overrides := []Option{
		configOverride("global", "mon allow pool delete", "true"),
		configOverride("global", "log file", ""),
		configOverride("global", "mon cluster log file", ""),
	}

	// We disable "bluestore warn on legacy statfs"
	// This setting appeared on 14.2.2, so if detected we disable the warning
	// As of 14.2.5 (https://github.com/rook/rook/issues/3539#issuecomment-531287051), Ceph will disable this flag by default so there is no need to apply it
	if cephVersion.IsAtLeast(version.CephVersion{Major: 14, Minor: 2, Extra: 2}) && version.IsInferior(cephVersion, version.CephVersion{Major: 14, Minor: 2, Extra: 5}) {
		overrides = append(overrides, []Option{
			configOverride("global", "bluestore warn on legacy statfs", "false"),
		}...)
	}

	// For Octopus
	if cephVersion.IsAtLeastOctopus() {
		overrides = append(overrides, []Option{
			configOverride("global", "mon allow pool size one", "true"),
		}...)
	}

	return overrides
}

// DefaultLegacyConfigs need to be added to the Ceph config file until the integration tests can be
// made to override these options for the Ceph clusters it creates.
func DefaultLegacyConfigs() []Option {
	overrides := []Option{
		// TODO: drop this when FlexVolume is no longer supported
		configOverride("global", "rbd_default_features", "3"),
	}
	return overrides
}
<commit_msg>ceph: allow size option on pacific only<commit_after>/*
Copyright 2019 The Rook Authors. All rights reserved.

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

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

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

// Package config provides default configurations which Rook will set in Ceph clusters.
package config

import (
	"github.com/rook/rook/pkg/operator/ceph/version"
)

// DefaultFlags returns the default configuration flags Rook will set on the command line for all
// calls to Ceph daemons and tools. Values specified here will not be able to be overridden using
// the mon's central KV store, and that is (and should be) by intent.
func DefaultFlags(fsid, mountedKeyringPath string, cephVersion version.CephVersion) []string {
	flags := []string{
		// fsid unnecessary but is a safety to make sure daemons can only connect to their cluster
		NewFlag("fsid", fsid),
		NewFlag("keyring", mountedKeyringPath),
		// For containers, we're expected to log everything to stderr
		NewFlag("log-to-stderr", "true"),
		NewFlag("err-to-stderr", "true"),
		NewFlag("mon-cluster-log-to-stderr", "true"),
		// differentiate debug text from audit text, and the space after 'debug' is critical
		NewFlag("log-stderr-prefix", "debug "),
		NewFlag("default-log-to-file", "false"),
		NewFlag("default-mon-cluster-log-to-file", "false"),
	}

	flags = append(flags, StoredMonHostEnvVarFlags()...)

	return flags
}

// makes it possible to be slightly less verbose to create a ConfigOverride here
func configOverride(who, option, value string) Option {
	return Option{Who: who, Option: option, Value: value}
}

// DefaultCentralizedConfigs returns the default configuration options Rook will set in Ceph's
// centralized config store.
func DefaultCentralizedConfigs(cephVersion version.CephVersion) []Option {
	overrides := []Option{
		configOverride("global", "mon allow pool delete", "true"),
		configOverride("global", "log file", ""),
		configOverride("global", "mon cluster log file", ""),
	}

	// We disable "bluestore warn on legacy statfs"
	// This setting appeared on 14.2.2, so if detected we disable the warning
	// As of 14.2.5 (https://github.com/rook/rook/issues/3539#issuecomment-531287051), Ceph will disable this flag by default so there is no need to apply it
	if cephVersion.IsAtLeast(version.CephVersion{Major: 14, Minor: 2, Extra: 2}) && version.IsInferior(cephVersion, version.CephVersion{Major: 14, Minor: 2, Extra: 5}) {
		overrides = append(overrides, []Option{
			configOverride("global", "bluestore warn on legacy statfs", "false"),
		}...)
	}

	// For Pacific
	if cephVersion.IsAtLeastPacific() {
		overrides = append(overrides, []Option{
			configOverride("global", "mon allow pool size one", "true"),
		}...)
	}

	return overrides
}

// DefaultLegacyConfigs need to be added to the Ceph config file until the integration tests can be
// made to override these options for the Ceph clusters it creates.
func DefaultLegacyConfigs() []Option {
	overrides := []Option{
		// TODO: drop this when FlexVolume is no longer supported
		configOverride("global", "rbd_default_features", "3"),
	}
	return overrides
}
<|endoftext|>


<commit_before>package postgresql

import (
	"database/sql"
	"strings"
	"strconv"
	"time"
        "github.com/influxdb-zabbix/input"
        //"fmt"

	_"github.com/lib/pq"
)

///
/// utils
///
const (
	millisPerSecond     = int64(time.Second / time.Millisecond)
	nanosPerMillisecond = int64(time.Millisecond / time.Nanosecond)
)

func msToTime(ms string) (time.Time, error) {
	msInt, err := strconv.ParseInt(ms, 10, 64)
	if err != nil {
		   return time.Time{}, err
	}
	return time.Unix(msInt/millisPerSecond,
	  (msInt%millisPerSecond)*nanosPerMillisecond), nil
}

///
/// main methods
///
type Input input.Input

func NewExtracter(address string, tablename string, intputrowsperbatch int, startdate time.Time, enddate time.Time) Input {
    i := Input{}
    i.Address = address
    i.Tablename = tablename
    i.Rowsperbatch = intputrowsperbatch
    i.Startdate = startdate
    i.Enddate = enddate
    return i
}

func (i *Input) Extract() error {

	if i.Address == "" || i.Address == "localhost" {
		i.Address = localhost
	}

	conn, err := sql.Open("postgres", i.Address)
	if err != nil {
		return err
	}
	defer conn.Close()

	query := strings.Replace(
            strings.Replace(
             strings.Replace(
                queries[i.Tablename],
                "##STARTDATE##", strconv.FormatInt(i.Startdate.Unix(), 10), -1),
                "##ENDDATE##", strconv.FormatInt(i.Enddate.Unix(), 10), -1),
                "##ROWSPERBATCH##", strconv.Itoa(i.Rowsperbatch), -1)
  // fmt.Println(
	//  		fmt.Sprintf(
	//  			"------ Query used for %s: %s",
  //        i.Tablename,
	//  			query))

	rows, err := conn.Query(query)
	if err != nil {
		return err
	}
	defer rows.Close()

	// fetch result
  resulttmp := []string{}
  var clock string

	for rows.Next() {
		var result string
		if err := rows.Scan(&result, &clock); err != nil {
			return err
		}
		resulttmp= append(resulttmp, result)
	}

	if err := rows.Err(); err != nil {
		return err
	}

  rows.Close()

	i.Result = resulttmp

  // last clock will be the start date for the next run
  if len(clock) > 0 {
      lastclock, err := msToTime(strings.Trim(clock, " "))
    if err != nil {
      return err
    }
    i.Startdate = lastclock
  }

	return nil
}

var localhost = "host=localhost sslmode=disable"

type MapQuery map[string]string

var queries MapQuery

func init() {
	queries = make(MapQuery)
	queries["history"] = sqlHistory
	queries["history_uint"] = sqlHistoryUInt
	queries["trends"] = sqlTrends
	queries["trends_uint"] = sqlTrendsUInt
}

const sqlOne string = `SELECT 1;`

const sqlTrends string = `SELECT
-- measurement
replace(replace(CASE
    WHEN (position('$2' in ite.name) > 0) AND (position('$4' in ite.name) > 0)
      THEN replace(replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2)), '$4', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 4))
    WHEN (position('$1' in ite.name) > 0) AND (position('$2' in ite.name) > 0)
      THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$1' in ite.name) > 0)
       THEN replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1))
    WHEN (position('$2' in ite.name) > 0)
       THEN replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    ELSE ite.name
  END, ',', ''), ' ', '\ ')
-- tags
|| ',host_name=' || replace(hos.name, ' ', '\ ')
|| ',group_name=' || replace(grp.name, ' ', '\ ')
|| ',applications=' || coalesce(replace((SELECT string_agg(app.name, ' | ')
    FROM public.items_applications iap
    INNER JOIN public.applications app on app.applicationid = iap.applicationid
    WHERE iap.itemid = ite.itemid), ' ', '\ '), 'No application')
|| ' value_min=' || CAST(tre.value_min as varchar(32))
|| ',value_avg=' || CAST(tre.value_avg as varchar(32))
|| ',value_max=' || CAST(tre.value_max as varchar(32))
-- timestamp (in ms)
|| ' ' || CAST((tre.clock * 1000.) as char(14)) as result
,  CAST((tre.clock * 1000.) as char(14)) as clock
FROM public.trends tre
INNER JOIN public.items ite on ite.itemid = tre.itemid
INNER JOIN public.hosts hos on hos.hostid = ite.hostid
INNER JOIN public.hosts_groups hg on hg.hostid = hos.hostid
INNER JOIN public.groups grp on grp.groupid = hg.groupid
WHERE grp.internal=0
   AND tre.clock > ##STARTDATE##
   AND tre.clock <= ##ENDDATE##
ORDER BY tre.clock ASC, tre.itemid ASC
LIMIT ##ROWSPERBATCH##;
`
const sqlTrendsUInt string = `SELECT
-- measurement
replace(replace(CASE
    WHEN (position('$2' in ite.name) > 0) AND (position('$4' in ite.name) > 0)
      THEN replace(replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2)), '$4', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 4))
    WHEN (position('$1' in ite.name) > 0) AND (position('$2' in ite.name) > 0)
      THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$1' in ite.name) > 0)
       THEN replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1))
    WHEN (position('$2' in ite.name) > 0)
       THEN replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$3' in ite.name) > 0)
       THEN replace(ite.name, '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    WHEN (position('$1' in ite.name) > 0) AND (position('$3' in ite.name) > 0)
       THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    ELSE ite.name
  END, ',', ''), ' ', '\ ')
-- tags
|| ',host_name=' || replace(hos.name, ' ', '\ ')
|| ',group_name=' || replace(grp.name, ' ', '\ ')
|| ',applications=' || coalesce(replace((SELECT string_agg(app.name, ' | ')
    FROM public.items_applications iap
    INNER JOIN public.applications app on app.applicationid = iap.applicationid
    WHERE iap.itemid = ite.itemid), ' ', '\ '), 'No application')
|| ' value_min=' || CAST(tre.value_min as varchar(32))
|| ',value_avg=' || CAST(tre.value_avg as varchar(32))
|| ',value_max=' || CAST(tre.value_max as varchar(32))
-- timestamp (in ms)
|| ' ' || CAST((tre.clock * 1000.) as char(14)) as result
,  CAST((tre.clock * 1000.) as char(14)) as clock
FROM public.trends_uint tre
INNER JOIN public.items ite on ite.itemid = tre.itemid
INNER JOIN public.hosts hos on hos.hostid = ite.hostid
INNER JOIN public.hosts_groups hg on hg.hostid = hos.hostid
INNER JOIN public.groups grp on grp.groupid = hg.groupid
WHERE grp.internal=0
   AND tre.clock > ##STARTDATE##
   AND tre.clock <= ##ENDDATE##
ORDER BY tre.clock ASC, tre.itemid ASC
LIMIT ##ROWSPERBATCH##;
`

const sqlHistory string = `SELECT
-- measurement
replace(replace(CASE
    WHEN (position('$2' in ite.name) > 0) AND (position('$4' in ite.name) > 0)
      THEN replace(replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2)), '$4', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 4))
    WHEN (position('$1' in ite.name) > 0) AND (position('$2' in ite.name) > 0)
      THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$1' in ite.name) > 0)
       THEN replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1))
    WHEN (position('$2' in ite.name) > 0)
       THEN replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$3' in ite.name) > 0)
       THEN replace(ite.name, '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    WHEN (position('$1' in ite.name) > 0) AND (position('$3' in ite.name) > 0)
       THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    ELSE ite.name
  END, ',', ''), ' ', '\ ')
-- tags
|| ',host_name=' || replace(hos.name, ' ', '\ ')
|| ',group_name=' || replace(grp.name, ' ', '\ ')
|| ',applications=' || coalesce(replace((SELECT string_agg(app.name, ' | ')
    FROM public.items_applications iap
    INNER JOIN public.applications app on app.applicationid = iap.applicationid
    WHERE iap.itemid = ite.itemid), ' ', '\ '), 'No application')
|| ' value=' || CAST(his.value as varchar(32))
-- timestamp (in ms)
|| ' ' || CAST((his.clock * 1000.) + round(his.ns / 1000000., 0) as char(14)) as result
,  CAST((his.clock * 1000.) + round(his.ns / 1000000., 0) as char(14)) as clock
FROM public.history his
INNER JOIN public.items ite on ite.itemid = his.itemid
INNER JOIN public.hosts hos on hos.hostid = ite.hostid
INNER JOIN public.hosts_groups hg on hg.hostid = hos.hostid
INNER JOIN public.groups grp on grp.groupid = hg.groupid
WHERE grp.internal=0
   AND his.clock > ##STARTDATE##
   AND his.clock <= ##ENDDATE##
ORDER BY his.clock ASC, his.ns ASC, his.itemid ASC
LIMIT ##ROWSPERBATCH##;
`

const sqlHistoryUInt string = `SELECT
-- measurement
replace(replace(CASE
    WHEN (position('$2' in ite.name) > 0) AND (position('$4' in ite.name) > 0)
      THEN replace(replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2)), '$4', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 4))
    WHEN (position('$1' in ite.name) > 0) AND (position('$2' in ite.name) > 0)
      THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$1' in ite.name) > 0)
       THEN replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1))
	  WHEN (position('$2' in ite.name) > 0)
       THEN replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$3' in ite.name) > 0)
       THEN replace(ite.name, '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    WHEN (position('$1' in ite.name) > 0) AND (position('$3' in ite.name) > 0)
       THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    ELSE ite.name
  END, ',', ''), ' ', '\ ')
-- tags
|| ',host_name=' || replace(hos.name, ' ', '\ ')
|| ',group_name=' || replace(grp.name, ' ', '\ ')
|| ',applications=' || coalesce(replace((SELECT string_agg(app.name, ' | ')
    FROM public.items_applications iap
    INNER JOIN public.applications app on app.applicationid = iap.applicationid
    WHERE iap.itemid = ite.itemid), ' ', '\ '), 'No application')
|| ' value=' || CAST(his.value as varchar(32))
-- timestamp (in ms)
|| ' ' || CAST((his.clock * 1000.) + round(his.ns / 1000000., 0) as char(14)) as result
,  CAST((his.clock * 1000.) + round(his.ns / 1000000., 0) as char(14)) as clock
FROM public.history_uint his
INNER JOIN public.items ite on ite.itemid = his.itemid
INNER JOIN public.hosts hos on hos.hostid = ite.hostid
INNER JOIN public.hosts_groups hg on hg.hostid = hos.hostid
INNER JOIN public.groups grp on grp.groupid = hg.groupid
WHERE grp.internal=0
   AND his.clock > ##STARTDATE##
   AND his.clock <= ##ENDDATE##
ORDER BY his.clock ASC, his.ns ASC, his.itemid ASC
LIMIT ##ROWSPERBATCH##;
`
<commit_msg>Update postgresql.go<commit_after>package postgresql

import (
	"database/sql"
	"strings"
	"strconv"
	"time"
        "github.com/zensqlmonitor/influxdb-zabbix/input"
        //"fmt"

	_"github.com/lib/pq"
)

///
/// utils
///
const (
	millisPerSecond     = int64(time.Second / time.Millisecond)
	nanosPerMillisecond = int64(time.Millisecond / time.Nanosecond)
)

func msToTime(ms string) (time.Time, error) {
	msInt, err := strconv.ParseInt(ms, 10, 64)
	if err != nil {
		   return time.Time{}, err
	}
	return time.Unix(msInt/millisPerSecond,
	  (msInt%millisPerSecond)*nanosPerMillisecond), nil
}

///
/// main methods
///
type Input input.Input

func NewExtracter(address string, tablename string, intputrowsperbatch int, startdate time.Time, enddate time.Time) Input {
    i := Input{}
    i.Address = address
    i.Tablename = tablename
    i.Rowsperbatch = intputrowsperbatch
    i.Startdate = startdate
    i.Enddate = enddate
    return i
}

func (i *Input) Extract() error {

	if i.Address == "" || i.Address == "localhost" {
		i.Address = localhost
	}

	conn, err := sql.Open("postgres", i.Address)
	if err != nil {
		return err
	}
	defer conn.Close()

	query := strings.Replace(
            strings.Replace(
             strings.Replace(
                queries[i.Tablename],
                "##STARTDATE##", strconv.FormatInt(i.Startdate.Unix(), 10), -1),
                "##ENDDATE##", strconv.FormatInt(i.Enddate.Unix(), 10), -1),
                "##ROWSPERBATCH##", strconv.Itoa(i.Rowsperbatch), -1)
  // fmt.Println(
	//  		fmt.Sprintf(
	//  			"------ Query used for %s: %s",
  //        i.Tablename,
	//  			query))

	rows, err := conn.Query(query)
	if err != nil {
		return err
	}
	defer rows.Close()

	// fetch result
  resulttmp := []string{}
  var clock string

	for rows.Next() {
		var result string
		if err := rows.Scan(&result, &clock); err != nil {
			return err
		}
		resulttmp= append(resulttmp, result)
	}

	if err := rows.Err(); err != nil {
		return err
	}

  rows.Close()

	i.Result = resulttmp

  // last clock will be the start date for the next run
  if len(clock) > 0 {
      lastclock, err := msToTime(strings.Trim(clock, " "))
    if err != nil {
      return err
    }
    i.Startdate = lastclock
  }

	return nil
}

var localhost = "host=localhost sslmode=disable"

type MapQuery map[string]string

var queries MapQuery

func init() {
	queries = make(MapQuery)
	queries["history"] = sqlHistory
	queries["history_uint"] = sqlHistoryUInt
	queries["trends"] = sqlTrends
	queries["trends_uint"] = sqlTrendsUInt
}

const sqlOne string = `SELECT 1;`

const sqlTrends string = `SELECT
-- measurement
replace(replace(CASE
    WHEN (position('$2' in ite.name) > 0) AND (position('$4' in ite.name) > 0)
      THEN replace(replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2)), '$4', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 4))
    WHEN (position('$1' in ite.name) > 0) AND (position('$2' in ite.name) > 0)
      THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$1' in ite.name) > 0)
       THEN replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1))
    WHEN (position('$2' in ite.name) > 0)
       THEN replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    ELSE ite.name
  END, ',', ''), ' ', '\ ')
-- tags
|| ',host_name=' || replace(hos.name, ' ', '\ ')
|| ',group_name=' || replace(grp.name, ' ', '\ ')
|| ',applications=' || coalesce(replace((SELECT string_agg(app.name, ' | ')
    FROM public.items_applications iap
    INNER JOIN public.applications app on app.applicationid = iap.applicationid
    WHERE iap.itemid = ite.itemid), ' ', '\ '), 'No application')
|| ' value_min=' || CAST(tre.value_min as varchar(32))
|| ',value_avg=' || CAST(tre.value_avg as varchar(32))
|| ',value_max=' || CAST(tre.value_max as varchar(32))
-- timestamp (in ms)
|| ' ' || CAST((tre.clock * 1000.) as char(14)) as result
,  CAST((tre.clock * 1000.) as char(14)) as clock
FROM public.trends tre
INNER JOIN public.items ite on ite.itemid = tre.itemid
INNER JOIN public.hosts hos on hos.hostid = ite.hostid
INNER JOIN public.hosts_groups hg on hg.hostid = hos.hostid
INNER JOIN public.groups grp on grp.groupid = hg.groupid
WHERE grp.internal=0
   AND tre.clock > ##STARTDATE##
   AND tre.clock <= ##ENDDATE##
ORDER BY tre.clock ASC, tre.itemid ASC
LIMIT ##ROWSPERBATCH##;
`
const sqlTrendsUInt string = `SELECT
-- measurement
replace(replace(CASE
    WHEN (position('$2' in ite.name) > 0) AND (position('$4' in ite.name) > 0)
      THEN replace(replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2)), '$4', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 4))
    WHEN (position('$1' in ite.name) > 0) AND (position('$2' in ite.name) > 0)
      THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$1' in ite.name) > 0)
       THEN replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1))
    WHEN (position('$2' in ite.name) > 0)
       THEN replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$3' in ite.name) > 0)
       THEN replace(ite.name, '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    WHEN (position('$1' in ite.name) > 0) AND (position('$3' in ite.name) > 0)
       THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    ELSE ite.name
  END, ',', ''), ' ', '\ ')
-- tags
|| ',host_name=' || replace(hos.name, ' ', '\ ')
|| ',group_name=' || replace(grp.name, ' ', '\ ')
|| ',applications=' || coalesce(replace((SELECT string_agg(app.name, ' | ')
    FROM public.items_applications iap
    INNER JOIN public.applications app on app.applicationid = iap.applicationid
    WHERE iap.itemid = ite.itemid), ' ', '\ '), 'No application')
|| ' value_min=' || CAST(tre.value_min as varchar(32))
|| ',value_avg=' || CAST(tre.value_avg as varchar(32))
|| ',value_max=' || CAST(tre.value_max as varchar(32))
-- timestamp (in ms)
|| ' ' || CAST((tre.clock * 1000.) as char(14)) as result
,  CAST((tre.clock * 1000.) as char(14)) as clock
FROM public.trends_uint tre
INNER JOIN public.items ite on ite.itemid = tre.itemid
INNER JOIN public.hosts hos on hos.hostid = ite.hostid
INNER JOIN public.hosts_groups hg on hg.hostid = hos.hostid
INNER JOIN public.groups grp on grp.groupid = hg.groupid
WHERE grp.internal=0
   AND tre.clock > ##STARTDATE##
   AND tre.clock <= ##ENDDATE##
ORDER BY tre.clock ASC, tre.itemid ASC
LIMIT ##ROWSPERBATCH##;
`

const sqlHistory string = `SELECT
-- measurement
replace(replace(CASE
    WHEN (position('$2' in ite.name) > 0) AND (position('$4' in ite.name) > 0)
      THEN replace(replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2)), '$4', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 4))
    WHEN (position('$1' in ite.name) > 0) AND (position('$2' in ite.name) > 0)
      THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$1' in ite.name) > 0)
       THEN replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1))
    WHEN (position('$2' in ite.name) > 0)
       THEN replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$3' in ite.name) > 0)
       THEN replace(ite.name, '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    WHEN (position('$1' in ite.name) > 0) AND (position('$3' in ite.name) > 0)
       THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    ELSE ite.name
  END, ',', ''), ' ', '\ ')
-- tags
|| ',host_name=' || replace(hos.name, ' ', '\ ')
|| ',group_name=' || replace(grp.name, ' ', '\ ')
|| ',applications=' || coalesce(replace((SELECT string_agg(app.name, ' | ')
    FROM public.items_applications iap
    INNER JOIN public.applications app on app.applicationid = iap.applicationid
    WHERE iap.itemid = ite.itemid), ' ', '\ '), 'No application')
|| ' value=' || CAST(his.value as varchar(32))
-- timestamp (in ms)
|| ' ' || CAST((his.clock * 1000.) + round(his.ns / 1000000., 0) as char(14)) as result
,  CAST((his.clock * 1000.) + round(his.ns / 1000000., 0) as char(14)) as clock
FROM public.history his
INNER JOIN public.items ite on ite.itemid = his.itemid
INNER JOIN public.hosts hos on hos.hostid = ite.hostid
INNER JOIN public.hosts_groups hg on hg.hostid = hos.hostid
INNER JOIN public.groups grp on grp.groupid = hg.groupid
WHERE grp.internal=0
   AND his.clock > ##STARTDATE##
   AND his.clock <= ##ENDDATE##
ORDER BY his.clock ASC, his.ns ASC, his.itemid ASC
LIMIT ##ROWSPERBATCH##;
`

const sqlHistoryUInt string = `SELECT
-- measurement
replace(replace(CASE
    WHEN (position('$2' in ite.name) > 0) AND (position('$4' in ite.name) > 0)
      THEN replace(replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2)), '$4', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 4))
    WHEN (position('$1' in ite.name) > 0) AND (position('$2' in ite.name) > 0)
      THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$1' in ite.name) > 0)
       THEN replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1))
	  WHEN (position('$2' in ite.name) > 0)
       THEN replace(ite.name, '$2', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 2))
    WHEN (position('$3' in ite.name) > 0)
       THEN replace(ite.name, '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    WHEN (position('$1' in ite.name) > 0) AND (position('$3' in ite.name) > 0)
       THEN replace(replace(ite.name, '$1', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 1)), '$3', split_part(substring(ite.key_ FROM '\[(.+)\]'), ',', 3))
    ELSE ite.name
  END, ',', ''), ' ', '\ ')
-- tags
|| ',host_name=' || replace(hos.name, ' ', '\ ')
|| ',group_name=' || replace(grp.name, ' ', '\ ')
|| ',applications=' || coalesce(replace((SELECT string_agg(app.name, ' | ')
    FROM public.items_applications iap
    INNER JOIN public.applications app on app.applicationid = iap.applicationid
    WHERE iap.itemid = ite.itemid), ' ', '\ '), 'No application')
|| ' value=' || CAST(his.value as varchar(32))
-- timestamp (in ms)
|| ' ' || CAST((his.clock * 1000.) + round(his.ns / 1000000., 0) as char(14)) as result
,  CAST((his.clock * 1000.) + round(his.ns / 1000000., 0) as char(14)) as clock
FROM public.history_uint his
INNER JOIN public.items ite on ite.itemid = his.itemid
INNER JOIN public.hosts hos on hos.hostid = ite.hostid
INNER JOIN public.hosts_groups hg on hg.hostid = hos.hostid
INNER JOIN public.groups grp on grp.groupid = hg.groupid
WHERE grp.internal=0
   AND his.clock > ##STARTDATE##
   AND his.clock <= ##ENDDATE##
ORDER BY his.clock ASC, his.ns ASC, his.itemid ASC
LIMIT ##ROWSPERBATCH##;
`
<|endoftext|>


<commit_before>package cli

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"time"

	"github.com/alecthomas/kingpin"
	"github.com/prometheus/alertmanager/cli/format"
	"github.com/prometheus/alertmanager/types"
)

type getResponse struct {
	Status    string        `json:"status"`
	Data      types.Silence `json:"data,omitempty"`
	ErrorType string        `json:"errorType,omitempty"`
	Error     string        `json:"error,omitempty"`
}

var (
	updateCmd       = silenceCmd.Command("update", "Update silences")
	updateExpires   = updateCmd.Flag("expires", "Duration of silence").Short('e').Duration()
	updateExpiresOn = updateCmd.Flag("expire-on", "Expire at a certain time (Overwrites expires) RFC3339 format 2006-01-02T15:04:05Z07:00").Time(time.RFC3339)
	updateComment   = updateCmd.Flag("comment", "A comment to help describe the silence").Short('c').String()
	updateIds       = updateCmd.Arg("update-ids", "Silence IDs to update").Strings()
)

func init() {
	updateCmd.Action(update)
	longHelpText["silence update"] = `Extend or update existing silence in Alertmanager.`
}

func update(element *kingpin.ParseElement, ctx *kingpin.ParseContext) error {
	if len(*updateIds) < 1 {
		return fmt.Errorf("no silence IDs specified")
	}

	alertmanagerUrl := GetAlertmanagerURL("/api/v1/silence")
	var updatedSilences []types.Silence
	for _, silenceId := range *updateIds {
		silence, err := getSilenceById(silenceId, alertmanagerUrl)
		if err != nil {
			return err
		}
		silence, err = updateSilence(silence)
		if err != nil {
			return err
		}
		updatedSilences = append(updatedSilences, *silence)
	}

	if *silenceQuiet {
		for _, silence := range updatedSilences {
			fmt.Println(silence.ID)
		}
	} else {
		formatter, found := format.Formatters[*output]
		if !found {
			return fmt.Errorf("unknown output formatter")
		}
		formatter.FormatSilences(updatedSilences)
	}
	return nil
}

// This takes an url.URL and not a pointer as we will modify it for our API call.
func getSilenceById(silenceId string, baseUrl url.URL) (*types.Silence, error) {
	baseUrl.Path = path.Join(baseUrl.Path, silenceId)
	res, err := http.Get(baseUrl.String())
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, fmt.Errorf("couldn't read response body: %v", err)
	}

	if res.StatusCode == 404 {
		return nil, fmt.Errorf("no silence found with id: %v", silenceId)
	}
	if res.StatusCode != 200 {
		return nil, fmt.Errorf("received %d response from Alertmanager: %v", res.StatusCode, body)
	}

	var response getResponse
	err = json.Unmarshal(body, &response)
	return &response.Data, nil
}

func updateSilence(silence *types.Silence) (*types.Silence, error) {
	if *updateExpires != 0 {
		silence.EndsAt = time.Now().UTC().Add(*updateExpires)
	}

	// expire-on will override expires value if both are specified
	if !(*updateExpiresOn).IsZero() {
		silence.EndsAt = *updateExpiresOn
	}

	if *updateComment != "" {
		silence.Comment = *updateComment
	}

	// addSilence can also be used to update an existing silence
	_, err := addSilence(silence)
	if err != nil {
		return nil, err
	}
	return silence, nil
}
<commit_msg>amtool silence update support dwy suffixes to expire flag (#1197)<commit_after>package cli

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"time"

	"github.com/alecthomas/kingpin"
	"github.com/prometheus/alertmanager/cli/format"
	"github.com/prometheus/alertmanager/types"
	"github.com/prometheus/common/model"
)

type getResponse struct {
	Status    string        `json:"status"`
	Data      types.Silence `json:"data,omitempty"`
	ErrorType string        `json:"errorType,omitempty"`
	Error     string        `json:"error,omitempty"`
}

var (
	updateCmd       = silenceCmd.Command("update", "Update silences")
	updateExpires   = updateCmd.Flag("expires", "Duration of silence").Short('e').Default("1h").String()
	updateExpiresOn = updateCmd.Flag("expire-on", "Expire at a certain time (Overwrites expires) RFC3339 format 2006-01-02T15:04:05Z07:00").Time(time.RFC3339)
	updateComment   = updateCmd.Flag("comment", "A comment to help describe the silence").Short('c').String()
	updateIds       = updateCmd.Arg("update-ids", "Silence IDs to update").Strings()
)

func init() {
	updateCmd.Action(update)
	longHelpText["silence update"] = `Extend or update existing silence in Alertmanager.`
}

func update(element *kingpin.ParseElement, ctx *kingpin.ParseContext) error {
	if len(*updateIds) < 1 {
		return fmt.Errorf("no silence IDs specified")
	}

	alertmanagerUrl := GetAlertmanagerURL("/api/v1/silence")
	var updatedSilences []types.Silence
	for _, silenceId := range *updateIds {
		silence, err := getSilenceById(silenceId, alertmanagerUrl)
		if err != nil {
			return err
		}
		silence, err = updateSilence(silence)
		if err != nil {
			return err
		}
		updatedSilences = append(updatedSilences, *silence)
	}

	if *silenceQuiet {
		for _, silence := range updatedSilences {
			fmt.Println(silence.ID)
		}
	} else {
		formatter, found := format.Formatters[*output]
		if !found {
			return fmt.Errorf("unknown output formatter")
		}
		formatter.FormatSilences(updatedSilences)
	}
	return nil
}

// This takes an url.URL and not a pointer as we will modify it for our API call.
func getSilenceById(silenceId string, baseUrl url.URL) (*types.Silence, error) {
	baseUrl.Path = path.Join(baseUrl.Path, silenceId)
	res, err := http.Get(baseUrl.String())
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, fmt.Errorf("couldn't read response body: %v", err)
	}

	if res.StatusCode == 404 {
		return nil, fmt.Errorf("no silence found with id: %v", silenceId)
	}
	if res.StatusCode != 200 {
		return nil, fmt.Errorf("received %d response from Alertmanager: %v", res.StatusCode, body)
	}

	var response getResponse
	err = json.Unmarshal(body, &response)
	return &response.Data, nil
}

func updateSilence(silence *types.Silence) (*types.Silence, error) {
	if *updateExpires != "" {
		duration, err := model.ParseDuration(*updateExpires)
		if err != nil {
			return nil, err
		}
		if duration == 0 {
			return nil, fmt.Errorf("silence duration must be greater than 0")
		}
		silence.EndsAt = time.Now().UTC().Add(time.Duration(duration))
	}

	// expire-on will override expires value if both are specified
	if !(*updateExpiresOn).IsZero() {
		silence.EndsAt = *updateExpiresOn
	}

	if *updateComment != "" {
		silence.Comment = *updateComment
	}

	// addSilence can also be used to update an existing silence
	_, err := addSilence(silence)
	if err != nil {
		return nil, err
	}
	return silence, nil
}
<|endoftext|>