<commit_before>package util

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/openshift/origin/pkg/client"
	"github.com/openshift/origin/pkg/cmd/cli"
	cmdapi "github.com/openshift/origin/pkg/cmd/cli/cmd"
	"github.com/openshift/origin/pkg/cmd/cli/config"
	configapi "github.com/openshift/origin/pkg/cmd/server/api"
	testutil "github.com/openshift/origin/test/util"
	"github.com/spf13/cobra"
	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client"
	clientcmd "k8s.io/kubernetes/pkg/client/clientcmd"
	"k8s.io/kubernetes/test/e2e"
)

// CLI provides function to call the OpenShift CLI and Kubernetes and OpenShift
// REST clients.
type CLI struct {
	verb            string
	configPath      string
	adminConfigPath string
	username        string
	outputDir       string
	globalArgs      []string
	commandArgs     []string
	finalArgs       []string
	stdout          io.Writer
	verbose         bool
	cmd             *cobra.Command
	kubeFramework   *e2e.Framework
}

// NewCLI initialize the upstream E2E framework and set the namespace to match
// with the project name. Note that this function does not initialize the project
// role bindings for the namespace.
func NewCLI(project, adminConfigPath string) *CLI {
	client := &CLI{}
	client.kubeFramework = e2e.InitializeFramework(project, client.SetupProject)
	client.outputDir = os.TempDir()
	client.username = "admin"
	if len(adminConfigPath) == 0 {
		FatalErr(fmt.Errorf("You must set the KUBECONFIG variable to admin kubeconfig."))
	}
	client.adminConfigPath = adminConfigPath
	return client
}

// KubeFramework returns Kubernetes framework which contains helper functions
// specific for Kubernetes resources
func (c *CLI) KubeFramework() *e2e.Framework {
	return c.kubeFramework
}

// Username returns the name of currently logged user. If there is no user assigned
// for the current session, it returns 'admin'.
func (c *CLI) Username() string {
	return c.username
}

// ChangeUser changes the user used by the current CLI session.
func (c *CLI) ChangeUser(name string) *CLI {
	adminClientConfig, err := testutil.GetClusterAdminClientConfig(c.adminConfigPath)
	if err != nil {
		FatalErr(err)
	}
	_, _, clientConfig, err := testutil.GetClientForUser(*adminClientConfig, name)
	if err != nil {
		FatalErr(err)
	}

	kubeConfig, err := config.CreateConfig(c.Namespace(), clientConfig)
	if err != nil {
		FatalErr(err)
	}

	c.configPath = filepath.Join(c.outputDir, name+".kubeconfig")
	err = clientcmd.WriteToFile(*kubeConfig, c.configPath)
	if err != nil {
		FatalErr(err)
	}

	c.username = name
	fmt.Printf("INFO: configPath is now %q\n", c.configPath)
	return c
}

// SetNamespace sets a new namespace
func (c *CLI) SetNamespace(ns string) *CLI {
	c.kubeFramework.Namespace = &kapi.Namespace{
		ObjectMeta: kapi.ObjectMeta{
			Name: ns,
		},
	}
	return c
}

// SetOutputDir change the default output directory for temporary files
func (c *CLI) SetOutputDir(dir string) *CLI {
	c.outputDir = dir
	return c
}

// SetupProject creates a new project and assign a random user to the project.
// All resources will be then created within this project and Kubernetes E2E
// suite will destroy the project after test case finish.
// Note that the kubeClient is not used and serves just to make this function
// compatible with upstream function.
func (c *CLI) SetupProject(name string, kubeClient *kclient.Client) (*kapi.Namespace, error) {
	newNamespace := kapi.SimpleNameGenerator.GenerateName(fmt.Sprintf("extended-test-%s-", name))
	c.SetNamespace(newNamespace).ChangeUser(fmt.Sprintf("%s-user", c.Namespace()))
	e2e.Logf("The user is now %q", c.Username())

	projectOpts := cmdapi.NewProjectOptions{
		ProjectName: c.Namespace(),
		Client:      c.REST(),
		Out:         c.stdout,
	}
	e2e.Logf("Creating project %q", c.Namespace())
	return c.kubeFramework.Namespace, projectOpts.Run()
}

// Verbose turns on printing verbose messages when executing OpenShift commands
func (c *CLI) Verbose() *CLI {
	c.verbose = true
	return c
}

// REST provides an OpenShift REST client for the current user. If the user is not
// set, then it provides REST client for the cluster admin user
func (c *CLI) REST() *client.Client {
	_, clientConfig, err := configapi.GetKubeClient(c.configPath)
	osClient, err := client.New(clientConfig)
	if err != nil {
		FatalErr(err)
	}
	return osClient
}

// AdminREST provides an OpenShift REST client for the cluster admin user.
func (c *CLI) AdminREST() *client.Client {
	_, clientConfig, err := configapi.GetKubeClient(c.adminConfigPath)
	osClient, err := client.New(clientConfig)
	if err != nil {
		FatalErr(err)
	}
	return osClient
}

// KubeREST provides a Kubernetes REST client for the current namespace
func (c *CLI) KubeREST() *kclient.Client {
	kubeClient, _, err := configapi.GetKubeClient(c.configPath)
	if err != nil {
		FatalErr(err)
	}
	return kubeClient
}

// Namespace returns the name of the namespace used in the current test case.
// If the namespace is not set, an empty string is returned.
func (c *CLI) Namespace() string {
	if c.kubeFramework.Namespace == nil {
		return ""
	}
	return c.kubeFramework.Namespace.Name
}

// SetOutput allows to override the default command output
func (c *CLI) SetOutput(out io.Writer) *CLI {
	c.stdout = out
	for _, subCmd := range c.cmd.Commands() {
		subCmd.SetOutput(c.stdout)
	}
	c.cmd.SetOutput(c.stdout)
	return c
}

// Run executes given OpenShift CLI command verb (iow. "oc <verb>").
// This function also override the default 'stdout' to redirect all output
// to a buffer and prepare the global flags such as namespace and config path.
func (c *CLI) Run(verb string) *CLI {
	out := new(bytes.Buffer)
	nc := &CLI{
		verb:            verb,
		kubeFramework:   c.KubeFramework(),
		adminConfigPath: c.adminConfigPath,
		configPath:      c.configPath,
		username:        c.username,
		outputDir:       c.outputDir,
		cmd:             cli.NewCommandCLI("oc", "openshift", out),
		globalArgs: []string{
			verb,
			fmt.Sprintf("--namespace=%s", c.Namespace()),
			fmt.Sprintf("--config=%s", c.configPath),
		},
	}
	return nc.SetOutput(out)
}

// Template sets a Go template for the OpenShift CLI command.
// This is equivalent of running "oc get foo -o template -t '{{ .spec }}'"
func (c *CLI) Template(t string) *CLI {
	if c.verb != "get" {
		FatalErr("Cannot use Template() for non-get verbs.")
	}
	templateArgs := []string{"--output=template", fmt.Sprintf("--template=%s", t)}
	commandArgs := append(c.commandArgs, templateArgs...)
	c.finalArgs = append(c.globalArgs, commandArgs...)
	c.cmd.SetArgs(c.finalArgs)
	return c
}

// Args sets the additional arguments for the OpenShift CLI command
func (c *CLI) Args(args ...string) *CLI {
	c.commandArgs = args
	c.finalArgs = append(c.globalArgs, c.commandArgs...)
	c.cmd.SetArgs(c.finalArgs)
	return c
}

func (c *CLI) printCmd() string {
	return strings.Join(c.finalArgs, " ")
}

// Output executes the command and return the output as string
func (c *CLI) Output() (string, error) {
	if c.verbose {
		fmt.Printf("DEBUG: oc %s\n", c.printCmd())
	}
	err := c.cmd.Execute()
	if err != nil {
		FatalErr(err)
	}
	if c.verbose {
		fmt.Printf("DEBUG: %q\n", trimmedOutput(c.stdout))
	}
	return trimmedOutput(c.stdout), err
}

// OutputToFile executes the command and store output to a file
func (c *CLI) OutputToFile(filename string) (string, error) {
	content, err := c.Output()
	if err != nil {
		return "", err
	}
	path := filepath.Join(c.outputDir, c.Namespace()+"-"+filename)
	return path, ioutil.WriteFile(path, []byte(content), 0644)
}

// Execute executes the current command and return error if the execution failed
// This function will set the default output to stdout.
func (c *CLI) Execute() error {
	out, err := c.Output()
	if err != nil {
		FatalErr(fmt.Errorf("%v", err))
	}
	if _, err := io.Copy(os.Stdout, strings.NewReader(out+"\n")); err != nil {
		fmt.Printf("ERROR: Unable to copy the output to stdout")
	}
	os.Stdout.Sync()
	return err
}

// trimmedOutput converts the stdout to a string and trims the trailing whitespaces
func trimmedOutput(stdout io.Writer) string {
	output, ok := stdout.(*bytes.Buffer)
	if !ok {
		fmt.Printf("WARNING: Unable to convert output to a buffer\n")
		return ""
	}
	return strings.TrimSpace(output.String())
}

// FatalErr exits the test in case a fatal error has occured.
func FatalErr(msg interface{}) {
	e2e.Failf("%v", msg)
}
<commit_msg>Capture panic in extended CLI and return them as Go errors<commit_after>package util

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/openshift/origin/pkg/client"
	"github.com/openshift/origin/pkg/cmd/cli"
	cmdapi "github.com/openshift/origin/pkg/cmd/cli/cmd"
	"github.com/openshift/origin/pkg/cmd/cli/config"
	configapi "github.com/openshift/origin/pkg/cmd/server/api"
	testutil "github.com/openshift/origin/test/util"
	"github.com/spf13/cobra"
	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client"
	clientcmd "k8s.io/kubernetes/pkg/client/clientcmd"
	kcmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/test/e2e"
)

// CLI provides function to call the OpenShift CLI and Kubernetes and OpenShift
// REST clients.
type CLI struct {
	verb            string
	configPath      string
	adminConfigPath string
	username        string
	outputDir       string
	globalArgs      []string
	commandArgs     []string
	finalArgs       []string
	stdout          io.Writer
	verbose         bool
	cmd             *cobra.Command
	kubeFramework   *e2e.Framework
}

// NewCLI initialize the upstream E2E framework and set the namespace to match
// with the project name. Note that this function does not initialize the project
// role bindings for the namespace.
func NewCLI(project, adminConfigPath string) *CLI {
	client := &CLI{}
	client.kubeFramework = e2e.InitializeFramework(project, client.SetupProject)
	client.outputDir = os.TempDir()
	client.username = "admin"
	if len(adminConfigPath) == 0 {
		FatalErr(fmt.Errorf("You must set the KUBECONFIG variable to admin kubeconfig."))
	}
	client.adminConfigPath = adminConfigPath
	kcmdutil.BehaviorOnFatal(func(msg string) { panic(msg) })
	return client
}

// KubeFramework returns Kubernetes framework which contains helper functions
// specific for Kubernetes resources
func (c *CLI) KubeFramework() *e2e.Framework {
	return c.kubeFramework
}

// Username returns the name of currently logged user. If there is no user assigned
// for the current session, it returns 'admin'.
func (c *CLI) Username() string {
	return c.username
}

// ChangeUser changes the user used by the current CLI session.
func (c *CLI) ChangeUser(name string) *CLI {
	adminClientConfig, err := testutil.GetClusterAdminClientConfig(c.adminConfigPath)
	if err != nil {
		FatalErr(err)
	}
	_, _, clientConfig, err := testutil.GetClientForUser(*adminClientConfig, name)
	if err != nil {
		FatalErr(err)
	}

	kubeConfig, err := config.CreateConfig(c.Namespace(), clientConfig)
	if err != nil {
		FatalErr(err)
	}

	c.configPath = filepath.Join(c.outputDir, name+".kubeconfig")
	err = clientcmd.WriteToFile(*kubeConfig, c.configPath)
	if err != nil {
		FatalErr(err)
	}

	c.username = name
	fmt.Printf("INFO: configPath is now %q\n", c.configPath)
	return c
}

// SetNamespace sets a new namespace
func (c *CLI) SetNamespace(ns string) *CLI {
	c.kubeFramework.Namespace = &kapi.Namespace{
		ObjectMeta: kapi.ObjectMeta{
			Name: ns,
		},
	}
	return c
}

// SetOutputDir change the default output directory for temporary files
func (c *CLI) SetOutputDir(dir string) *CLI {
	c.outputDir = dir
	return c
}

// SetupProject creates a new project and assign a random user to the project.
// All resources will be then created within this project and Kubernetes E2E
// suite will destroy the project after test case finish.
// Note that the kubeClient is not used and serves just to make this function
// compatible with upstream function.
func (c *CLI) SetupProject(name string, kubeClient *kclient.Client) (*kapi.Namespace, error) {
	newNamespace := kapi.SimpleNameGenerator.GenerateName(fmt.Sprintf("extended-test-%s-", name))
	c.SetNamespace(newNamespace).ChangeUser(fmt.Sprintf("%s-user", c.Namespace()))
	e2e.Logf("The user is now %q", c.Username())

	projectOpts := cmdapi.NewProjectOptions{
		ProjectName: c.Namespace(),
		Client:      c.REST(),
		Out:         c.stdout,
	}
	e2e.Logf("Creating project %q", c.Namespace())
	return c.kubeFramework.Namespace, projectOpts.Run()
}

// Verbose turns on printing verbose messages when executing OpenShift commands
func (c *CLI) Verbose() *CLI {
	c.verbose = true
	return c
}

// REST provides an OpenShift REST client for the current user. If the user is not
// set, then it provides REST client for the cluster admin user
func (c *CLI) REST() *client.Client {
	_, clientConfig, err := configapi.GetKubeClient(c.configPath)
	osClient, err := client.New(clientConfig)
	if err != nil {
		FatalErr(err)
	}
	return osClient
}

// AdminREST provides an OpenShift REST client for the cluster admin user.
func (c *CLI) AdminREST() *client.Client {
	_, clientConfig, err := configapi.GetKubeClient(c.adminConfigPath)
	osClient, err := client.New(clientConfig)
	if err != nil {
		FatalErr(err)
	}
	return osClient
}

// KubeREST provides a Kubernetes REST client for the current namespace
func (c *CLI) KubeREST() *kclient.Client {
	kubeClient, _, err := configapi.GetKubeClient(c.configPath)
	if err != nil {
		FatalErr(err)
	}
	return kubeClient
}

// Namespace returns the name of the namespace used in the current test case.
// If the namespace is not set, an empty string is returned.
func (c *CLI) Namespace() string {
	if c.kubeFramework.Namespace == nil {
		return ""
	}
	return c.kubeFramework.Namespace.Name
}

// SetOutput allows to override the default command output
func (c *CLI) SetOutput(out io.Writer) *CLI {
	c.stdout = out
	for _, subCmd := range c.cmd.Commands() {
		subCmd.SetOutput(c.stdout)
	}
	c.cmd.SetOutput(c.stdout)
	return c
}

// Run executes given OpenShift CLI command verb (iow. "oc <verb>").
// This function also override the default 'stdout' to redirect all output
// to a buffer and prepare the global flags such as namespace and config path.
func (c *CLI) Run(verb string) *CLI {
	out := new(bytes.Buffer)
	nc := &CLI{
		verb:            verb,
		kubeFramework:   c.KubeFramework(),
		adminConfigPath: c.adminConfigPath,
		configPath:      c.configPath,
		username:        c.username,
		outputDir:       c.outputDir,
		cmd:             cli.NewCommandCLI("oc", "openshift", out),
		globalArgs: []string{
			verb,
			fmt.Sprintf("--namespace=%s", c.Namespace()),
			fmt.Sprintf("--config=%s", c.configPath),
		},
	}
	return nc.SetOutput(out)
}

// Template sets a Go template for the OpenShift CLI command.
// This is equivalent of running "oc get foo -o template -t '{{ .spec }}'"
func (c *CLI) Template(t string) *CLI {
	if c.verb != "get" {
		FatalErr("Cannot use Template() for non-get verbs.")
	}
	templateArgs := []string{"--output=template", fmt.Sprintf("--template=%s", t)}
	commandArgs := append(c.commandArgs, templateArgs...)
	c.finalArgs = append(c.globalArgs, commandArgs...)
	c.cmd.SetArgs(c.finalArgs)
	return c
}

// Args sets the additional arguments for the OpenShift CLI command
func (c *CLI) Args(args ...string) *CLI {
	c.commandArgs = args
	c.finalArgs = append(c.globalArgs, c.commandArgs...)
	c.cmd.SetArgs(c.finalArgs)
	return c
}

func (c *CLI) printCmd() string {
	return strings.Join(c.finalArgs, " ")
}

// Output executes the command and return the output as string
func (c *CLI) Output() (out string, err error) {
	if c.verbose {
		fmt.Printf("DEBUG: oc %s\n", c.printCmd())
	}
	// Capture the panic and convert it to a regular error
	defer func() {
		if r := recover(); r != nil {
			out = fmt.Sprintf("%s", r)
			err = fmt.Errorf("PANIC: %s", out)
		}
	}()
	err = c.cmd.Execute()
	if c.verbose {
		fmt.Printf("DEBUG: %q\n", trimmedOutput(c.stdout))
	}
	return trimmedOutput(c.stdout), err
}

// Stdout returns the current stdout writer
func (c *CLI) Stdout() io.Writer {
	return c.stdout
}

// OutputToFile executes the command and store output to a file
func (c *CLI) OutputToFile(filename string) (string, error) {
	content, err := c.Output()
	if err != nil {
		return "", err
	}
	path := filepath.Join(c.outputDir, c.Namespace()+"-"+filename)
	return path, ioutil.WriteFile(path, []byte(content), 0644)
}

// Execute executes the current command and return error if the execution failed
// This function will set the default output to stdout.
func (c *CLI) Execute() error {
	out, err := c.Output()
	if err != nil {
		FatalErr(fmt.Errorf("%v", err))
	}
	if _, err := io.Copy(os.Stdout, strings.NewReader(out+"\n")); err != nil {
		fmt.Printf("ERROR: Unable to copy the output to stdout")
	}
	os.Stdout.Sync()
	return err
}

// trimmedOutput converts the stdout to a string and trims the trailing whitespaces
func trimmedOutput(stdout io.Writer) string {
	output, ok := stdout.(*bytes.Buffer)
	if !ok {
		fmt.Printf("WARNING: Unable to convert output to a buffer\n")
		return ""
	}
	return strings.TrimSpace(output.String())
}

// FatalErr exits the test in case a fatal error has occured.
func FatalErr(msg interface{}) {
	e2e.Failf("%v", msg)
}
<|endoftext|>


<commit_before>package sqldb

import (
	"bytes"
	"database/sql"
	"fmt"
	"reflect"
	"strings"
	"unicode"
	"unicode/utf8"
)

type Column struct {
	Name         string
	Type         string
	Precision    string
	DBType       string
	Primary      bool
	AutoIncr     bool
	Notnull      bool
	Default      bool
	DefaultVal   string
	Unique       bool
	UniqueName   string
	ForeignTable string
	ForeignCol   string
}

type Table struct {
	Name string
	Cols []Column
}

type Parser struct {
	DBDialect       DBDialect
	FieldTag        string
	Default         bool
	Notnull         bool
	TablenamePrefix string
	NameMapper      NameMapper
}

func (p *Parser) initDefault() {
	if p.DBDialect == nil {
		p.DBDialect = Postgres{}
	}
	if p.FieldTag == "" {
		p.FieldTag = "sqldb"
	}
	if p.NameMapper == nil {
		p.NameMapper = SnakeCase
	}
}

func (p *Parser) CreateTables(db *sql.DB, models ...interface{}) error {
	for _, mod := range models {
		table, err := p.StructTable(mod)
		if err != nil {
			return err
		}
		s, err := p.SQLCreate(table)
		if err != nil {
			return fmt.Errorf("%s: %s", table.Name, err.Error())
		}
		_, err = db.Exec(s)
		if err != nil {
			return fmt.Errorf("%s: %s", table.Name, err.Error())
		}
	}
	return nil
}

func (p *Parser) SQLCreate(table Table) (string, error) {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "CREATE TABLE IF NOT EXISTS %s (\n", table.Name)
	var (
		uniques   map[string][]string
		primaries []string
		foreigns  []int
		lastQuite string
	)
	for i, col := range table.Cols {
		dbTyp, defaultVal, err := p.DBDialect.Type(col.Type, col.Precision, col.DefaultVal)
		if err != nil {
			return "", err
		}
		if col.Primary {
			primaries = append(primaries, col.Name)
		}
		if col.ForeignTable != "" {
			foreigns = append(foreigns, i)
		}
		if col.DBType != "" {
			dbTyp = col.DBType
		}
		var constraints string
		if col.Unique {
			if col.UniqueName == "" {
				constraints += " UNIQUE"
			} else {
				if uniques == nil {
					uniques = make(map[string][]string)
				}
				uniques[col.UniqueName] = append(uniques[col.UniqueName], col.Name)
			}
		}
		if col.AutoIncr {
			constraints += " AUTO INCREAMENT"
		}
		if !col.Notnull {
			constraints += " NOT NULL"
		}
		if col.Default {
			constraints += " DEFAULT " + defaultVal
		}
		lastQuite = ""
		if i != len(table.Cols)-1 || len(primaries) != 0 || len(uniques) != 0 || len(foreigns) != 0 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    %s %s %s%s\n", col.Name, dbTyp, constraints, lastQuite)
	}
	if len(primaries) > 0 {
		lastQuite = ""
		if len(uniques) != 0 || len(foreigns) != 0 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    PRIMARY KEY (%s)%s\n", strings.Join(primaries, ","), lastQuite)
	}
	for name, keys := range uniques {
		lastQuite = ""
		if len(foreigns) != 0 || len(uniques) != 1 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    CONSTRAINT %s UNIQUE (%s)%s\n", name, strings.Join(keys, ","), lastQuite)
		delete(uniques, name)
	}
	for i, index := range foreigns {
		col := table.Cols[index]
		lastQuite = ""
		if i != len(foreigns)-1 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    FOREIGN KEY(%s) REFERENCES %s(%s)%s\n", col.Name, col.ForeignTable, col.ForeignCol, lastQuite)
	}
	fmt.Fprintf(&buf, ");\n")
	return buf.String(), nil
}

func (p *Parser) parseColumn(t *Table, f *reflect.StructField) (Column, error) {
	col := Column{
		Name:    p.NameMapper(f.Name),
		Type:    f.Type.Kind().String(),
		Default: p.Default,
		Notnull: !p.Notnull,
	}

	var conds []string
	tag := strings.TrimSpace(f.Tag.Get(p.FieldTag))
	if tag != "" {
		conds = strings.Split(tag, " ")
	}
	for _, sec := range conds {
		sec = strings.TrimSpace(sec)
		var (
			keyCond  = strings.SplitN(sec, ":", 2)
			condName = keyCond[0]
			condVal  string
		)
		if len(keyCond) > 1 {
			condVal = keyCond[1]
		}

		switch condName {
		case "table":
			t.Name = condVal
		case "col":
			col.Name = condVal
			if condVal == "" {
				return col, fmt.Errorf("invalid column name")
			}
			if condVal == "-" {
				col.Name = ""
				return col, nil
			}
			col.Name = condVal
		case "type":
			if condVal == "" {
				return col, fmt.Errorf("invalid column type: %s", col.Name)
			}
			col.Type = condVal
		case "precision":
			if condVal == "" {
				return col, fmt.Errorf("invalid column precision: %s", col.Name)
			}
			col.Precision = condVal
		case "dbtype":
			if condVal == "" {
				return col, fmt.Errorf("invalid column db type: %s", col.Name)
			}
			col.DBType = condVal
		case "pk":
			col.Primary = condVal == "" || condVal == "true"
		case "autoincr":
			col.AutoIncr = condVal == "" || condVal == "true"
		case "notnull":
			col.Notnull = condVal == "" || condVal == "true"
		case "default":
			col.Default = condVal != "-"
			if p.Default {
				col.DefaultVal = condVal
			}
		case "unique":
			col.Unique = true
			col.UniqueName = condVal
		case "fk":
			fkConds := strings.SplitN(condVal, ".", 2)
			if len(fkConds) != 2 || fkConds[0] == "" || fkConds[1] == "" {
				return col, fmt.Errorf("invalid foreign key: %s", condVal)
			}
			col.ForeignTable = fkConds[0]
			col.ForeignCol = fkConds[1]
		default:
			return col, fmt.Errorf("unsupported tag: %s", condName)
		}
	}
	return col, nil
}

func (p *Parser) StructTable(v interface{}) (Table, error) {
	p.initDefault()

	refv := reflect.ValueOf(v)
	if refv.Kind() == reflect.Ptr {
		refv = refv.Elem()
	}
	if refv.Kind() != reflect.Struct {
		return Table{}, fmt.Errorf("invalid artument type, expect (pointer of) structure")
	}
	reft := refv.Type()

	t := Table{
		Name: p.TablenamePrefix + p.NameMapper(reft.Name()),
	}
	n := reft.NumField()
	for i := 0; i < n; i++ {
		f := reft.Field(i)
		col, err := p.parseColumn(&t, &f)
		if err != nil {
			return t, err
		}
		if col.Name != "" {
			t.Cols = append(t.Cols, col)
		}
	}
	return t, nil
}

func SnakeCase(s string) string {
	var (
		hasUpper  bool
		size      = utf8.RuneCountInString(s)
		prevUpper bool
	)
	for i, r := range s {
		if unicode.IsUpper(r) {
			hasUpper = true
			if i != 0 && !prevUpper {
				size++
			}
			prevUpper = true
		} else {
			prevUpper = false
		}
	}
	if !hasUpper {
		return s
	}
	var (
		buf = make([]rune, 0, size)
	)
	prevUpper = false
	for i, r := range s {
		isUpper := unicode.IsUpper(r)
		if isUpper && i != 0 && !prevUpper {
			buf = append(buf, '_')
		}
		if prevUpper = isUpper; isUpper {
			buf = append(buf, unicode.ToLower(r))
		} else {
			buf = append(buf, r)
		}
	}
	return string(buf)
}
<commit_msg>escape table and column name<commit_after>package sqldb

import (
	"bytes"
	"database/sql"
	"fmt"
	"reflect"
	"strings"
	"unicode"
	"unicode/utf8"
)

type Column struct {
	Name         string
	Type         string
	Precision    string
	DBType       string
	Primary      bool
	AutoIncr     bool
	Notnull      bool
	Default      bool
	DefaultVal   string
	Unique       bool
	UniqueName   string
	ForeignTable string
	ForeignCol   string
}

type Table struct {
	Name string
	Cols []Column
}

type Parser struct {
	DBDialect       DBDialect
	FieldTag        string
	Default         bool
	Notnull         bool
	TablenamePrefix string
	NameMapper      NameMapper
}

func (p *Parser) initDefault() {
	if p.DBDialect == nil {
		p.DBDialect = Postgres{}
	}
	if p.FieldTag == "" {
		p.FieldTag = "sqldb"
	}
	if p.NameMapper == nil {
		p.NameMapper = SnakeCase
	}
}

func (p *Parser) CreateTables(db *sql.DB, models ...interface{}) error {
	for _, mod := range models {
		table, err := p.StructTable(mod)
		if err != nil {
			return err
		}
		s, err := p.SQLCreate(table)
		if err != nil {
			return fmt.Errorf("%s: %s", table.Name, err.Error())
		}
		_, err = db.Exec(s)
		if err != nil {
			return fmt.Errorf("%s: %s", table.Name, err.Error())
		}
	}
	return nil
}

func (p *Parser) EscapeName(name string) string {
	return `"` + name + `"`
}

func (p *Parser) SQLCreate(table Table) (string, error) {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "CREATE TABLE IF NOT EXISTS %s (\n", p.EscapeName(table.Name))
	var (
		uniques   map[string][]string
		primaries []string
		foreigns  []int
		lastQuite string
	)
	for i, col := range table.Cols {
		dbTyp, defaultVal, err := p.DBDialect.Type(col.Type, col.Precision, col.DefaultVal)
		if err != nil {
			return "", err
		}
		if col.Primary {
			primaries = append(primaries, p.EscapeName(col.Name))
		}
		if col.ForeignTable != "" {
			foreigns = append(foreigns, i)
		}
		if col.DBType != "" {
			dbTyp = col.DBType
		}
		var constraints string
		if col.Unique {
			if col.UniqueName == "" {
				constraints += " UNIQUE"
			} else {
				if uniques == nil {
					uniques = make(map[string][]string)
				}
				uniques[col.UniqueName] = append(uniques[col.UniqueName], p.EscapeName(col.Name))
			}
		}
		if col.AutoIncr {
			constraints += " AUTO INCREAMENT"
		}
		if !col.Notnull {
			constraints += " NOT NULL"
		}
		if col.Default {
			constraints += " DEFAULT " + defaultVal
		}
		lastQuite = ""
		if i != len(table.Cols)-1 || len(primaries) != 0 || len(uniques) != 0 || len(foreigns) != 0 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    %s %s %s%s\n", p.EscapeName(col.Name), dbTyp, constraints, lastQuite)
	}
	if len(primaries) > 0 {
		lastQuite = ""
		if len(uniques) != 0 || len(foreigns) != 0 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    PRIMARY KEY (%s)%s\n", strings.Join(primaries, ","), lastQuite)
	}
	for name, keys := range uniques {
		lastQuite = ""
		if len(foreigns) != 0 || len(uniques) != 1 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    CONSTRAINT %s UNIQUE (%s)%s\n", name, strings.Join(keys, ","), lastQuite)
		delete(uniques, name)
	}
	for i, index := range foreigns {
		col := table.Cols[index]
		lastQuite = ""
		if i != len(foreigns)-1 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    FOREIGN KEY(%s) REFERENCES %s(%s)%s\n", p.EscapeName(col.Name), col.ForeignTable, col.ForeignCol, lastQuite)
	}
	fmt.Fprintf(&buf, ");\n")
	return buf.String(), nil
}

func (p *Parser) parseColumn(t *Table, f *reflect.StructField) (Column, error) {
	col := Column{
		Name:    p.NameMapper(f.Name),
		Type:    f.Type.Kind().String(),
		Default: p.Default,
		Notnull: !p.Notnull,
	}

	var conds []string
	tag := strings.TrimSpace(f.Tag.Get(p.FieldTag))
	if tag != "" {
		conds = strings.Split(tag, " ")
	}
	for _, sec := range conds {
		sec = strings.TrimSpace(sec)
		var (
			keyCond  = strings.SplitN(sec, ":", 2)
			condName = keyCond[0]
			condVal  string
		)
		if len(keyCond) > 1 {
			condVal = keyCond[1]
		}

		switch condName {
		case "table":
			t.Name = condVal
		case "col":
			col.Name = condVal
			if condVal == "" {
				return col, fmt.Errorf("invalid column name")
			}
			if condVal == "-" {
				col.Name = ""
				return col, nil
			}
			col.Name = condVal
		case "type":
			if condVal == "" {
				return col, fmt.Errorf("invalid column type: %s", col.Name)
			}
			col.Type = condVal
		case "precision":
			if condVal == "" {
				return col, fmt.Errorf("invalid column precision: %s", col.Name)
			}
			col.Precision = condVal
		case "dbtype":
			if condVal == "" {
				return col, fmt.Errorf("invalid column db type: %s", col.Name)
			}
			col.DBType = condVal
		case "pk":
			col.Primary = condVal == "" || condVal == "true"
		case "autoincr":
			col.AutoIncr = condVal == "" || condVal == "true"
		case "notnull":
			col.Notnull = condVal == "" || condVal == "true"
		case "default":
			col.Default = condVal != "-"
			if p.Default {
				col.DefaultVal = condVal
			}
		case "unique":
			col.Unique = true
			col.UniqueName = condVal
		case "fk":
			fkConds := strings.SplitN(condVal, ".", 2)
			if len(fkConds) != 2 || fkConds[0] == "" || fkConds[1] == "" {
				return col, fmt.Errorf("invalid foreign key: %s", condVal)
			}
			col.ForeignTable = fkConds[0]
			col.ForeignCol = fkConds[1]
		default:
			return col, fmt.Errorf("unsupported tag: %s", condName)
		}
	}
	return col, nil
}

func (p *Parser) StructTable(v interface{}) (Table, error) {
	p.initDefault()

	refv := reflect.ValueOf(v)
	if refv.Kind() == reflect.Ptr {
		refv = refv.Elem()
	}
	if refv.Kind() != reflect.Struct {
		return Table{}, fmt.Errorf("invalid artument type, expect (pointer of) structure")
	}
	reft := refv.Type()

	t := Table{
		Name: p.TablenamePrefix + p.NameMapper(reft.Name()),
	}
	n := reft.NumField()
	for i := 0; i < n; i++ {
		f := reft.Field(i)
		col, err := p.parseColumn(&t, &f)
		if err != nil {
			return t, err
		}
		if col.Name != "" {
			t.Cols = append(t.Cols, col)
		}
	}
	return t, nil
}

func SnakeCase(s string) string {
	var (
		hasUpper  bool
		size      = utf8.RuneCountInString(s)
		prevUpper bool
	)
	for i, r := range s {
		if unicode.IsUpper(r) {
			hasUpper = true
			if i != 0 && !prevUpper {
				size++
			}
			prevUpper = true
		} else {
			prevUpper = false
		}
	}
	if !hasUpper {
		return s
	}
	var (
		buf = make([]rune, 0, size)
	)
	prevUpper = false
	for i, r := range s {
		isUpper := unicode.IsUpper(r)
		if isUpper && i != 0 && !prevUpper {
			buf = append(buf, '_')
		}
		if prevUpper = isUpper; isUpper {
			buf = append(buf, unicode.ToLower(r))
		} else {
			buf = append(buf, r)
		}
	}
	return string(buf)
}
<|endoftext|>


<commit_before>// Copyright 2014 Oleku Konko All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

// This module is a Table Writer  API for the Go Programming Language.
// The protocols were written in pure Go and works on windows and unix systems

// Create & Generate text based table
package tablewriter

import (
	"fmt"
	"io"
	"regexp"
	"strings"
)

const (
	MAX_ROW_WIDTH = 30
)

const (
	CENTRE = "+"
	ROW    = "-"
	COLUMN = "|"
	SPACE  = " "
)

const (
	ALIGN_DEFAULT = iota
	ALIGN_CENTRE
	ALIGN_RIGHT
	ALIGN_LEFT
)

var (
	decimal = regexp.MustCompile(`^[0-9]+(.[0-9]+?)$`)
	percent = regexp.MustCompile(`^[0-9]+(.[0-9]+?)%$`)
)

type Table struct {
	out      io.Writer
	rows     [][]string
	lines    [][][]string
	cs       map[int]int
	rs       map[int]int
	headers  []string
	footers  []string
	autoFmt  bool
	autoWrap bool
	mW       int
	pCenter  string
	pRow     string
	pColumn  string
	tColumn  int
	tRow     int
	align    int
	rowLine  bool
	border   bool
	colSize  int
}

// Start New Table
// Take io.Writer Directly
func NewWriter(writer io.Writer) *Table {
	t := &Table{
		out:      writer,
		rows:     [][]string{},
		lines:    [][][]string{},
		cs:       make(map[int]int),
		rs:       make(map[int]int),
		headers:  []string{},
		footers:  []string{},
		autoFmt:  true,
		autoWrap: true,
		mW:       MAX_ROW_WIDTH,
		pCenter:  CENTRE,
		pRow:     ROW,
		pColumn:  COLUMN,
		tColumn:  -1,
		tRow:     -1,
		align:    ALIGN_DEFAULT,
		rowLine:  false,
		border:   true,
		colSize:  -1}
	return t
}

// Render table output
func (t Table) Render() {
	if t.border {
		t.printLine(true)
	}
	t.printHeading()
	t.printRows()

	if !t.rowLine && t.border {
		t.printLine(true)
	}
	t.printFooter()

}

// Set table header
func (t *Table) SetHeader(keys []string) {
	t.colSize = len(keys)
	for i, v := range keys {
		t.parseDimension(v, i, -1)
		t.headers = append(t.headers, v)
	}
}

// Set table Footer
func (t *Table) SetFooter(keys []string) {
	//t.colSize = len(keys)
	for i, v := range keys {
		t.parseDimension(v, i, -1)
		t.footers = append(t.footers, v)
	}
}

// Turn header autoformatting on/off. Default is on (true).
func (t *Table) SetAutoFormatHeaders(auto bool) {
	t.autoFmt = auto
}

// Turn automatic multiline text adjustment on/off. Default is on (true).
func (t *Table) SetAutoWrapText(auto bool) {
	t.autoWrap = auto
}

// Set the Default column width
func (t *Table) SetColWidth(width int) {
	t.mW = width
}

// Set the Column Separator
func (t *Table) SetColumnSeparator(sep string) {
	t.pColumn = sep
}

// Set the Row Separator
func (t *Table) SetRowSeparator(sep string) {
	t.pRow = sep
}

// Set the center Separator
func (t *Table) SetCenterSeparator(sep string) {
	t.pCenter = sep
}

// Set Table Alignment
func (t *Table) SetAlignment(align int) {
	t.align = align
}

// Set Row Line
// This would enable / disable a line on each row of the table
func (t *Table) SetRowLine(line bool) {
	t.rowLine = line
}

// Set Table Border
// This would enable / disable line around the table
func (t *Table) SetBorder(border bool) {
	t.border = border
}

// Append row to table
func (t *Table) Append(row []string) error {
	rowSize := len(t.headers)
	if rowSize > t.colSize {
		t.colSize = rowSize
	}

	n := len(t.lines)
	line := [][]string{}
	for i, v := range row {

		// Detect string  width
		// Detect String height
		// Break strings into words
		out := t.parseDimension(v, i, n)

		// Append broken words
		line = append(line, out)
	}
	t.lines = append(t.lines, line)
	return nil
}

// Allow Support for Bulk Append
// Eliminates repeated for loops
func (t *Table) AppendBulk(rows [][]string) (err error) {
	for _, row := range rows {
		err = t.Append(row)
		if err != nil {
			return err
		}
	}
	return nil
}

// Print line based on row width
func (t Table) printLine(nl bool) {
	fmt.Fprint(t.out, t.pCenter)
	for i := 0; i < len(t.cs); i++ {
		v := t.cs[i]
		fmt.Fprintf(t.out, "%s%s%s%s",
			t.pRow,
			strings.Repeat(string(t.pRow), v),
			t.pRow,
			t.pCenter)
	}
	if nl {
		fmt.Fprintln(t.out)
	}
}

// Print heading information
func (t Table) printHeading() {
	// Check if headers is available
	if len(t.headers) < 1 {
		return
	}

	// Check if border is set
	// Replace with space if not set
	fmt.Fprint(t.out, ConditionString(t.border, t.pColumn, SPACE))

	// Identify last column
	end := len(t.cs) - 1

	// Print Heading column
	for i := 0; i <= end; i++ {
		v := t.cs[i]
		h := t.headers[i]
		if t.autoFmt {
			h = Title(h)
		}
		pad := ConditionString((i == end && !t.border), SPACE, t.pColumn)
		fmt.Fprintf(t.out, " %s %s",
			Pad(h, SPACE, v),
			pad)
	}
	// Next line
	fmt.Fprintln(t.out)
	t.printLine(true)
}

// Print heading information
func (t Table) printFooter() {
	// Check if headers is available
	if len(t.footers) < 1 {
		return
	}

	// Only print line if border is not set
	if !t.border {
		t.printLine(true)
	}
	// Check if border is set
	// Replace with space if not set
	fmt.Fprint(t.out, ConditionString(t.border, t.pColumn, SPACE))

	// Identify last column
	end := len(t.cs) - 1

	// Print Heading column
	for i := 0; i <= end; i++ {
		v := t.cs[i]
		f := t.footers[i]
		if t.autoFmt {
			f = Title(f)
		}
		pad := ConditionString((i == end && !t.border), SPACE, t.pColumn)

		if len(t.footers[i]) == 0 {
			pad = SPACE
		}
		fmt.Fprintf(t.out, " %s %s",
			Pad(f, SPACE, v),
			pad)
	}
	// Next line
	fmt.Fprintln(t.out)
	//t.printLine(true)

	hasPrinted := false

	for i := 0; i <= end; i++ {
		v := t.cs[i]
		pad := t.pRow
		center := t.pCenter
		length := len(t.footers[i])

		if length > 0 {
			hasPrinted = true
		}

		// Set center to be space if length is 0
		if length == 0 && !t.border {
			center = SPACE
		}

		// Print first junction
		if i == 0 {
			fmt.Fprint(t.out, center)
		}

		// Pad With space of length is 0
		if length == 0 {
			pad = SPACE
		}
		// Ignore left space of it has printed before
		if hasPrinted || t.border {
			pad = t.pRow
			center = t.pCenter
		}

		// Change Center start position
		if center == SPACE {
			if i < end && len(t.footers[i+1]) != 0 {
				center = t.pCenter
			}
		}

		// Print the footer
		fmt.Fprintf(t.out, "%s%s%s%s",
			pad,
			strings.Repeat(string(pad), v),
			pad,
			center)

	}

	fmt.Fprintln(t.out)

}

func (t Table) printRows() {
	for i, lines := range t.lines {
		t.printRow(lines, i)
	}

}

// Print Row Information
// Adjust column alignment based on type

func (t Table) printRow(columns [][]string, colKey int) {
	// Get Maximum Height
	max := t.rs[colKey]
	total := len(columns)

	// TODO Fix uneven col size
	// if total < t.colSize {
	//	for n := t.colSize - total; n < t.colSize ; n++ {
	//		columns = append(columns, []string{SPACE})
	//		t.cs[n] = t.mW
	//	}
	//}

	// Pad Each Height
	// pads := []int{}
	pads := []int{}

	for i, line := range columns {
		length := len(line)
		pad := max - length
		pads = append(pads, pad)
		for n := 0; n < pad; n++ {
			columns[i] = append(columns[i], "  ")
		}
	}
	//fmt.Println(max, "\n")
	for x := 0; x < max; x++ {
		for y := 0; y < total; y++ {

			// Check if border is set
			fmt.Fprint(t.out, ConditionString((!t.border && y == 0), SPACE, t.pColumn))

			fmt.Fprintf(t.out, SPACE)
			str := columns[y][x]

			// This would print alignment
			// Default alignment  would use multiple configuration
			switch t.align {
			case ALIGN_CENTRE: //
				fmt.Fprintf(t.out, "%s", Pad(str, SPACE, t.cs[y]))
			case ALIGN_RIGHT:
				fmt.Fprintf(t.out, "%s", PadLeft(str, SPACE, t.cs[y]))
			case ALIGN_LEFT:
				fmt.Fprintf(t.out, "%s", PadRight(str, SPACE, t.cs[y]))
			default:
				if decimal.MatchString(strings.TrimSpace(str)) || percent.MatchString(strings.TrimSpace(str)) {
					fmt.Fprintf(t.out, "%s", PadLeft(str, SPACE, t.cs[y]))
				} else {
					fmt.Fprintf(t.out, "%s", PadRight(str, SPACE, t.cs[y]))

					// TODO Custom alignment per column
					//if max == 1 || pads[y] > 0 {
					//	fmt.Fprintf(t.out, "%s", Pad(str, SPACE, t.cs[y]))
					//} else {
					//	fmt.Fprintf(t.out, "%s", PadRight(str, SPACE, t.cs[y]))
					//}

				}
			}
			fmt.Fprintf(t.out, SPACE)
		}
		// Check if border is set
		// Replace with space if not set
		fmt.Fprint(t.out, ConditionString(t.border, t.pColumn, SPACE))
		fmt.Fprintln(t.out)
	}

	if t.rowLine {
		t.printLine(true)
	}

}

func (t *Table) parseDimension(str string, colKey, rowKey int) []string {
	var (
		raw []string
		max int
	)
	w := DisplayWidth(str)
	// Calculate Width
	// Check if with is grater than maximum width
	if w > t.mW {
		w = t.mW
	}

	// Check if width exists
	v, ok := t.cs[colKey]
	if !ok || v < w || v == 0 {
		t.cs[colKey] = w
	}

	if rowKey == -1 {
		return raw
	}
	// Calculate Height
	if t.autoWrap {
		raw, _ = WrapString(str, t.cs[colKey])
	} else {
		raw = getLines(str)
	}

	for _, line := range raw {
		if w := DisplayWidth(line); w > max {
			max = w
		}
	}

	// Make sure the with is the same length as maximum word
	// Important for cases where the width is smaller than maxu word
	if max > t.cs[colKey] {
		t.cs[colKey] = max
	}

	h := len(raw)
	v, ok = t.rs[rowKey]

	if !ok || v < h || v == 0 {
		t.rs[rowKey] = h
	}
	//fmt.Printf("Raw %+v %d\n", raw, len(raw))
	return raw
}
<commit_msg>fix decimal and percent regex to allow numbers without '.'<commit_after>// Copyright 2014 Oleku Konko All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

// This module is a Table Writer  API for the Go Programming Language.
// The protocols were written in pure Go and works on windows and unix systems

// Create & Generate text based table
package tablewriter

import (
	"fmt"
	"io"
	"regexp"
	"strings"
)

const (
	MAX_ROW_WIDTH = 30
)

const (
	CENTRE = "+"
	ROW    = "-"
	COLUMN = "|"
	SPACE  = " "
)

const (
	ALIGN_DEFAULT = iota
	ALIGN_CENTRE
	ALIGN_RIGHT
	ALIGN_LEFT
)

var (
	decimal = regexp.MustCompile(`^\d*\.?\d*$`)
	percent = regexp.MustCompile(`^\d*\.?\d*$%$`)
)

type Table struct {
	out      io.Writer
	rows     [][]string
	lines    [][][]string
	cs       map[int]int
	rs       map[int]int
	headers  []string
	footers  []string
	autoFmt  bool
	autoWrap bool
	mW       int
	pCenter  string
	pRow     string
	pColumn  string
	tColumn  int
	tRow     int
	align    int
	rowLine  bool
	border   bool
	colSize  int
}

// Start New Table
// Take io.Writer Directly
func NewWriter(writer io.Writer) *Table {
	t := &Table{
		out:      writer,
		rows:     [][]string{},
		lines:    [][][]string{},
		cs:       make(map[int]int),
		rs:       make(map[int]int),
		headers:  []string{},
		footers:  []string{},
		autoFmt:  true,
		autoWrap: true,
		mW:       MAX_ROW_WIDTH,
		pCenter:  CENTRE,
		pRow:     ROW,
		pColumn:  COLUMN,
		tColumn:  -1,
		tRow:     -1,
		align:    ALIGN_DEFAULT,
		rowLine:  false,
		border:   true,
		colSize:  -1}
	return t
}

// Render table output
func (t Table) Render() {
	if t.border {
		t.printLine(true)
	}
	t.printHeading()
	t.printRows()

	if !t.rowLine && t.border {
		t.printLine(true)
	}
	t.printFooter()

}

// Set table header
func (t *Table) SetHeader(keys []string) {
	t.colSize = len(keys)
	for i, v := range keys {
		t.parseDimension(v, i, -1)
		t.headers = append(t.headers, v)
	}
}

// Set table Footer
func (t *Table) SetFooter(keys []string) {
	//t.colSize = len(keys)
	for i, v := range keys {
		t.parseDimension(v, i, -1)
		t.footers = append(t.footers, v)
	}
}

// Turn header autoformatting on/off. Default is on (true).
func (t *Table) SetAutoFormatHeaders(auto bool) {
	t.autoFmt = auto
}

// Turn automatic multiline text adjustment on/off. Default is on (true).
func (t *Table) SetAutoWrapText(auto bool) {
	t.autoWrap = auto
}

// Set the Default column width
func (t *Table) SetColWidth(width int) {
	t.mW = width
}

// Set the Column Separator
func (t *Table) SetColumnSeparator(sep string) {
	t.pColumn = sep
}

// Set the Row Separator
func (t *Table) SetRowSeparator(sep string) {
	t.pRow = sep
}

// Set the center Separator
func (t *Table) SetCenterSeparator(sep string) {
	t.pCenter = sep
}

// Set Table Alignment
func (t *Table) SetAlignment(align int) {
	t.align = align
}

// Set Row Line
// This would enable / disable a line on each row of the table
func (t *Table) SetRowLine(line bool) {
	t.rowLine = line
}

// Set Table Border
// This would enable / disable line around the table
func (t *Table) SetBorder(border bool) {
	t.border = border
}

// Append row to table
func (t *Table) Append(row []string) error {
	rowSize := len(t.headers)
	if rowSize > t.colSize {
		t.colSize = rowSize
	}

	n := len(t.lines)
	line := [][]string{}
	for i, v := range row {

		// Detect string  width
		// Detect String height
		// Break strings into words
		out := t.parseDimension(v, i, n)

		// Append broken words
		line = append(line, out)
	}
	t.lines = append(t.lines, line)
	return nil
}

// Allow Support for Bulk Append
// Eliminates repeated for loops
func (t *Table) AppendBulk(rows [][]string) (err error) {
	for _, row := range rows {
		err = t.Append(row)
		if err != nil {
			return err
		}
	}
	return nil
}

// Print line based on row width
func (t Table) printLine(nl bool) {
	fmt.Fprint(t.out, t.pCenter)
	for i := 0; i < len(t.cs); i++ {
		v := t.cs[i]
		fmt.Fprintf(t.out, "%s%s%s%s",
			t.pRow,
			strings.Repeat(string(t.pRow), v),
			t.pRow,
			t.pCenter)
	}
	if nl {
		fmt.Fprintln(t.out)
	}
}

// Print heading information
func (t Table) printHeading() {
	// Check if headers is available
	if len(t.headers) < 1 {
		return
	}

	// Check if border is set
	// Replace with space if not set
	fmt.Fprint(t.out, ConditionString(t.border, t.pColumn, SPACE))

	// Identify last column
	end := len(t.cs) - 1

	// Print Heading column
	for i := 0; i <= end; i++ {
		v := t.cs[i]
		h := t.headers[i]
		if t.autoFmt {
			h = Title(h)
		}
		pad := ConditionString((i == end && !t.border), SPACE, t.pColumn)
		fmt.Fprintf(t.out, " %s %s",
			Pad(h, SPACE, v),
			pad)
	}
	// Next line
	fmt.Fprintln(t.out)
	t.printLine(true)
}

// Print heading information
func (t Table) printFooter() {
	// Check if headers is available
	if len(t.footers) < 1 {
		return
	}

	// Only print line if border is not set
	if !t.border {
		t.printLine(true)
	}
	// Check if border is set
	// Replace with space if not set
	fmt.Fprint(t.out, ConditionString(t.border, t.pColumn, SPACE))

	// Identify last column
	end := len(t.cs) - 1

	// Print Heading column
	for i := 0; i <= end; i++ {
		v := t.cs[i]
		f := t.footers[i]
		if t.autoFmt {
			f = Title(f)
		}
		pad := ConditionString((i == end && !t.border), SPACE, t.pColumn)

		if len(t.footers[i]) == 0 {
			pad = SPACE
		}
		fmt.Fprintf(t.out, " %s %s",
			Pad(f, SPACE, v),
			pad)
	}
	// Next line
	fmt.Fprintln(t.out)
	//t.printLine(true)

	hasPrinted := false

	for i := 0; i <= end; i++ {
		v := t.cs[i]
		pad := t.pRow
		center := t.pCenter
		length := len(t.footers[i])

		if length > 0 {
			hasPrinted = true
		}

		// Set center to be space if length is 0
		if length == 0 && !t.border {
			center = SPACE
		}

		// Print first junction
		if i == 0 {
			fmt.Fprint(t.out, center)
		}

		// Pad With space of length is 0
		if length == 0 {
			pad = SPACE
		}
		// Ignore left space of it has printed before
		if hasPrinted || t.border {
			pad = t.pRow
			center = t.pCenter
		}

		// Change Center start position
		if center == SPACE {
			if i < end && len(t.footers[i+1]) != 0 {
				center = t.pCenter
			}
		}

		// Print the footer
		fmt.Fprintf(t.out, "%s%s%s%s",
			pad,
			strings.Repeat(string(pad), v),
			pad,
			center)

	}

	fmt.Fprintln(t.out)

}

func (t Table) printRows() {
	for i, lines := range t.lines {
		t.printRow(lines, i)
	}

}

// Print Row Information
// Adjust column alignment based on type

func (t Table) printRow(columns [][]string, colKey int) {
	// Get Maximum Height
	max := t.rs[colKey]
	total := len(columns)

	// TODO Fix uneven col size
	// if total < t.colSize {
	//	for n := t.colSize - total; n < t.colSize ; n++ {
	//		columns = append(columns, []string{SPACE})
	//		t.cs[n] = t.mW
	//	}
	//}

	// Pad Each Height
	// pads := []int{}
	pads := []int{}

	for i, line := range columns {
		length := len(line)
		pad := max - length
		pads = append(pads, pad)
		for n := 0; n < pad; n++ {
			columns[i] = append(columns[i], "  ")
		}
	}
	//fmt.Println(max, "\n")
	for x := 0; x < max; x++ {
		for y := 0; y < total; y++ {

			// Check if border is set
			fmt.Fprint(t.out, ConditionString((!t.border && y == 0), SPACE, t.pColumn))

			fmt.Fprintf(t.out, SPACE)
			str := columns[y][x]

			// This would print alignment
			// Default alignment  would use multiple configuration
			switch t.align {
			case ALIGN_CENTRE: //
				fmt.Fprintf(t.out, "%s", Pad(str, SPACE, t.cs[y]))
			case ALIGN_RIGHT:
				fmt.Fprintf(t.out, "%s", PadLeft(str, SPACE, t.cs[y]))
			case ALIGN_LEFT:
				fmt.Fprintf(t.out, "%s", PadRight(str, SPACE, t.cs[y]))
			default:
				if decimal.MatchString(strings.TrimSpace(str)) || percent.MatchString(strings.TrimSpace(str)) {
					fmt.Fprintf(t.out, "%s", PadLeft(str, SPACE, t.cs[y]))
				} else {
					fmt.Fprintf(t.out, "%s", PadRight(str, SPACE, t.cs[y]))

					// TODO Custom alignment per column
					//if max == 1 || pads[y] > 0 {
					//	fmt.Fprintf(t.out, "%s", Pad(str, SPACE, t.cs[y]))
					//} else {
					//	fmt.Fprintf(t.out, "%s", PadRight(str, SPACE, t.cs[y]))
					//}

				}
			}
			fmt.Fprintf(t.out, SPACE)
		}
		// Check if border is set
		// Replace with space if not set
		fmt.Fprint(t.out, ConditionString(t.border, t.pColumn, SPACE))
		fmt.Fprintln(t.out)
	}

	if t.rowLine {
		t.printLine(true)
	}

}

func (t *Table) parseDimension(str string, colKey, rowKey int) []string {
	var (
		raw []string
		max int
	)
	w := DisplayWidth(str)
	// Calculate Width
	// Check if with is grater than maximum width
	if w > t.mW {
		w = t.mW
	}

	// Check if width exists
	v, ok := t.cs[colKey]
	if !ok || v < w || v == 0 {
		t.cs[colKey] = w
	}

	if rowKey == -1 {
		return raw
	}
	// Calculate Height
	if t.autoWrap {
		raw, _ = WrapString(str, t.cs[colKey])
	} else {
		raw = getLines(str)
	}

	for _, line := range raw {
		if w := DisplayWidth(line); w > max {
			max = w
		}
	}

	// Make sure the with is the same length as maximum word
	// Important for cases where the width is smaller than maxu word
	if max > t.cs[colKey] {
		t.cs[colKey] = max
	}

	h := len(raw)
	v, ok = t.rs[rowKey]

	if !ok || v < h || v == 0 {
		t.rs[rowKey] = h
	}
	//fmt.Printf("Raw %+v %d\n", raw, len(raw))
	return raw
}
<|endoftext|>


<commit_before>package fitsio

import (
	"fmt"
	"reflect"
)

type Table struct {
	hdr    Header
	binary bool

	data []byte // main data table
	heap []byte // heap data table (for variable length arrays)

	rowsz  int   // size of each row in bytes (ie: NAXIS1)
	nrows  int64 // number of rows (ie: NAXIS2)
	cols   []Column
	colidx map[string]int // associates a column name to its index
}

// Close closes this HDU, cleaning up cycles (if any) for garbage collection
func (t *Table) Close() error {
	return nil
}

// Header returns the Header part of this HDU block.
func (t *Table) Header() *Header {
	return &t.hdr
}

// Type returns the Type of this HDU
func (t *Table) Type() HDUType {
	return t.hdr.Type()
}

// Name returns the value of the 'EXTNAME' Card.
func (t *Table) Name() string {
	card := t.hdr.Get("EXTNAME")
	if card == nil {
		return ""
	}
	return card.Value.(string)
}

// Version returns the value of the 'EXTVER' Card (or 1 if none)
func (t *Table) Version() int {
	card := t.hdr.Get("EXTVER")
	if card == nil {
		return 1
	}
	return card.Value.(int)
}

// Data returns the image payload
func (t *Table) Data() (Value, error) {
	panic("not implemented")
}

func (t *Table) NumRows() int64 {
	return t.nrows
}

func (t *Table) NumCols() int {
	return len(t.cols)
}

func (t *Table) Cols() []Column {
	return t.cols
}

func (t *Table) Col(i int) *Column {
	return &t.cols[i]
}

// Index returns the index of the first column with name `n` or -1
func (t *Table) Index(n string) int {
	idx, ok := t.colidx[n]
	if !ok {
		return -1
	}
	return idx
}

// ReadRange reads rows over the range [beg, end) and returns the corresponding iterator.
// if end > maxrows, the iteration will stop at maxrows
// ReadRange has the same semantics than a `for i=0; i < max; i+=inc {...}` loop
func (t *Table) ReadRange(beg, end, inc int64) (*Rows, error) {
	var err error
	var rows *Rows

	maxrows := t.NumRows()
	if end > maxrows {
		end = maxrows
	}

	if beg < 0 {
		beg = 0
	}

	cols := make([]int, len(t.cols))
	for i := range t.cols {
		cols[i] = i
	}

	rows = &Rows{
		table: t,
		cols:  cols,
		i:     beg,
		n:     end,
		inc:   inc,
		cur:   beg - inc,
		err:   nil,
		icols: make(map[reflect.Type][][2]int),
	}
	return rows, err
}

// Read reads rows over the range [beg, end) and returns the corresponding iterator.
// if end > maxrows, the iteration will stop at maxrows
// ReadRange has the same semantics than a `for i=0; i < max; i++ {...}` loop
func (t *Table) Read(beg, end int64) (*Rows, error) {
	return t.ReadRange(beg, end, 1)
}

// NewTable creates a new table in the given FITS file
func NewTable(name string, cols []Column, hdutype HDUType) (*Table, error) {
	var err error

	isbinary := true
	switch hdutype {
	case ASCII_TBL:
		isbinary = false
	case BINARY_TBL:
		isbinary = true
	default:
		return nil, fmt.Errorf("fitsio: invalid HDUType (%v)", hdutype)
	}

	ncols := len(cols)
	table := &Table{
		hdr:    Header{},
		binary: isbinary,
		data:   make([]byte, 0),
		heap:   make([]byte, 0),
		rowsz:  0, // NAXIS1 in bytes
		nrows:  0, // NAXIS2
		cols:   make([]Column, ncols),
		colidx: make(map[string]int, ncols),
	}

	copy(table.cols, cols)

	cards := make([]Card, 0, len(cols)+2)
	cards = append(
		cards,
		Card{
			Name:    "TFIELDS",
			Value:   ncols,
			Comment: "number of fields in each row",
		},
	)

	offset := 0
	for i := 0; i < ncols; i++ {
		col := &table.cols[i]
		col.offset = offset
		switch hdutype {
		case BINARY_TBL:
			col.write = col.writeBin
			col.read = col.readBin
		case ASCII_TBL:
			col.write = col.writeTxt
			col.read = col.readTxt
		default:
			return nil, fmt.Errorf("fitsio: invalid HDUType (%v)", hdutype)
		}

		table.colidx[col.Name] = i

		if col.Format == "" {
			return nil, fmt.Errorf("fitsio: column (col=%s) has NO valid format", col.Name)
		}

		cards = append(cards,
			Card{
				Name:    fmt.Sprintf("TTYPE%d", i+1),
				Value:   col.Name,
				Comment: fmt.Sprintf("label for column %d", i+1),
			},
			Card{
				Name:    fmt.Sprintf("TFORM%d", i+1),
				Value:   col.Format,
				Comment: fmt.Sprintf("data format for column %d", i+1),
			},
		)

		col.dtype, err = typeFromForm(col.Format, hdutype)
		if err != nil {
			return nil, err
		}

		offset += col.dtype.dsize * col.dtype.len
		col.txtfmt = txtfmtFromForm(col.Format)

		if offset == 0 && i > 0 {
			return nil, fmt.Errorf("fitsio: invalid data-layout")
		}

		if col.Unit != "" {
			cards = append(cards,
				Card{
					Name:    fmt.Sprintf("TUNIT%d", i+1),
					Value:   col.Unit,
					Comment: fmt.Sprintf("unit for column %d", i+1),
				},
			)
		}

		if col.Null != "" {
			cards = append(cards,
				Card{
					Name:    fmt.Sprintf("TNULL%d", i+1),
					Value:   col.Null,
					Comment: fmt.Sprintf("default value for column %d", i+1),
				},
			)
		}

		cards = append(cards,
			Card{
				Name:    fmt.Sprintf("TSCAL%d", i+1),
				Value:   col.Bscale,
				Comment: fmt.Sprintf("scaling offset for column %d", i+1),
			},
		)

		cards = append(cards,
			Card{
				Name:    fmt.Sprintf("TZERO%d", i+1),
				Value:   col.Bzero,
				Comment: fmt.Sprintf("zero value for column %d", i+1),
			},
		)

		cards = append(cards,
			Card{
				Name:  fmt.Sprintf("TBCOL%d", i+1),
				Value: offset - col.dtype.dsize*col.dtype.len + 1,
			},
		)

		if col.Display != "" {
			cards = append(cards,
				Card{
					Name:    fmt.Sprintf("TDISP%d", i+1),
					Value:   col.Display,
					Comment: fmt.Sprintf("display format for column %d", i+1),
				},
			)
		}

		if len(col.Dim) > 0 {
			str := "("
			for idim, dim := range col.Dim {
				str += fmt.Sprintf("%d", dim)
				if idim+1 < len(col.Dim) {
					str += ","
				}
			}
			str += ")"
			cards = append(cards,
				Card{
					Name:  fmt.Sprintf("TDIM%d", i+1),
					Value: str,
				},
			)
		}

		if col.Start != 0 {
			cards = append(cards,
				Card{
					Name:  fmt.Sprintf("TBCOL%d", i+1),
					Value: int(col.Start),
				},
			)
		}
	}

	cards = append(
		cards,
		Card{
			Name:    "EXTNAME",
			Value:   name,
			Comment: "name of this table extension",
		},
	)

	bitpix := 8
	hdr := newHeader(cards, hdutype, bitpix, []int{offset, 0})
	table.hdr = *hdr
	table.rowsz = offset

	return table, err
}

// NewTableFrom creates a new table in the given FITS file, using the struct v as schema
func NewTableFrom(name string, v Value, hdutype HDUType) (*Table, error) {
	rv := reflect.Indirect(reflect.ValueOf(v))
	rt := rv.Type()
	if rt.Kind() != reflect.Struct {
		return nil, fmt.Errorf("fitsio: NewTableFrom takes a struct value. got: %#T", v)
	}
	nmax := rt.NumField()
	cols := make([]Column, 0, nmax)
	for i := 0; i < nmax; i++ {
		ft := rt.Field(i)
		name := ft.Tag.Get("fits")
		if name == "" {
			name = ft.Name
		}
		field := rv.Field(i)
		form := formFromGoType(field.Type(), hdutype)
		if form == "" {
			return nil, fmt.Errorf("fitsio: no FITS TFORM for field [%d] %#v", i, field.Interface())
		}
		cols = append(cols,
			Column{
				Name:   name,
				Format: form,
			},
		)
	}
	return NewTable(name, cols, hdutype)
}

// Write writes the data into the columns at the current row.
func (t *Table) Write(args ...interface{}) error {
	var err error

	t.data = append(t.data, make([]byte, t.rowsz)...)

	switch len(args) {
	case 0:
		return fmt.Errorf("fitsio: Rows.Scan needs at least one argument")

	case 1:
		// maybe special case: map? struct?
		rt := reflect.TypeOf(args[0]).Elem()
		switch rt.Kind() {
		case reflect.Map:
			err = t.writeMap(*args[0].(*map[string]interface{}))
		case reflect.Struct:
			err = t.writeStruct(args[0])
		default:
			err = t.write(args[0])
		}
	default:
		err = t.write(args...)
	}

	if err != nil {
		return err
	}

	t.nrows += 1
	t.hdr.axes[1] += 1
	return err
}

func (t *Table) write(args ...interface{}) error {
	var err error
	if len(args) != len(t.cols) {
		return fmt.Errorf(
			"fitsio.Table.Write: invalid number of arguments (got %d. expected %d)",
			len(args),
			len(t.cols),
		)
	}

	for i := range t.cols {
		err = t.cols[i].write(t, i, t.nrows, args[i])
		if err != nil {
			return err
		}
	}

	return err
}

func (t *Table) writeMap(data map[string]interface{}) error {
	var err error
	icols := make([]int, 0, len(data))
	switch len(data) {
	case 0:
		icols = make([]int, len(t.cols))
		for i := range t.cols {
			icols[i] = i
		}
	default:
		for k := range data {
			icol := t.Index(k)
			if icol >= 0 {
				icols = append(icols, icol)
			}
		}
	}

	for _, icol := range icols {
		col := t.Col(icol)
		val := reflect.New(col.Type())
		err = col.write(t, icol, t.nrows, val.Interface())
		if err != nil {
			return err
		}
		data[col.Name] = val.Elem().Interface()
	}
	return err
}

func (t *Table) writeStruct(data interface{}) error {
	var err error
	rt := reflect.TypeOf(data).Elem()
	rv := reflect.ValueOf(data).Elem()
	icols := make([][2]int, 0, rt.NumField())

	if true { // fixme: devise a cache ?
		for i := 0; i < rt.NumField(); i++ {
			f := rt.Field(i)
			n := f.Tag.Get("fits")
			if n == "" {
				n = f.Name
			}
			icol := t.Index(n)
			if icol >= 0 {
				icols = append(icols, [2]int{i, icol})
			}
		}
	}

	for _, icol := range icols {
		col := &t.cols[icol[1]]
		field := rv.Field(icol[0])
		value := field.Addr().Interface()
		err = col.write(t, icol[1], t.nrows, value)
		if err != nil {
			return err
		}
	}
	return err
}

// freeze freezes a Table before writing, calculating offsets and finalizing header values.
func (t *Table) freeze() error {
	var err error
	nrows := t.nrows
	t.hdr.axes[1] = int(nrows)

	hduext := ""
	if t.binary {
		hduext = "BINTABLE"
	} else {
		hduext = "TABLE   "
	}
	cards := []Card{
		{
			Name:    "XTENSION",
			Value:   hduext,
			Comment: "table extension",
		},
		{
			Name:    "BITPIX",
			Value:   t.Header().Bitpix(),
			Comment: "number of bits per data pixel",
		},
		{
			Name:    "NAXIS",
			Value:   len(t.Header().Axes()),
			Comment: "number of data axes",
		},
		{
			Name:    "NAXIS1",
			Value:   t.Header().Axes()[0],
			Comment: "length of data axis 1",
		},
		{
			Name:    "NAXIS2",
			Value:   t.Header().Axes()[1],
			Comment: "length of data axis 2",
		},
		{
			Name:    "PCOUNT",
			Value:   len(t.heap),
			Comment: "heap area size (bytes)",
		},
		{
			Name:    "GCOUNT",
			Value:   1,
			Comment: "one data group",
		},
	}

	err = t.hdr.prepend(cards...)
	if err != nil {
		return err
	}

	err = t.hdr.Append([]Card{
		{
			Name:    "THEAP",
			Value:   0,
			Comment: "gap size (bytes)",
		},
	}...)

	return err
}
<commit_msg>table: cleverer freeze function - dont add duplicates<commit_after>package fitsio

import (
	"fmt"
	"reflect"
)

type Table struct {
	hdr    Header
	binary bool

	data []byte // main data table
	heap []byte // heap data table (for variable length arrays)

	rowsz  int   // size of each row in bytes (ie: NAXIS1)
	nrows  int64 // number of rows (ie: NAXIS2)
	cols   []Column
	colidx map[string]int // associates a column name to its index
}

// Close closes this HDU, cleaning up cycles (if any) for garbage collection
func (t *Table) Close() error {
	return nil
}

// Header returns the Header part of this HDU block.
func (t *Table) Header() *Header {
	return &t.hdr
}

// Type returns the Type of this HDU
func (t *Table) Type() HDUType {
	return t.hdr.Type()
}

// Name returns the value of the 'EXTNAME' Card.
func (t *Table) Name() string {
	card := t.hdr.Get("EXTNAME")
	if card == nil {
		return ""
	}
	return card.Value.(string)
}

// Version returns the value of the 'EXTVER' Card (or 1 if none)
func (t *Table) Version() int {
	card := t.hdr.Get("EXTVER")
	if card == nil {
		return 1
	}
	return card.Value.(int)
}

// Data returns the image payload
func (t *Table) Data() (Value, error) {
	panic("not implemented")
}

func (t *Table) NumRows() int64 {
	return t.nrows
}

func (t *Table) NumCols() int {
	return len(t.cols)
}

func (t *Table) Cols() []Column {
	return t.cols
}

func (t *Table) Col(i int) *Column {
	return &t.cols[i]
}

// Index returns the index of the first column with name `n` or -1
func (t *Table) Index(n string) int {
	idx, ok := t.colidx[n]
	if !ok {
		return -1
	}
	return idx
}

// ReadRange reads rows over the range [beg, end) and returns the corresponding iterator.
// if end > maxrows, the iteration will stop at maxrows
// ReadRange has the same semantics than a `for i=0; i < max; i+=inc {...}` loop
func (t *Table) ReadRange(beg, end, inc int64) (*Rows, error) {
	var err error
	var rows *Rows

	maxrows := t.NumRows()
	if end > maxrows {
		end = maxrows
	}

	if beg < 0 {
		beg = 0
	}

	cols := make([]int, len(t.cols))
	for i := range t.cols {
		cols[i] = i
	}

	rows = &Rows{
		table: t,
		cols:  cols,
		i:     beg,
		n:     end,
		inc:   inc,
		cur:   beg - inc,
		err:   nil,
		icols: make(map[reflect.Type][][2]int),
	}
	return rows, err
}

// Read reads rows over the range [beg, end) and returns the corresponding iterator.
// if end > maxrows, the iteration will stop at maxrows
// ReadRange has the same semantics than a `for i=0; i < max; i++ {...}` loop
func (t *Table) Read(beg, end int64) (*Rows, error) {
	return t.ReadRange(beg, end, 1)
}

// NewTable creates a new table in the given FITS file
func NewTable(name string, cols []Column, hdutype HDUType) (*Table, error) {
	var err error

	isbinary := true
	switch hdutype {
	case ASCII_TBL:
		isbinary = false
	case BINARY_TBL:
		isbinary = true
	default:
		return nil, fmt.Errorf("fitsio: invalid HDUType (%v)", hdutype)
	}

	ncols := len(cols)
	table := &Table{
		hdr:    Header{},
		binary: isbinary,
		data:   make([]byte, 0),
		heap:   make([]byte, 0),
		rowsz:  0, // NAXIS1 in bytes
		nrows:  0, // NAXIS2
		cols:   make([]Column, ncols),
		colidx: make(map[string]int, ncols),
	}

	copy(table.cols, cols)

	cards := make([]Card, 0, len(cols)+2)
	cards = append(
		cards,
		Card{
			Name:    "TFIELDS",
			Value:   ncols,
			Comment: "number of fields in each row",
		},
	)

	offset := 0
	for i := 0; i < ncols; i++ {
		col := &table.cols[i]
		col.offset = offset
		switch hdutype {
		case BINARY_TBL:
			col.write = col.writeBin
			col.read = col.readBin
		case ASCII_TBL:
			col.write = col.writeTxt
			col.read = col.readTxt
		default:
			return nil, fmt.Errorf("fitsio: invalid HDUType (%v)", hdutype)
		}

		table.colidx[col.Name] = i

		if col.Format == "" {
			return nil, fmt.Errorf("fitsio: column (col=%s) has NO valid format", col.Name)
		}

		cards = append(cards,
			Card{
				Name:    fmt.Sprintf("TTYPE%d", i+1),
				Value:   col.Name,
				Comment: fmt.Sprintf("label for column %d", i+1),
			},
			Card{
				Name:    fmt.Sprintf("TFORM%d", i+1),
				Value:   col.Format,
				Comment: fmt.Sprintf("data format for column %d", i+1),
			},
		)

		col.dtype, err = typeFromForm(col.Format, hdutype)
		if err != nil {
			return nil, err
		}

		offset += col.dtype.dsize * col.dtype.len
		col.txtfmt = txtfmtFromForm(col.Format)

		if offset == 0 && i > 0 {
			return nil, fmt.Errorf("fitsio: invalid data-layout")
		}

		if col.Unit != "" {
			cards = append(cards,
				Card{
					Name:    fmt.Sprintf("TUNIT%d", i+1),
					Value:   col.Unit,
					Comment: fmt.Sprintf("unit for column %d", i+1),
				},
			)
		}

		if col.Null != "" {
			cards = append(cards,
				Card{
					Name:    fmt.Sprintf("TNULL%d", i+1),
					Value:   col.Null,
					Comment: fmt.Sprintf("default value for column %d", i+1),
				},
			)
		}

		cards = append(cards,
			Card{
				Name:    fmt.Sprintf("TSCAL%d", i+1),
				Value:   col.Bscale,
				Comment: fmt.Sprintf("scaling offset for column %d", i+1),
			},
		)

		cards = append(cards,
			Card{
				Name:    fmt.Sprintf("TZERO%d", i+1),
				Value:   col.Bzero,
				Comment: fmt.Sprintf("zero value for column %d", i+1),
			},
		)

		cards = append(cards,
			Card{
				Name:  fmt.Sprintf("TBCOL%d", i+1),
				Value: offset - col.dtype.dsize*col.dtype.len + 1,
			},
		)

		if col.Display != "" {
			cards = append(cards,
				Card{
					Name:    fmt.Sprintf("TDISP%d", i+1),
					Value:   col.Display,
					Comment: fmt.Sprintf("display format for column %d", i+1),
				},
			)
		}

		if len(col.Dim) > 0 {
			str := "("
			for idim, dim := range col.Dim {
				str += fmt.Sprintf("%d", dim)
				if idim+1 < len(col.Dim) {
					str += ","
				}
			}
			str += ")"
			cards = append(cards,
				Card{
					Name:  fmt.Sprintf("TDIM%d", i+1),
					Value: str,
				},
			)
		}

		if col.Start != 0 {
			cards = append(cards,
				Card{
					Name:  fmt.Sprintf("TBCOL%d", i+1),
					Value: int(col.Start),
				},
			)
		}
	}

	cards = append(
		cards,
		Card{
			Name:    "EXTNAME",
			Value:   name,
			Comment: "name of this table extension",
		},
	)

	bitpix := 8
	hdr := newHeader(cards, hdutype, bitpix, []int{offset, 0})
	table.hdr = *hdr
	table.rowsz = offset

	return table, err
}

// NewTableFrom creates a new table in the given FITS file, using the struct v as schema
func NewTableFrom(name string, v Value, hdutype HDUType) (*Table, error) {
	rv := reflect.Indirect(reflect.ValueOf(v))
	rt := rv.Type()
	if rt.Kind() != reflect.Struct {
		return nil, fmt.Errorf("fitsio: NewTableFrom takes a struct value. got: %#T", v)
	}
	nmax := rt.NumField()
	cols := make([]Column, 0, nmax)
	for i := 0; i < nmax; i++ {
		ft := rt.Field(i)
		name := ft.Tag.Get("fits")
		if name == "" {
			name = ft.Name
		}
		field := rv.Field(i)
		form := formFromGoType(field.Type(), hdutype)
		if form == "" {
			return nil, fmt.Errorf("fitsio: no FITS TFORM for field [%d] %#v", i, field.Interface())
		}
		cols = append(cols,
			Column{
				Name:   name,
				Format: form,
			},
		)
	}
	return NewTable(name, cols, hdutype)
}

// Write writes the data into the columns at the current row.
func (t *Table) Write(args ...interface{}) error {
	var err error

	t.data = append(t.data, make([]byte, t.rowsz)...)

	switch len(args) {
	case 0:
		return fmt.Errorf("fitsio: Rows.Scan needs at least one argument")

	case 1:
		// maybe special case: map? struct?
		rt := reflect.TypeOf(args[0]).Elem()
		switch rt.Kind() {
		case reflect.Map:
			err = t.writeMap(*args[0].(*map[string]interface{}))
		case reflect.Struct:
			err = t.writeStruct(args[0])
		default:
			err = t.write(args[0])
		}
	default:
		err = t.write(args...)
	}

	if err != nil {
		return err
	}

	t.nrows += 1
	t.hdr.axes[1] += 1
	return err
}

func (t *Table) write(args ...interface{}) error {
	var err error
	if len(args) != len(t.cols) {
		return fmt.Errorf(
			"fitsio.Table.Write: invalid number of arguments (got %d. expected %d)",
			len(args),
			len(t.cols),
		)
	}

	for i := range t.cols {
		err = t.cols[i].write(t, i, t.nrows, args[i])
		if err != nil {
			return err
		}
	}

	return err
}

func (t *Table) writeMap(data map[string]interface{}) error {
	var err error
	icols := make([]int, 0, len(data))
	switch len(data) {
	case 0:
		icols = make([]int, len(t.cols))
		for i := range t.cols {
			icols[i] = i
		}
	default:
		for k := range data {
			icol := t.Index(k)
			if icol >= 0 {
				icols = append(icols, icol)
			}
		}
	}

	for _, icol := range icols {
		col := t.Col(icol)
		val := reflect.New(col.Type())
		err = col.write(t, icol, t.nrows, val.Interface())
		if err != nil {
			return err
		}
		data[col.Name] = val.Elem().Interface()
	}
	return err
}

func (t *Table) writeStruct(data interface{}) error {
	var err error
	rt := reflect.TypeOf(data).Elem()
	rv := reflect.ValueOf(data).Elem()
	icols := make([][2]int, 0, rt.NumField())

	if true { // fixme: devise a cache ?
		for i := 0; i < rt.NumField(); i++ {
			f := rt.Field(i)
			n := f.Tag.Get("fits")
			if n == "" {
				n = f.Name
			}
			icol := t.Index(n)
			if icol >= 0 {
				icols = append(icols, [2]int{i, icol})
			}
		}
	}

	for _, icol := range icols {
		col := &t.cols[icol[1]]
		field := rv.Field(icol[0])
		value := field.Addr().Interface()
		err = col.write(t, icol[1], t.nrows, value)
		if err != nil {
			return err
		}
	}
	return err
}

// freeze freezes a Table before writing, calculating offsets and finalizing header values.
func (t *Table) freeze() error {
	var err error
	nrows := t.nrows
	t.hdr.axes[1] = int(nrows)

	if card := t.Header().Get("XTENSION"); card == nil {
		hduext := ""
		if t.binary {
			hduext = "BINTABLE"
		} else {
			hduext = "TABLE   "
		}
		cards := []Card{
			{
				Name:    "XTENSION",
				Value:   hduext,
				Comment: "table extension",
			},
			{
				Name:    "BITPIX",
				Value:   t.Header().Bitpix(),
				Comment: "number of bits per data pixel",
			},
			{
				Name:    "NAXIS",
				Value:   len(t.Header().Axes()),
				Comment: "number of data axes",
			},
			{
				Name:    "NAXIS1",
				Value:   t.Header().Axes()[0],
				Comment: "length of data axis 1",
			},
			{
				Name:    "NAXIS2",
				Value:   t.Header().Axes()[1],
				Comment: "length of data axis 2",
			},
			{
				Name:    "PCOUNT",
				Value:   len(t.heap),
				Comment: "heap area size (bytes)",
			},
			{
				Name:    "GCOUNT",
				Value:   1,
				Comment: "one data group",
			},
		}

		err = t.hdr.prepend(cards...)
		if err != nil {
			return err
		}
	}

	if card := t.Header().Get("THEAP"); card == nil {
		err = t.hdr.Append([]Card{
			{
				Name:    "THEAP",
				Value:   0,
				Comment: "gap size (bytes)",
			},
		}...)
	}

	return err
}
<|endoftext|>


<commit_before>package mktmpio

import (
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

// Instance represents a server that has been created on the mktmpio service.
type Instance struct {
	ID          string
	Host        string
	Port        int
	Error       string
	RemoteShell shell
	Type        string
	Username    string
	Password    string
	client      Client
}

type shell struct {
	Cmd []string
	Env map[string]string
}

// Destroy the server on the mktmpio service
func (i *Instance) Destroy() {
	i.client.Destroy(i.ID)
}

// Cmd returns an exec.Cmd that is pre-populated with the command, arguments,
// and environment variables required for spawning a local shell connected to
// the remote server.
func (i *Instance) Cmd() *exec.Cmd {
	cmd := exec.Command(i.RemoteShell.Cmd[0], i.RemoteShell.Cmd[1:]...)
	if len(i.RemoteShell.Env) > 0 {
		cmd.Env = append(os.Environ(), envList(i.RemoteShell.Env)...)
	}
	return cmd
}

// LoadEnv modifies the current environment by setting environment variables
// that contain the host, port and credentials required for connecting to the
// remote server represented by the Instance.
func (i *Instance) LoadEnv() error {
	var err error
	setEnv := func(key, val string) {
		if err == nil {
			err = os.Setenv(envKey(i, key), val)
		}
	}
	setEnv("host", i.Host)
	setEnv("port", strconv.Itoa(i.Port))
	setEnv("username", i.Username)
	setEnv("password", i.Password)
	return err
}

func envKey(i *Instance, field interface{}) string {
	return strings.ToUpper(fmt.Sprintf("%s_%s", i.Type, field))
}

func envList(kv map[string]string) []string {
	env := make([]string, len(kv))
	for k, v := range kv {
		env = append(env, k+"="+v)
	}
	return env
}
<commit_msg>return error from Instance.Destroy()<commit_after>package mktmpio

import (
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

// Instance represents a server that has been created on the mktmpio service.
type Instance struct {
	ID          string
	Host        string
	Port        int
	Error       string
	RemoteShell shell
	Type        string
	Username    string
	Password    string
	client      Client
}

type shell struct {
	Cmd []string
	Env map[string]string
}

// Destroy the server on the mktmpio service
func (i *Instance) Destroy() error {
	return i.client.Destroy(i.ID)
}

// Cmd returns an exec.Cmd that is pre-populated with the command, arguments,
// and environment variables required for spawning a local shell connected to
// the remote server.
func (i *Instance) Cmd() *exec.Cmd {
	cmd := exec.Command(i.RemoteShell.Cmd[0], i.RemoteShell.Cmd[1:]...)
	if len(i.RemoteShell.Env) > 0 {
		cmd.Env = append(os.Environ(), envList(i.RemoteShell.Env)...)
	}
	return cmd
}

// LoadEnv modifies the current environment by setting environment variables
// that contain the host, port and credentials required for connecting to the
// remote server represented by the Instance.
func (i *Instance) LoadEnv() error {
	var err error
	setEnv := func(key, val string) {
		if err == nil {
			err = os.Setenv(envKey(i, key), val)
		}
	}
	setEnv("host", i.Host)
	setEnv("port", strconv.Itoa(i.Port))
	setEnv("username", i.Username)
	setEnv("password", i.Password)
	return err
}

func envKey(i *Instance, field interface{}) string {
	return strings.ToUpper(fmt.Sprintf("%s_%s", i.Type, field))
}

func envList(kv map[string]string) []string {
	env := make([]string, len(kv))
	for k, v := range kv {
		env = append(env, k+"="+v)
	}
	return env
}
<|endoftext|>


<commit_before>// Copyright 2018 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

package discovery

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"reflect"
	"testing"
	"time"

	"github.com/open-policy-agent/opa/ast"
	bundleApi "github.com/open-policy-agent/opa/bundle"
	"github.com/open-policy-agent/opa/download"
	"github.com/open-policy-agent/opa/plugins"
	"github.com/open-policy-agent/opa/plugins/bundle"
	"github.com/open-policy-agent/opa/plugins/status"
	"github.com/open-policy-agent/opa/storage/inmem"
	"github.com/open-policy-agent/opa/util"
	"github.com/open-policy-agent/opa/version"
)

func TestMain(m *testing.M) {
	if version.Version == "" {
		version.Version = "unit-test"
	}
	os.Exit(m.Run())
}

func TestEvaluateBundle(t *testing.T) {

	sampleModule := `
		package foo.bar

		bundle = {
			"name": rt.name,
			"service": "example"
		} {
			rt := opa.runtime()
		}
	`

	b := &bundleApi.Bundle{
		Manifest: bundleApi.Manifest{
			Revision: "quickbrownfaux",
		},
		Data: map[string]interface{}{
			"foo": map[string]interface{}{
				"bar": map[string]interface{}{
					"status": map[string]interface{}{},
				},
			},
		},
		Modules: []bundleApi.ModuleFile{
			{
				Path:   `/example.rego`,
				Raw:    []byte(sampleModule),
				Parsed: ast.MustParseModule(sampleModule),
			},
		},
	}

	info := ast.MustParseTerm(`{"name": "test/bundle1"}`)

	config, err := evaluateBundle(context.Background(), "test-id", info, b, "data.foo.bar")
	if err != nil {
		t.Fatal(err)
	}

	if config.Bundle == nil {
		t.Fatal("Expected a bundle configuration")
	}

	var parsedConfig bundle.Config

	if err := util.Unmarshal(config.Bundle, &parsedConfig); err != nil {
		t.Fatal("Unexpected error:", err)
	}

	expectedBundleConfig := bundle.Config{
		Name:    "test/bundle1",
		Service: "example",
	}

	if !reflect.DeepEqual(expectedBundleConfig, parsedConfig) {
		t.Fatalf("Expected bundle config %v, but got %v", expectedBundleConfig, parsedConfig)
	}

}

func TestProcessBundle(t *testing.T) {

	ctx := context.Background()

	manager, err := plugins.New([]byte(`{
		"services": {
			"default": {
				"url": "http://localhost:8181"
			}
		}
	}`), "test-id", inmem.New())
	if err != nil {
		t.Fatal(err)
	}

	initialBundle := makeDataBundle(1, `
		{
			"config": {
				"bundle": {"name": "test1"},
				"status": {},
				"decision_logs": {}
			}
		}
	`)

	_, ps, err := processBundle(ctx, manager, nil, initialBundle, "data.config")
	if err != nil {
		t.Fatal(err)
	}

	if len(ps.Start) != 3 || len(ps.Reconfig) != 0 {
		t.Fatalf("Expected exactly three start events but got %v", ps)
	}

	updatedBundle := makeDataBundle(1, `
		{
			"config": {
				"bundle": {"name": "test2"},
				"status": {"partition_name": "foo"},
				"decision_logs": {"partition_name": "bar"}
			}
		}
	`)

	_, ps, err = processBundle(ctx, manager, nil, updatedBundle, "data.config")
	if err != nil {
		t.Fatal(err)
	}

	if len(ps.Start) != 0 || len(ps.Reconfig) != 3 {
		t.Fatalf("Expected exactly three start events but got %v", ps)
	}

	updatedBundle = makeDataBundle(2, `
		{
			"config": {
				"bundle": {"service": "missing service name", "name": "test2"}
			}
		}
	`)

	_, _, err = processBundle(ctx, manager, nil, updatedBundle, "data.config")
	if err == nil {
		t.Fatal("Expected error but got success")
	}

}

type testFactory struct {
	p *reconfigureTestPlugin
}

func (f testFactory) Validate(*plugins.Manager, []byte) (interface{}, error) {
	return nil, nil
}

func (f testFactory) New(*plugins.Manager, interface{}) plugins.Plugin {
	return f.p
}

type reconfigureTestPlugin struct {
	counts map[string]int
}

func (r *reconfigureTestPlugin) Start(context.Context) error {
	r.counts["start"]++
	return nil
}

func (r *reconfigureTestPlugin) Stop(context.Context) {
}

func (r *reconfigureTestPlugin) Reconfigure(_ context.Context, config interface{}) {
	r.counts["reconfig"]++
}

func TestReconfigure(t *testing.T) {

	manager, err := plugins.New([]byte(`{
		"labels": {"x": "y"},
		"services": {
			"localhost": {
				"url": "http://localhost:9999"
			}
		},
		"discovery": {"name": "config"},
	}`), "test-id", inmem.New())
	if err != nil {
		t.Fatal(err)
	}

	testPlugin := &reconfigureTestPlugin{counts: map[string]int{}}
	testFactory := testFactory{p: testPlugin}

	disco, err := New(manager, Factories(map[string]plugins.Factory{"test_plugin": testFactory}))
	if err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()

	initialBundle := makeDataBundle(1, `
		{
			"config": {
				"labels": {"x": "label value changed"},
				"default_decision": "bar/baz",
				"default_authorization_decision": "baz/qux",
				"plugins": {
					"test_plugin": {"a": "b"}
				}
			}
		}
	`)

	disco.oneShot(ctx, download.Update{Bundle: initialBundle})

	// Verify labels are unchanged
	exp := map[string]string{"x": "y", "id": "test-id", "version": version.Version}
	if !reflect.DeepEqual(manager.Labels(), exp) {
		t.Errorf("Expected labels to be unchanged (%v) but got %v", exp, manager.Labels())
	}

	// Verify decision ids set
	expDecision := ast.MustParseTerm("data.bar.baz")
	expAuthzDecision := ast.MustParseTerm("data.baz.qux")
	if !manager.Config.DefaultDecisionRef().Equal(expDecision.Value) {
		t.Errorf("Expected default decision to be %v but got %v", expDecision, manager.Config.DefaultDecisionRef())
	}
	if !manager.Config.DefaultAuthorizationDecisionRef().Equal(expAuthzDecision.Value) {
		t.Errorf("Expected default authz decision to be %v but got %v", expAuthzDecision, manager.Config.DefaultAuthorizationDecisionRef())
	}

	// Verify plugins started
	if !reflect.DeepEqual(testPlugin.counts, map[string]int{"start": 1}) {
		t.Errorf("Expected exactly one plugin start but got %v", testPlugin)
	}

	// Verify plugins reconfigured
	updatedBundle := makeDataBundle(2, `
		{
			"config": {
				"labels": {"x": "label value changed"},
				"default_decision": "bar/baz",
				"default_authorization_decision": "baz/qux",
				"plugins": {
					"test_plugin": {"a": "plugin parameter value changed"}
				}
			}
		}
	`)

	disco.oneShot(ctx, download.Update{Bundle: updatedBundle})

	if !reflect.DeepEqual(testPlugin.counts, map[string]int{"start": 1, "reconfig": 1}) {
		t.Errorf("Expected one plugin start and one reconfig but got %v", testPlugin)
	}

}

type testServer struct {
	t       *testing.T
	server  *httptest.Server
	updates []status.UpdateRequestV1
}

func (ts *testServer) Start() {
	ts.server = httptest.NewServer(http.HandlerFunc(ts.handle))
}

func (ts *testServer) Stop() {
	ts.server.Close()
}

func (ts *testServer) handle(w http.ResponseWriter, r *http.Request) {

	var update status.UpdateRequestV1

	if err := json.NewDecoder(r.Body).Decode(&update); err != nil {
		ts.t.Fatal(err)
	}

	ts.updates = append(ts.updates, update)

	w.WriteHeader(200)
}

func TestStatusUpdates(t *testing.T) {

	ts := testServer{t: t}
	ts.Start()
	defer ts.Stop()

	manager, err := plugins.New([]byte(fmt.Sprintf(`{
			"labels": {"x": "y"},
			"services": {
				"localhost": {
					"url": %q
				}
			},
			"discovery": {"name": "config"},
		}`, ts.server.URL)), "test-id", inmem.New())
	if err != nil {
		t.Fatal(err)
	}

	disco, err := New(manager)
	if err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()

	// Enable status plugin which sends initial update.
	disco.oneShot(ctx, download.Update{ETag: "etag-1", Bundle: makeDataBundle(1, `{
		"config": {
			"status": {}
		}
	}`)})

	// Downloader error.
	disco.oneShot(ctx, download.Update{Error: fmt.Errorf("unknown error")})

	// Clear error.
	disco.oneShot(ctx, download.Update{ETag: "etag-2", Bundle: makeDataBundle(2, `{
		"config": {
			"status": {}
		}
	}`)})

	// Configuration error.
	disco.oneShot(ctx, download.Update{ETag: "etag-3", Bundle: makeDataBundle(3, `{
		"config": {
			"status": {"service": "missing service"}
		}
	}`)})

	// Clear error (last successful reconfigure).
	disco.oneShot(ctx, download.Update{ETag: "etag-2"})

	// Check that all updates were received and active revisions are expected.
	var ok bool
	t0 := time.Now()

	for !ok && time.Since(t0) < time.Second {
		ok = len(ts.updates) == 5 &&
			ts.updates[0].Discovery.ActiveRevision == "test-revision-1" && ts.updates[0].Discovery.Code == "" &&
			ts.updates[1].Discovery.ActiveRevision == "test-revision-1" && ts.updates[1].Discovery.Code == "bundle_error" &&
			ts.updates[2].Discovery.ActiveRevision == "test-revision-2" && ts.updates[2].Discovery.Code == "" &&
			ts.updates[3].Discovery.ActiveRevision == "test-revision-2" && ts.updates[3].Discovery.Code == "bundle_error" &&
			ts.updates[4].Discovery.ActiveRevision == "test-revision-2" && ts.updates[4].Discovery.Code == ""
	}

	if !ok {
		t.Fatalf("Did not receive expected updates before timeout expired. Received: %+v", ts.updates)
	}
}

func makeDataBundle(n int, s string) *bundleApi.Bundle {
	return &bundleApi.Bundle{
		Manifest: bundleApi.Manifest{Revision: fmt.Sprintf("test-revision-%v", n)},
		Data:     util.MustUnmarshalJSON([]byte(s)).(map[string]interface{}),
	}
}
<commit_msg>discovery: Fix race condition in test case<commit_after>// Copyright 2018 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

package discovery

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"reflect"
	"sync"
	"testing"
	"time"

	"github.com/open-policy-agent/opa/ast"
	bundleApi "github.com/open-policy-agent/opa/bundle"
	"github.com/open-policy-agent/opa/download"
	"github.com/open-policy-agent/opa/plugins"
	"github.com/open-policy-agent/opa/plugins/bundle"
	"github.com/open-policy-agent/opa/plugins/status"
	"github.com/open-policy-agent/opa/storage/inmem"
	"github.com/open-policy-agent/opa/util"
	"github.com/open-policy-agent/opa/version"
)

func TestMain(m *testing.M) {
	if version.Version == "" {
		version.Version = "unit-test"
	}
	os.Exit(m.Run())
}

func TestEvaluateBundle(t *testing.T) {

	sampleModule := `
		package foo.bar

		bundle = {
			"name": rt.name,
			"service": "example"
		} {
			rt := opa.runtime()
		}
	`

	b := &bundleApi.Bundle{
		Manifest: bundleApi.Manifest{
			Revision: "quickbrownfaux",
		},
		Data: map[string]interface{}{
			"foo": map[string]interface{}{
				"bar": map[string]interface{}{
					"status": map[string]interface{}{},
				},
			},
		},
		Modules: []bundleApi.ModuleFile{
			{
				Path:   `/example.rego`,
				Raw:    []byte(sampleModule),
				Parsed: ast.MustParseModule(sampleModule),
			},
		},
	}

	info := ast.MustParseTerm(`{"name": "test/bundle1"}`)

	config, err := evaluateBundle(context.Background(), "test-id", info, b, "data.foo.bar")
	if err != nil {
		t.Fatal(err)
	}

	if config.Bundle == nil {
		t.Fatal("Expected a bundle configuration")
	}

	var parsedConfig bundle.Config

	if err := util.Unmarshal(config.Bundle, &parsedConfig); err != nil {
		t.Fatal("Unexpected error:", err)
	}

	expectedBundleConfig := bundle.Config{
		Name:    "test/bundle1",
		Service: "example",
	}

	if !reflect.DeepEqual(expectedBundleConfig, parsedConfig) {
		t.Fatalf("Expected bundle config %v, but got %v", expectedBundleConfig, parsedConfig)
	}

}

func TestProcessBundle(t *testing.T) {

	ctx := context.Background()

	manager, err := plugins.New([]byte(`{
		"services": {
			"default": {
				"url": "http://localhost:8181"
			}
		}
	}`), "test-id", inmem.New())
	if err != nil {
		t.Fatal(err)
	}

	initialBundle := makeDataBundle(1, `
		{
			"config": {
				"bundle": {"name": "test1"},
				"status": {},
				"decision_logs": {}
			}
		}
	`)

	_, ps, err := processBundle(ctx, manager, nil, initialBundle, "data.config")
	if err != nil {
		t.Fatal(err)
	}

	if len(ps.Start) != 3 || len(ps.Reconfig) != 0 {
		t.Fatalf("Expected exactly three start events but got %v", ps)
	}

	updatedBundle := makeDataBundle(1, `
		{
			"config": {
				"bundle": {"name": "test2"},
				"status": {"partition_name": "foo"},
				"decision_logs": {"partition_name": "bar"}
			}
		}
	`)

	_, ps, err = processBundle(ctx, manager, nil, updatedBundle, "data.config")
	if err != nil {
		t.Fatal(err)
	}

	if len(ps.Start) != 0 || len(ps.Reconfig) != 3 {
		t.Fatalf("Expected exactly three start events but got %v", ps)
	}

	updatedBundle = makeDataBundle(2, `
		{
			"config": {
				"bundle": {"service": "missing service name", "name": "test2"}
			}
		}
	`)

	_, _, err = processBundle(ctx, manager, nil, updatedBundle, "data.config")
	if err == nil {
		t.Fatal("Expected error but got success")
	}

}

type testFactory struct {
	p *reconfigureTestPlugin
}

func (f testFactory) Validate(*plugins.Manager, []byte) (interface{}, error) {
	return nil, nil
}

func (f testFactory) New(*plugins.Manager, interface{}) plugins.Plugin {
	return f.p
}

type reconfigureTestPlugin struct {
	counts map[string]int
}

func (r *reconfigureTestPlugin) Start(context.Context) error {
	r.counts["start"]++
	return nil
}

func (r *reconfigureTestPlugin) Stop(context.Context) {
}

func (r *reconfigureTestPlugin) Reconfigure(_ context.Context, config interface{}) {
	r.counts["reconfig"]++
}

func TestReconfigure(t *testing.T) {

	manager, err := plugins.New([]byte(`{
		"labels": {"x": "y"},
		"services": {
			"localhost": {
				"url": "http://localhost:9999"
			}
		},
		"discovery": {"name": "config"},
	}`), "test-id", inmem.New())
	if err != nil {
		t.Fatal(err)
	}

	testPlugin := &reconfigureTestPlugin{counts: map[string]int{}}
	testFactory := testFactory{p: testPlugin}

	disco, err := New(manager, Factories(map[string]plugins.Factory{"test_plugin": testFactory}))
	if err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()

	initialBundle := makeDataBundle(1, `
		{
			"config": {
				"labels": {"x": "label value changed"},
				"default_decision": "bar/baz",
				"default_authorization_decision": "baz/qux",
				"plugins": {
					"test_plugin": {"a": "b"}
				}
			}
		}
	`)

	disco.oneShot(ctx, download.Update{Bundle: initialBundle})

	// Verify labels are unchanged
	exp := map[string]string{"x": "y", "id": "test-id", "version": version.Version}
	if !reflect.DeepEqual(manager.Labels(), exp) {
		t.Errorf("Expected labels to be unchanged (%v) but got %v", exp, manager.Labels())
	}

	// Verify decision ids set
	expDecision := ast.MustParseTerm("data.bar.baz")
	expAuthzDecision := ast.MustParseTerm("data.baz.qux")
	if !manager.Config.DefaultDecisionRef().Equal(expDecision.Value) {
		t.Errorf("Expected default decision to be %v but got %v", expDecision, manager.Config.DefaultDecisionRef())
	}
	if !manager.Config.DefaultAuthorizationDecisionRef().Equal(expAuthzDecision.Value) {
		t.Errorf("Expected default authz decision to be %v but got %v", expAuthzDecision, manager.Config.DefaultAuthorizationDecisionRef())
	}

	// Verify plugins started
	if !reflect.DeepEqual(testPlugin.counts, map[string]int{"start": 1}) {
		t.Errorf("Expected exactly one plugin start but got %v", testPlugin)
	}

	// Verify plugins reconfigured
	updatedBundle := makeDataBundle(2, `
		{
			"config": {
				"labels": {"x": "label value changed"},
				"default_decision": "bar/baz",
				"default_authorization_decision": "baz/qux",
				"plugins": {
					"test_plugin": {"a": "plugin parameter value changed"}
				}
			}
		}
	`)

	disco.oneShot(ctx, download.Update{Bundle: updatedBundle})

	if !reflect.DeepEqual(testPlugin.counts, map[string]int{"start": 1, "reconfig": 1}) {
		t.Errorf("Expected one plugin start and one reconfig but got %v", testPlugin)
	}

}

type testServer struct {
	t       *testing.T
	mtx     sync.Mutex
	server  *httptest.Server
	updates []status.UpdateRequestV1
}

func (ts *testServer) Start() {
	ts.server = httptest.NewServer(http.HandlerFunc(ts.handle))
}

func (ts *testServer) Stop() {
	ts.server.Close()
}

func (ts *testServer) Updates() []status.UpdateRequestV1 {
	ts.mtx.Lock()
	defer ts.mtx.Unlock()
	return ts.updates
}

func (ts *testServer) handle(w http.ResponseWriter, r *http.Request) {

	var update status.UpdateRequestV1

	if err := json.NewDecoder(r.Body).Decode(&update); err != nil {
		ts.t.Fatal(err)
	}

	func() {
		ts.mtx.Lock()
		defer ts.mtx.Unlock()
		ts.updates = append(ts.updates, update)
	}()

	w.WriteHeader(200)
}

func TestStatusUpdates(t *testing.T) {

	ts := testServer{t: t}
	ts.Start()
	defer ts.Stop()

	manager, err := plugins.New([]byte(fmt.Sprintf(`{
			"labels": {"x": "y"},
			"services": {
				"localhost": {
					"url": %q
				}
			},
			"discovery": {"name": "config"},
		}`, ts.server.URL)), "test-id", inmem.New())
	if err != nil {
		t.Fatal(err)
	}

	disco, err := New(manager)
	if err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()

	// Enable status plugin which sends initial update.
	disco.oneShot(ctx, download.Update{ETag: "etag-1", Bundle: makeDataBundle(1, `{
		"config": {
			"status": {}
		}
	}`)})

	// Downloader error.
	disco.oneShot(ctx, download.Update{Error: fmt.Errorf("unknown error")})

	// Clear error.
	disco.oneShot(ctx, download.Update{ETag: "etag-2", Bundle: makeDataBundle(2, `{
		"config": {
			"status": {}
		}
	}`)})

	// Configuration error.
	disco.oneShot(ctx, download.Update{ETag: "etag-3", Bundle: makeDataBundle(3, `{
		"config": {
			"status": {"service": "missing service"}
		}
	}`)})

	// Clear error (last successful reconfigure).
	disco.oneShot(ctx, download.Update{ETag: "etag-2"})

	// Check that all updates were received and active revisions are expected.
	var ok bool
	var updates []status.UpdateRequestV1
	t0 := time.Now()

	for !ok && time.Since(t0) < time.Second {
		updates = ts.Updates()
		ok = len(updates) == 5 &&
			updates[0].Discovery.ActiveRevision == "test-revision-1" && updates[0].Discovery.Code == "" &&
			updates[1].Discovery.ActiveRevision == "test-revision-1" && updates[1].Discovery.Code == "bundle_error" &&
			updates[2].Discovery.ActiveRevision == "test-revision-2" && updates[2].Discovery.Code == "" &&
			updates[3].Discovery.ActiveRevision == "test-revision-2" && updates[3].Discovery.Code == "bundle_error" &&
			updates[4].Discovery.ActiveRevision == "test-revision-2" && updates[4].Discovery.Code == ""
	}

	if !ok {
		t.Fatalf("Did not receive expected updates before timeout expired. Received: %+v", updates)
	}
}

func makeDataBundle(n int, s string) *bundleApi.Bundle {
	return &bundleApi.Bundle{
		Manifest: bundleApi.Manifest{Revision: fmt.Sprintf("test-revision-%v", n)},
		Data:     util.MustUnmarshalJSON([]byte(s)).(map[string]interface{}),
	}
}
<|endoftext|>