<commit_before>package main

import (
	"database/sql"
	"fmt"
	"io"
	"os"
	"strings"
	"text/template"

	"github.com/Masterminds/squirrel"
	"github.com/codegangsta/cli"

	_ "github.com/lib/pq"
)

const version = "DEV"

const Usage = `Read a schema and generate Structable structs.

This utility generates Structable structs be reading your database table and
generating the appropriate code.
`

const fileHeader = `package %s

// This file is automatically generated by schema2struct.

import (
	"time"

	"github.com/Masterminds/squirrel"
	"github.com/Masterminds/structable"
	_ "github.com/lib/pq"
)

// QueryFunc modifies a SelectBuilder prior to execution.
//
// The SelectBuilder is modified in place. An error is returned under any
// conditions where the query should not be executed.
type QueryFunc func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error)

`

const structTemplate = `// {{.StructName}} maps to database table {{.TableName}}
type {{.StructName}} struct {
	tableName string {{ann "tablename" .TableName}}
	structable.Recorder
	builder squirrel.StatementBuilderType
	{{range .Fields}}{{.}}
	{{end}}
}

// New{{.StructName}} creates a new {{.StructName}} wired to structable.
func New{{.StructName}}(db squirrel.DBProxyBeginner, flavor string) *{{.StructName}} {
	o := new({{.StructName}})
	o.Recorder = structable.New(db, flavor).Bind("{{.TableName}}", o)
	return o
}

// List{{.StructName}} returns a list of {{.StructName}} objects.
//
// Limit is the max number of items. Offset is the offset the results will
// begin with.
func List{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, limit, offset uint64) ([]*{{.StructName}}, error) {
	fn := func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error) {
		return q.Limit(limit).Offset(offset), nil
	}
	return Query{{.StructName}}(db, flavor, fn)
}

// Query{{.StructName}} builds a base query, but allows the query to be modified before execution.
//
// This creates a new Select, settings the columns and table name, and then calling QueryFunc with the
// query. The QueryFunc can then add a Where clause, etc. Provided QueryFunc does not exit with an
// error, Query{{.StructName}} will then execute the query, extract the results into a slice of
// {{.StructName}} structs, and then return.
//
// The QueryFunc should not modify the list of fields returned or the table name,
// as the intent is to construct a complete {{.StructName}} from each result.
// More sophisticated queries should be written directly.
func Query{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, fn QueryFunc) ([]*{{.StructName}}, error){
	var tn string = "{{.TableName}}"

	// We need a prototype structable to learn about the table structure.
	ps := New{{.StructName}}(db, flavor)
	cols := ps.Columns(true)

	q := ps.Builder().Select(cols...).From(tn)
	var err error
	if q, err = fn(q); err != nil {
		return []*{{.StructName}}{}, err
	}
	rows, err := q.Query()
	if err != nil || rows == nil {
		return []*{{.StructName}}{}, err
	}
	defer rows.Close()

	buf := []*{{.StructName}}{}
	for rows.Next() {
		o := New{{.StructName}}(db, flavor)
		dest := o.FieldReferences(true)
		if err := rows.Scan(dest...); err != nil {
			return buf, err
		}
		buf = append(buf, o)
	}
	return buf, rows.Err()
}

// Len{{.StructName}} returns the number of {{.StructName}} objects in the database.
func Len{{.StructName}}(db squirrel.DBProxyBeginner, flavor string) (int, error) {
	fn := func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error) {return q, nil}
	return QueryLen{{.StructName}}(db, flavor, fn)
}

// QueryLen{{.StructName}} returns the length of a table.
//
// The QueryFunc can be used to modify the query. For a simple length call, you
// may prefer to use Len{{.StructName}}.
func QueryLen{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, fn QueryFunc) (int, error) {
	tn := "{{.TableName}}"
	ps := New{{.StructName}}(db, flavor)
	q := ps.Builder().Select("COUNT(*)").From(tn)
	var err error
	if q, err = fn(q); err != nil {
		return 0, err
	}
	var count int
	err = q.Scan(&count)
	return count, err
}

`

type structDesc struct {
	StructName string
	TableName  string
	Fields     []string
}

func main() {
	app := cli.NewApp()
	app.Name = "schema2struct"
	app.Version = "version"
	app.Usage = Usage
	app.Action = importTables
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "driver,d",
			Value: "postgres",
			Usage: "The name of the SQL driver to use.",
		},
		cli.StringFlag{
			Name:  "connection,c",
			Value: "user=$USER dbname=$USER sslmode=disable",
			Usage: "The database connection string. Environment variables are expanded.",
		},
		cli.StringFlag{
			Name:  "tables,t",
			Value: "",
			Usage: "The list of tables to generate, comma separated. If none specified, the entire schema is used.",
		},
		cli.StringFlag{
			Name:  "file,f",
			Value: "",
			Usage: "The file to send the output.",
		},
		cli.StringFlag{
			Name:   "package,p",
			Value:  "main",
			Usage:  "The name of the destination package.",
			EnvVar: "GOPACKAGE",
		},
	}

	app.Run(os.Args)
}

func driver(c *cli.Context) string {
	return c.String("driver")
}
func conn(c *cli.Context) string {
	return os.ExpandEnv(c.String("connection"))
}

// dest gets the destination output writer.
func dest(c *cli.Context) io.Writer {
	if out := c.String("file"); out != "" {
		f, err := os.Create(out)
		if err != nil {
			panic(f)
		}
		return f
	}
	return os.Stdout
}

func tableList(c *cli.Context) []string {
	z := c.String("tables")
	if z != "" {
		return strings.Split(z, ",")
	}
	return []string{}
}

func cxdie(c *cli.Context, err error) {
	fmt.Fprintf(os.Stderr, "Failed to connect to %s (type %s): %s", conn(c), driver(c), err)
	os.Exit(1)
}

var funcMap = map[string]interface{}{
	"ann": func(tag, val string) string {
		return fmt.Sprintf("`%s:\"%s\"`", tag, val)
	},
}

func importTables(c *cli.Context) {
	ttt := template.Must(template.New("st").Funcs(funcMap).Parse(structTemplate))
	cxn, err := sql.Open(driver(c), conn(c))
	if err != nil {
		cxdie(c, err)
	}
	// Many drivers defer connections until the first statement. We test
	// that here.
	if err := cxn.Ping(); err != nil {
		cxdie(c, err)
	}
	defer cxn.Close()

	// Set up Squirrel
	stmts := squirrel.NewStmtCacher(cxn)
	bldr := squirrel.StatementBuilder.RunWith(stmts)
	if driver(c) == "postgres" {
		bldr = bldr.PlaceholderFormat(squirrel.Dollar)
	}

	// Set up destination
	out := dest(c)
	fmt.Fprintf(out, fileHeader, c.String("package"))

	tables := tableList(c)

	if len(tables) == 0 {
		tables, err = publicTables(bldr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Cannot fetch list of tables: %s\n", err)
			os.Exit(2)
		}
	}

	for _, t := range tables {
		f, err := importTable(t, bldr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to import table %s: %s", t, err)
		}

		//fmt.Fprintf(out, "%s %s %s\n", f.StructName, f.TableName, f.Fields)
		ttt.Execute(out, f)
	}
}

type column struct {
	Name, DataType string
	Max            int64
}

func publicTables(b squirrel.StatementBuilderType) ([]string, error) {
	rows, err := b.Select("table_name").From("INFORMATION_SCHEMA.TABLES").
		Where("table_schema = 'public'").Query()

	res := []string{}
	if err != nil {
		return res, err
	}

	for rows.Next() {
		var s string
		rows.Scan(&s)
		res = append(res, s)
	}

	return res, nil
}

// importTable reads a table definition and writes a corresponding struct.
// SELECT table_name, column_name, data_type, character_maximum_length
//   FROM INFORMATION_SCHEMA.COLUMNS WHERE table_name = 'goose_db_version'
func importTable(tbl string, b squirrel.StatementBuilderType) (*structDesc, error) {

	pks, err := primaryKeyField(tbl, b)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error getting primary keys: %s", err)
	}

	q := b.Select("column_name, data_type, character_maximum_length").
		From("INFORMATION_SCHEMA.COLUMNS").
		Where("table_name = ?", tbl)

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

	ff := []string{}
	for rows.Next() {
		c := &column{}
		var length sql.NullInt64
		if err := rows.Scan(&c.Name, &c.DataType, &length); err != nil {
			return nil, err
		}
		c.Max = length.Int64
		ff = append(ff, structField(c, pks, tbl, b))
	}
	sd := &structDesc{
		StructName: goName(tbl),
		TableName:  tbl,
		Fields:     ff,
	}

	return sd, nil
}

func primaryKeyField(tbl string, b squirrel.StatementBuilderType) ([]string, error) {
	q := b.Select("column_name").
		From("INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS c").
		LeftJoin("INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS t USING(constraint_name)").
		Where("t.table_name = ? AND t.constraint_type = 'PRIMARY KEY'", tbl).
		OrderBy("ordinal_position")

	rows, err := q.Query()
	if err != nil {
		return []string{}, err
	}

	res := []string{}
	for rows.Next() {
		var s string
		rows.Scan(&s)
		res = append(res, s)
	}
	return res, nil
}

func sequentialKey(tbl, pk string, b squirrel.StatementBuilderType) bool {

	tlen := 58

	stbl := tbl
	if len(tbl) > 29 {
		stbl = tbl[0:29]
	}

	left := tlen - len(stbl)
	spk := pk
	if len(pk) > left {
		spk = pk[0:left]
	}
	seq := fmt.Sprintf("%s_%s_seq", stbl, spk)

	q := b.Select("COUNT(*)").
		From("INFORMATION_SCHEMA.SEQUENCES").
		Where("sequence_name = ?", seq)

	var num int
	if err := q.Scan(&num); err != nil {
		panic(err)
	}
	return num > 0
}

func structField(c *column, pks []string, tbl string, b squirrel.StatementBuilderType) string {
	tpl := "%s %s `stbl:\"%s\"`"
	gn := destutter(goName(c.Name), goName(tbl))
	tt := goType(c.DataType)

	tag := c.Name
	for _, p := range pks {
		if c.Name == p {
			tag += ",PRIMARY_KEY"
			if sequentialKey(tbl, c.Name, b) {
				tag += ",SERIAL"
			}
		}
	}

	return fmt.Sprintf(tpl, gn, tt, tag)
}

// goType takes a SQL type and returns a string containin the name of a Go type.
//
// The goal is not to provide an exact match for every type, but to provide a
// safe Go representation of a SQL type.
//
// For some floating point SQL types, for example, we store them as strings
// so as not to lose precision while also not adding new types.
//
// The default type is string.
func goType(sqlType string) string {
	switch sqlType {
	case "smallint", "smallserial":
		return "int16"
	case "integer", "serial":
		return "int32"
	case "bigint", "bigserial":
		return "int"
	case "real":
		return "float32"
	case "double precision":
		return "float64"
	// Because we need to preserve base-10 precision.
	case "money":
		return "string"
	case "text", "varchar", "char", "character", "character varying", "uuid":
		return "string"
	case "bytea":
		return "[]byte"
	case "boolean":
		return "bool"
	case "timezone", "timezonetz", "date", "time":
		return "time.Time"
	case "interval":
		return "time.Duration"
	}
	return "string"
}

// Convert a SQL name to a Go name.
func goName(sqlName string) string {
	// This can definitely be done better.
	goName := strings.Replace(sqlName, "_", " ", -1)
	goName = strings.Replace(goName, ".", " ", -1)
	goName = strings.Title(goName)
	goName = strings.Replace(goName, " ", "", -1)

	return goName
}

// destutter removes a stutter prefix.
func destutter(str, prefix string) string {
	return strings.TrimPrefix(str, prefix)
}
<commit_msg>Add prixy beginner and flavor<commit_after>package main

import (
	"database/sql"
	"fmt"
	"io"
	"os"
	"strings"
	"text/template"

	"github.com/Masterminds/squirrel"
	"github.com/codegangsta/cli"

	_ "github.com/lib/pq"
)

const version = "DEV"

const Usage = `Read a schema and generate Structable structs.

This utility generates Structable structs be reading your database table and
generating the appropriate code.
`

const fileHeader = `package %s

// This file is automatically generated by schema2struct.

import (
	"time"

	"github.com/Masterminds/squirrel"
	"github.com/Masterminds/structable"
	_ "github.com/lib/pq"
)

// QueryFunc modifies a SelectBuilder prior to execution.
//
// The SelectBuilder is modified in place. An error is returned under any
// conditions where the query should not be executed.
type QueryFunc func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error)

`

const structTemplate = `// {{.StructName}} maps to database table {{.TableName}}
type {{.StructName}} struct {
	tableName string {{ann "tablename" .TableName}}
	structable.Recorder
	builder squirrel.StatementBuilderType
	{{range .Fields}}{{.}}
	{{end}}db squirrel.DBProxyBeginner
	flavor string
}

// New{{.StructName}} creates a new {{.StructName}} wired to structable.
func New{{.StructName}}(db squirrel.DBProxyBeginner, flavor string) *{{.StructName}} {
	o := &{{.StructName}}{db: db, flavor: flavor}
	o.Recorder = structable.New(db, flavor).Bind("{{.TableName}}", o)
	return o
}

// List{{.StructName}} returns a list of {{.StructName}} objects.
//
// Limit is the max number of items. Offset is the offset the results will
// begin with.
func List{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, limit, offset uint64) ([]*{{.StructName}}, error) {
	fn := func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error) {
		return q.Limit(limit).Offset(offset), nil
	}
	return Query{{.StructName}}(db, flavor, fn)
}

// Query{{.StructName}} builds a base query, but allows the query to be modified before execution.
//
// This creates a new Select, settings the columns and table name, and then calling QueryFunc with the
// query. The QueryFunc can then add a Where clause, etc. Provided QueryFunc does not exit with an
// error, Query{{.StructName}} will then execute the query, extract the results into a slice of
// {{.StructName}} structs, and then return.
//
// The QueryFunc should not modify the list of fields returned or the table name,
// as the intent is to construct a complete {{.StructName}} from each result.
// More sophisticated queries should be written directly.
func Query{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, fn QueryFunc) ([]*{{.StructName}}, error){
	var tn string = "{{.TableName}}"

	// We need a prototype structable to learn about the table structure.
	ps := New{{.StructName}}(db, flavor)
	cols := ps.Columns(true)

	q := ps.Builder().Select(cols...).From(tn)
	var err error
	if q, err = fn(q); err != nil {
		return []*{{.StructName}}{}, err
	}
	rows, err := q.Query()
	if err != nil || rows == nil {
		return []*{{.StructName}}{}, err
	}
	defer rows.Close()

	buf := []*{{.StructName}}{}
	for rows.Next() {
		o := New{{.StructName}}(db, flavor)
		dest := o.FieldReferences(true)
		if err := rows.Scan(dest...); err != nil {
			return buf, err
		}
		buf = append(buf, o)
	}
	return buf, rows.Err()
}

// Len{{.StructName}} returns the number of {{.StructName}} objects in the database.
func Len{{.StructName}}(db squirrel.DBProxyBeginner, flavor string) (int, error) {
	fn := func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error) {return q, nil}
	return QueryLen{{.StructName}}(db, flavor, fn)
}

// QueryLen{{.StructName}} returns the length of a table.
//
// The QueryFunc can be used to modify the query. For a simple length call, you
// may prefer to use Len{{.StructName}}.
func QueryLen{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, fn QueryFunc) (int, error) {
	tn := "{{.TableName}}"
	ps := New{{.StructName}}(db, flavor)
	q := ps.Builder().Select("COUNT(*)").From(tn)
	var err error
	if q, err = fn(q); err != nil {
		return 0, err
	}
	var count int
	err = q.Scan(&count)
	return count, err
}

`

type structDesc struct {
	StructName string
	TableName  string
	Fields     []string
}

func main() {
	app := cli.NewApp()
	app.Name = "schema2struct"
	app.Version = "version"
	app.Usage = Usage
	app.Action = importTables
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "driver,d",
			Value: "postgres",
			Usage: "The name of the SQL driver to use.",
		},
		cli.StringFlag{
			Name:  "connection,c",
			Value: "user=$USER dbname=$USER sslmode=disable",
			Usage: "The database connection string. Environment variables are expanded.",
		},
		cli.StringFlag{
			Name:  "tables,t",
			Value: "",
			Usage: "The list of tables to generate, comma separated. If none specified, the entire schema is used.",
		},
		cli.StringFlag{
			Name:  "file,f",
			Value: "",
			Usage: "The file to send the output.",
		},
		cli.StringFlag{
			Name:   "package,p",
			Value:  "main",
			Usage:  "The name of the destination package.",
			EnvVar: "GOPACKAGE",
		},
	}

	app.Run(os.Args)
}

func driver(c *cli.Context) string {
	return c.String("driver")
}
func conn(c *cli.Context) string {
	return os.ExpandEnv(c.String("connection"))
}

// dest gets the destination output writer.
func dest(c *cli.Context) io.Writer {
	if out := c.String("file"); out != "" {
		f, err := os.Create(out)
		if err != nil {
			panic(f)
		}
		return f
	}
	return os.Stdout
}

func tableList(c *cli.Context) []string {
	z := c.String("tables")
	if z != "" {
		return strings.Split(z, ",")
	}
	return []string{}
}

func cxdie(c *cli.Context, err error) {
	fmt.Fprintf(os.Stderr, "Failed to connect to %s (type %s): %s", conn(c), driver(c), err)
	os.Exit(1)
}

var funcMap = map[string]interface{}{
	"ann": func(tag, val string) string {
		return fmt.Sprintf("`%s:\"%s\"`", tag, val)
	},
}

func importTables(c *cli.Context) {
	ttt := template.Must(template.New("st").Funcs(funcMap).Parse(structTemplate))
	cxn, err := sql.Open(driver(c), conn(c))
	if err != nil {
		cxdie(c, err)
	}
	// Many drivers defer connections until the first statement. We test
	// that here.
	if err := cxn.Ping(); err != nil {
		cxdie(c, err)
	}
	defer cxn.Close()

	// Set up Squirrel
	stmts := squirrel.NewStmtCacher(cxn)
	bldr := squirrel.StatementBuilder.RunWith(stmts)
	if driver(c) == "postgres" {
		bldr = bldr.PlaceholderFormat(squirrel.Dollar)
	}

	// Set up destination
	out := dest(c)
	fmt.Fprintf(out, fileHeader, c.String("package"))

	tables := tableList(c)

	if len(tables) == 0 {
		tables, err = publicTables(bldr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Cannot fetch list of tables: %s\n", err)
			os.Exit(2)
		}
	}

	for _, t := range tables {
		f, err := importTable(t, bldr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to import table %s: %s", t, err)
		}

		//fmt.Fprintf(out, "%s %s %s\n", f.StructName, f.TableName, f.Fields)
		ttt.Execute(out, f)
	}
}

type column struct {
	Name, DataType string
	Max            int64
}

func publicTables(b squirrel.StatementBuilderType) ([]string, error) {
	rows, err := b.Select("table_name").From("INFORMATION_SCHEMA.TABLES").
		Where("table_schema = 'public'").Query()

	res := []string{}
	if err != nil {
		return res, err
	}

	for rows.Next() {
		var s string
		rows.Scan(&s)
		res = append(res, s)
	}

	return res, nil
}

// importTable reads a table definition and writes a corresponding struct.
// SELECT table_name, column_name, data_type, character_maximum_length
//   FROM INFORMATION_SCHEMA.COLUMNS WHERE table_name = 'goose_db_version'
func importTable(tbl string, b squirrel.StatementBuilderType) (*structDesc, error) {

	pks, err := primaryKeyField(tbl, b)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error getting primary keys: %s", err)
	}

	q := b.Select("column_name, data_type, character_maximum_length").
		From("INFORMATION_SCHEMA.COLUMNS").
		Where("table_name = ?", tbl)

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

	ff := []string{}
	for rows.Next() {
		c := &column{}
		var length sql.NullInt64
		if err := rows.Scan(&c.Name, &c.DataType, &length); err != nil {
			return nil, err
		}
		c.Max = length.Int64
		ff = append(ff, structField(c, pks, tbl, b))
	}
	sd := &structDesc{
		StructName: goName(tbl),
		TableName:  tbl,
		Fields:     ff,
	}

	return sd, nil
}

func primaryKeyField(tbl string, b squirrel.StatementBuilderType) ([]string, error) {
	q := b.Select("column_name").
		From("INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS c").
		LeftJoin("INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS t USING(constraint_name)").
		Where("t.table_name = ? AND t.constraint_type = 'PRIMARY KEY'", tbl).
		OrderBy("ordinal_position")

	rows, err := q.Query()
	if err != nil {
		return []string{}, err
	}

	res := []string{}
	for rows.Next() {
		var s string
		rows.Scan(&s)
		res = append(res, s)
	}
	return res, nil
}

func sequentialKey(tbl, pk string, b squirrel.StatementBuilderType) bool {

	tlen := 58

	stbl := tbl
	if len(tbl) > 29 {
		stbl = tbl[0:29]
	}

	left := tlen - len(stbl)
	spk := pk
	if len(pk) > left {
		spk = pk[0:left]
	}
	seq := fmt.Sprintf("%s_%s_seq", stbl, spk)

	q := b.Select("COUNT(*)").
		From("INFORMATION_SCHEMA.SEQUENCES").
		Where("sequence_name = ?", seq)

	var num int
	if err := q.Scan(&num); err != nil {
		panic(err)
	}
	return num > 0
}

func structField(c *column, pks []string, tbl string, b squirrel.StatementBuilderType) string {
	tpl := "%s %s `stbl:\"%s\"`"
	gn := destutter(goName(c.Name), goName(tbl))
	tt := goType(c.DataType)

	tag := c.Name
	for _, p := range pks {
		if c.Name == p {
			tag += ",PRIMARY_KEY"
			if sequentialKey(tbl, c.Name, b) {
				tag += ",SERIAL"
			}
		}
	}

	return fmt.Sprintf(tpl, gn, tt, tag)
}

// goType takes a SQL type and returns a string containin the name of a Go type.
//
// The goal is not to provide an exact match for every type, but to provide a
// safe Go representation of a SQL type.
//
// For some floating point SQL types, for example, we store them as strings
// so as not to lose precision while also not adding new types.
//
// The default type is string.
func goType(sqlType string) string {
	switch sqlType {
	case "smallint", "smallserial":
		return "int16"
	case "integer", "serial":
		return "int32"
	case "bigint", "bigserial":
		return "int"
	case "real":
		return "float32"
	case "double precision":
		return "float64"
	// Because we need to preserve base-10 precision.
	case "money":
		return "string"
	case "text", "varchar", "char", "character", "character varying", "uuid":
		return "string"
	case "bytea":
		return "[]byte"
	case "boolean":
		return "bool"
	case "timezone", "timezonetz", "date", "time":
		return "time.Time"
	case "interval":
		return "time.Duration"
	}
	return "string"
}

// Convert a SQL name to a Go name.
func goName(sqlName string) string {
	// This can definitely be done better.
	goName := strings.Replace(sqlName, "_", " ", -1)
	goName = strings.Replace(goName, ".", " ", -1)
	goName = strings.Title(goName)
	goName = strings.Replace(goName, " ", "", -1)

	return goName
}

// destutter removes a stutter prefix.
func destutter(str, prefix string) string {
	return strings.TrimPrefix(str, prefix)
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"io"
	"os"
	"strings"
	"text/template"

	"github.com/Masterminds/squirrel"
	"github.com/codegangsta/cli"

	_ "github.com/lib/pq"
)

const version = "DEV"

const Usage = `Read a schema and generate Structable structs.

This utility generates Structable structs be reading your database table and
generating the appropriate code.
`

const fileHeader = `package %s

// This file is automatically generated by schema2struct.

import (
	"time"

	"github.com/Masterminds/squirrel"
	"github.com/Masterminds/structable"
	_ "github.com/lib/pq"
)

// QueryFunc modifies a SelectBuilder prior to execution.
//
// The SelectBuilder is modified in place. An error is returned under any
// conditions where the query should not be executed.
type QueryFunc func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error)

`

const structTemplate = `// {{.StructName}} maps to database table {{.TableName}}
type {{.StructName}} struct {
	tableName string {{ann "tablename" .TableName}}
	structable.Recorder
	builder squirrel.StatementBuilderType
	{{range .Fields}}{{.}}
	{{end}}db squirrel.DBProxyBeginner
	flavor string
}

// New{{.StructName}} creates a new {{.StructName}} wired to structable.
func New{{.StructName}}(db squirrel.DBProxyBeginner, flavor string) *{{.StructName}} {
	o := &{{.StructName}}{db: db, flavor: flavor}
	o.Recorder = structable.New(db, flavor).Bind("{{.TableName}}", o)
	return o
}

// List{{.StructName}} returns a list of {{.StructName}} objects.
//
// Limit is the max number of items. Offset is the offset the results will
// begin with.
func List{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, limit, offset uint64) ([]*{{.StructName}}, error) {
	fn := func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error) {
		return q.Limit(limit).Offset(offset), nil
	}
	return Query{{.StructName}}(db, flavor, fn)
}

// Query{{.StructName}} builds a base query, but allows the query to be modified before execution.
//
// This creates a new Select, settings the columns and table name, and then calling QueryFunc with the
// query. The QueryFunc can then add a Where clause, etc. Provided QueryFunc does not exit with an
// error, Query{{.StructName}} will then execute the query, extract the results into a slice of
// {{.StructName}} structs, and then return.
//
// The QueryFunc should not modify the list of fields returned or the table name,
// as the intent is to construct a complete {{.StructName}} from each result.
// More sophisticated queries should be written directly.
func Query{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, fn QueryFunc) ([]*{{.StructName}}, error){
	var tn string = "{{.TableName}}"

	// We need a prototype structable to learn about the table structure.
	ps := New{{.StructName}}(db, flavor)
	cols := ps.Columns(true)

	q := ps.Builder().Select(cols...).From(tn)
	var err error
	if q, err = fn(q); err != nil {
		return []*{{.StructName}}{}, err
	}
	rows, err := q.Query()
	if err != nil || rows == nil {
		return []*{{.StructName}}{}, err
	}
	defer rows.Close()

	buf := []*{{.StructName}}{}
	for rows.Next() {
		o := New{{.StructName}}(db, flavor)
		dest := o.FieldReferences(true)
		if err := rows.Scan(dest...); err != nil {
			return buf, err
		}
		buf = append(buf, o)
	}
	return buf, rows.Err()
}

// Len{{.StructName}} returns the number of {{.StructName}} objects in the database.
func Len{{.StructName}}(db squirrel.DBProxyBeginner, flavor string) (int, error) {
	fn := func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error) {return q, nil}
	return QueryLen{{.StructName}}(db, flavor, fn)
}

// QueryLen{{.StructName}} returns the length of a table.
//
// The QueryFunc can be used to modify the query. For a simple length call, you
// may prefer to use Len{{.StructName}}.
func QueryLen{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, fn QueryFunc) (int, error) {
	tn := "{{.TableName}}"
	ps := New{{.StructName}}(db, flavor)
	q := ps.Builder().Select("COUNT(*)").From(tn)
	var err error
	if q, err = fn(q); err != nil {
		return 0, err
	}
	var count int
	err = q.Scan(&count)
	return count, err
}

`

type structDesc struct {
	StructName string
	TableName  string
	Fields     []string
}

func main() {
	app := cli.NewApp()
	app.Name = "schema2struct"
	app.Version = "version"
	app.Usage = Usage
	app.Action = importTables
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "driver,d",
			Value: "postgres",
			Usage: "The name of the SQL driver to use.",
		},
		cli.StringFlag{
			Name:  "connection,c",
			Value: "user=$USER dbname=$USER sslmode=disable",
			Usage: "The database connection string. Environment variables are expanded.",
		},
		cli.StringFlag{
			Name:  "tables,t",
			Value: "",
			Usage: "The list of tables to generate, comma separated. If none specified, the entire schema is used.",
		},
		cli.StringFlag{
			Name:  "file,f",
			Value: "",
			Usage: "The file to send the output.",
		},
		cli.StringFlag{
			Name:   "package,p",
			Value:  "main",
			Usage:  "The name of the destination package.",
			EnvVar: "GOPACKAGE",
		},
	}

	app.Run(os.Args)
}

func driver(c *cli.Context) string {
	return c.String("driver")
}
func conn(c *cli.Context) string {
	return os.ExpandEnv(c.String("connection"))
}

// dest gets the destination output writer.
func dest(c *cli.Context) io.Writer {
	if out := c.String("file"); out != "" {
		f, err := os.Create(out)
		if err != nil {
			panic(f)
		}
		return f
	}
	return os.Stdout
}

func tableList(c *cli.Context) []string {
	z := c.String("tables")
	if z != "" {
		return strings.Split(z, ",")
	}
	return []string{}
}

func cxdie(c *cli.Context, err error) {
	fmt.Fprintf(os.Stderr, "Failed to connect to %s (type %s): %s", conn(c), driver(c), err)
	os.Exit(1)
}

var funcMap = map[string]interface{}{
	"ann": func(tag, val string) string {
		return fmt.Sprintf("`%s:\"%s\"`", tag, val)
	},
}

func importTables(c *cli.Context) {
	ttt := template.Must(template.New("st").Funcs(funcMap).Parse(structTemplate))
	cxn, err := sql.Open(driver(c), conn(c))
	if err != nil {
		cxdie(c, err)
	}
	// Many drivers defer connections until the first statement. We test
	// that here.
	if err := cxn.Ping(); err != nil {
		cxdie(c, err)
	}
	defer cxn.Close()

	// Set up Squirrel
	stmts := squirrel.NewStmtCacher(cxn)
	bldr := squirrel.StatementBuilder.RunWith(stmts)
	if driver(c) == "postgres" {
		bldr = bldr.PlaceholderFormat(squirrel.Dollar)
	}

	// Set up destination
	out := dest(c)
	fmt.Fprintf(out, fileHeader, c.String("package"))

	tables := tableList(c)

	if len(tables) == 0 {
		tables, err = publicTables(bldr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Cannot fetch list of tables: %s\n", err)
			os.Exit(2)
		}
	}

	for _, t := range tables {
		f, err := importTable(t, bldr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to import table %s: %s", t, err)
		}

		//fmt.Fprintf(out, "%s %s %s\n", f.StructName, f.TableName, f.Fields)
		ttt.Execute(out, f)
	}
}

type column struct {
	Name, DataType string
	Max            int64
}

func publicTables(b squirrel.StatementBuilderType) ([]string, error) {
	rows, err := b.Select("table_name").From("INFORMATION_SCHEMA.TABLES").
		Where("table_schema = 'public'").Query()

	res := []string{}
	if err != nil {
		return res, err
	}

	for rows.Next() {
		var s string
		rows.Scan(&s)
		res = append(res, s)
	}

	return res, nil
}

// importTable reads a table definition and writes a corresponding struct.
// SELECT table_name, column_name, data_type, character_maximum_length
//   FROM INFORMATION_SCHEMA.COLUMNS WHERE table_name = 'goose_db_version'
func importTable(tbl string, b squirrel.StatementBuilderType) (*structDesc, error) {

	pks, err := primaryKeyField(tbl, b)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error getting primary keys: %s", err)
	}

	q := b.Select("column_name, data_type, character_maximum_length").
		From("INFORMATION_SCHEMA.COLUMNS").
		Where("table_name = ?", tbl)

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

	ff := []string{}
	for rows.Next() {
		c := &column{}
		var length sql.NullInt64
		if err := rows.Scan(&c.Name, &c.DataType, &length); err != nil {
			return nil, err
		}
		c.Max = length.Int64
		ff = append(ff, structField(c, pks, tbl, b))
	}
	sd := &structDesc{
		StructName: goName(tbl),
		TableName:  tbl,
		Fields:     ff,
	}

	return sd, nil
}

func primaryKeyField(tbl string, b squirrel.StatementBuilderType) ([]string, error) {
	q := b.Select("column_name").
		From("INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS c").
		LeftJoin("INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS t USING(constraint_name)").
		Where("t.table_name = ? AND t.constraint_type = 'PRIMARY KEY'", tbl).
		OrderBy("ordinal_position")

	rows, err := q.Query()
	if err != nil {
		return []string{}, err
	}

	res := []string{}
	for rows.Next() {
		var s string
		rows.Scan(&s)
		res = append(res, s)
	}
	return res, nil
}

func sequentialKey(tbl, pk string, b squirrel.StatementBuilderType) bool {

	tlen := 58

	stbl := tbl
	if len(tbl) > 29 {
		stbl = tbl[0:29]
	}

	left := tlen - len(stbl)
	spk := pk
	if len(pk) > left {
		spk = pk[0:left]
	}
	seq := fmt.Sprintf("%s_%s_seq", stbl, spk)

	q := b.Select("COUNT(*)").
		From("INFORMATION_SCHEMA.SEQUENCES").
		Where("sequence_name = ?", seq)

	var num int
	if err := q.Scan(&num); err != nil {
		panic(err)
	}
	return num > 0
}

func structField(c *column, pks []string, tbl string, b squirrel.StatementBuilderType) string {
	tpl := "%s %s `stbl:\"%s\"`"
	gn := destutter(goName(c.Name), goName(tbl))
	tt := goType(c.DataType)

	tag := c.Name
	for _, p := range pks {
		if c.Name == p {
			tag += ",PRIMARY_KEY"
			if sequentialKey(tbl, c.Name, b) {
				tag += ",SERIAL"
			}
		}
	}

	return fmt.Sprintf(tpl, gn, tt, tag)
}

// goType takes a SQL type and returns a string containin the name of a Go type.
//
// The goal is not to provide an exact match for every type, but to provide a
// safe Go representation of a SQL type.
//
// For some floating point SQL types, for example, we store them as strings
// so as not to lose precision while also not adding new types.
//
// The default type is string.
func goType(sqlType string) string {
	switch sqlType {
	case "smallint", "smallserial":
		return "int16"
	case "integer", "serial":
		return "int32"
	case "bigint", "bigserial":
		return "int"
	case "real":
		return "float32"
	case "double precision":
		return "float64"
	// Because we need to preserve base-10 precision.
	case "money":
		return "string"
	case "text", "varchar", "char", "character", "character varying", "uuid":
		return "string"
	case "bytea":
		return "[]byte"
	case "boolean":
		return "bool"
	case "timezone", "timezonetz", "date", "time":
		return "time.Time"
	case "interval":
		return "time.Duration"
	}
	return "string"
}

// Convert a SQL name to a Go name.
func goName(sqlName string) string {
	// This can definitely be done better.
	goName := strings.Replace(sqlName, "_", " ", -1)
	goName = strings.Replace(goName, ".", " ", -1)
	goName = strings.Title(goName)
	goName = strings.Replace(goName, " ", "", -1)

	return goName
}

// destutter removes a stutter prefix.
func destutter(str, prefix string) string {
	return strings.TrimPrefix(str, prefix)
}
<commit_msg>add support of mysql driver in schema2struct utility<commit_after>package main

import (
	"database/sql"
	"fmt"
	"io"
	"os"
	"strings"
	"text/template"

	"github.com/Masterminds/squirrel"
	"github.com/codegangsta/cli"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"
)

const version = "DEV"

// Usage : exported const Usage
const Usage = `Read a schema and generate Structable structs.

This utility generates Structable structs be reading your database table and
generating the appropriate code.
`

const fileHeader = `package %s

// This file is automatically generated by schema2struct.

import (
	"time"

	"github.com/Masterminds/squirrel"
	"github.com/Masterminds/structable"
	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"
)

// QueryFunc modifies a SelectBuilder prior to execution.
//
// The SelectBuilder is modified in place. An error is returned under any
// conditions where the query should not be executed.
type QueryFunc func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error)

`

const structTemplate = `// {{.StructName}} maps to database table {{.TableName}}
type {{.StructName}} struct {
	tableName string {{ann "tablename" .TableName}}
	structable.Recorder
	builder squirrel.StatementBuilderType
	{{range .Fields}}{{.}}
	{{end}}db squirrel.DBProxyBeginner
	flavor string
}

// New{{.StructName}} creates a new {{.StructName}} wired to structable.
func New{{.StructName}}(db squirrel.DBProxyBeginner, flavor string) *{{.StructName}} {
	o := &{{.StructName}}{db: db, flavor: flavor}
	o.Recorder = structable.New(db, flavor).Bind("{{.TableName}}", o)
	return o
}

// List{{.StructName}} returns a list of {{.StructName}} objects.
//
// Limit is the max number of items. Offset is the offset the results will
// begin with.
func List{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, limit, offset uint64) ([]*{{.StructName}}, error) {
	fn := func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error) {
		return q.Limit(limit).Offset(offset), nil
	}
	return Query{{.StructName}}(db, flavor, fn)
}

// Query{{.StructName}} builds a base query, but allows the query to be modified before execution.
//
// This creates a new Select, settings the columns and table name, and then calling QueryFunc with the
// query. The QueryFunc can then add a Where clause, etc. Provided QueryFunc does not exit with an
// error, Query{{.StructName}} will then execute the query, extract the results into a slice of
// {{.StructName}} structs, and then return.
//
// The QueryFunc should not modify the list of fields returned or the table name,
// as the intent is to construct a complete {{.StructName}} from each result.
// More sophisticated queries should be written directly.
func Query{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, fn QueryFunc) ([]*{{.StructName}}, error){
	var tn string = "{{.TableName}}"

	// We need a prototype structable to learn about the table structure.
	ps := New{{.StructName}}(db, flavor)
	cols := ps.Columns(true)

	q := ps.Builder().Select(cols...).From(tn)
	var err error
	if q, err = fn(q); err != nil {
		return []*{{.StructName}}{}, err
	}
	rows, err := q.Query()
	if err != nil || rows == nil {
		return []*{{.StructName}}{}, err
	}
	defer rows.Close()

	buf := []*{{.StructName}}{}
	for rows.Next() {
		o := New{{.StructName}}(db, flavor)
		dest := o.FieldReferences(true)
		if err := rows.Scan(dest...); err != nil {
			return buf, err
		}
		buf = append(buf, o)
	}
	return buf, rows.Err()
}

// Len{{.StructName}} returns the number of {{.StructName}} objects in the database.
func Len{{.StructName}}(db squirrel.DBProxyBeginner, flavor string) (int, error) {
	fn := func(q squirrel.SelectBuilder) (squirrel.SelectBuilder, error) {return q, nil}
	return QueryLen{{.StructName}}(db, flavor, fn)
}

// QueryLen{{.StructName}} returns the length of a table.
//
// The QueryFunc can be used to modify the query. For a simple length call, you
// may prefer to use Len{{.StructName}}.
func QueryLen{{.StructName}}(db squirrel.DBProxyBeginner, flavor string, fn QueryFunc) (int, error) {
	tn := "{{.TableName}}"
	ps := New{{.StructName}}(db, flavor)
	q := ps.Builder().Select("COUNT(*)").From(tn)
	var err error
	if q, err = fn(q); err != nil {
		return 0, err
	}
	var count int
	err = q.Scan(&count)
	return count, err
}

`

type structDesc struct {
	StructName string
	TableName  string
	Fields     []string
}

func main() {
	app := cli.NewApp()
	app.Name = "schema2struct"
	app.Version = "version"
	app.Usage = Usage
	app.Action = importTables
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "driver,d",
			Value: "postgres",
			Usage: "The name of the SQL driver to use.",
		},
		cli.StringFlag{
			Name:  "connection,c",
			Value: "user=$USER dbname=$USER sslmode=disable",
			Usage: "The database connection string. Environment variables are expanded.",
		},
		cli.StringFlag{
			Name:  "tables,t",
			Value: "",
			Usage: "The list of tables to generate, comma separated. If none specified, the entire schema is used.",
		},
		cli.StringFlag{
			Name:  "file,f",
			Value: "",
			Usage: "The file to send the output.",
		},
		cli.StringFlag{
			Name:   "package,p",
			Value:  "main",
			Usage:  "The name of the destination package.",
			EnvVar: "GOPACKAGE",
		},
	}

	app.Run(os.Args)
}

func driver(c *cli.Context) string {
	return c.String("driver")
}
func conn(c *cli.Context) string {
	return os.ExpandEnv(c.String("connection"))
}

// dest gets the destination output writer.
func dest(c *cli.Context) io.Writer {
	if out := c.String("file"); out != "" {
		f, err := os.Create(out)
		if err != nil {
			panic(f)
		}
		return f
	}
	return os.Stdout
}

func tableList(c *cli.Context) []string {
	z := c.String("tables")
	if z != "" {
		return strings.Split(z, ",")
	}
	return []string{}
}

func cxdie(c *cli.Context, err error) {
	fmt.Fprintf(os.Stderr, "Failed to connect to %s (type %s): %s", conn(c), driver(c), err)
	os.Exit(1)
}

var funcMap = map[string]interface{}{
	"ann": func(tag, val string) string {
		return fmt.Sprintf("`%s:\"%s\"`", tag, val)
	},
}

func importTables(c *cli.Context) {
	ttt := template.Must(template.New("st").Funcs(funcMap).Parse(structTemplate))
	cxn, err := sql.Open(driver(c), conn(c))
	if err != nil {
		cxdie(c, err)
	}
	// Many drivers defer connections until the first statement. We test
	// that here.
	if err := cxn.Ping(); err != nil {
		cxdie(c, err)
	}
	defer cxn.Close()

	// Set up Squirrel
	stmts := squirrel.NewStmtCacher(cxn)
	bldr := squirrel.StatementBuilder.RunWith(stmts)
	if driver(c) == "postgres" {
		bldr = bldr.PlaceholderFormat(squirrel.Dollar)
	}

	// Set up destination
	out := dest(c)
	fmt.Fprintf(out, fileHeader, c.String("package"))

	tables := tableList(c)

	if len(tables) == 0 {
		tables, err = publicTables(bldr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Cannot fetch list of tables: %s\n", err)
			os.Exit(2)
		}
	}

	for _, t := range tables {
		f, err := importTable(t, bldr, driver(c))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to import table %s: %s", t, err)
		}

		//fmt.Fprintf(out, "%s %s %s\n", f.StructName, f.TableName, f.Fields)
		ttt.Execute(out, f)
	}
}

type column struct {
	Name, DataType string
	Max            int64
}

func publicTables(b squirrel.StatementBuilderType) ([]string, error) {
	rows, err := b.Select("table_name").From("INFORMATION_SCHEMA.TABLES").
		Where("table_schema = 'public'").Query()

	res := []string{}
	if err != nil {
		return res, err
	}

	for rows.Next() {
		var s string
		rows.Scan(&s)
		res = append(res, s)
	}

	return res, nil
}

// importTable reads a table definition and writes a corresponding struct.
// SELECT table_name, column_name, data_type, character_maximum_length
//   FROM INFORMATION_SCHEMA.COLUMNS WHERE table_name = 'goose_db_version'
func importTable(tbl string, b squirrel.StatementBuilderType, driver string) (*structDesc, error) {

	pks, err := primaryKeyField(tbl, b)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error getting primary keys: %s", err)
	}

	q := b.Select("column_name, data_type, character_maximum_length").
		From("INFORMATION_SCHEMA.COLUMNS").
		Where("table_name = ?", tbl)

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

	ff := []string{}
	for rows.Next() {
		c := &column{}
		var length sql.NullInt64
		if err := rows.Scan(&c.Name, &c.DataType, &length); err != nil {
			return nil, err
		}
		c.Max = length.Int64
		switch driver {
		case "mysql":
			ff = append(ff, structFieldMySQL(c, pks, tbl, b))
		case "postgres":
			ff = append(ff, structField(c, pks, tbl, b))
		}
	}
	sd := &structDesc{
		StructName: goName(tbl),
		TableName:  tbl,
		Fields:     ff,
	}

	return sd, nil
}

func primaryKeyField(tbl string, b squirrel.StatementBuilderType) ([]string, error) {
	q := b.Select("column_name").
		From("INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS c").
		LeftJoin("INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS t USING(constraint_name)").
		Where("t.table_name = ? AND t.constraint_type = 'PRIMARY KEY'", tbl).
		OrderBy("ordinal_position")

	rows, err := q.Query()
	if err != nil {
		return []string{}, err
	}

	res := []string{}
	for rows.Next() {
		var s string
		rows.Scan(&s)
		res = append(res, s)
	}
	return res, nil
}

func autoincrementKey(tbl, pk string, b squirrel.StatementBuilderType) bool {
	q := b.Select("COUNT(*)").
		From("INFORMATION_SCHEMA.COLUMNS").
		Where("TABLE_NAME = ? AND COLUMN_NAME = ? AND EXTRA = 'auto_increment'", tbl, pk)
	var num int
	if err := q.Scan(&num); err != nil {
		panic(err)
	}
	return num > 0
}

func sequentialKey(tbl, pk string, b squirrel.StatementBuilderType) bool {
	tlen := 58

	stbl := tbl
	if len(tbl) > 29 {
		stbl = tbl[0:29]
	}

	left := tlen - len(stbl)
	spk := pk
	if len(pk) > left {
		spk = pk[0:left]
	}
	seq := fmt.Sprintf("%s_%s_seq", stbl, spk)

	q := b.Select("COUNT(*)").
		From("INFORMATION_SCHEMA.SEQUENCES").
		Where("sequence_name = ?", seq)

	var num int
	if err := q.Scan(&num); err != nil {
		panic(err)
	}
	return num > 0
}

func structFieldMySQL(c *column, pks []string, tbl string, b squirrel.StatementBuilderType) string {
	tpl := "%s %s `stbl:\"%s\"`"
	gn := destutter(goName(c.Name), goName(tbl))
	tt := goType(c.DataType)

	tag := c.Name
	for _, p := range pks {
		if c.Name == p {
			tag += ",PRIMARY_KEY"
			if autoincrementKey(tbl, c.Name, b) {
				tag += ",AUTO_INCREMENT"
			}
		}
	}

	return fmt.Sprintf(tpl, gn, tt, tag)
}

func structField(c *column, pks []string, tbl string, b squirrel.StatementBuilderType) string {
	tpl := "%s %s `stbl:\"%s\"`"
	gn := destutter(goName(c.Name), goName(tbl))
	tt := goType(c.DataType)

	tag := c.Name
	for _, p := range pks {
		if c.Name == p {
			tag += ",PRIMARY_KEY"
			if sequentialKey(tbl, c.Name, b) {
				tag += ",SERIAL"
			}
		}
	}

	return fmt.Sprintf(tpl, gn, tt, tag)
}

// goType takes a SQL type and returns a string containin the name of a Go type.
//
// The goal is not to provide an exact match for every type, but to provide a
// safe Go representation of a SQL type.
//
// For some floating point SQL types, for example, we store them as strings
// so as not to lose precision while also not adding new types.
//
// The default type is string.
func goType(sqlType string) string {
	switch sqlType {
	case "smallint", "smallserial":
		return "int16"
	case "integer", "serial":
		return "int32"
	case "bigint", "bigserial":
		return "int"
	case "real":
		return "float32"
	case "double precision":
		return "float64"
	// Because we need to preserve base-10 precision.
	case "money":
		return "string"
	case "text", "varchar", "char", "character", "character varying", "uuid":
		return "string"
	case "bytea":
		return "[]byte"
	case "boolean":
		return "bool"
	case "timezone", "timezonetz", "date", "time":
		return "time.Time"
	case "interval":
		return "time.Duration"
	}
	return "string"
}

// Convert a SQL name to a Go name.
func goName(sqlName string) string {
	// This can definitely be done better.
	goName := strings.Replace(sqlName, "_", " ", -1)
	goName = strings.Replace(goName, ".", " ", -1)
	goName = strings.Title(goName)
	goName = strings.Replace(goName, " ", "", -1)

	return goName
}

// destutter removes a stutter prefix.
func destutter(str, prefix string) string {
	return strings.TrimPrefix(str, prefix)
}
<|endoftext|>


<commit_before>// Copyright 2012-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// testramfs tests things, badly
package main

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"syscall"

	flag "github.com/spf13/pflag"

	"github.com/u-root/u-root/pkg/cpio"
	"github.com/u-root/u-root/pkg/pty"
)

const (
	unshareFlags = syscall.CLONE_NEWNS
	cloneFlags   = syscall.CLONE_NEWIPC |
		syscall.CLONE_NEWNET |
		// making newpid work will be more tricky,
		// since none of my CLs to fix go runtime for
		// it ever got in.
		//syscall.CLONE_NEWPID |
		syscall.CLONE_NEWUTS |
		0
)

var (
	noremove    = flag.BoolP("noremove", "n", false, "remove tempdir when done")
	interactive = flag.BoolP("interactive", "i", false, "interactive mode")
)

func main() {
	flag.Parse()

	if flag.NArg() != 1 {
		log.Fatalln("usage: %s <cpio-path>", os.Args[0])
	}

	c := flag.Args()[0]

	f, err := os.Open(c)
	if err != nil {
		log.Fatal(err)
	}

	// So, what's the plan here?
	//
	// - new mount namespace
	//   - root mount is a tmpfs mount filled with the archive.
	//
	// - new PID namespace
	//   - archive/init actually runs as PID 1.

	tempDir, err := ioutil.TempDir("", "u-root")
	if err != nil {
		log.Fatal(err)
	}
	// Don't do a RemoveAll. This should be empty and
	// an error can tell us we got something wrong.
	if !*noremove {
		defer func(n string) {
			log.Printf("Removing %v", n)
			if err := os.Remove(n); err != nil {
				log.Fatal(err)
			}
		}(tempDir)
	}
	if err := syscall.Mount("", tempDir, "tmpfs", 0, ""); err != nil {
		log.Fatal(err)
	}
	if !*noremove {
		defer func(n string) {
			log.Printf("Unmounting %v", n)
			if err := syscall.Unmount(n, syscall.MNT_DETACH); err != nil {
				log.Fatal(err)
			}
		}(tempDir)
	}

	archiver, err := cpio.Format("newc")
	if err != nil {
		log.Fatal(err)
	}

	r := archiver.Reader(f)
	for {
		rec, err := r.ReadRecord()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}
		cpio.CreateFileInRoot(rec, tempDir)
	}

	cmd, err := pty.New()
	if err != nil {
		log.Fatal(err)
	}
	cmd.Command("/init")
	cmd.C.SysProcAttr.Chroot = tempDir
	cmd.C.SysProcAttr.Cloneflags = cloneFlags
	cmd.C.SysProcAttr.Unshareflags = cloneFlags
	if *interactive {
		if err := cmd.Run(); err != nil {
			log.Fatal(err)
		}
		return
	}
	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}
	go io.Copy(cmd.TTY, cmd.Ptm)

	// At this point you could use an array of commands/output templates to
	// drive the test, and end with the exit command shown nere.
	if n, err := cmd.Ptm.Write([]byte("exit\n")); err != nil {
		log.Printf("Writing exit: want (5, nil); got (%d, %v)\n", n, err)
	}
	if err := cmd.Wait(); err != nil {
		log.Fatal(err)
	}

}
<commit_msg>testramfs: set umask to 0, restore it on exit<commit_after>// Copyright 2012-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// testramfs tests things, badly
package main

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"syscall"

	flag "github.com/spf13/pflag"
	"golang.org/x/sys/unix"

	"github.com/u-root/u-root/pkg/cpio"
	"github.com/u-root/u-root/pkg/pty"
)

const (
	unshareFlags = syscall.CLONE_NEWNS
	cloneFlags   = syscall.CLONE_NEWIPC |
		syscall.CLONE_NEWNET |
		// making newpid work will be more tricky,
		// since none of my CLs to fix go runtime for
		// it ever got in.
		//syscall.CLONE_NEWPID |
		syscall.CLONE_NEWUTS |
		0
)

var (
	noremove    = flag.BoolP("noremove", "n", false, "remove tempdir when done")
	interactive = flag.BoolP("interactive", "i", false, "interactive mode")
)

func main() {
	flag.Parse()

	if flag.NArg() != 1 {
		log.Fatalln("usage: %s <cpio-path>", os.Args[0])
	}

	c := flag.Args()[0]

	f, err := os.Open(c)
	if err != nil {
		log.Fatal(err)
	}

	// So, what's the plan here?
	//
	// - new mount namespace
	//   - root mount is a tmpfs mount filled with the archive.
	//
	// - new PID namespace
	//   - archive/init actually runs as PID 1.

	// Note this is basically a chroot and umask is inherited.
	// The umask has to be zero else some creation will end
	// up with incorrect permissions, a particular problem
	// in device creation.
	u := unix.Umask(0)
	defer unix.Umask(u)

	tempDir, err := ioutil.TempDir("", "u-root")
	if err != nil {
		log.Fatal(err)
	}
	// Don't do a RemoveAll. This should be empty and
	// an error can tell us we got something wrong.
	if !*noremove {
		defer func(n string) {
			log.Printf("Removing %v", n)
			if err := os.Remove(n); err != nil {
				log.Fatal(err)
			}
		}(tempDir)
	}
	if err := syscall.Mount("", tempDir, "tmpfs", 0, ""); err != nil {
		log.Fatal(err)
	}
	if !*noremove {
		defer func(n string) {
			log.Printf("Unmounting %v", n)
			if err := syscall.Unmount(n, syscall.MNT_DETACH); err != nil {
				log.Fatal(err)
			}
		}(tempDir)
	}

	archiver, err := cpio.Format("newc")
	if err != nil {
		log.Fatal(err)
	}

	r := archiver.Reader(f)
	for {
		rec, err := r.ReadRecord()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}
		cpio.CreateFileInRoot(rec, tempDir)
	}

	cmd, err := pty.New()
	if err != nil {
		log.Fatal(err)
	}
	cmd.Command("/init")
	cmd.C.SysProcAttr.Chroot = tempDir
	cmd.C.SysProcAttr.Cloneflags = cloneFlags
	cmd.C.SysProcAttr.Unshareflags = cloneFlags
	if *interactive {
		if err := cmd.Run(); err != nil {
			log.Fatal(err)
		}
		return
	}
	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}
	go io.Copy(cmd.TTY, cmd.Ptm)

	// At this point you could use an array of commands/output templates to
	// drive the test, and end with the exit command shown nere.
	if n, err := cmd.Ptm.Write([]byte("exit\n")); err != nil {
		log.Printf("Writing exit: want (5, nil); got (%d, %v)\n", n, err)
	}
	if err := cmd.Wait(); err != nil {
		log.Fatal(err)
	}

}
<|endoftext|>


<commit_before>package discovery

import (
	version "github.com/hashicorp/go-version"
)

// A ConstraintStr is a string containing a possibly-invalid representation
// of a version constraint provided in configuration. Call Parse on it to
// obtain a real Constraint object, or discover that it is invalid.
type ConstraintStr string

// Parse transforms a ConstraintStr into a VersionSet if it is
// syntactically valid. If it isn't then an error is returned instead.
func (s ConstraintStr) Parse() (VersionSet, error) {
	raw, err := version.NewConstraint(string(s))
	if err != nil {
		return VersionSet{}, err
	}
	return VersionSet{raw}, nil
}

// MustParse is like Parse but it panics if the constraint string is invalid.
func (s ConstraintStr) MustParse() VersionSet {
	ret, err := s.Parse()
	if err != nil {
		panic(err)
	}
	return ret
}

// VersionSet represents a set of versions which any given Version is either
// a member of or not.
type VersionSet struct {
	// Internally a version set is actually a list of constraints that
	// *remove* versions from the set. Thus a VersionSet with an empty
	// Constraints list would be one that contains *all* versions.
	raw version.Constraints
}

// Has returns true if the given version is in the receiving set.
func (s VersionSet) Has(v Version) bool {
	return s.raw.Check(v.raw)
}

// Intersection combines the receving set with the given other set to produce a
// set that is the intersection of both sets, which is to say that it contains
// only the versions that are members of both sets.
func (s VersionSet) Intersection(other VersionSet) VersionSet {
	raw := make(version.Constraints, 0, len(s.raw)+len(other.raw))

	// Since "raw" is a list of constraints that remove versions from the set,
	// "Intersection" is implemented by concatenating together those lists,
	// thus leaving behind only the versions not removed by either list.
	raw = append(raw, s.raw...)
	raw = append(raw, other.raw...)

	return VersionSet{raw}
}

// String returns a string representation of the set members as a set
// of range constraints.
func (s VersionSet) String() string {
	return s.raw.String()
}
<commit_msg>plugin/discovery: provide an AllVersions set<commit_after>package discovery

import (
	version "github.com/hashicorp/go-version"
)

// A ConstraintStr is a string containing a possibly-invalid representation
// of a version constraint provided in configuration. Call Parse on it to
// obtain a real Constraint object, or discover that it is invalid.
type ConstraintStr string

// Parse transforms a ConstraintStr into a VersionSet if it is
// syntactically valid. If it isn't then an error is returned instead.
func (s ConstraintStr) Parse() (VersionSet, error) {
	raw, err := version.NewConstraint(string(s))
	if err != nil {
		return VersionSet{}, err
	}
	return VersionSet{raw}, nil
}

// MustParse is like Parse but it panics if the constraint string is invalid.
func (s ConstraintStr) MustParse() VersionSet {
	ret, err := s.Parse()
	if err != nil {
		panic(err)
	}
	return ret
}

// VersionSet represents a set of versions which any given Version is either
// a member of or not.
type VersionSet struct {
	// Internally a version set is actually a list of constraints that
	// *remove* versions from the set. Thus a VersionSet with an empty
	// Constraints list would be one that contains *all* versions.
	raw version.Constraints
}

// AllVersions is a VersionSet containing all versions
var AllVersions VersionSet

func init() {
	AllVersions = VersionSet{
		raw: make(version.Constraints, 0),
	}
}

// Has returns true if the given version is in the receiving set.
func (s VersionSet) Has(v Version) bool {
	return s.raw.Check(v.raw)
}

// Intersection combines the receving set with the given other set to produce a
// set that is the intersection of both sets, which is to say that it contains
// only the versions that are members of both sets.
func (s VersionSet) Intersection(other VersionSet) VersionSet {
	raw := make(version.Constraints, 0, len(s.raw)+len(other.raw))

	// Since "raw" is a list of constraints that remove versions from the set,
	// "Intersection" is implemented by concatenating together those lists,
	// thus leaving behind only the versions not removed by either list.
	raw = append(raw, s.raw...)
	raw = append(raw, other.raw...)

	return VersionSet{raw}
}

// String returns a string representation of the set members as a set
// of range constraints.
func (s VersionSet) String() string {
	return s.raw.String()
}
<|endoftext|>


<commit_before>package plugins

import (
	"fmt"

	"github.com/InnovaCo/serve/manifest"
)

func init() {
	manifest.PluginRegestry.Add("db.create.postgresql", DBCreatePostgresql{})
}

type DBCreatePostgresql struct{}

func (p DBCreatePostgresql) Run(data manifest.Manifest) error {
	if data.GetBool("purge") {
		return p.Drop(data)
	} else {
		return p.Create(data)
	}
}

func (p DBCreatePostgresql) Create(data manifest.Manifest) error {
	var cmd string

	if data.Has("source") {
		s := data.GetString("source")
		t := data.GetString("target")
		cmd = fmt.Sprintf("sudo -Hu postgres createdb -O "+
			"`sudo -Hu postgres psql postgres -c \"SELECT d.datname, pg_catalog.pg_get_userbyid(d.datdba) " +
			"FROM pg_catalog.pg_database d "+
			"WHERE d.datname='%s' ORDER BY 1;\" | grep %s | awk '{print $3}'` \"%s\" && pg_dump \"%s\" | " +
			"psql \"%s\"", s, s, t, s, t)

	} else {
		cmd = fmt.Sprintf("sudo -EHu postgres createdb -O %s \"%s\"", data.GetStringOr("db-user", "postgres"),
																	  data.GetString("target"))
	}

	return runSshCmd(data.GetString("host"), data.GetString("ssh-user"), cmd)
}

func (p DBCreatePostgresql) Drop(data manifest.Manifest) error {
	return runSshCmd(
		data.GetString("host"),
		data.GetString("ssh-user"),
		fmt.Sprintf("sudo -EHu postgres dropdb \"%s\"", data.GetString("target")),
	)
}
<commit_msg>fix db.create.postgresql plugin<commit_after>package plugins

import (
	"fmt"

	"github.com/InnovaCo/serve/manifest"
)

func init() {
	manifest.PluginRegestry.Add("db.create.postgresql", DBCreatePostgresql{})
}

type DBCreatePostgresql struct{}

func (p DBCreatePostgresql) Run(data manifest.Manifest) error {
	if data.GetBool("purge") {
		return p.Drop(data)
	} else {
		return p.Create(data)
	}
}

func (p DBCreatePostgresql) Create(data manifest.Manifest) error {
	var cmd string

	if data.Has("source") {
		t := data.GetString("target")
		cmd = fmt.Sprintf("sudo -Hu postgres createdb -O %s \"%s\" && pg_dump \"%s\" | psql \"%s\"",
			              data.GetStringOr("db-user", "postgres"), t, data.GetString("source"), t)

	} else {
		cmd = fmt.Sprintf("sudo -Hu postgres createdb -O %s \"%s\"",
			              data.GetStringOr("db-user", "postgres"), data.GetString("target"))
	}

	return runSshCmd(data.GetString("host"), data.GetString("ssh-user"), cmd)
}

func (p DBCreatePostgresql) Drop(data manifest.Manifest) error {
	return runSshCmd(
		data.GetString("host"),
		data.GetString("ssh-user"),
		fmt.Sprintf("sudo -EHu postgres dropdb \"%s\"", data.GetString("target")),
	)
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Cisco and/or its affiliates.
//
// 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 cmdimpl

import (
	"encoding/json"
	"fmt"
	nodeinfomodel "github.com/contiv/vpp/plugins/contiv/model/node"

	"github.com/contiv/vpp/plugins/crd/cache/telemetrymodel"
	"github.com/contiv/vpp/plugins/netctl/http"
	"github.com/coreos/etcd/clientv3"
	"github.com/ligato/cn-infra/db/keyval/etcd"
	"github.com/ligato/cn-infra/logging"
	"github.com/ligato/cn-infra/logging/logrus"
	"os"
	"strings"
	"text/tabwriter"
	"time"
)

//PrintNodes will print out all of the cmdimpl in a network in a table format.
func PrintNodes() {
	cfg := &etcd.ClientConfig{
		Config: &clientv3.Config{
			Endpoints: []string{"127.0.0.1:32379"},
		},
		OpTimeout: 1 * time.Second,
	}
	logger := logrus.DefaultLogger()
	logger.SetLevel(logging.FatalLevel)
	w := tabwriter.NewWriter(os.Stdout, 0, 8, 2, ' ', 0)
	// w := tabwriter.NewWriter(os.Stdout, 0, 8, 4, '\t', 0)
	// Create connection to etcd.
	db, err := etcd.NewEtcdConnectionWithBytes(*cfg, logger)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	itr, err := db.ListValues("/vnf-agent/contiv-ksr/allocatedIDs/")
	if err != nil {
		fmt.Printf("Error getting values")
		return
	}
	fmt.Fprintf(w, "ID\tNAME\tVPP-IP\tHOST-IP\tBUILD-DATE\tBUILD-VERSION\tSTART-TIME\tSTATE\n")
	for {
		kv, stop := itr.GetNext()
		if stop {
			break
		}
		buf := kv.GetValue()
		nodeInfo := &nodeinfomodel.NodeInfo{}
		err = json.Unmarshal(buf, nodeInfo)
		//fmt.Printf("NodeInfo: %+v\n", nodeInfo)
		// Do whatever processing we need to do
		bytes := http.GetNodeInfo(nodeInfo.ManagementIpAddress, "liveness")
		var liveness telemetrymodel.NodeLiveness
		err = json.Unmarshal(bytes, &liveness)
		if err != nil {
			fmt.Println(err)
			liveness.BuildDate = "Not Available"
		}

		fmt.Fprintf(w, "%d\t%s\t%s\t%s\t%s\t%s\t%s\t%d\n",
			nodeInfo.Id,
			nodeInfo.Name,
			strings.Split(nodeInfo.IpAddress, "/")[0],
			nodeInfo.ManagementIpAddress,
			liveness.BuildDate,
			liveness.BuildVersion,
			time.Unix(int64(liveness.StartTime), 0),
			liveness.State)
	}
	w.Flush()
	db.Close()
}
<commit_msg>Improve readability of the ipam command output<commit_after>// Copyright (c) 2018 Cisco and/or its affiliates.
//
// 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 cmdimpl

import (
	"encoding/json"
	"fmt"
	nodeinfomodel "github.com/contiv/vpp/plugins/contiv/model/node"

	"github.com/contiv/vpp/plugins/crd/cache/telemetrymodel"
	"github.com/contiv/vpp/plugins/netctl/http"
	"github.com/coreos/etcd/clientv3"
	"github.com/ligato/cn-infra/db/keyval/etcd"
	"github.com/ligato/cn-infra/logging"
	"github.com/ligato/cn-infra/logging/logrus"
	"os"
	"strings"
	"text/tabwriter"
	"time"
)

const timeLayout = "Mon Jan _2 15:04:05 2006"

//PrintNodes will print out all of the cmdimpl in a network in a table format.
func PrintNodes() {
	cfg := &etcd.ClientConfig{
		Config: &clientv3.Config{
			Endpoints: []string{"127.0.0.1:32379"},
		},
		OpTimeout: 1 * time.Second,
	}
	logger := logrus.DefaultLogger()
	logger.SetLevel(logging.FatalLevel)
	w := tabwriter.NewWriter(os.Stdout, 0, 8, 2, ' ', 0)
	// w := tabwriter.NewWriter(os.Stdout, 0, 8, 4, '\t', 0)
	// Create connection to etcd.
	db, err := etcd.NewEtcdConnectionWithBytes(*cfg, logger)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	itr, err := db.ListValues("/vnf-agent/contiv-ksr/allocatedIDs/")
	if err != nil {
		fmt.Printf("Error getting values")
		return
	}
	fmt.Fprintf(w, "ID\tNAME\tVPP-IP\tHOST-IP\tSTART-TIME\tSTATE\tBUILD-VERSION\tBUILD-DATE\n")
	for {
		kv, stop := itr.GetNext()
		if stop {
			break
		}

		buf := kv.GetValue()
		nodeInfo := &nodeinfomodel.NodeInfo{}
		err = json.Unmarshal(buf, nodeInfo)
		//fmt.Printf("NodeInfo: %+v\n", nodeInfo)
		// Do whatever processing we need to do
		bytes := http.GetNodeInfo(nodeInfo.ManagementIpAddress, "liveness")
		var liveness telemetrymodel.NodeLiveness
		err = json.Unmarshal(bytes, &liveness)

		if err != nil {
			fmt.Println(err)
			liveness.BuildDate = "Not Available"
		}

		buildDate := liveness.BuildDate
		bd, err1 := time.Parse("2006-01-02T15:04+00:00", buildDate)
		if err1 == nil {
			buildDate = bd.Format(timeLayout)
		}

		fmt.Fprintf(w, "%d\t%s\t%s\t%s\t%s\t%d\t%s\t%s\n",
			nodeInfo.Id,
			nodeInfo.Name,
			strings.Split(nodeInfo.IpAddress, "/")[0],
			nodeInfo.ManagementIpAddress,
			time.Unix(int64(liveness.StartTime), 0).Format(timeLayout),
			liveness.State,
			liveness.BuildVersion,
			buildDate)
	}
	w.Flush()
	db.Close()
}
<|endoftext|>


<commit_before>// Copyright 2015 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package policy

import (
	"database/sql"
	"fmt"
	"strings"

	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"

	gorp "github.com/letsencrypt/boulder/Godeps/_workspace/src/gopkg.in/gorp.v1"
)

var errDBFailure = core.InternalServerError("Error checking policy DB.")

const whitelisted = "whitelist"
const blacklisted = "blacklist"

type domainRule struct {
	Host string `db:"host"`
}

// BlacklistRule is used to hold rules blacklisting a DNS name
type BlacklistRule domainRule

// WhitelistRule is used to hold rules whitelisting a DNS name
type WhitelistRule domainRule

// RawRuleSet describes the rule set file format
type RawRuleSet struct {
	Blacklist []string
	Whitelist []string
}

// RuleSet describes the rules to load into the policy database
type RuleSet struct {
	Blacklist []BlacklistRule
	Whitelist []WhitelistRule
}

type gorpDbMap interface {
	AddTableWithName(interface{}, string) *gorp.TableMap
	Begin() (*gorp.Transaction, error)
	SelectOne(interface{}, string, ...interface{}) error
	Select(interface{}, string, ...interface{}) ([]interface{}, error)
}

// PolicyAuthorityDatabaseImpl enforces policy decisions based on various rule
// lists
type PolicyAuthorityDatabaseImpl struct {
	log   *blog.AuditLogger
	dbMap gorpDbMap
}

// NewPolicyAuthorityDatabaseImpl constructs a Policy Authority Database (and
// creates tables if they are non-existent)
func NewPolicyAuthorityDatabaseImpl(dbMap gorpDbMap) (padb *PolicyAuthorityDatabaseImpl, err error) {
	logger := blog.GetAuditLogger()

	dbMap.AddTableWithName(BlacklistRule{}, "blacklist")
	dbMap.AddTableWithName(WhitelistRule{}, "whitelist")

	padb = &PolicyAuthorityDatabaseImpl{
		dbMap: dbMap,
		log:   logger,
	}

	return padb, nil
}

// LoadRules loads the whitelist and blacklist into the database in a transaction
// deleting any previous content
func (padb *PolicyAuthorityDatabaseImpl) LoadRules(rs RuleSet) error {
	tx, err := padb.dbMap.Begin()
	if err != nil {
		tx.Rollback()
		return err
	}
	_, err = tx.Exec("DELETE FROM blacklist")
	if err != nil {
		tx.Rollback()
		return err
	}
	for _, r := range rs.Blacklist {
		r.Host = core.ReverseName(r.Host)
		tx.Insert(&r)
	}
	_, err = tx.Exec("DELETE FROM whitelist")
	if err != nil {
		tx.Rollback()
		return err
	}
	for _, r := range rs.Whitelist {
		tx.Insert(&r)
	}

	err = tx.Commit()
	return err
}

// DumpRules retrieves all domainRules in the database so they can be written to
// disk
func (padb *PolicyAuthorityDatabaseImpl) DumpRules() (rs RuleSet, err error) {
	var bList []BlacklistRule
	_, err = padb.dbMap.Select(&bList, "SELECT * FROM blacklist")
	if err != nil {
		return
	}
	for _, r := range bList {
		r.Host = core.ReverseName(r.Host)
		rs.Blacklist = append(rs.Blacklist, r)
	}
	var wList []WhitelistRule
	_, err = padb.dbMap.Select(&wList, "SELECT * FROM whitelist")
	if err != nil {
		return
	}
	rs.Whitelist = wList
	return rs, err
}

// allowedByBlacklist returns nil if the host is allowed, errBlacklisted if the
// host is disallowed, or an InternalServerError if there was another problem
// checking the database.
func (padb *PolicyAuthorityDatabaseImpl) allowedByBlacklist(host string) error {
	var rule BlacklistRule
	// Use lexical ordering to quickly find blacklisted root domains
	err := padb.dbMap.SelectOne(
		&rule,
		`SELECT * FROM blacklist WHERE :host >= host ORDER BY host DESC LIMIT 1`,
		map[string]interface{}{"host": host},
	)
	if err != nil {
		// No rows means not blacklisted, so no error.
		if err == sql.ErrNoRows {
			return nil
		}
		padb.log.Err(fmt.Sprintf("Error checking policy DB: %s", err))
		return errDBFailure
	}
	if host == rule.Host || strings.HasPrefix(host, rule.Host+".") {
		return errBlacklisted
	}
	// If we got a result but it's not a match, that means the host is not
	// blacklisted.
	return nil
}

func (padb *PolicyAuthorityDatabaseImpl) allowedByWhitelist(host string) bool {
	var rule WhitelistRule
	err := padb.dbMap.SelectOne(
		&rule,
		`SELECT * FROM whitelist WHERE :host = host LIMIT 1`,
		map[string]interface{}{"host": host},
	)
	if err != nil {
		if err == sql.ErrNoRows {
			return false
		}
		return false
	}
	return true
}

// CheckHostLists will query the database for white/blacklist rules that match host,
// if both whitelist and blacklist rules are found the blacklist will always win
func (padb *PolicyAuthorityDatabaseImpl) CheckHostLists(host string, requireWhitelisted bool) error {
	if requireWhitelisted {
		if !padb.allowedByWhitelist(host) {
			return errNotWhitelisted
		}
	}
	// Overrides the whitelist if a blacklist rule is found
	host = core.ReverseName(host)
	return padb.allowedByBlacklist(host)
}
<commit_msg>Update comment.<commit_after>// Copyright 2015 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package policy

import (
	"database/sql"
	"fmt"
	"strings"

	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"

	gorp "github.com/letsencrypt/boulder/Godeps/_workspace/src/gopkg.in/gorp.v1"
)

var errDBFailure = core.InternalServerError("Error checking policy DB.")

const whitelisted = "whitelist"
const blacklisted = "blacklist"

type domainRule struct {
	Host string `db:"host"`
}

// BlacklistRule is used to hold rules blacklisting a DNS name
type BlacklistRule domainRule

// WhitelistRule is used to hold rules whitelisting a DNS name
type WhitelistRule domainRule

// RawRuleSet describes the rule set file format
type RawRuleSet struct {
	Blacklist []string
	Whitelist []string
}

// RuleSet describes the rules to load into the policy database
type RuleSet struct {
	Blacklist []BlacklistRule
	Whitelist []WhitelistRule
}

type gorpDbMap interface {
	AddTableWithName(interface{}, string) *gorp.TableMap
	Begin() (*gorp.Transaction, error)
	SelectOne(interface{}, string, ...interface{}) error
	Select(interface{}, string, ...interface{}) ([]interface{}, error)
}

// PolicyAuthorityDatabaseImpl enforces policy decisions based on various rule
// lists
type PolicyAuthorityDatabaseImpl struct {
	log   *blog.AuditLogger
	dbMap gorpDbMap
}

// NewPolicyAuthorityDatabaseImpl constructs a Policy Authority Database (and
// creates tables if they are non-existent)
func NewPolicyAuthorityDatabaseImpl(dbMap gorpDbMap) (padb *PolicyAuthorityDatabaseImpl, err error) {
	logger := blog.GetAuditLogger()

	dbMap.AddTableWithName(BlacklistRule{}, "blacklist")
	dbMap.AddTableWithName(WhitelistRule{}, "whitelist")

	padb = &PolicyAuthorityDatabaseImpl{
		dbMap: dbMap,
		log:   logger,
	}

	return padb, nil
}

// LoadRules loads the whitelist and blacklist into the database in a transaction
// deleting any previous content
func (padb *PolicyAuthorityDatabaseImpl) LoadRules(rs RuleSet) error {
	tx, err := padb.dbMap.Begin()
	if err != nil {
		tx.Rollback()
		return err
	}
	_, err = tx.Exec("DELETE FROM blacklist")
	if err != nil {
		tx.Rollback()
		return err
	}
	for _, r := range rs.Blacklist {
		r.Host = core.ReverseName(r.Host)
		tx.Insert(&r)
	}
	_, err = tx.Exec("DELETE FROM whitelist")
	if err != nil {
		tx.Rollback()
		return err
	}
	for _, r := range rs.Whitelist {
		tx.Insert(&r)
	}

	err = tx.Commit()
	return err
}

// DumpRules retrieves all domainRules in the database so they can be written to
// disk
func (padb *PolicyAuthorityDatabaseImpl) DumpRules() (rs RuleSet, err error) {
	var bList []BlacklistRule
	_, err = padb.dbMap.Select(&bList, "SELECT * FROM blacklist")
	if err != nil {
		return
	}
	for _, r := range bList {
		r.Host = core.ReverseName(r.Host)
		rs.Blacklist = append(rs.Blacklist, r)
	}
	var wList []WhitelistRule
	_, err = padb.dbMap.Select(&wList, "SELECT * FROM whitelist")
	if err != nil {
		return
	}
	rs.Whitelist = wList
	return rs, err
}

// allowedByBlacklist returns nil if the host is allowed, errBlacklisted if the
// host is disallowed, or an InternalServerError if there was another problem
// checking the database.
func (padb *PolicyAuthorityDatabaseImpl) allowedByBlacklist(host string) error {
	var rule BlacklistRule
	// Use lexical ordering to quickly find blacklisted root domains
	err := padb.dbMap.SelectOne(
		&rule,
		`SELECT * FROM blacklist WHERE :host >= host ORDER BY host DESC LIMIT 1`,
		map[string]interface{}{"host": host},
	)
	if err != nil {
		// No rows means not blacklisted, so no error.
		if err == sql.ErrNoRows {
			return nil
		}
		padb.log.Err(fmt.Sprintf("Error checking policy DB: %s", err))
		return errDBFailure
	}
	if host == rule.Host || strings.HasPrefix(host, rule.Host+".") {
		return errBlacklisted
	}
	// If we got a result but it's not a match, that means the host is not
	// blacklisted.
	return nil
}

func (padb *PolicyAuthorityDatabaseImpl) allowedByWhitelist(host string) bool {
	var rule WhitelistRule
	err := padb.dbMap.SelectOne(
		&rule,
		`SELECT * FROM whitelist WHERE :host = host LIMIT 1`,
		map[string]interface{}{"host": host},
	)
	if err != nil {
		if err == sql.ErrNoRows {
			return false
		}
		return false
	}
	return true
}

// CheckHostLists will query the database for white/blacklist rules that match host,
// if both whitelist and blacklist rules are found the blacklist will always win
// Returns errNotWhitelisted, errBlacklisted, or errDBFailure for the
// appropriate problems, or nil if the host is allowable.
func (padb *PolicyAuthorityDatabaseImpl) CheckHostLists(host string, requireWhitelisted bool) error {
	if requireWhitelisted {
		if !padb.allowedByWhitelist(host) {
			return errNotWhitelisted
		}
	}
	// Overrides the whitelist if a blacklist rule is found
	host = core.ReverseName(host)
	return padb.allowedByBlacklist(host)
}
<|endoftext|>


<commit_before>// Copyright 2013 Joyent Inc.
// Licensed under the AGPLv3, see LICENCE file for details.

package joyent_test

import (
	stdtesting "testing"

	"github.com/juju/testing"

	gc "gopkg.in/check.v1"
)

func TestPackage(t *stdtesting.T) {
	if testing.RaceEnabled {
		t.Skip("skipping package under -race, see LP 1497801")
	}
	registerLocalTests()
	gc.TestingT(t)
}
<commit_msg>provider/joyent: re-enable data race tests<commit_after>// Copyright 2013 Joyent Inc.
// Licensed under the AGPLv3, see LICENCE file for details.

package joyent_test

import (
	stdtesting "testing"

	gc "gopkg.in/check.v1"
)

func TestPackage(t *stdtesting.T) {
	registerLocalTests()
	gc.TestingT(t)
}
<|endoftext|>


<commit_before>package twitter

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
	"github.com/gin-gonic/gin"
	"github.com/kurrik/oauth1a"
	"github.com/kurrik/twittergo"
	"net/http"
	"net/url"
	"os"
	"time"
)

type Twitter struct{}

var (
	TWITTER_COSUMER_KEY     = os.Getenv("TWITTER_CONSUMER_KEY")
	TWITTER_CONSUMER_SECRET = os.Getenv("TWITTER_CONSUMER_SECRET")
	conn, redisErr          = redis.Dial("tcp", os.Getenv("REDISTOGO"))
)

func (tw Twitter) LoadCredentials() (client *twittergo.Client, err error) {
	config := &oauth1a.ClientConfig{
		ConsumerKey:    TWITTER_COSUMER_KEY,
		ConsumerSecret: TWITTER_CONSUMER_SECRET,
	}
	user := oauth1a.NewAuthorizedConfig("", "")
	client = twittergo.NewClient(config, user)
	return
}

func (tw Twitter) Memoize(resp map[string]string, tweetId uint64, key string, value string) map[string]string {
	cache, _ := redis.String(conn.Do("HGET", fmt.Sprintf("tweet:%v", tweetId), key))
    if cache == "" {
		conn.Do("HSET", fmt.Sprintf("tweet:%v", tweetId), key, value)
		cache = value
	}
	resp[key] = cache
	return resp
}

func (tw Twitter) TweetsFromResults(c *gin.Context, results *twittergo.SearchResults, t chan map[string]string, done chan bool) {

	fmt.Printf("start TweetsFromResults \n")

	for _, tweet := range results.Statuses() {

		user := tweet.User()
		entities := tweet["entities"].(map[string]interface{})
		urls := entities["urls"].([]interface{})

		if len(urls) > 0 {
			url := urls[0].(map[string]interface{})

			if redisErr != nil {
				fmt.Printf("Redis error: %v\n", redisErr)
			}

			resp := make(map[string]string)

			reply, _ := redis.Values(conn.Do("KEYS", fmt.Sprintf("tweet:%v", tweet.Id())))
			fmt.Printf("values: %v", reply)
			if len(reply) == 0 {
				conn.Do("RPUSH", "tweets", tweet.Id())
			}

			resp = tw.Memoize(resp, tweet.Id(), "tweet", fmt.Sprintf("%v", tweet.Text()))
			resp = tw.Memoize(resp, tweet.Id(), "image", fmt.Sprintf("%v", url["expanded_url"]))
			resp = tw.Memoize(resp, tweet.Id(), "user", fmt.Sprintf("%v (@%v) ", user.Name(), user.ScreenName()))
			resp = tw.Memoize(resp, tweet.Id(), "created", fmt.Sprintf("%v", tweet.CreatedAt().Format(time.RFC1123)))
			t <- resp
		}
	}
	fmt.Printf("end TweetsFromResults \n")
}

func (tw Twitter) RetrieveMaxId() (string, bool){
	l, _ := redis.Int(conn.Do("LLEN", "tweets"))
	reply, _ := redis.String( conn.Do("LINDEX", "tweets", l-1) )
	if reply == "" {
		return reply, false
	} else {
		return reply, true
	}
}

func (tw Twitter) ConstructParams() url.Values {
	query := url.Values{}
	query.Set("q", "#beautifulsky")
	query.Set("result_type", "mixed")
	query.Set("count", "100")

	if maxId, present := tw.RetrieveMaxId(); present {
		fmt.Printf("maxId: %v", maxId)
		query.Set("max_id", maxId)
	}
	return query
}

func (tw Twitter) SearchTweets( k chan *twittergo.SearchResults, r chan *twittergo.APIResponse) {
	var (
		err     error
		results *twittergo.SearchResults
		req     *http.Request
		resp    *twittergo.APIResponse
		client *twittergo.Client
	)
	client, _ = tw.LoadCredentials()

	query := tw.ConstructParams()

	url := fmt.Sprintf("/1.1/search/tweets.json?%v", query.Encode())
	req, err = http.NewRequest("GET", url, nil)
	if err != nil {
		fmt.Printf("Could not parse request: %v\n", err)
	}

	resp, err = client.SendRequest(req)
	if err != nil {
		fmt.Printf("Could not send request: %v\n", err)
	}

	results = &twittergo.SearchResults{}
	err = resp.Parse(results)
	if err != nil {
		fmt.Printf("Problem parsing response: %v\n", err)
	}

	k <- results
	close(k)
	r <- resp
}

func (tw Twitter) TweetsFromCache(t chan map[string]string, cacheDone chan bool) {
    conn, redisErr = redis.Dial("tcp", os.Getenv("REDISTOGO"))
	l, _ := redis.Int(conn.Do("LLEN", "tweets"))
	for i:=0; i< l; i++ {
		tweetId, _ := redis.String(conn.Do("LINDEX", "tweets", i))
		var resp map[string]string = make(map[string]string)
		for _, key := range []string{"tweet","image","user","created"} {
			reply,_ := redis.String( conn.Do("HGET", "tweet:"+tweetId, key) )
			resp[key] = reply
		}
		t <- resp
	}
	cacheDone <- true
}

func (tw Twitter) PrintRateLimit(r chan *twittergo.APIResponse) {
	fmt.Printf("start PrintRateLimit \n")
	resp := <-r
	if resp.HasRateLimit() {
		fmt.Printf("Rate limit:           %v\n", resp.RateLimit())
		fmt.Printf("Rate limit remaining: %v\n", resp.RateLimitRemaining())
		fmt.Printf("Rate limit reset:     %v\n", resp.RateLimitReset())
	} else {
		fmt.Printf("Could not parse rate limit from response.\n")
	}
	fmt.Printf("end PrintRateLimit \n")
}
<commit_msg>add url parsing for redistogo<commit_after>package twitter

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
	"github.com/gin-gonic/gin"
	"github.com/kurrik/oauth1a"
	"github.com/kurrik/twittergo"
	"net/http"
	"net/url"
	"os"
	"time"
)

type Twitter struct{}

var (
	TWITTER_COSUMER_KEY     = os.Getenv("TWITTER_CONSUMER_KEY")
	TWITTER_CONSUMER_SECRET = os.Getenv("TWITTER_CONSUMER_SECRET")
	REDISTOGO, _ = url.Parse(os.Getenv("REDISTOGO"))
	conn, redisErr = redis.Dial("tcp", REDISTOGO.Host)
)

func (tw Twitter) LoadCredentials() (client *twittergo.Client, err error) {
	config := &oauth1a.ClientConfig{
		ConsumerKey:    TWITTER_COSUMER_KEY,
		ConsumerSecret: TWITTER_CONSUMER_SECRET,
	}
	user := oauth1a.NewAuthorizedConfig("", "")
	client = twittergo.NewClient(config, user)
	return
}

func (tw Twitter) Memoize(resp map[string]string, tweetId uint64, key string, value string) map[string]string {
	cache, _ := redis.String(conn.Do("HGET", fmt.Sprintf("tweet:%v", tweetId), key))
    if cache == "" {
		conn.Do("HSET", fmt.Sprintf("tweet:%v", tweetId), key, value)
		cache = value
	}
	resp[key] = cache
	return resp
}

func (tw Twitter) TweetsFromResults(c *gin.Context, results *twittergo.SearchResults, t chan map[string]string, done chan bool) {

	fmt.Printf("start TweetsFromResults \n")

	for _, tweet := range results.Statuses() {

		user := tweet.User()
		entities := tweet["entities"].(map[string]interface{})
		urls := entities["urls"].([]interface{})

		if len(urls) > 0 {
			url := urls[0].(map[string]interface{})

			if redisErr != nil {
				fmt.Printf("Redis error: %v\n", redisErr)
			}

			resp := make(map[string]string)

			reply, _ := redis.Values(conn.Do("KEYS", fmt.Sprintf("tweet:%v", tweet.Id())))
			fmt.Printf("values: %v", reply)
			if len(reply) == 0 {
				conn.Do("RPUSH", "tweets", tweet.Id())
			}

			resp = tw.Memoize(resp, tweet.Id(), "tweet", fmt.Sprintf("%v", tweet.Text()))
			resp = tw.Memoize(resp, tweet.Id(), "image", fmt.Sprintf("%v", url["expanded_url"]))
			resp = tw.Memoize(resp, tweet.Id(), "user", fmt.Sprintf("%v (@%v) ", user.Name(), user.ScreenName()))
			resp = tw.Memoize(resp, tweet.Id(), "created", fmt.Sprintf("%v", tweet.CreatedAt().Format(time.RFC1123)))
			t <- resp
		}
	}
	fmt.Printf("end TweetsFromResults \n")
}

func (tw Twitter) RetrieveMaxId() (string, bool){
	l, _ := redis.Int(conn.Do("LLEN", "tweets"))
	reply, _ := redis.String( conn.Do("LINDEX", "tweets", l-1) )
	if reply == "" {
		return reply, false
	} else {
		return reply, true
	}
}

func (tw Twitter) ConstructParams() url.Values {
	query := url.Values{}
	query.Set("q", "#beautifulsky")
	query.Set("result_type", "mixed")
	query.Set("count", "100")

	if maxId, present := tw.RetrieveMaxId(); present {
		fmt.Printf("maxId: %v", maxId)
		query.Set("max_id", maxId)
	}
	return query
}

func (tw Twitter) SearchTweets( k chan *twittergo.SearchResults, r chan *twittergo.APIResponse) {
	var (
		err     error
		results *twittergo.SearchResults
		req     *http.Request
		resp    *twittergo.APIResponse
		client *twittergo.Client
	)
	client, _ = tw.LoadCredentials()

	query := tw.ConstructParams()

	url := fmt.Sprintf("/1.1/search/tweets.json?%v", query.Encode())
	req, err = http.NewRequest("GET", url, nil)
	if err != nil {
		fmt.Printf("Could not parse request: %v\n", err)
	}

	resp, err = client.SendRequest(req)
	if err != nil {
		fmt.Printf("Could not send request: %v\n", err)
	}

	results = &twittergo.SearchResults{}
	err = resp.Parse(results)
	if err != nil {
		fmt.Printf("Problem parsing response: %v\n", err)
	}

	k <- results
	close(k)
	r <- resp
}

func (tw Twitter) TweetsFromCache(t chan map[string]string, cacheDone chan bool) {
    conn, redisErr = redis.Dial("tcp", REDISTOGO.Host)
	l, _ := redis.Int(conn.Do("LLEN", "tweets"))
	for i:=0; i< l; i++ {
		tweetId, _ := redis.String(conn.Do("LINDEX", "tweets", i))
		var resp map[string]string = make(map[string]string)
		for _, key := range []string{"tweet","image","user","created"} {
			reply,_ := redis.String( conn.Do("HGET", "tweet:"+tweetId, key) )
			resp[key] = reply
		}
		t <- resp
	}
	cacheDone <- true
}

func (tw Twitter) PrintRateLimit(r chan *twittergo.APIResponse) {
	fmt.Printf("start PrintRateLimit \n")
	resp := <-r
	if resp.HasRateLimit() {
		fmt.Printf("Rate limit:           %v\n", resp.RateLimit())
		fmt.Printf("Rate limit remaining: %v\n", resp.RateLimitRemaining())
		fmt.Printf("Rate limit reset:     %v\n", resp.RateLimitReset())
	} else {
		fmt.Printf("Could not parse rate limit from response.\n")
	}
	fmt.Printf("end PrintRateLimit \n")
}
<|endoftext|>


<commit_before>package udp

import (
	"testing"
	"time"

	"github.com/hsheth2/gonet/ipv4"
)

const rwport = 20102

func TestReadWriteLocal(t *testing.T) {
	readWriteTest(t, ipv4.LoopbackIPAddress, 0)
}

func TestReadWriteLocalFragmentation(t *testing.T) {
	readWriteTest(t, ipv4.LoopbackIPAddress, 10)
}

func TestReadWriteExternal(t *testing.T) {
	t.Skip("External tests actually don't work")
	readWriteTest(t, ipv4.ExternalIPAddress, 0)
}

func readWriteTest(t *testing.T, ip *ipv4.Address, exp int) {
	success := make(chan bool, 1)

	r, err := NewReader(rwport, ip)
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()

	data := []byte{'h', 'e', 'l', 'l', 'o'}
	for i := 0; i < exp; i++ {
		data = append(data, data...)
	}

	go func() {
		w, err := NewWriter(20000, rwport, ip)
		if err != nil {
			t.Fatal(err)
		}

		_, err = w.Write(data)
		if err != nil {
			t.Fatal(err)
		} else {
			t.Log("Wrote the data:", data)
		}

		w.Close()
	}()

	go func() {
		//time.Sleep(10*time.Second)
		p, err := r.Read(maxUDPPacketLength)
		if err != nil {
			t.Fatal(err)
		}
		t.Log("Output:", string(p))

		if string(p) == string(data) {
			t.Log("Got correct output:", p)
			success <- true
		} else {
			t.Error("Got Wrong Output:", p)
		}
	}()

	select {
	case <-success:
		t.Log("Success")
	case <-time.After(5 * time.Second):
		t.Error("Timed out")
	}
}
<commit_msg>Passed in data into goroutine as a parameter<commit_after>package udp

import (
	"testing"
	"time"

	"github.com/hsheth2/gonet/ipv4"
)

const rwport = 20102

func TestReadWriteLocal(t *testing.T) {
	readWriteTest(t, ipv4.LoopbackIPAddress, 0)
}

func TestReadWriteLocalFragmentation(t *testing.T) {
	readWriteTest(t, ipv4.LoopbackIPAddress, 10)
}

func TestReadWriteExternal(t *testing.T) {
	t.Skip("External tests actually don't work")
	readWriteTest(t, ipv4.ExternalIPAddress, 0)
}

func readWriteTest(t *testing.T, ip *ipv4.Address, exp int) {
	success := make(chan bool, 1)

	r, err := NewReader(rwport, ip)
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()

	data := []byte{'h', 'e', 'l', 'l', 'o'}
	for i := 0; i < exp; i++ {
		data = append(data, data...)
	}

	go func(data []byte) {
		w, err := NewWriter(20000, rwport, ip)
		if err != nil {
			t.Fatal(err)
		}

		_, err = w.Write(data)
		if err != nil {
			t.Fatal(err)
		} else {
			t.Log("Wrote the data:", data)
		}

		w.Close()
	}(data)

	go func() {
		//time.Sleep(10*time.Second)
		p, err := r.Read(maxUDPPacketLength)
		if err != nil {
			t.Fatal(err)
		}
		t.Log("Output:", string(p))

		if string(p) == string(data) {
			t.Log("Got correct output:", p)
			success <- true
		} else {
			t.Error("Got Wrong Output:", p)
		}
	}()

	select {
	case <-success:
		t.Log("Success")
	case <-time.After(5 * time.Second):
		t.Error("Timed out")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"./help"
)

type MaxIntHeap struct {
	data     []int
	count    int
	capacity int
}

func (heap *MaxIntHeap) InsertSort(d int) {
	if heap.capacity < heap.count+1 {
		fmt.Println("out of bounds")
		return
	}

	heap.data[heap.count+1] = d
	heap.count++
	heap.shfitUp(heap.count)

}

func (heap *MaxIntHeap) shfitUp(k int) {
	for k > 1 && heap.data[k] > heap.data[k/2] {
		heap.data[k], heap.data[k/2] = heap.data[k/2], heap.data[k]
		k /= 2
	}

}

func (heap *MaxIntHeap) Size() int {
	return heap.count

}

func (heap *MaxIntHeap) Init(arr []int) {

	n := len(arr)
	heap.capacity = n
	heap.count = n

	heap.data = make([]int, n+1)
	for i := 0; i < n; i++ {
		heap.data[i+1] = arr[i]
	}

	for i := len(arr) / 2; i > 0; i-- {
		heap.shiftDown(i)
	}

}

func (heap *MaxIntHeap) shiftDown(k int) {
	for 2*k <= heap.count {
		j := k * 2
		if j+1 <= heap.count && heap.data[j+1] > heap.data[j] {
			j++
		}
		if heap.data[k] > heap.data[j] {
			break
		}
		heap.data[j], heap.data[k] = heap.data[k], heap.data[j]
		k = j
	}

}

func (heap *MaxIntHeap) ExtractMax() int {
	item := heap.data[1]
	heap.data[1], heap.data[heap.count] = heap.data[heap.count], heap.data[1]
	heap.data = heap.data[0:heap.count]
	heap.count--
	heap.shiftDown(1)
	return item
}

func HeapSort(arr []int, n int) []int {
	maxIntHeap := MaxIntHeap{}
	maxIntHeap.Init(arr)
	for i := n - 1; i >= 0; i-- {
		arr[i] = maxIntHeap.ExtractMax()

	}
	return arr

}

func main() {

	arr := help.GenRanArray(15, 1, 9)
	fmt.Println(arr)
	help.TestSort(HeapSort)(arr, len(arr))
	fmt.Println(arr)
	// arr := make([]int, 6)
	// fmt.Println(arr)

}
<commit_msg>调整最大堆排序参数<commit_after>package main

import (
	"fmt"

	"./help"
)

type MaxIntHeap struct {
	data     []int
	count    int
	capacity int
}

func (heap *MaxIntHeap) InsertSort(d int) {
	if heap.capacity < heap.count+1 {
		fmt.Println("out of bounds")
		return
	}

	heap.data[heap.count+1] = d
	heap.count++
	heap.shfitUp(heap.count)

}

func (heap *MaxIntHeap) shfitUp(k int) {
	for k > 1 && heap.data[k] > heap.data[k/2] {
		heap.data[k], heap.data[k/2] = heap.data[k/2], heap.data[k]
		k /= 2
	}

}

func (heap *MaxIntHeap) Size() int {
	return heap.count

}

func (heap *MaxIntHeap) Init(arr []int) {

	n := len(arr)
	heap.capacity = n
	heap.count = n

	heap.data = make([]int, n+1)
	for i := 0; i < n; i++ {
		heap.data[i+1] = arr[i]
	}

	for i := len(arr) / 2; i > 0; i-- {
		heap.shiftDown(i)
	}

}

func (heap *MaxIntHeap) shiftDown(k int) {
	for 2*k <= heap.count {
		j := k * 2
		if j+1 <= heap.count && heap.data[j+1] > heap.data[j] {
			j++
		}
		if heap.data[k] > heap.data[j] {
			break
		}
		heap.data[j], heap.data[k] = heap.data[k], heap.data[j]
		k = j
	}

}

func (heap *MaxIntHeap) ExtractMax() int {
	item := heap.data[1]
	heap.data[1], heap.data[heap.count] = heap.data[heap.count], heap.data[1]
	heap.data = heap.data[0:heap.count]
	heap.count--
	heap.shiftDown(1)
	return item
}

func HeapSort(arr []int, n int) []int {
	maxIntHeap := MaxIntHeap{}
	maxIntHeap.Init(arr)
	for i := n - 1; i >= 0; i-- {
		arr[i] = maxIntHeap.ExtractMax()

	}
	return arr

}

func main() {
	arr := help.GenRanArray(10000000, 1, 999999)
	help.TestSort(HeapSort)(arr, len(arr))
}
<|endoftext|>


<commit_before>package bitbucket

import (
	"encoding/json"
	"os"

	"github.com/k0kubun/pp"
)

type BranchRestrictions struct {
	c *Client
}

func (b *BranchRestrictions) Gets(bo *BranchRestrictionsOptions) (interface{}, error) {
	urlStr := b.c.requestUrl("/repositories/%s/%s/branch-restrictions", bo.Owner, bo.RepoSlug)
	return b.c.execute("GET", urlStr, "")
}

func (b *BranchRestrictions) Create(bo *BranchRestrictionsOptions) (interface{}, error) {
	data := b.buildBranchRestrictionsBody(bo)
	urlStr := b.c.requestUrl("/repositories/%s/%s/branch-restrictions", bo.Owner, bo.RepoSlug)
	return b.c.execute("POST", urlStr, data)
}

func (b *BranchRestrictions) Get(bo *BranchRestrictionsOptions) (interface{}, error) {
	urlStr := b.c.requestUrl("/repositories/%s/%s/branch-restrictions/%s", bo.Owner, bo.RepoSlug, bo.ID)
	return b.c.execute("GET", urlStr, "")
}

func (b *BranchRestrictions) Update(bo *BranchRestrictionsOptions) (interface{}, error) {
	data := b.buildBranchRestrictionsBody(bo)
	urlStr := b.c.requestUrl("/repositories/%s/%s/branch-restrictions/%s", bo.Owner, bo.RepoSlug, bo.ID)
	return b.c.execute("PUT", urlStr, data)
}

func (b *BranchRestrictions) Delete(bo *BranchRestrictionsOptions) (interface{}, error) {
	urlStr := b.c.requestUrl("/repositories/%s/%s/branch-restrictions/%s", bo.Owner, bo.RepoSlug, bo.ID)
	return b.c.execute("DELETE", urlStr, "")
}

type branchRestrictionsBody struct {
	Kind    string `json:"kind"`
	Pattern string `json:"pattern"`
	Links   struct {
		Self struct {
			Href string `json:"href"`
		} `json:"self"`
	} `json:"links"`
	Value  interface{}                   `json:"value"`
	ID     int                           `json:"id"`
	Users  []branchRestrictionsBodyUser  `json:"users"`
	Groups []branchRestrictionsBodyGroup `json:"groups"`
}

type branchRestrictionsBodyGroup struct {
	Name  string `json:"name"`
	Links struct {
		Self struct {
			Href string `json:"href"`
		} `json:"self"`
		Html struct {
			Href string `json:"href"`
		} `json:"html"`
		FullSlug string `json:"full_slug"`
		Members  int    `json:"members"`
		Slug     string `json:"slug"`
	} `json:"links"`
}

type branchRestrictionsBodyUser struct {
	Username     string `json:"username"`
	Website      string `json:"website"`
	Display_name string `json:"display_name"`
	UUID         string `json:"uuid"`
	Created_on   string `json:"created_on"`
	Links        struct {
		Self struct {
			Href string `json:"href"`
		} `json:"self"`
		Repositories struct {
			Href string `json:"href"`
		} `json:"repositories"`
		Html struct {
			Href string `json:"href"`
		} `json:"html"`
		Followers struct {
			Href string `json:"href"`
		} `json:"followers"`
		Avatar struct {
			Href string `json:"href"`
		} `json:"avatar"`
		Following struct {
			Href string `json:"href"`
		} `json:"following"`
	} `json:"links"`
}

func (b *BranchRestrictions) buildBranchRestrictionsBody(bo *BranchRestrictionsOptions) string {

	var users []branchRestrictionsBodyUser
	var groups []branchRestrictionsBodyGroup
	for _, u := range bo.Users {
		user := branchRestrictionsBodyUser{
			Username: u,
		}
		users = append(users, user)
	}
	for _, g := range bo.Groups {
		group := branchRestrictionsBodyGroup{
			Name: g,
		}
		groups = append(groups, group)
	}

	body := branchRestrictionsBody{
		Kind:    bo.Kind,
		Pattern: bo.Pattern,
		Users:   users,
		Groups:  groups,
		Value:   bo.Value,
	}

	data, err := json.Marshal(body)
	if err != nil {
		pp.Println(err)
		os.Exit(9)
	}

	return string(data)
}
<commit_msg>add decode to BranchRestrictions (#129)<commit_after>package bitbucket

import (
	"encoding/json"
	"os"

	"github.com/k0kubun/pp"
	"github.com/mitchellh/mapstructure"
)

type BranchRestrictions struct {
	c *Client

	ID      int
	Pattern string
	Kind    string
	Value   *int
}

func (b *BranchRestrictions) Gets(bo *BranchRestrictionsOptions) (interface{}, error) {
	urlStr := b.c.requestUrl("/repositories/%s/%s/branch-restrictions", bo.Owner, bo.RepoSlug)
	return b.c.execute("GET", urlStr, "")
}

func (b *BranchRestrictions) Create(bo *BranchRestrictionsOptions) (*BranchRestrictions, error) {
	data := b.buildBranchRestrictionsBody(bo)
	urlStr := b.c.requestUrl("/repositories/%s/%s/branch-restrictions", bo.Owner, bo.RepoSlug)
	response, err := b.c.execute("POST", urlStr, data)
	if err != nil {
		return nil, err
	}

	return decodeBranchRestriction(response)
}

func (b *BranchRestrictions) Get(bo *BranchRestrictionsOptions) (*BranchRestrictions, error) {
	urlStr := b.c.requestUrl("/repositories/%s/%s/branch-restrictions/%s", bo.Owner, bo.RepoSlug, bo.ID)
	response, err := b.c.execute("GET", urlStr, "")
	if err != nil {
		return nil, err
	}

	return decodeBranchRestriction(response)
}

func (b *BranchRestrictions) Update(bo *BranchRestrictionsOptions) (interface{}, error) {
	data := b.buildBranchRestrictionsBody(bo)
	urlStr := b.c.requestUrl("/repositories/%s/%s/branch-restrictions/%s", bo.Owner, bo.RepoSlug, bo.ID)
	response, err := b.c.execute("PUT", urlStr, data)
	if err != nil {
		return nil, err
	}

	return decodeBranchRestriction(response)
}

func (b *BranchRestrictions) Delete(bo *BranchRestrictionsOptions) (interface{}, error) {
	urlStr := b.c.requestUrl("/repositories/%s/%s/branch-restrictions/%s", bo.Owner, bo.RepoSlug, bo.ID)
	return b.c.execute("DELETE", urlStr, "")
}

type branchRestrictionsBody struct {
	Kind    string `json:"kind"`
	Pattern string `json:"pattern"`
	Links   struct {
		Self struct {
			Href string `json:"href"`
		} `json:"self"`
	} `json:"links"`
	Value  interface{}                   `json:"value"`
	ID     int                           `json:"id"`
	Users  []branchRestrictionsBodyUser  `json:"users"`
	Groups []branchRestrictionsBodyGroup `json:"groups"`
}

type branchRestrictionsBodyGroup struct {
	Name  string `json:"name"`
	Links struct {
		Self struct {
			Href string `json:"href"`
		} `json:"self"`
		Html struct {
			Href string `json:"href"`
		} `json:"html"`
		FullSlug string `json:"full_slug"`
		Members  int    `json:"members"`
		Slug     string `json:"slug"`
	} `json:"links"`
}

type branchRestrictionsBodyUser struct {
	Username     string `json:"username"`
	Website      string `json:"website"`
	Display_name string `json:"display_name"`
	UUID         string `json:"uuid"`
	Created_on   string `json:"created_on"`
	Links        struct {
		Self struct {
			Href string `json:"href"`
		} `json:"self"`
		Repositories struct {
			Href string `json:"href"`
		} `json:"repositories"`
		Html struct {
			Href string `json:"href"`
		} `json:"html"`
		Followers struct {
			Href string `json:"href"`
		} `json:"followers"`
		Avatar struct {
			Href string `json:"href"`
		} `json:"avatar"`
		Following struct {
			Href string `json:"href"`
		} `json:"following"`
	} `json:"links"`
}

func (b *BranchRestrictions) buildBranchRestrictionsBody(bo *BranchRestrictionsOptions) string {

	var users []branchRestrictionsBodyUser
	var groups []branchRestrictionsBodyGroup
	for _, u := range bo.Users {
		user := branchRestrictionsBodyUser{
			Username: u,
		}
		users = append(users, user)
	}
	for _, g := range bo.Groups {
		group := branchRestrictionsBodyGroup{
			Name: g,
		}
		groups = append(groups, group)
	}

	body := branchRestrictionsBody{
		Kind:    bo.Kind,
		Pattern: bo.Pattern,
		Users:   users,
		Groups:  groups,
		Value:   bo.Value,
	}

	data, err := json.Marshal(body)
	if err != nil {
		pp.Println(err)
		os.Exit(9)
	}

	return string(data)
}

func decodeBranchRestriction(branchResponse interface{}) (*BranchRestrictions, error) {
	branchMap := branchResponse.(map[string]interface{})

	if branchMap["type"] == "error" {
		return nil, DecodeError(branchMap)
	}

	var branchRestriction = new(BranchRestrictions)
	err := mapstructure.Decode(branchMap, branchRestriction)
	if err != nil {
		return nil, err
	}
	return branchRestriction, nil
}
<|endoftext|>