<commit_before>package sqlite

import (
	"database/sql"
	"os"
	"strings"
	"testing"
	"time"
)

func TestOpen(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("drop table foo")
	_, err = db.Exec("create table foo (id integer)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	if stat, err := os.Stat("./foo.db"); err != nil || stat.IsDir() {
		t.Error("Failed to create ./foo.db")
	}
}

func TestInsert(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("drop table foo")
	_, err = db.Exec("create table foo (id integer)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	res, err := db.Exec("insert into foo(id) values(123)")
	if err != nil {
		t.Fatal("Failed to insert record:", err)
	}
	affected, _ := res.RowsAffected()
	if affected != 1 {
		t.Fatalf("Expected %d for affected rows, but %d:", 1, affected)
	}

	rows, err := db.Query("select id from foo")
	if err != nil {
		t.Fatal("Failed to select records:", err)
	}
	defer rows.Close()

	rows.Next()

	var result int
	rows.Scan(&result)
	if result != 123 {
		t.Errorf("Fetched %q; expected %q", 123, result)
	}
}

func TestUpdate(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("drop table foo")
	_, err = db.Exec("create table foo (id integer)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	res, err := db.Exec("insert into foo(id) values(123)")
	if err != nil {
		t.Fatal("Failed to insert record:", err)
	}
	expected, err := res.LastInsertId()
	if err != nil {
		t.Fatal("Failed to get LastInsertId:", err)
	}
	affected, _ := res.RowsAffected()
	if err != nil {
		t.Fatal("Failed to get RowsAffected:", err)
	}
	if affected != 1 {
		t.Fatalf("Expected %d for affected rows, but %d:", 1, affected)
	}

	res, err = db.Exec("update foo set id = 234")
	if err != nil {
		t.Fatal("Failed to update record:", err)
	}
	lastId, err := res.LastInsertId()
	if err != nil {
		t.Fatal("Failed to get LastInsertId:", err)
	}
	if expected != lastId {
		t.Errorf("Expected %q for last Id, but %q:", expected, lastId)
	}
	affected, _ = res.RowsAffected()
	if err != nil {
		t.Fatal("Failed to get RowsAffected:", err)
	}
	if affected != 1 {
		t.Fatalf("Expected %d for affected rows, but %d:", 1, affected)
	}

	rows, err := db.Query("select id from foo")
	if err != nil {
		t.Fatal("Failed to select records:", err)
	}
	defer rows.Close()

	rows.Next()

	var result int
	rows.Scan(&result)
	if result != 234 {
		t.Errorf("Fetched %q; expected %q", 234, result)
	}
}

func TestDelete(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("drop table foo")
	_, err = db.Exec("create table foo (id integer)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	res, err := db.Exec("insert into foo(id) values(123)")
	if err != nil {
		t.Fatal("Failed to insert record:", err)
	}
	expected, err := res.LastInsertId()
	if err != nil {
		t.Fatal("Failed to get LastInsertId:", err)
	}
	affected, err := res.RowsAffected()
	if err != nil {
		t.Fatal("Failed to get RowsAffected:", err)
	}
	if affected != 1 {
		t.Errorf("Expected %d for cout of affected rows, but %q:", 1, affected)
	}

	res, err = db.Exec("delete from foo where id = 123")
	if err != nil {
		t.Fatal("Failed to delete record:", err)
	}
	lastId, err := res.LastInsertId()
	if err != nil {
		t.Fatal("Failed to get LastInsertId:", err)
	}
	if expected != lastId {
		t.Errorf("Expected %q for last Id, but %q:", expected, lastId)
	}
	affected, err = res.RowsAffected()
	if err != nil {
		t.Fatal("Failed to get RowsAffected:", err)
	}
	if affected != 1 {
		t.Errorf("Expected %d for cout of affected rows, but %q:", 1, affected)
	}

	rows, err := db.Query("select id from foo")
	if err != nil {
		t.Fatal("Failed to select records:", err)
	}
	defer rows.Close()

	if rows.Next() {
		t.Error("Fetched row but expected not rows")
	}
}

func TestBooleanRoundtrip(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("DROP TABLE foo")
	_, err = db.Exec("CREATE TABLE foo(id INTEGER, value BOOL)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	_, err = db.Exec("INSERT INTO foo(id, value) VALUES(1, ?)", true)
	if err != nil {
		t.Fatal("Failed to insert true value:", err)
	}

	_, err = db.Exec("INSERT INTO foo(id, value) VALUES(2, ?)", false)
	if err != nil {
		t.Fatal("Failed to insert false value:", err)
	}

	rows, err := db.Query("SELECT id, value FROM foo")
	if err != nil {
		t.Fatal("Unable to query foo table:", err)
	}
	defer rows.Close()

	for rows.Next() {
		var id int
		var value bool

		if err := rows.Scan(&id, &value); err != nil {
			t.Error("Unable to scan results:", err)
			continue
		}

		if id == 1 && !value {
			t.Error("Value for id 1 should be true, not false")

		} else if id == 2 && value {
			t.Error("Value for id 2 should be false, not true")
		}
	}
}

func TestTimestamp(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("DROP TABLE foo")
	_, err = db.Exec("CREATE TABLE foo(id INTEGER, ts timeSTAMP)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	timestamp1 := time.Date(2012, time.April, 6, 22, 50, 0, 0, time.UTC)
	_, err = db.Exec("INSERT INTO foo(id, ts) VALUES(1, ?)", timestamp1)
	if err != nil {
		t.Fatal("Failed to insert timestamp:", err)
	}

	timestamp2 := time.Date(2012, time.April, 6, 23, 22, 0, 0, time.UTC)
	_, err = db.Exec("INSERT INTO foo(id, ts) VALUES(2, ?)", timestamp2.Unix())
	if err != nil {
		t.Fatal("Failed to insert timestamp:", err)
	}

	_, err = db.Exec("INSERT INTO foo(id, ts) VALUES(3, ?)", "nonsense")
	if err != nil {
		t.Fatal("Failed to insert nonsense:", err)
	}

	rows, err := db.Query("SELECT id, ts FROM foo ORDER BY id ASC")
	if err != nil {
		t.Fatal("Unable to query foo table:", err)
	}
	defer rows.Close()

	seen := 0
	for rows.Next() {
		var id int
		var ts time.Time

		if err := rows.Scan(&id, &ts); err != nil {
			t.Error("Unable to scan results:", err)
			continue
		}

		if id == 1 {
			seen += 1
			if !timestamp1.Equal(ts) {
				t.Errorf("Value for id 1 should be %v, not %v", timestamp1, ts)
			}
		}

		if id == 2 {
			seen += 1
			if !timestamp2.Equal(ts) {
				t.Errorf("Value for id 2 should be %v, not %v", timestamp2, ts)
			}
		}
	}

	if seen != 2 {
		t.Error("Expected to see two valid timestamps")
	}

	// make sure "nonsense" triggered an error
	err = rows.Err()
	if err == nil || !strings.Contains(err.Error(), "cannot parse \"nonsense\"") {
		t.Error("Expected error from \"nonsense\" timestamp")
	}
}

func TestBoolean(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}

	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("CREATE TABLE foo(id INTEGER, fbool BOOLEAN)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	bool1 := true
	_, err = db.Exec("INSERT INTO foo(id, fbool) VALUES(1, ?)", bool1)
	if err != nil {
		t.Fatal("Failed to insert boolean:", err)
	}

	bool2 := false
	_, err = db.Exec("INSERT INTO foo(id, fbool) VALUES(2, ?)", bool2)
	if err != nil {
		t.Fatal("Failed to insert boolean:", err)
	}

	bool3 := "nonsense"
	_, err = db.Exec("INSERT INTO foo(id, fbool) VALUES(3, ?)", bool3)
	if err != nil {
		t.Fatal("Failed to insert nonsense:", err)
	}

	rows, err := db.Query("SELECT id, fbool FROM foo where fbool is ?", bool1)
	if err != nil {
		t.Fatal("Unable to query foo table:", err)
	}
	counter := 0

	var id int
	var fbool bool

	for rows.Next() {
		if err := rows.Scan(&id, &fbool); err != nil {
			t.Fatal("Unable to scan results:", err)
		}
		counter++
	}

	if counter != 1 {
		t.Fatalf("Expected 1 row but %v", counter)
	}

	if id != 1 && fbool != true {
		t.Fatalf("Value for id 1 should be %v, not %v", bool1, fbool)
	}

	rows, err = db.Query("SELECT id, fbool FROM foo where fbool is ?", bool2)
	if err != nil {
		t.Fatal("Unable to query foo table:", err)
	}

	counter = 0

	for rows.Next() {
		if err := rows.Scan(&id, &fbool); err != nil {
			t.Fatal("Unable to scan results:", err)
		}
		counter++
	}

	if counter != 1 {
		t.Fatalf("Expected 1 row but %v", counter)
	}

	if id != 2 && fbool != false {
		t.Fatalf("Value for id 2 should be %v, not %v", bool2, fbool)
	}

	// make sure "nonsense" triggered an error
	rows, err = db.Query("SELECT id, fbool FROM foo where id=?;", 3)
	if err != nil {
		t.Fatal("Unable to query foo table:", err)
	}

	rows.Next()
	err = rows.Scan(&id, &fbool)
	if err == nil {
		t.Error("Expected error from \"nonsense\" bool")
	}
}
<commit_msg>'is' does not work in some older sqlite3.<commit_after>package sqlite

import (
	"database/sql"
	"os"
	"strings"
	"testing"
	"time"
)

func TestOpen(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("drop table foo")
	_, err = db.Exec("create table foo (id integer)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	if stat, err := os.Stat("./foo.db"); err != nil || stat.IsDir() {
		t.Error("Failed to create ./foo.db")
	}
}

func TestInsert(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("drop table foo")
	_, err = db.Exec("create table foo (id integer)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	res, err := db.Exec("insert into foo(id) values(123)")
	if err != nil {
		t.Fatal("Failed to insert record:", err)
	}
	affected, _ := res.RowsAffected()
	if affected != 1 {
		t.Fatalf("Expected %d for affected rows, but %d:", 1, affected)
	}

	rows, err := db.Query("select id from foo")
	if err != nil {
		t.Fatal("Failed to select records:", err)
	}
	defer rows.Close()

	rows.Next()

	var result int
	rows.Scan(&result)
	if result != 123 {
		t.Errorf("Fetched %q; expected %q", 123, result)
	}
}

func TestUpdate(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("drop table foo")
	_, err = db.Exec("create table foo (id integer)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	res, err := db.Exec("insert into foo(id) values(123)")
	if err != nil {
		t.Fatal("Failed to insert record:", err)
	}
	expected, err := res.LastInsertId()
	if err != nil {
		t.Fatal("Failed to get LastInsertId:", err)
	}
	affected, _ := res.RowsAffected()
	if err != nil {
		t.Fatal("Failed to get RowsAffected:", err)
	}
	if affected != 1 {
		t.Fatalf("Expected %d for affected rows, but %d:", 1, affected)
	}

	res, err = db.Exec("update foo set id = 234")
	if err != nil {
		t.Fatal("Failed to update record:", err)
	}
	lastId, err := res.LastInsertId()
	if err != nil {
		t.Fatal("Failed to get LastInsertId:", err)
	}
	if expected != lastId {
		t.Errorf("Expected %q for last Id, but %q:", expected, lastId)
	}
	affected, _ = res.RowsAffected()
	if err != nil {
		t.Fatal("Failed to get RowsAffected:", err)
	}
	if affected != 1 {
		t.Fatalf("Expected %d for affected rows, but %d:", 1, affected)
	}

	rows, err := db.Query("select id from foo")
	if err != nil {
		t.Fatal("Failed to select records:", err)
	}
	defer rows.Close()

	rows.Next()

	var result int
	rows.Scan(&result)
	if result != 234 {
		t.Errorf("Fetched %q; expected %q", 234, result)
	}
}

func TestDelete(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("drop table foo")
	_, err = db.Exec("create table foo (id integer)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	res, err := db.Exec("insert into foo(id) values(123)")
	if err != nil {
		t.Fatal("Failed to insert record:", err)
	}
	expected, err := res.LastInsertId()
	if err != nil {
		t.Fatal("Failed to get LastInsertId:", err)
	}
	affected, err := res.RowsAffected()
	if err != nil {
		t.Fatal("Failed to get RowsAffected:", err)
	}
	if affected != 1 {
		t.Errorf("Expected %d for cout of affected rows, but %q:", 1, affected)
	}

	res, err = db.Exec("delete from foo where id = 123")
	if err != nil {
		t.Fatal("Failed to delete record:", err)
	}
	lastId, err := res.LastInsertId()
	if err != nil {
		t.Fatal("Failed to get LastInsertId:", err)
	}
	if expected != lastId {
		t.Errorf("Expected %q for last Id, but %q:", expected, lastId)
	}
	affected, err = res.RowsAffected()
	if err != nil {
		t.Fatal("Failed to get RowsAffected:", err)
	}
	if affected != 1 {
		t.Errorf("Expected %d for cout of affected rows, but %q:", 1, affected)
	}

	rows, err := db.Query("select id from foo")
	if err != nil {
		t.Fatal("Failed to select records:", err)
	}
	defer rows.Close()

	if rows.Next() {
		t.Error("Fetched row but expected not rows")
	}
}

func TestBooleanRoundtrip(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("DROP TABLE foo")
	_, err = db.Exec("CREATE TABLE foo(id INTEGER, value BOOL)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	_, err = db.Exec("INSERT INTO foo(id, value) VALUES(1, ?)", true)
	if err != nil {
		t.Fatal("Failed to insert true value:", err)
	}

	_, err = db.Exec("INSERT INTO foo(id, value) VALUES(2, ?)", false)
	if err != nil {
		t.Fatal("Failed to insert false value:", err)
	}

	rows, err := db.Query("SELECT id, value FROM foo")
	if err != nil {
		t.Fatal("Unable to query foo table:", err)
	}
	defer rows.Close()

	for rows.Next() {
		var id int
		var value bool

		if err := rows.Scan(&id, &value); err != nil {
			t.Error("Unable to scan results:", err)
			continue
		}

		if id == 1 && !value {
			t.Error("Value for id 1 should be true, not false")

		} else if id == 2 && value {
			t.Error("Value for id 2 should be false, not true")
		}
	}
}

func TestTimestamp(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}
	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("DROP TABLE foo")
	_, err = db.Exec("CREATE TABLE foo(id INTEGER, ts timeSTAMP)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	timestamp1 := time.Date(2012, time.April, 6, 22, 50, 0, 0, time.UTC)
	_, err = db.Exec("INSERT INTO foo(id, ts) VALUES(1, ?)", timestamp1)
	if err != nil {
		t.Fatal("Failed to insert timestamp:", err)
	}

	timestamp2 := time.Date(2012, time.April, 6, 23, 22, 0, 0, time.UTC)
	_, err = db.Exec("INSERT INTO foo(id, ts) VALUES(2, ?)", timestamp2.Unix())
	if err != nil {
		t.Fatal("Failed to insert timestamp:", err)
	}

	_, err = db.Exec("INSERT INTO foo(id, ts) VALUES(3, ?)", "nonsense")
	if err != nil {
		t.Fatal("Failed to insert nonsense:", err)
	}

	rows, err := db.Query("SELECT id, ts FROM foo ORDER BY id ASC")
	if err != nil {
		t.Fatal("Unable to query foo table:", err)
	}
	defer rows.Close()

	seen := 0
	for rows.Next() {
		var id int
		var ts time.Time

		if err := rows.Scan(&id, &ts); err != nil {
			t.Error("Unable to scan results:", err)
			continue
		}

		if id == 1 {
			seen += 1
			if !timestamp1.Equal(ts) {
				t.Errorf("Value for id 1 should be %v, not %v", timestamp1, ts)
			}
		}

		if id == 2 {
			seen += 1
			if !timestamp2.Equal(ts) {
				t.Errorf("Value for id 2 should be %v, not %v", timestamp2, ts)
			}
		}
	}

	if seen != 2 {
		t.Error("Expected to see two valid timestamps")
	}

	// make sure "nonsense" triggered an error
	err = rows.Err()
	if err == nil || !strings.Contains(err.Error(), "cannot parse \"nonsense\"") {
		t.Error("Expected error from \"nonsense\" timestamp")
	}
}

func TestBoolean(t *testing.T) {
	db, err := sql.Open("sqlite3", "./foo.db")
	if err != nil {
		t.Fatal("Failed to open database:", err)
	}

	defer os.Remove("./foo.db")
	defer db.Close()

	_, err = db.Exec("CREATE TABLE foo(id INTEGER, fbool BOOLEAN)")
	if err != nil {
		t.Fatal("Failed to create table:", err)
	}

	bool1 := true
	_, err = db.Exec("INSERT INTO foo(id, fbool) VALUES(1, ?)", bool1)
	if err != nil {
		t.Fatal("Failed to insert boolean:", err)
	}

	bool2 := false
	_, err = db.Exec("INSERT INTO foo(id, fbool) VALUES(2, ?)", bool2)
	if err != nil {
		t.Fatal("Failed to insert boolean:", err)
	}

	bool3 := "nonsense"
	_, err = db.Exec("INSERT INTO foo(id, fbool) VALUES(3, ?)", bool3)
	if err != nil {
		t.Fatal("Failed to insert nonsense:", err)
	}

	rows, err := db.Query("SELECT id, fbool FROM foo where fbool = ?", bool1)
	if err != nil {
		t.Fatal("Unable to query foo table:", err)
	}
	counter := 0

	var id int
	var fbool bool

	for rows.Next() {
		if err := rows.Scan(&id, &fbool); err != nil {
			t.Fatal("Unable to scan results:", err)
		}
		counter++
	}

	if counter != 1 {
		t.Fatalf("Expected 1 row but %v", counter)
	}

	if id != 1 && fbool != true {
		t.Fatalf("Value for id 1 should be %v, not %v", bool1, fbool)
	}

	rows, err = db.Query("SELECT id, fbool FROM foo where fbool = ?", bool2)
	if err != nil {
		t.Fatal("Unable to query foo table:", err)
	}

	counter = 0

	for rows.Next() {
		if err := rows.Scan(&id, &fbool); err != nil {
			t.Fatal("Unable to scan results:", err)
		}
		counter++
	}

	if counter != 1 {
		t.Fatalf("Expected 1 row but %v", counter)
	}

	if id != 2 && fbool != false {
		t.Fatalf("Value for id 2 should be %v, not %v", bool2, fbool)
	}

	// make sure "nonsense" triggered an error
	rows, err = db.Query("SELECT id, fbool FROM foo where id=?;", 3)
	if err != nil {
		t.Fatal("Unable to query foo table:", err)
	}

	rows.Next()
	err = rows.Scan(&id, &fbool)
	if err == nil {
		t.Error("Expected error from \"nonsense\" bool")
	}
}
<|endoftext|>


<commit_before>// Package minitrace contains primitives to support
// propagation of tracing information.
package minitrace

// TODO(jbd): Rename this package to trace and propose it to the standard library.

import (
	"context"
	"net/http"
)

// Span identifies a span globally in a tracing system
// and may contain annotations the span is propagating.
//
// The identification and the formatting of propagated
// annotations might be tracing backend-specific.
type Span []byte

// NewContext returns a context with the given span.
func NewContext(ctx context.Context, s Span) context.Context {
	return context.WithValue(ctx, spanKey, s)
}

// FromContext returns the span from the context.
func FromContext(ctx context.Context) Span {
	return ctx.Value(spanKey).(Span)
}

// HTTPCarrier allows spans to be propagated via HTTP requests.
//
// SpanFromReq returns a span from the incoming HTTP request.
// If the incoming request doesn't contain trace information,
// a nil span is returned with no errors.
//
// SpanToReq mutates the outgoing request with the span, and
// returns a shallow copy of the request. If span is nil,
// req is not mutated.
type HTTPCarrier interface {
	SpanFromReq(req *http.Request) (Span, error)
	SpanToReq(req *http.Request, s Span) (*http.Request, error)
}

type contextKey struct{}

var spanKey = contextKey{}
<commit_msg>add another TODO<commit_after>// Package minitrace contains primitives to support
// propagation of tracing information.
package minitrace

// TODO(jbd): Rename this package to trace and propose it to the standard library.

// TODO(jbd): Rename Span to SpanID? Then, does it represent propagated labels?

import (
	"context"
	"net/http"
)

// Span identifies a span globally in a tracing system
// and may contain annotations the span is propagating.
//
// The identification and the formatting of propagated
// annotations might be tracing backend-specific.
type Span []byte

// NewContext returns a context with the given span.
func NewContext(ctx context.Context, s Span) context.Context {
	return context.WithValue(ctx, spanKey, s)
}

// FromContext returns the span from the context.
func FromContext(ctx context.Context) Span {
	return ctx.Value(spanKey).(Span)
}

// HTTPCarrier allows spans to be propagated via HTTP requests.
//
// SpanFromReq returns a span from the incoming HTTP request.
// If the incoming request doesn't contain trace information,
// a nil span is returned with no errors.
//
// SpanToReq mutates the outgoing request with the span, and
// returns a shallow copy of the request. If span is nil,
// req is not mutated.
type HTTPCarrier interface {
	SpanFromReq(req *http.Request) (Span, error)
	SpanToReq(req *http.Request, s Span) (*http.Request, error)
}

type contextKey struct{}

var spanKey = contextKey{}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"os"
	"regexp"

	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/packer/plugin"
	"github.com/hashicorp/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	MitamaeVersion string `mapstructure:"mitamae_version"`

	BinDir string `mapstructure:"bin_dir"`

	Option string

	RecipePath string `mapstructure:"recipe_path"`

	ctx interpolate.Context
}

type MitamaeProvisioner struct {
	config Config
}

func (mp *MitamaeProvisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&mp.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &mp.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if mp.config.RecipePath == "" {
		return errors.New("recipe_path is required")
	}

	if mp.config.BinDir == "" {
		mp.config.BinDir = "/usr/local/bin"
	}

	if mp.config.MitamaeVersion == "" {
		mp.config.MitamaeVersion = "v1.4.5"
	}

	return nil
}

func (mp *MitamaeProvisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	ui.Say("MItamae provisioning start")

	filename, err := mp.getMItamaeFileName(comm)
	if err != nil {
		return err
	}

	err = mp.downloadMItamae(ui, comm, filename)
	if err != nil {
		return err
	}

	err = mp.execRecipe(ui, comm, filename)
	if err != nil {
		return err
	}

	ui.Say("MItamae provisioning end")

	return nil
}

func (mp *MitamaeProvisioner) getMItamaeFileName(comm packer.Communicator) (string, error) {
	var cmd packer.RemoteCmd
	cmd.Command = "uname -s -m | awk '{printf(\"%s-%s\",$2,tolower($1))}'"

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := comm.Start(&cmd); err != nil {
		return "", err
	}
	cmd.Wait()
	if stderr.String() != "" {
		return "", errors.New(stderr.String())
	}
	os_arch := stdout.String()

	os := regexp.MustCompile(`^[^-]+-`).ReplaceAllString(os_arch, "")
	switch os {
	case "linux", "darwin":
		return fmt.Sprintf("mitamae-%s", os_arch), nil
	default:
		return "", fmt.Errorf("%s is not support.", os_arch)
	}
}

func (mp *MitamaeProvisioner) downloadMItamae(ui packer.Ui, comm packer.Communicator, filename string) error {
	download_url := fmt.Sprintf("https://github.com/itamae-kitchen/mitamae/releases/download/%s/%s", mp.config.MitamaeVersion, filename)
	bin_path := fmt.Sprintf("%s/%s", mp.config.BinDir, filename)

	log.Printf("Start MItamae Download from %s to %s", download_url, bin_path)

	var cmd packer.RemoteCmd
	cmd.Command = fmt.Sprintf("wget %s -q -P %s -O %s && chmod +x %s", download_url, mp.config.BinDir, filename, bin_path)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	if err := comm.Start(&cmd); err != nil {
		return err
	}
	cmd.Wait()
	if stderr.String() != "" {
		return errors.New(stderr.String())
	}

	log.Printf("Success MItamae Download")

	return nil
}

func (mp *MitamaeProvisioner) execRecipe(ui packer.Ui, comm packer.Communicator, filename string) error {
	bin_path := fmt.Sprintf("%s/%s", mp.config.BinDir, filename)

	var cmd packer.RemoteCmd
	cmd.Command = fmt.Sprintf("%s local %s %s", bin_path, mp.config.Option, mp.config.RecipePath)
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("MItamae execution failed")
	}

	return nil
}

func (mp *MitamaeProvisioner) Cancel() {
	os.Exit(0)
}

func main() {
	server, err := plugin.Server()
	if err != nil {
		panic(err)
	}
	server.RegisterProvisioner(new(MitamaeProvisioner))
	server.Serve()
}
<commit_msg>variable name change snake_case to camelCase<commit_after>package main

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"os"
	"regexp"

	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/packer/plugin"
	"github.com/hashicorp/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	MitamaeVersion string `mapstructure:"mitamae_version"`

	BinDir string `mapstructure:"bin_dir"`

	Option string

	RecipePath string `mapstructure:"recipe_path"`

	ctx interpolate.Context
}

type MitamaeProvisioner struct {
	config Config
}

func (mp *MitamaeProvisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&mp.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &mp.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if mp.config.RecipePath == "" {
		return errors.New("recipe_path is required")
	}

	if mp.config.BinDir == "" {
		mp.config.BinDir = "/usr/local/bin"
	}

	if mp.config.MitamaeVersion == "" {
		mp.config.MitamaeVersion = "v1.4.5"
	}

	return nil
}

func (mp *MitamaeProvisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	ui.Say("MItamae provisioning start")

	filename, err := mp.getMItamaeFileName(comm)
	if err != nil {
		return err
	}

	err = mp.downloadMItamae(ui, comm, filename)
	if err != nil {
		return err
	}

	err = mp.execRecipe(ui, comm, filename)
	if err != nil {
		return err
	}

	ui.Say("MItamae provisioning end")

	return nil
}

func (mp *MitamaeProvisioner) getMItamaeFileName(comm packer.Communicator) (string, error) {
	var cmd packer.RemoteCmd
	cmd.Command = "uname -s -m | awk '{printf(\"%s-%s\",$2,tolower($1))}'"

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := comm.Start(&cmd); err != nil {
		return "", err
	}
	cmd.Wait()
	if stderr.String() != "" {
		return "", errors.New(stderr.String())
	}
	osArch := stdout.String()

	os := regexp.MustCompile(`^[^-]+-`).ReplaceAllString(osArch, "")
	switch os {
	case "linux", "darwin":
		return fmt.Sprintf("mitamae-%s", osArch), nil
	default:
		return "", fmt.Errorf("%s is not support.", osArch)
	}
}

func (mp *MitamaeProvisioner) downloadMItamae(ui packer.Ui, comm packer.Communicator, filename string) error {
	downloadUrl := fmt.Sprintf("https://github.com/itamae-kitchen/mitamae/releases/download/%s/%s", mp.config.MitamaeVersion, filename)
	binPath := fmt.Sprintf("%s/%s", mp.config.BinDir, filename)

	log.Printf("Start MItamae Download from %s to %s", downloadUrl, binPath)

	var cmd packer.RemoteCmd
	cmd.Command = fmt.Sprintf("wget %s -q -P %s -O %s && chmod +x %s", downloadUrl, mp.config.BinDir, filename, binPath)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	if err := comm.Start(&cmd); err != nil {
		return err
	}
	cmd.Wait()
	if stderr.String() != "" {
		return errors.New(stderr.String())
	}

	log.Printf("Success MItamae Download")

	return nil
}

func (mp *MitamaeProvisioner) execRecipe(ui packer.Ui, comm packer.Communicator, filename string) error {
	binPath := fmt.Sprintf("%s/%s", mp.config.BinDir, filename)

	var cmd packer.RemoteCmd
	cmd.Command = fmt.Sprintf("%s local %s %s", binPath, mp.config.Option, mp.config.RecipePath)
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("MItamae execution failed")
	}

	return nil
}

func (mp *MitamaeProvisioner) Cancel() {
	os.Exit(0)
}

func main() {
	server, err := plugin.Server()
	if err != nil {
		panic(err)
	}
	server.RegisterProvisioner(new(MitamaeProvisioner))
	server.Serve()
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

package config

import (
	"encoding"
	"encoding/json"
	"encoding/xml"
	"testing"
	"time"

	"github.com/issue9/assert"
	"gopkg.in/yaml.v2"
)

var (
	dur time.Duration

	_ encoding.TextMarshaler   = Duration(1)
	_ encoding.TextUnmarshaler = (*Duration)(&dur)
)

type testDuration struct {
	Duration Duration `xml:"dur" json:"dur" yaml:"dur"`
}

func TestDuration_Duration(t *testing.T) {
	a := assert.New(t)

	dur := time.Second * 2

	a.Equal(dur, Duration(dur).Duration())
}

func TestDuration_YAML(t *testing.T) {
	a := assert.New(t)

	m := &testDuration{
		Duration: Duration(time.Nanosecond * 5),
	}

	bs, err := yaml.Marshal(m)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `dur: 5ns
`)

	rm := &testDuration{}
	a.NotError(yaml.Unmarshal(bs, rm))
	a.Equal(rm, m)
}

func TestDuration_XML(t *testing.T) {
	a := assert.New(t)

	m := &testDuration{
		Duration: Duration(time.Nanosecond * 5),
	}

	bs, err := xml.MarshalIndent(m, "", "  ")
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `<testDuration>
  <dur>5ns</dur>
</testDuration>`)

	rm := &testDuration{}
	a.NotError(xml.Unmarshal(bs, rm))
	a.Equal(rm, m)
}

func TestDuration_XMLAttr(t *testing.T) {
	a := assert.New(t)

	type obj struct {
		D Duration `xml:"d,attr"`
	}
	m := &obj{
		D: Duration(time.Nanosecond * 5),
	}

	bs, err := xml.MarshalIndent(m, "", "  ")
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `<obj d="5ns"></obj>`)

	rm := &obj{}
	a.NotError(xml.Unmarshal(bs, rm))
	a.Equal(rm, m)
}

func TestDuration_JSON(t *testing.T) {
	a := assert.New(t)

	m := &testDuration{
		Duration: Duration(time.Nanosecond * 5),
	}

	bs, err := json.Marshal(m)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `{"dur":"5ns"}`)

	rm := &testDuration{}
	a.NotError(json.Unmarshal(bs, rm))
	a.Equal(rm, m)
}
<commit_msg>test(config): 添加测试用例<commit_after>// SPDX-License-Identifier: MIT

package config

import (
	"encoding"
	"encoding/json"
	"encoding/xml"
	"io/fs"
	"os"
	"testing"
	"time"

	"github.com/issue9/assert"
	"gopkg.in/yaml.v2"
)

var (
	dur time.Duration

	_ encoding.TextMarshaler   = Duration(1)
	_ encoding.TextUnmarshaler = (*Duration)(&dur)
)

type testDuration struct {
	Duration Duration `xml:"dur" json:"dur" yaml:"dur"`
}

func TestDuration_Duration(t *testing.T) {
	a := assert.New(t)

	dur := time.Second * 2

	a.Equal(dur, Duration(dur).Duration())
}

func TestDuration_YAML(t *testing.T) {
	a := assert.New(t)

	m := &testDuration{
		Duration: Duration(time.Nanosecond * 5),
	}

	bs, err := yaml.Marshal(m)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `dur: 5ns
`)

	rm := &testDuration{}
	a.NotError(yaml.Unmarshal(bs, rm))
	a.Equal(rm, m)
}

func TestDuration_XML(t *testing.T) {
	a := assert.New(t)

	m := &testDuration{
		Duration: Duration(time.Nanosecond * 5),
	}

	bs, err := xml.MarshalIndent(m, "", "  ")
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `<testDuration>
  <dur>5ns</dur>
</testDuration>`)

	rm := &testDuration{}
	a.NotError(xml.Unmarshal(bs, rm))
	a.Equal(rm, m)
}

func TestDuration_XMLAttr(t *testing.T) {
	a := assert.New(t)

	type obj struct {
		D Duration `xml:"d,attr"`
	}
	m := &obj{
		D: Duration(time.Nanosecond * 5),
	}

	bs, err := xml.MarshalIndent(m, "", "  ")
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `<obj d="5ns"></obj>`)

	rm := &obj{}
	a.NotError(xml.Unmarshal(bs, rm))
	a.Equal(rm, m)
}

func TestDuration_JSON(t *testing.T) {
	a := assert.New(t)

	m := &testDuration{
		Duration: Duration(time.Nanosecond * 5),
	}

	bs, err := json.Marshal(m)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `{"dur":"5ns"}`)

	rm := &testDuration{}
	a.NotError(json.Unmarshal(bs, rm))
	a.Equal(rm, m)
}

func TestNewServer(t *testing.T) {
	a := assert.New(t)

	srv, err := NewServer("app", "1.1.1", os.DirFS("./testdata"), nil, nil)
	a.NotError(err).NotNil(srv)

	srv, err = NewServer("app", "1.1.1", os.DirFS("./testdata/not-exists"), nil, nil)
	a.ErrorIs(err, fs.ErrNotExist).Nil(srv)
}
<|endoftext|>


<commit_before>package mod_test

import (
	"io/ioutil"
	"net/http"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/prometheus/common/expfmt"
	"github.com/robustirc/bridge/robustsession"
	"github.com/robustirc/internal/health"
	"github.com/robustirc/robustirc/internal/localnet"
	"github.com/robustirc/robustirc/internal/robust"
)

func TestMessageOfDeath(t *testing.T) {
	tempdir, err := ioutil.TempDir("", "robustirc-message-of-death-")
	if err != nil {
		t.Fatalf("Could not create tempdir: %v", err)
	}

	l, err := localnet.NewLocalnet(-1, tempdir)
	if err != nil {
		t.Fatalf("Could not start local RobustIRC network: %v", err)
	}
	defer l.Kill(true)

	l.EnablePanicCommand = "1"

	// For each of the nodes, start a goroutine that verifies that the node crashes, then start it again
	var wg sync.WaitGroup

	for i := 0; i < 3; i++ {
		cmd, tempdir, addr := l.StartIRCServer(i == 0)
		wg.Add(1)
		go func(cmd *exec.Cmd, tempdir string, addr string) {
			defer wg.Done()

			terminated := make(chan error)
			skipped := make(chan bool)

			go func() {
				terminated <- cmd.Wait()
			}()

			go func() {
				// Poll messages of death counter.
				parser := &expfmt.TextParser{}
				for {
					time.Sleep(50 * time.Millisecond)
					req, err := http.NewRequest("GET", addr+"metrics", nil)
					if err != nil {
						continue
					}
					resp, err := l.Httpclient.Do(req)
					if err != nil {
						continue
					}
					if resp.StatusCode != http.StatusOK {
						continue
					}
					metrics, err := parser.TextToMetricFamilies(resp.Body)
					if err != nil {
						continue
					}
					applied, ok := metrics["applied_messages"]
					if !ok {
						continue
					}
					for _, m := range applied.GetMetric() {
						for _, labelpair := range m.GetLabel() {
							if labelpair.GetName() == "type" &&
								labelpair.GetValue() == robust.Type(robust.MessageOfDeath).String() {
								if m.GetCounter().GetValue() > 0 {
									skipped <- true
								}
							}
						}
					}
				}
			}()

			// Wait for the server to either crash or skip a message of death.
			select {
			case <-terminated:
				t.Logf("Node terminated (as expected)")
			case <-skipped:
				t.Logf("Node skipped message of death")
			}

			// Run restart.sh for that node.
			rcmd := exec.Command(filepath.Join(tempdir, "restart.sh"))
			if err := rcmd.Start(); err != nil {
				t.Errorf("Cannot restart node: %v", err)
				return
			}
			l.RecordResource("pid", strconv.Itoa(rcmd.Process.Pid))

			// Ensure the node comes back up.
			started := time.Now()
			for time.Since(started) < 10*time.Second {
				if _, err := health.GetServerStatus(addr, l.NetworkPassword); err != nil {
					t.Logf("Node %s unhealthy: %v", addr, err)
					time.Sleep(1 * time.Second)
					continue
				}
				t.Logf("Node %s became healthy", addr)
				return
			}
			t.Errorf("Node did not become healthy within 10s")
		}(cmd, tempdir, addr)
	}

	// Connect and send the PANIC message.
	session, err := robustsession.Create(strings.Join(l.Servers(), ","), filepath.Join(tempdir, "cert.pem"))
	if err != nil {
		t.Fatalf("Could not create robustsession: %v", err)
	}
	foundjoin := make(chan bool)
	go func() {
		for msg := range session.Messages {
			if !strings.HasPrefix(msg, ":mod!1@") ||
				!strings.HasSuffix(msg, " JOIN #mod") {
				continue
			}
			select {
			case foundjoin <- true:
			default:
				t.Errorf("Found JOIN too early or too late (channel write did not block)")
			}
		}
	}()
	go func() {
		for err := range session.Errors {
			t.Errorf("RobustSession error: %v", err)
		}
	}()

	session.PostMessage("NICK mod")
	session.PostMessage("USER 1 2 3 4")
	session.PostMessage("PANIC")

	t.Logf("Message of death sent")

	wg.Wait()

	healthy := false
	for try := 0; try < 5; try++ {
		if l.Healthy() {
			healthy = true
			break
		}
		time.Sleep(1 * time.Second)
	}

	if !healthy {
		t.Fatalf("Expected recovery, but not all nodes are healthy")
	}

	// Verify sending a JOIN now results in an output message.
	session.PostMessage("JOIN #mod")
	select {
	case <-foundjoin:
		t.Logf("JOIN reply received, network progressing")
	case <-time.After(10 * time.Second):
		t.Errorf("Timeout waiting for JOIN message")
	}
}
<commit_msg>mod_test: use state var+mutex instead of relying on scheduling<commit_after>package mod_test

import (
	"io/ioutil"
	"net/http"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/prometheus/common/expfmt"
	"github.com/robustirc/bridge/robustsession"
	"github.com/robustirc/internal/health"
	"github.com/robustirc/robustirc/internal/localnet"
	"github.com/robustirc/robustirc/internal/robust"
)

func TestMessageOfDeath(t *testing.T) {
	tempdir, err := ioutil.TempDir("", "robustirc-message-of-death-")
	if err != nil {
		t.Fatalf("Could not create tempdir: %v", err)
	}

	l, err := localnet.NewLocalnet(-1, tempdir)
	if err != nil {
		t.Fatalf("Could not start local RobustIRC network: %v", err)
	}
	defer l.Kill(true)

	l.EnablePanicCommand = "1"

	// For each of the nodes, start a goroutine that verifies that the node crashes, then start it again
	var wg sync.WaitGroup

	for i := 0; i < 3; i++ {
		cmd, tempdir, addr := l.StartIRCServer(i == 0)
		wg.Add(1)
		go func(cmd *exec.Cmd, tempdir string, addr string) {
			defer wg.Done()

			terminated := make(chan error)
			skipped := make(chan bool)

			go func() {
				terminated <- cmd.Wait()
			}()

			go func() {
				// Poll messages of death counter.
				parser := &expfmt.TextParser{}
				for {
					time.Sleep(50 * time.Millisecond)
					req, err := http.NewRequest("GET", addr+"metrics", nil)
					if err != nil {
						continue
					}
					resp, err := l.Httpclient.Do(req)
					if err != nil {
						continue
					}
					if resp.StatusCode != http.StatusOK {
						continue
					}
					metrics, err := parser.TextToMetricFamilies(resp.Body)
					if err != nil {
						continue
					}
					applied, ok := metrics["applied_messages"]
					if !ok {
						continue
					}
					for _, m := range applied.GetMetric() {
						for _, labelpair := range m.GetLabel() {
							if labelpair.GetName() == "type" &&
								labelpair.GetValue() == robust.Type(robust.MessageOfDeath).String() {
								if m.GetCounter().GetValue() > 0 {
									skipped <- true
								}
							}
						}
					}
				}
			}()

			// Wait for the server to either crash or skip a message of death.
			select {
			case <-terminated:
				t.Logf("Node terminated (as expected)")
			case <-skipped:
				t.Logf("Node skipped message of death")
			}

			// Run restart.sh for that node.
			rcmd := exec.Command(filepath.Join(tempdir, "restart.sh"))
			if err := rcmd.Start(); err != nil {
				t.Errorf("Cannot restart node: %v", err)
				return
			}
			l.RecordResource("pid", strconv.Itoa(rcmd.Process.Pid))

			// Ensure the node comes back up.
			started := time.Now()
			for time.Since(started) < 10*time.Second {
				if _, err := health.GetServerStatus(addr, l.NetworkPassword); err != nil {
					t.Logf("Node %s unhealthy: %v", addr, err)
					time.Sleep(1 * time.Second)
					continue
				}
				t.Logf("Node %s became healthy", addr)
				return
			}
			t.Errorf("Node did not become healthy within 10s")
		}(cmd, tempdir, addr)
	}

	// Connect and send the PANIC message.
	session, err := robustsession.Create(strings.Join(l.Servers(), ","), filepath.Join(tempdir, "cert.pem"))
	if err != nil {
		t.Fatalf("Could not create robustsession: %v", err)
	}
	var (
		mu    sync.Mutex
		state = "no-join-yet"
	)
	foundjoin := make(chan bool)
	go func() {
		for msg := range session.Messages {
			if !strings.HasPrefix(msg, ":mod!1@") ||
				!strings.HasSuffix(msg, " JOIN #mod") {
				continue
			}
			mu.Lock()
			if state == "no-join-yet" {
				t.Errorf("Found JOIN too early")
			}
			foundjoin <- true
			mu.Unlock()
		}
	}()
	go func() {
		for err := range session.Errors {
			t.Errorf("RobustSession error: %v", err)
		}
	}()

	session.PostMessage("NICK mod")
	session.PostMessage("USER 1 2 3 4")
	session.PostMessage("PANIC")

	t.Logf("Message of death sent")

	wg.Wait()

	healthy := false
	for try := 0; try < 5; try++ {
		if l.Healthy() {
			healthy = true
			break
		}
		time.Sleep(1 * time.Second)
	}

	if !healthy {
		t.Fatalf("Expected recovery, but not all nodes are healthy")
	}

	// Verify sending a JOIN now results in an output message.
	mu.Lock()
	state = "expect-join"
	mu.Unlock()
	session.PostMessage("JOIN #mod")
	select {
	case <-foundjoin:
		t.Logf("JOIN reply received, network progressing")
	case <-time.After(10 * time.Second):
		t.Errorf("Timeout waiting for JOIN message")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"time"

	"io/ioutil"

	"github.com/docker/libkv"
	"github.com/docker/libkv/store"
)

const (
	appName = "kviator"
	version = "0.0.5"

	helpText = `
	kviator is a cli client for accessing consul, etcd, or zookeper KV.

	Syntax:

	    kviator --kvstore [consul|etcd|zzookeper] --client <kv_addr> <command> <key> [<val>]

	Options:
	    --kvstore     The kvstore to connect to. Can be consul, etcd, or zookeper.
	    --client      The url of the kvstore. (eg. localhost:8500)
	    --show-value  Show the value when using the list command.

	Commands:
	    put           put a key value pair in the kvstore
	    get           retrieve a key value pair from the kvstore
	    del           removes a key value pair from the kvstore
	    deltree       removes an entire tree structure in the kvstore
	    list          list all kv of a given subtree/key.
	    cas           put a key value pair in the keystore only when it's empty
	    exists        returns true when key value pair exists

	Arguments:
	    key           The key. Required for all commands.
	    val           The value. required for put and cas.

	Note:

	    kviator can also read the value from Stdin. The syntax would look like this:

	        cmd | kviator ... put -
	        kviator ... put - < val.file

	    The - character is necessary to force kviator to read from Stdin. Without the -, Stdin
	    is ignored.
	`
)

var (
	kvstore   string
	client    string
	showValue bool
)

func init() {
	flag.StringVar(&kvstore, "kvstore", "", "the kvstore to connect to. Can be consul, etcd, or zookeper.")
	flag.StringVar(&client, "client", "", "the client IP address")
	flag.BoolVar(&showValue, "show-value", false, "show the value of the listed keys")
	flag.Usage = help
	flag.Parse()
}

func main() {
	switch flag.Arg(0) {
	case "put":
		key := flag.Arg(1)
		val := strings.Join(flag.Args()[2:], " ")
		val = parseVal(val)
		save(key, val)
	case "get":
		key := flag.Arg(1)
		retrieve(key)
	case "del":
		key := flag.Arg(1)
		delete(key)
	case "deltree":
		key := flag.Arg(1)
		deleteTree(key)
	case "list":
		key := flag.Arg(1)
		list(key)
	case "cas":
		key := flag.Arg(1)
		val := strings.Join(flag.Args()[2:], " ")
		val = parseVal(val)
		checkAndSave(key, val)
	case "exists":
		key := flag.Arg(1)
		keyExists(key)
	default:
		help()
		os.Exit(8)
	}
}

func parseVal(arg string) string {
	arg = strings.TrimSpace(arg)
	if arg == "-" {
		stat, err := os.Stdin.Stat()
		if err == nil {
			if (stat.Mode() & os.ModeCharDevice) == 0 {
				in, _ := ioutil.ReadAll(os.Stdin)
				inStr := strings.TrimSuffix(string(in), "\n")
				inStr = strings.TrimSpace(inStr)
				return inStr
			}
		}
	}
	return arg
}

func kvstoreConn(kvstore, client string) store.Store {
	var backend store.Backend
	switch kvstore {
	case "consul":
		backend = store.CONSUL
		consul.Register()
	case "etcd":
		backend = store.ETCD
		etcd.Register()
	case "zookeper":
		backend = store.ZK
		zookeeper.Register()
	}
	kv, err := libkv.NewStore(
		backend,
		[]string{client},
		&store.Config{
			ConnectionTimeout: 10 * time.Second,
		},
	)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	return kv
}

func save(key, val string) {
	kv := kvstoreConn(kvstore, client)
	err := kv.Put(key, []byte(val), nil)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}
}

func retrieve(key string) {
	kv := kvstoreConn(kvstore, client)
	kvPair, err := kv.Get(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	} else {
		fmt.Fprintln(os.Stdout, string(kvPair.Value))
	}
}

func delete(key string) {
	kv := kvstoreConn(kvstore, client)
	err := kv.Delete(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(4)
	}
}

func deleteTree(key string) {
	if key == "" {
		fmt.Fprintln(os.Stderr, "Please specify subtree. To delete all, use /.")
		os.Exit(11)
	}
	if key == "/" {
		key = ""
	}
	kv := kvstoreConn(kvstore, client)
	err := kv.DeleteTree(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(10)
	}
}

func list(key string) {
	kv := kvstoreConn(kvstore, client)
	kvPairs, err := kv.List(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(11)
	}
	for _, kvPair := range kvPairs {
		k := kvPair.Key
		v := string(kvPair.Value)
		if showValue {
			fmt.Fprintf(os.Stdout, "%s=%s\n", k, v)
		} else {
			fmt.Fprintln(os.Stdout, k)
		}

	}
}

func checkAndSave(key, val string) {
	kv := kvstoreConn(kvstore, client)
	_, err := kv.Get(key)
	if err != nil {
		err := kv.Put(key, []byte(val), nil)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(6)
		}
	} else {
		fmt.Fprintln(os.Stderr, "key is already set")
		os.Exit(5)
	}
}

func keyExists(key string) {
	kv := kvstoreConn(kvstore, client)
	_, err := kv.Get(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, "false")
		os.Exit(7)
	} else {
		fmt.Fprintln(os.Stdout, "true")
	}
}

func help() {
	fmt.Fprintf(os.Stdout, "\n\t%s %s\n\n%s\n", appName, version, helpText)
}
<commit_msg>dependencies<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"time"

	"io/ioutil"

	"github.com/docker/libkv"
	"github.com/docker/libkv/store"
	"github.com/docker/libkv/store/etcd"
	"github.com/docker/libkv/store/consul"
	"github.com/docker/libkv/store/zookeeper"
)

const (
	appName = "kviator"
	version = "0.0.5"

	helpText = `
	kviator is a cli client for accessing consul, etcd, or zookeper KV.

	Syntax:

	    kviator --kvstore [consul|etcd|zzookeper] --client <kv_addr> <command> <key> [<val>]

	Options:
	    --kvstore     The kvstore to connect to. Can be consul, etcd, or zookeper.
	    --client      The url of the kvstore. (eg. localhost:8500)
	    --show-value  Show the value when using the list command.

	Commands:
	    put           put a key value pair in the kvstore
	    get           retrieve a key value pair from the kvstore
	    del           removes a key value pair from the kvstore
	    deltree       removes an entire tree structure in the kvstore
	    list          list all kv of a given subtree/key.
	    cas           put a key value pair in the keystore only when it's empty
	    exists        returns true when key value pair exists

	Arguments:
	    key           The key. Required for all commands.
	    val           The value. required for put and cas.

	Note:

	    kviator can also read the value from Stdin. The syntax would look like this:

	        cmd | kviator ... put -
	        kviator ... put - < val.file

	    The - character is necessary to force kviator to read from Stdin. Without the -, Stdin
	    is ignored.
	`
)

var (
	kvstore   string
	client    string
	showValue bool
)

func init() {
	flag.StringVar(&kvstore, "kvstore", "", "the kvstore to connect to. Can be consul, etcd, or zookeper.")
	flag.StringVar(&client, "client", "", "the client IP address")
	flag.BoolVar(&showValue, "show-value", false, "show the value of the listed keys")
	flag.Usage = help
	flag.Parse()
}

func main() {
	switch flag.Arg(0) {
	case "put":
		key := flag.Arg(1)
		val := strings.Join(flag.Args()[2:], " ")
		val = parseVal(val)
		save(key, val)
	case "get":
		key := flag.Arg(1)
		retrieve(key)
	case "del":
		key := flag.Arg(1)
		delete(key)
	case "deltree":
		key := flag.Arg(1)
		deleteTree(key)
	case "list":
		key := flag.Arg(1)
		list(key)
	case "cas":
		key := flag.Arg(1)
		val := strings.Join(flag.Args()[2:], " ")
		val = parseVal(val)
		checkAndSave(key, val)
	case "exists":
		key := flag.Arg(1)
		keyExists(key)
	default:
		help()
		os.Exit(8)
	}
}

func parseVal(arg string) string {
	arg = strings.TrimSpace(arg)
	if arg == "-" {
		stat, err := os.Stdin.Stat()
		if err == nil {
			if (stat.Mode() & os.ModeCharDevice) == 0 {
				in, _ := ioutil.ReadAll(os.Stdin)
				inStr := strings.TrimSuffix(string(in), "\n")
				inStr = strings.TrimSpace(inStr)
				return inStr
			}
		}
	}
	return arg
}

func kvstoreConn(kvstore, client string) store.Store {
	var backend store.Backend
	switch kvstore {
	case "consul":
		backend = store.CONSUL
		consul.Register()
	case "etcd":
		backend = store.ETCD
		etcd.Register()
	case "zookeper":
		backend = store.ZK
		zookeeper.Register()
	}
	kv, err := libkv.NewStore(
		backend,
		[]string{client},
		&store.Config{
			ConnectionTimeout: 10 * time.Second,
		},
	)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	return kv
}

func save(key, val string) {
	kv := kvstoreConn(kvstore, client)
	err := kv.Put(key, []byte(val), nil)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}
}

func retrieve(key string) {
	kv := kvstoreConn(kvstore, client)
	kvPair, err := kv.Get(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	} else {
		fmt.Fprintln(os.Stdout, string(kvPair.Value))
	}
}

func delete(key string) {
	kv := kvstoreConn(kvstore, client)
	err := kv.Delete(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(4)
	}
}

func deleteTree(key string) {
	if key == "" {
		fmt.Fprintln(os.Stderr, "Please specify subtree. To delete all, use /.")
		os.Exit(11)
	}
	if key == "/" {
		key = ""
	}
	kv := kvstoreConn(kvstore, client)
	err := kv.DeleteTree(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(10)
	}
}

func list(key string) {
	kv := kvstoreConn(kvstore, client)
	kvPairs, err := kv.List(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(11)
	}
	for _, kvPair := range kvPairs {
		k := kvPair.Key
		v := string(kvPair.Value)
		if showValue {
			fmt.Fprintf(os.Stdout, "%s=%s\n", k, v)
		} else {
			fmt.Fprintln(os.Stdout, k)
		}

	}
}

func checkAndSave(key, val string) {
	kv := kvstoreConn(kvstore, client)
	_, err := kv.Get(key)
	if err != nil {
		err := kv.Put(key, []byte(val), nil)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(6)
		}
	} else {
		fmt.Fprintln(os.Stderr, "key is already set")
		os.Exit(5)
	}
}

func keyExists(key string) {
	kv := kvstoreConn(kvstore, client)
	_, err := kv.Get(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, "false")
		os.Exit(7)
	} else {
		fmt.Fprintln(os.Stdout, "true")
	}
}

func help() {
	fmt.Fprintf(os.Stdout, "\n\t%s %s\n\n%s\n", appName, version, helpText)
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 the original author or 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 main

import (
	"context"
	"fmt"
	"gopkg.in/Shopify/sarama.v1"
	"github.com/bsm/sarama-cluster"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"time"
	"github.com/satori/go.uuid"
	"github.com/projectriff/http-gateway/pkg/message"
	"syscall"
)

// Function messageHandler creates an http handler that posts the http body as a message to Kafka, replying
// immediately with a successful http response
func messageHandler(producer sarama.AsyncProducer) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		topic := r.URL.Path[len("/messages/"):]
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		scsMessage := message.Message{Payload: b, Headers: headersToPropagate(r)}

		bytesOut, err := message.EncodeMessage(scsMessage)
		if err != nil {
			log.Printf("Error encoding message: %v", err)
			return
		}
		kafkaMsg := &sarama.ProducerMessage{Topic: topic, Value: sarama.ByteEncoder(bytesOut)}

		select {
		case producer.Input() <- kafkaMsg:
			w.Write([]byte("message published to topic: " + topic + "\n"))
		}
	}
}

// Function replyHandler creates an http handler that posts the http body as a message to Kafka, then waits
// for a message on a go channel it creates for a reply (this is expected to be set by the main thread) and sends
// that as an http response.
func replyHandler(producer sarama.AsyncProducer, replies map[string]chan message.Message) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		topic := r.URL.Path[len("/requests/"):]
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		correlationId := uuid.NewV4().String()
		replyChan := make(chan message.Message)
		replies[correlationId] = replyChan
		headers := headersToPropagate(r)
		headers["correlationId"] = correlationId
		scsMessage := message.Message{Payload: b, Headers: headers}

		bytesOut, err := message.EncodeMessage(scsMessage)
		if err != nil {
			log.Printf("Error encoding message: %v", err)
			return
		}
		kafkaMsg := &sarama.ProducerMessage{Topic: topic, Value: sarama.ByteEncoder(bytesOut)}

		select {
		case producer.Input() <- kafkaMsg:
			select {
			case reply := <-replyChan:
				delete(replies, correlationId)
				p := reply.Payload
				w.Write(p.([]byte)) // TODO equivalent of Spring's HttpMessageConverter handling
			case <- time.After(time.Second * 60):
				delete(replies, correlationId)
				w.WriteHeader(404)
			}
		}
	}
}

func healthHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"status":"UP"}`))
	}
}

func startHttpServer(producer sarama.AsyncProducer, replies map[string]chan message.Message) *http.Server {
	srv := &http.Server{Addr: ":8080"}

	http.HandleFunc("/messages/", messageHandler(producer))
	http.HandleFunc("/requests/", replyHandler(producer, replies))
	http.HandleFunc("/application/status", healthHandler())

	go func() {
		if err := srv.ListenAndServe(); err != nil {
			log.Printf("Httpserver: ListenAndServe() error: %s", err)
		}
	}()

	log.Printf("Listening on %v", srv.Addr)
	return srv
}

func headersToPropagate(r *http.Request) map[string]interface{} {
	result := make(map[string]interface{})
	if ct, ok := r.Header["Content-Type"]; ok {
		result["Content-Type"] = ct[0]
	}
	if ct, ok := r.Header["Accept"]; ok {
		result["Accept"] = ct[0]
	}
	return result
}

func main() {
	// Trap signals to trigger a proper shutdown.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt, syscall.SIGTERM, os.Kill)

	// Key is correlationId, value is channel used to pass message received from main Kafka consumer loop
	replies := make(map[string]chan message.Message)

	brokers := []string{os.Getenv("SPRING_CLOUD_STREAM_KAFKA_BINDER_BROKERS")}
	producer, err := sarama.NewAsyncProducer(brokers, nil)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := producer.Close(); err != nil {
			log.Fatalln(err)
		}
	}()

	consumerConfig := makeConsumerConfig()
	consumer, err := cluster.NewConsumer(brokers, "gateway", []string{"replies"}, consumerConfig)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()
	if consumerConfig.Consumer.Return.Errors {
		go consumeErrors(consumer)
	}
	if consumerConfig.Group.Return.Notifications {
		go consumeNotifications(consumer)
	}

	srv := startHttpServer(producer, replies)

MainLoop:
	for {
		select {
		case <-signals:
			fmt.Println("Shutting Down...")
			timeout, c := context.WithTimeout(context.Background(), 1*time.Second)
			defer c()
			if err := srv.Shutdown(timeout); err != nil {
				panic(err) // failure/timeout shutting down the server gracefully
			}
			break MainLoop
		case msg, ok := <-consumer.Messages():
			if ok {
				messageWithHeaders, err := message.ExtractMessage(msg.Value)
				if err != nil {
					log.Println("Failed to extract message ", err)
					break
				}
				correlationId, ok := messageWithHeaders.Headers["correlationId"].(string)
				if ok {
					c := replies[correlationId]
					if c != nil {
						fmt.Printf("Sending %v\n", messageWithHeaders)
						c <- messageWithHeaders
						consumer.MarkOffset(msg, "") // mark message as processed
					} else {
						log.Printf("Did not find communication channel for correlationId %v. Timed out?", correlationId)
						consumer.MarkOffset(msg, "") // mark message as processed
					}
				}
			}

		case err := <-producer.Errors():
			log.Println("Failed to produce kafka message ", err)
		}
	}
}

func makeConsumerConfig() *cluster.Config {
	consumerConfig := cluster.NewConfig()
	consumerConfig.Consumer.Return.Errors = true
	consumerConfig.Group.Return.Notifications = true
	return consumerConfig
}

func consumeNotifications(consumer *cluster.Consumer) {
	for ntf := range consumer.Notifications() {
		log.Printf("Rebalanced: %+v\n", ntf)
	}
}
func consumeErrors(consumer *cluster.Consumer) {
	for err := range consumer.Errors() {
		log.Printf("Error: %s\n", err.Error())
	}
}
<commit_msg>Send back result headers to the caller<commit_after>/*
 * Copyright 2017 the original author or 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 main

import (
	"context"
	"fmt"
	"gopkg.in/Shopify/sarama.v1"
	"github.com/bsm/sarama-cluster"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"time"
	"github.com/satori/go.uuid"
	"github.com/projectriff/http-gateway/pkg/message"
	"syscall"
)

// Function messageHandler creates an http handler that posts the http body as a message to Kafka, replying
// immediately with a successful http response
func messageHandler(producer sarama.AsyncProducer) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		topic := r.URL.Path[len("/messages/"):]
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		scsMessage := message.Message{Payload: b, Headers: headersToPropagate(r)}

		bytesOut, err := message.EncodeMessage(scsMessage)
		if err != nil {
			log.Printf("Error encoding message: %v", err)
			return
		}
		kafkaMsg := &sarama.ProducerMessage{Topic: topic, Value: sarama.ByteEncoder(bytesOut)}

		select {
		case producer.Input() <- kafkaMsg:
			w.Write([]byte("message published to topic: " + topic + "\n"))
		}
	}
}

// Function replyHandler creates an http handler that posts the http body as a message to Kafka, then waits
// for a message on a go channel it creates for a reply (this is expected to be set by the main thread) and sends
// that as an http response.
func replyHandler(producer sarama.AsyncProducer, replies map[string]chan message.Message) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		topic := r.URL.Path[len("/requests/"):]
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		correlationId := uuid.NewV4().String()
		replyChan := make(chan message.Message)
		replies[correlationId] = replyChan
		headers := headersToPropagate(r)
		headers["correlationId"] = correlationId
		scsMessage := message.Message{Payload: b, Headers: headers}

		bytesOut, err := message.EncodeMessage(scsMessage)
		if err != nil {
			log.Printf("Error encoding message: %v", err)
			return
		}
		kafkaMsg := &sarama.ProducerMessage{Topic: topic, Value: sarama.ByteEncoder(bytesOut)}

		select {
		case producer.Input() <- kafkaMsg:
			select {
			case reply := <-replyChan:
				delete(replies, correlationId)
				p := reply.Payload
				for k, v := range reply.Headers {
					switch value := v.(type) {
					case string:
						w.Header()[k] = []string{value}
					case []string:
						w.Header()[k] = value
					}
				}
				w.Write(p.([]byte)) // TODO equivalent of Spring's HttpMessageConverter handling
			case <- time.After(time.Second * 60):
				delete(replies, correlationId)
				w.WriteHeader(404)
			}
		}
	}
}

func healthHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"status":"UP"}`))
	}
}

func startHttpServer(producer sarama.AsyncProducer, replies map[string]chan message.Message) *http.Server {
	srv := &http.Server{Addr: ":8080"}

	http.HandleFunc("/messages/", messageHandler(producer))
	http.HandleFunc("/requests/", replyHandler(producer, replies))
	http.HandleFunc("/application/status", healthHandler())

	go func() {
		if err := srv.ListenAndServe(); err != nil {
			log.Printf("Httpserver: ListenAndServe() error: %s", err)
		}
	}()

	log.Printf("Listening on %v", srv.Addr)
	return srv
}

func headersToPropagate(r *http.Request) map[string]interface{} {
	result := make(map[string]interface{})
	if ct, ok := r.Header["Content-Type"]; ok {
		result["Content-Type"] = ct[0]
	}
	if ct, ok := r.Header["Accept"]; ok {
		result["Accept"] = ct[0]
	}
	return result
}

func main() {
	// Trap signals to trigger a proper shutdown.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt, syscall.SIGTERM, os.Kill)

	// Key is correlationId, value is channel used to pass message received from main Kafka consumer loop
	replies := make(map[string]chan message.Message)

	brokers := []string{os.Getenv("SPRING_CLOUD_STREAM_KAFKA_BINDER_BROKERS")}
	producer, err := sarama.NewAsyncProducer(brokers, nil)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := producer.Close(); err != nil {
			log.Fatalln(err)
		}
	}()

	consumerConfig := makeConsumerConfig()
	consumer, err := cluster.NewConsumer(brokers, "gateway", []string{"replies"}, consumerConfig)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()
	if consumerConfig.Consumer.Return.Errors {
		go consumeErrors(consumer)
	}
	if consumerConfig.Group.Return.Notifications {
		go consumeNotifications(consumer)
	}

	srv := startHttpServer(producer, replies)

MainLoop:
	for {
		select {
		case <-signals:
			fmt.Println("Shutting Down...")
			timeout, c := context.WithTimeout(context.Background(), 1*time.Second)
			defer c()
			if err := srv.Shutdown(timeout); err != nil {
				panic(err) // failure/timeout shutting down the server gracefully
			}
			break MainLoop
		case msg, ok := <-consumer.Messages():
			if ok {
				messageWithHeaders, err := message.ExtractMessage(msg.Value)
				if err != nil {
					log.Println("Failed to extract message ", err)
					break
				}
				correlationId, ok := messageWithHeaders.Headers["correlationId"].(string)
				if ok {
					c := replies[correlationId]
					if c != nil {
						fmt.Printf("Sending %v\n", messageWithHeaders)
						c <- messageWithHeaders
						consumer.MarkOffset(msg, "") // mark message as processed
					} else {
						log.Printf("Did not find communication channel for correlationId %v. Timed out?", correlationId)
						consumer.MarkOffset(msg, "") // mark message as processed
					}
				}
			}

		case err := <-producer.Errors():
			log.Println("Failed to produce kafka message ", err)
		}
	}
}

func makeConsumerConfig() *cluster.Config {
	consumerConfig := cluster.NewConfig()
	consumerConfig.Consumer.Return.Errors = true
	consumerConfig.Group.Return.Notifications = true
	return consumerConfig
}

func consumeNotifications(consumer *cluster.Consumer) {
	for ntf := range consumer.Notifications() {
		log.Printf("Rebalanced: %+v\n", ntf)
	}
}
func consumeErrors(consumer *cluster.Consumer) {
	for err := range consumer.Errors() {
		log.Printf("Error: %s\n", err.Error())
	}
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

// Package tunny implements a simple pool for maintaining independant worker goroutines.
package tunny

import (
	"reflect"
	"errors"
	"time"
	"sync"
)

type TunnyWorker interface {
	Job(interface{}) (interface{})
	Ready() bool
}

type TunnyExtendedWorker interface {
	Job(interface{}) (interface{})
	Ready() bool
	Initialize()
	Terminate()
}

// Default implementation of worker

type tunnyDefaultWorker struct {
	job *func(interface{}) (interface{})
}

func (worker *tunnyDefaultWorker) Job(data interface{}) interface{} {
	return (*worker.job)(data)
}

func (worker *tunnyDefaultWorker) Ready() bool {
	return true
}

/*
WorkPool allows you to contain and send work to your worker pool.
You must first indicate that the pool should run by calling Open(), then send work to the workers
through SendWork.
*/
type WorkPool struct {
	workers []*workerWrapper
	selects []reflect.SelectCase
	mutex   sync.RWMutex
	running bool
}

/*
SendWorkTimed - Send a job to a worker and return the result, this is a blocking call with a timeout.
SendWorkTimed - Args:    milliTimeout time.Duration, jobData interface{}
SendWorkTimed - Summary: the timeout period in milliseconds, the input data for the worker to process
*/
func (pool *WorkPool) SendWorkTimed(milliTimeout time.Duration, jobData interface{}) (interface{}, error) {
	pool.mutex.RLock()
	defer pool.mutex.RUnlock()

	if pool.running {
		before := time.Now()

		// Create new selectcase[] and add time out case
		selectCases := append(pool.selects[:], reflect.SelectCase {
			Dir: reflect.SelectRecv,
			Chan: reflect.ValueOf(time.After(milliTimeout * time.Millisecond)),
		})

		// Wait for workers, or time out
		if chosen, _, ok := reflect.Select(selectCases); ok {
			if chosen < ( len(selectCases) - 1 ) {
				(*pool.workers[chosen]).jobChan <- jobData

				// Wait for response, or time out
				select {
				case data := <-(*pool.workers[chosen]).outputChan:
					return data, nil
				case <- time.After((milliTimeout * time.Millisecond) - time.Since(before)):
					/* If we time out here we also need to ensure that the output is still
					 * collected and that the worker can move on. Therefore, we fork the
					 * waiting process into a new goroutine.
					 */
					go func() {
						<-(*pool.workers[chosen]).outputChan
					}()
					return nil, errors.New("Request timed out whilst waiting for job to complete")
				}
			} else {
				return nil, errors.New("Request timed out whilst waiting for a worker")
			}
		} else {
			return nil, errors.New("Failed to find a worker")
		}
	} else {
		return nil, errors.New("Pool is not running! Call Open() before sending work")
	}
}

/*
SendWorkTimedAsync - Send a timed job to a worker without blocking, and send the result to a receiving closure.
SendWorkTimedAsync - Args:    milliTimeout time.Duration, jobData interface{}, after func(interface{}, error)
SendWorkTimedAsync - Summary: the timeout period in milliseconds
SendWorkTimedAsync - Summary: the input data for the worker to process
SendWorkTimedAsync - Summary: the closure to hand the result to
*/
func (pool *WorkPool) SendWorkTimedAsync(
	milliTimeout time.Duration,
	jobData interface{},
	after func(interface{}, error),
) {
	go func() {
		result, err := pool.SendWorkTimed(milliTimeout, jobData)
		if after != nil {
			after(result, err)
		}
	}()
}

/*
SendWork - Send a job to a worker and return the result, this is a blocking call.
SendWork - Args:    jobData interface{}
SendWork - Summary: the input data for the worker to process
*/
func (pool *WorkPool) SendWork(jobData interface{}) (interface{}, error) {
	pool.mutex.RLock()
	defer pool.mutex.RUnlock()

	if pool.running {

		if chosen, _, ok := reflect.Select(pool.selects); ok && chosen >= 0 {
			(*pool.workers[chosen]).jobChan <- jobData
			return <- (*pool.workers[chosen]).outputChan, nil
		}

		return nil, errors.New("Failed to find or wait for a worker")

	} else {
		return nil, errors.New("Pool is not running! Call Open() before sending work")
	}
}

/*
SendWorkAsync - Send a job to a worker without blocking, and send the result to a receiving closure.
SendWorkAsync - Args:    jobData interface{}, after func(interface{}, error)
SendWorkAsync - Summary: the input data for the worker to process, the closure to hand the result to
*/
func (pool *WorkPool) SendWorkAsync(jobData interface{}, after func(interface{}, error)) {
	go func() {
		result, err := pool.SendWork(jobData)
		if after != nil {
			after(result, err)
		}
	}()
}

/*
Open - Open all channels and launch the background goroutines managed by the pool.
*/
func (pool *WorkPool) Open() (*WorkPool, error) {
	pool.mutex.Lock()
	defer pool.mutex.Unlock()

	if !pool.running {

		pool.selects = make( []reflect.SelectCase, len(pool.workers) )

		for i, workerWrapper := range pool.workers {
			(*workerWrapper).Open()

			pool.selects[i] = reflect.SelectCase {
				Dir: reflect.SelectRecv,
				Chan: reflect.ValueOf((*workerWrapper).readyChan),
			}
		}

		pool.running = true
		return pool, nil

	} else {
		return nil, errors.New("Pool is already running!")
	}
}

/*
Close - Close all channels and goroutines managed by the pool.
*/
func (pool *WorkPool) Close() error {
	pool.mutex.Lock()
	defer pool.mutex.Unlock()

	if pool.running {

		for _, workerWrapper := range pool.workers {
			workerWrapper.Close()
		}

		for _, workerWrapper := range pool.workers {
			workerWrapper.Join()
		}

		pool.running = false
		return nil

	} else {
		return errors.New("Cannot close when the pool is not running!")
	}
}

/*
CreatePool - Creates a pool of workers.
CreatePool - Args:    numWorkers int,    job func(interface{}) (interface{})
CreatePool - Summary: number of workers, the closure to run for each job
*/
func CreatePool(numWorkers int, job func(interface{}) interface{}) *WorkPool {
	pool := WorkPool { running: false }

	pool.workers = make ([]*workerWrapper, numWorkers)
	for i, _ := range pool.workers {
		newWorker := workerWrapper {
			worker: &(tunnyDefaultWorker { &job }),
		}
		pool.workers[i] = &newWorker
	}

	return &pool
}

/*
CreatePoolGeneric - Creates a pool of generic workers, they take a func as their only argument and execute it.
CreatePoolGeneric - Args:    numWorkers int
CreatePoolGeneric - Summary: number of workers
*/
func CreatePoolGeneric(numWorkers int) *WorkPool {

	return CreatePool(numWorkers, func (jobCall interface{}) interface{} {
		if method, ok := jobCall.(func()); ok {
			method()
			return nil
		}
		return errors.New("Generic worker not given a func()")
	})

}

/*
CreateCustomPool - Creates a pool for an array of custom workers.
CreateCustomPool - Args:    customWorkers []TunnyWorker
CreateCustomPool - Summary: An array of workers to use in the pool, each worker gets its own goroutine
*/
func CreateCustomPool(customWorkers []TunnyWorker) *WorkPool {
	pool := WorkPool { running: false }

	pool.workers = make ([]*workerWrapper, len(customWorkers))
	for i, _ := range pool.workers {
		newWorker := workerWrapper {
			worker: customWorkers[i],
		}
		pool.workers[i] = &newWorker
	}

	return &pool
}
<commit_msg>Check worker output channel<commit_after>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

// Package tunny implements a simple pool for maintaining independant worker goroutines.
package tunny

import (
	"reflect"
	"errors"
	"time"
	"sync"
)

type TunnyWorker interface {
	Job(interface{}) (interface{})
	Ready() bool
}

type TunnyExtendedWorker interface {
	Job(interface{}) (interface{})
	Ready() bool
	Initialize()
	Terminate()
}

// Default implementation of worker

type tunnyDefaultWorker struct {
	job *func(interface{}) (interface{})
}

func (worker *tunnyDefaultWorker) Job(data interface{}) interface{} {
	return (*worker.job)(data)
}

func (worker *tunnyDefaultWorker) Ready() bool {
	return true
}

/*
WorkPool allows you to contain and send work to your worker pool.
You must first indicate that the pool should run by calling Open(), then send work to the workers
through SendWork.
*/
type WorkPool struct {
	workers []*workerWrapper
	selects []reflect.SelectCase
	mutex   sync.RWMutex
	running bool
}

/*
SendWorkTimed - Send a job to a worker and return the result, this is a blocking call with a timeout.
SendWorkTimed - Args:    milliTimeout time.Duration, jobData interface{}
SendWorkTimed - Summary: the timeout period in milliseconds, the input data for the worker to process
*/
func (pool *WorkPool) SendWorkTimed(milliTimeout time.Duration, jobData interface{}) (interface{}, error) {
	pool.mutex.RLock()
	defer pool.mutex.RUnlock()

	if pool.running {
		before := time.Now()

		// Create new selectcase[] and add time out case
		selectCases := append(pool.selects[:], reflect.SelectCase {
			Dir: reflect.SelectRecv,
			Chan: reflect.ValueOf(time.After(milliTimeout * time.Millisecond)),
		})

		// Wait for workers, or time out
		if chosen, _, ok := reflect.Select(selectCases); ok {
			if chosen < ( len(selectCases) - 1 ) {
				(*pool.workers[chosen]).jobChan <- jobData

				// Wait for response, or time out
				select {
				case data, open := <-(*pool.workers[chosen]).outputChan:
					if !open {
						return nil, errors.New("Worker was closed before reaching a result")
					}
					return data, nil
				case <- time.After((milliTimeout * time.Millisecond) - time.Since(before)):
					/* If we time out here we also need to ensure that the output is still
					 * collected and that the worker can move on. Therefore, we fork the
					 * waiting process into a new goroutine.
					 */
					go func() {
						<-(*pool.workers[chosen]).outputChan
					}()
					return nil, errors.New("Request timed out whilst waiting for job to complete")
				}
			} else {
				return nil, errors.New("Request timed out whilst waiting for a worker")
			}
		} else {
			return nil, errors.New("Failed to find a worker")
		}
	} else {
		return nil, errors.New("Pool is not running! Call Open() before sending work")
	}
}

/*
SendWorkTimedAsync - Send a timed job to a worker without blocking, and send the result to a receiving closure.
SendWorkTimedAsync - Args:    milliTimeout time.Duration, jobData interface{}, after func(interface{}, error)
SendWorkTimedAsync - Summary: the timeout period in milliseconds
SendWorkTimedAsync - Summary: the input data for the worker to process
SendWorkTimedAsync - Summary: the closure to hand the result to
*/
func (pool *WorkPool) SendWorkTimedAsync(
	milliTimeout time.Duration,
	jobData interface{},
	after func(interface{}, error),
) {
	go func() {
		result, err := pool.SendWorkTimed(milliTimeout, jobData)
		if after != nil {
			after(result, err)
		}
	}()
}

/*
SendWork - Send a job to a worker and return the result, this is a blocking call.
SendWork - Args:    jobData interface{}
SendWork - Summary: the input data for the worker to process
*/
func (pool *WorkPool) SendWork(jobData interface{}) (interface{}, error) {
	pool.mutex.RLock()
	defer pool.mutex.RUnlock()

	if pool.running {

		if chosen, _, ok := reflect.Select(pool.selects); ok && chosen >= 0 {
			(*pool.workers[chosen]).jobChan <- jobData
			result, open := <-(*pool.workers[chosen]).outputChan

			if !open {
				return nil, errors.New("Worker was closed before reaching a result")
			}
			return result, nil
		}

		return nil, errors.New("Failed to find or wait for a worker")

	} else {
		return nil, errors.New("Pool is not running! Call Open() before sending work")
	}
}

/*
SendWorkAsync - Send a job to a worker without blocking, and send the result to a receiving closure.
SendWorkAsync - Args:    jobData interface{}, after func(interface{}, error)
SendWorkAsync - Summary: the input data for the worker to process, the closure to hand the result to
*/
func (pool *WorkPool) SendWorkAsync(jobData interface{}, after func(interface{}, error)) {
	go func() {
		result, err := pool.SendWork(jobData)
		if after != nil {
			after(result, err)
		}
	}()
}

/*
Open - Open all channels and launch the background goroutines managed by the pool.
*/
func (pool *WorkPool) Open() (*WorkPool, error) {
	pool.mutex.Lock()
	defer pool.mutex.Unlock()

	if !pool.running {

		pool.selects = make( []reflect.SelectCase, len(pool.workers) )

		for i, workerWrapper := range pool.workers {
			(*workerWrapper).Open()

			pool.selects[i] = reflect.SelectCase {
				Dir: reflect.SelectRecv,
				Chan: reflect.ValueOf((*workerWrapper).readyChan),
			}
		}

		pool.running = true
		return pool, nil

	} else {
		return nil, errors.New("Pool is already running!")
	}
}

/*
Close - Close all channels and goroutines managed by the pool.
*/
func (pool *WorkPool) Close() error {
	pool.mutex.Lock()
	defer pool.mutex.Unlock()

	if pool.running {

		for _, workerWrapper := range pool.workers {
			workerWrapper.Close()
		}

		for _, workerWrapper := range pool.workers {
			workerWrapper.Join()
		}

		pool.running = false
		return nil

	} else {
		return errors.New("Cannot close when the pool is not running!")
	}
}

/*
CreatePool - Creates a pool of workers.
CreatePool - Args:    numWorkers int,    job func(interface{}) (interface{})
CreatePool - Summary: number of workers, the closure to run for each job
*/
func CreatePool(numWorkers int, job func(interface{}) interface{}) *WorkPool {
	pool := WorkPool { running: false }

	pool.workers = make ([]*workerWrapper, numWorkers)
	for i, _ := range pool.workers {
		newWorker := workerWrapper {
			worker: &(tunnyDefaultWorker { &job }),
		}
		pool.workers[i] = &newWorker
	}

	return &pool
}

/*
CreatePoolGeneric - Creates a pool of generic workers, they take a func as their only argument and execute it.
CreatePoolGeneric - Args:    numWorkers int
CreatePoolGeneric - Summary: number of workers
*/
func CreatePoolGeneric(numWorkers int) *WorkPool {

	return CreatePool(numWorkers, func (jobCall interface{}) interface{} {
		if method, ok := jobCall.(func()); ok {
			method()
			return nil
		}
		return errors.New("Generic worker not given a func()")
	})

}

/*
CreateCustomPool - Creates a pool for an array of custom workers.
CreateCustomPool - Args:    customWorkers []TunnyWorker
CreateCustomPool - Summary: An array of workers to use in the pool, each worker gets its own goroutine
*/
func CreateCustomPool(customWorkers []TunnyWorker) *WorkPool {
	pool := WorkPool { running: false }

	pool.workers = make ([]*workerWrapper, len(customWorkers))
	for i, _ := range pool.workers {
		newWorker := workerWrapper {
			worker: customWorkers[i],
		}
		pool.workers[i] = &newWorker
	}

	return &pool
}
<|endoftext|>


<commit_before>package db

import (
	"fmt"

	"github.com/jmoiron/sqlx"
)

type Executor struct {
	db *sqlx.DB
}

func (ex *Executor) Query(input string) (qr QueryResult, err error) {
	dbRows, err := ex.db.Queryx(input)
	if err != nil {
		return qr, err
	}

	dbCols, err := dbRows.Columns()
	if err != nil {
		return qr, err
	}

	qr.Columns = make([]Column, len(dbCols))
	for i, dbCol := range dbCols {
		qr.Columns[i].Name = dbCol
	}

	for dbRows.Next() {
		results, err := dbRows.SliceScan()
		if err != nil {
			return qr, err
		}

		cols := make([]Cell, len(results))
		for i, res := range results {
			var value string
			switch t := res.(type) {
			case []uint8: // this happens with MySQL result and we don't know why yet. To investigate.
				value = string(res.([]uint8))
			default:
				value = fmt.Sprint(t)
			}
			cols[i] = Cell{Value: value}
		}

		qr.Rows = append(qr.Rows, cols)
	}

	return qr, dbRows.Close()
}
<commit_msg>Better naming<commit_after>package db

import (
	"fmt"

	"github.com/jmoiron/sqlx"
)

type Executor struct {
	db *sqlx.DB
}

func (ex *Executor) Query(input string) (qr QueryResult, err error) {
	dbRows, err := ex.db.Queryx(input)
	if err != nil {
		return qr, err
	}

	dbCols, err := dbRows.Columns()
	if err != nil {
		return qr, err
	}

	qr.Columns = make([]Column, len(dbCols))
	for i, dbCol := range dbCols {
		qr.Columns[i].Name = dbCol
	}

	for dbRows.Next() {
		results, err := dbRows.SliceScan()
		if err != nil {
			return qr, err
		}

		cells := make([]Cell, len(results))
		for i, res := range results {
			var value string
			switch t := res.(type) {
			case []uint8: // this happens with MySQL result and we don't know why yet. To investigate.
				value = string(res.([]uint8))
			default:
				value = fmt.Sprint(t)
			}
			cells[i] = Cell{Value: value}
		}

		qr.Rows = append(qr.Rows, cells)
	}

	return qr, dbRows.Close()
}
<|endoftext|>


<commit_before>package dbservice

import (
	"fmt"

	"github.com/go-xorm/core"

	log "github.com/lessgo/lessgo/logs"
	"github.com/lessgo/lessgo/logs/logs"
)

type ILogger struct {
	BeeLogger *logs.BeeLogger
	level     core.LogLevel
	showSQL   bool
}

func NewILogger(channelLen int64, l int, filename string) *ILogger {
	tl := logs.NewLogger(channelLen)
	tl.SetLogFuncCallDepth(3)
	tl.AddAdapter("console", "")
	tl.AddAdapter("file", `{"filename":"Logger/`+filename+`.lessgo.log"}`)
	return &ILogger{
		BeeLogger: tl,
		level:     level(core.LogLevel(l)),
	}
}

func level(l core.LogLevel) core.LogLevel {
	switch int(l) {
	case log.DEBUG:
		return core.LOG_DEBUG
	case log.INFO:
		return core.LOG_INFO
	case log.WARN:
		return core.LOG_WARNING
	case log.ERROR:
		return core.LOG_ERR
	case log.OFF, log.FATAL:
		return core.LOG_OFF
	}
	return core.LOG_UNKNOWN
}

func (i *ILogger) Debug(v ...interface{}) {
	i.BeeLogger.Debug(fmt.Sprintln(v...))
	return
}

func (i *ILogger) Debugf(format string, v ...interface{}) {
	i.BeeLogger.Debug(format, v...)
	return
}

func (i *ILogger) Error(v ...interface{}) {
	i.BeeLogger.Error(fmt.Sprintln(v...))
}

func (i *ILogger) Errorf(format string, v ...interface{}) {
	i.BeeLogger.Error(format, v...)
}

func (i *ILogger) Info(v ...interface{}) {
	i.BeeLogger.Info(fmt.Sprintln(v...))
}

func (i *ILogger) Infof(format string, v ...interface{}) {
	i.BeeLogger.Info(format, v...)
}

func (i *ILogger) Warn(v ...interface{}) {
	i.BeeLogger.Warn(fmt.Sprintln(v...))
}
func (i *ILogger) Warnf(format string, v ...interface{}) {
	i.BeeLogger.Warn(format, v...)
}

func (i *ILogger) Level() core.LogLevel {
	return i.level
}

func (i *ILogger) SetLevel(l core.LogLevel) {
	i.level = level(l)
	i.BeeLogger.SetLevel(int(i.level))
}

func (i *ILogger) ShowSQL(show ...bool) {
	if len(show) == 0 {
		i.showSQL = true
		return
	}
	i.showSQL = show[0]
}

func (i *ILogger) IsShowSQL() bool {
	return i.showSQL
}
<commit_msg>修改数据库日志文件名<commit_after>package dbservice

import (
	"fmt"

	"github.com/go-xorm/core"

	log "github.com/lessgo/lessgo/logs"
	"github.com/lessgo/lessgo/logs/logs"
)

type ILogger struct {
	BeeLogger *logs.BeeLogger
	level     core.LogLevel
	showSQL   bool
}

func NewILogger(channelLen int64, l int, filename string) *ILogger {
	tl := logs.NewLogger(channelLen)
	tl.SetLogFuncCallDepth(3)
	tl.AddAdapter("console", "")
	tl.AddAdapter("file", `{"filename":"Logger/`+filename+`.db.log"}`)
	return &ILogger{
		BeeLogger: tl,
		level:     level(core.LogLevel(l)),
	}
}

func level(l core.LogLevel) core.LogLevel {
	switch int(l) {
	case log.DEBUG:
		return core.LOG_DEBUG
	case log.INFO:
		return core.LOG_INFO
	case log.WARN:
		return core.LOG_WARNING
	case log.ERROR:
		return core.LOG_ERR
	case log.OFF, log.FATAL:
		return core.LOG_OFF
	}
	return core.LOG_UNKNOWN
}

func (i *ILogger) Debug(v ...interface{}) {
	i.BeeLogger.Debug(fmt.Sprintln(v...))
	return
}

func (i *ILogger) Debugf(format string, v ...interface{}) {
	i.BeeLogger.Debug(format, v...)
	return
}

func (i *ILogger) Error(v ...interface{}) {
	i.BeeLogger.Error(fmt.Sprintln(v...))
}

func (i *ILogger) Errorf(format string, v ...interface{}) {
	i.BeeLogger.Error(format, v...)
}

func (i *ILogger) Info(v ...interface{}) {
	i.BeeLogger.Info(fmt.Sprintln(v...))
}

func (i *ILogger) Infof(format string, v ...interface{}) {
	i.BeeLogger.Info(format, v...)
}

func (i *ILogger) Warn(v ...interface{}) {
	i.BeeLogger.Warn(fmt.Sprintln(v...))
}
func (i *ILogger) Warnf(format string, v ...interface{}) {
	i.BeeLogger.Warn(format, v...)
}

func (i *ILogger) Level() core.LogLevel {
	return i.level
}

func (i *ILogger) SetLevel(l core.LogLevel) {
	i.level = level(l)
	i.BeeLogger.SetLevel(int(i.level))
}

func (i *ILogger) ShowSQL(show ...bool) {
	if len(show) == 0 {
		i.showSQL = true
		return
	}
	i.showSQL = show[0]
}

func (i *ILogger) IsShowSQL() bool {
	return i.showSQL
}
<|endoftext|>