<commit_before>package main

import (
	"fmt"
	"flag"
	"os"
	"io/ioutil"
	"strings"
	"regexp"
	"sort"
)

var cluster, command, instanceId, instanceName, service, sshPem, output, credentialsFile, awsKey, awsSecretKey, version, loadBalancer string
var recursive, verbose, moreVerbose bool
var region = "eu-west-1"
var auth *Auth
var verboseLevel = 0

type Auth struct {
	key    string `toml:"aws_access_key_id"`
	secret string `toml:"aws_secret_access_key"`
}

func init() {
	flag.StringVar(&command, "command", "", "The command to use. The command 'help' displays available commands.")
	flag.StringVar(&cluster, "cluster", "", "Specify cluster to use")
	flag.StringVar(&instanceId, "instanceId", "", "Specify the EC2 instance")
	flag.StringVar(&instanceName, "instanceName", "", "Specify the EC2 instance(s) name")
	flag.StringVar(&service, "service", "", "Specify ECS service")
	flag.StringVar(&sshPem, "pemfile", "", "Specify PEM file for SSH access")
	flag.BoolVar(&recursive, "recursive", false, "Specify recursive operation")
	flag.StringVar(&output, "output", "", "Specify output directory")
	flag.StringVar(&credentialsFile, "credentials", "", "Specify credentials used for accessing AWS. Should be of format: .aws/credentials")
	flag.StringVar(&awsKey, "awsKey", "", "AWS key used for authentication. Overrides credentials file")
	flag.StringVar(&awsSecretKey, "awsSecretKey", "", "AWS secret key used for authentication, used in conjunction with 'awsKey'")
	flag.StringVar(&version, "version", "", "The version to use for docker image in the task definition")
	flag.StringVar(&loadBalancer, "loadBalancer", "", "Specifies the load balancer name to use")
	flag.BoolVar(&verbose, "v", false, "Making output more verbose, where applicable")
	flag.BoolVar(&moreVerbose, "vv", false, "Making output more verbose, where applicable")
}

func printCommandHelp() {
	var m = map[string]string{
		"help" : "Prints this help.",
		"listClusters" : "List available clusters.",
		"listServices" : "List available services. Needs -cluster flag.",
		"listTasks" : "List tasks for a service. Needs -cluster, -service flags.",
		"describeService" : "Describes the service. Needs -cluster, -service flags. Optionaly -v and -vv may be used.",
		"updateService" : "Stop/start all running tasks for the specified service. Needs -cluster, -service flags.",
		"releaseService" : "Creates a new release for the service. Neews -cluster, -service, -version flags.",
		"listEc2Instances" : "List available EC2 instances.",
		"listLoadBalancers" : "List available Load Balancers and their contained EC2 instances.",
		"getEntity" : "Gets an entity from the writer load balancer\n" +
			"                      -loadBalancer : The load balancer fronting the writer instances    (required)\n" +
			"                      {entityId}    : The ID of the entity to fetch    (required)\n",
		"ssh" : "Executes a command over SSH for the specified service.\n" +
			"                      -instanceName : The aws instance(s) to use as source(s). Operation will occur on all instances with the specific name   (required if instanceId is not specified)\n" +
			"                      -instanceId   : The specific aws instance to use as source.   (required if instanceName is not specified)\n" +
			"                      -pemfile      : The SSH pem file used for authentication    (required)\n" +
			"                      {command}     : The command to execute (e.g. 'ls -l')   (required)\n" +
			"                      Example: -command ssh -instanceName writer -pemfile ~/.ssh/pem-files/im-dev tail -20 /var/log/writer/writer.log\n",

		"scp" : "Copies files from the specified instance(s). Needs -instanceName or -instanceId, -output and optionally -recursive flags.\n" +
			"                      -instanceName : The aws instance(s) to use as source(s). Operation will occur on all instances with the specific name   (required if instanceId is not specified)\n" +
			"                      -instanceId   : The specific aws instance to use as source.   (required if instanceName is not specified)\n" +
			"                      -pemfile      : The SSH pem file used for authentication    (required)\n" +
			"                      -output       : the target directory   (required)\n" +
			"                      -recursive    : copies from source recursively\n" +
			"                      Example: -command scp -instanceName writer -pemfile ~/.ssh/pem-files/im-dev -output Documents -recursive /var/log/writer\n",
	}

	k := sortKeys(m)

	for _, v := range k {
		fmt.Print(v);
		for j := 0; j < 20 - (len([]rune(v))); j++ {
			fmt.Print(" ")
		}
		fmt.Println(m[v])
	}

}

func sortKeys(m map[string]string) []string {
	keys := make([]string, 0, len(m))
	for key := range m {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	return keys
}

func errUsage(message string) {
	fmt.Println(message)
	os.Exit(1);
}

func errState(message string) {
	fmt.Println(message)
	os.Exit(2)
}

func errStatef(message string, a ...interface{}) {
	fmt.Printf(message, a)
	os.Exit(2)
}

func _getClusterArn() string {
	if (cluster == "") {
		errUsage("You must specify a cluster name with: -cluster");
	}

	arn := GetClusterArn(cluster)
	if (arn == "") {
		errUsage("Could not find cluster ARN for name: " + cluster);
	}

	return arn
}

func _getServiceArn() string {
	if (cluster == "") {
		errUsage("You must specify a cluster name with: -cluster");
	}

	if (service == "") {
		errUsage("You must specify a service name with: -service");
	}

	clusterArn := _getClusterArn();

	serviceArn := GetServiceArn(clusterArn, service)

	return serviceArn
}

func _getVersion() string {
	if (version == "") {
		errUsage("You must specify a version with: -version")
	}

	return version;
}

func getAwsCredentials(filepath string) (awsAccessKeyId, awsSecretKey string) {
	var awsAccessKeyIdRegexp = regexp.MustCompile("aws_access_key_id\\s*=\\s*(.*)")
	var awsSecretKeyRegexp = regexp.MustCompile("aws_secret_access_key\\s*=\\s*(.*)")

	file, err := ioutil.ReadFile(filepath)
	if err != nil {
		errUsage(err.Error())
	}

	matches := awsAccessKeyIdRegexp.FindStringSubmatch(string(file))
	if len(matches) > 1 {
		awsAccessKeyId = strings.TrimSpace(matches[1])
	} else {
		errUsage("Could not find aws access key id")
	}

	matches = awsSecretKeyRegexp.FindStringSubmatch(string(file))

	if len(matches) > 1 {
		awsSecretKey = strings.TrimSpace(matches[1])
	} else {
		errUsage("Could not find aws secret key")
	}

	return awsAccessKeyId, awsSecretKey
}

func main() {
	flag.Parse()

	if (verbose) {
		verboseLevel = 1
	}
	if (moreVerbose) {
		verboseLevel = 2
	}

	if (command == "") {
		flag.PrintDefaults();
		return
	}

	if (credentialsFile != "") {
		key, secret := getAwsCredentials(credentialsFile)
		auth = new(Auth)
		auth.key = key
		auth.secret = secret
	}

	if (awsKey != "" && awsSecretKey == "" ) {
		errUsage("Missing secretKey")
	}

	if (awsKey != "" && awsSecretKey != "" ) {
		auth = new(Auth)
		auth.key = awsKey
		auth.secret = awsSecretKey
	}

	switch command {
	case "listClusters":
		ListClusters();
	case "listServices":
		clusterArn := _getClusterArn()
		ListServices(clusterArn)
	case "listTasks":
		clusterArn := _getClusterArn()
		serviceArn := _getServiceArn()
		ListTasks(clusterArn, serviceArn)
	case "describeService":
		clusterArn := _getClusterArn()
		serviceArn := _getServiceArn()
		DescribeService(clusterArn, serviceArn)
	case "updateService":
		clusterArn := _getClusterArn()
		serviceArn := _getServiceArn()
		UpdateService(clusterArn, serviceArn)
	case "releaseService":
		clusterArn := _getClusterArn()
		serviceArn := _getServiceArn()
		version := _getVersion()
		ReleaseService(clusterArn, serviceArn, version)
	case "listEc2Instances":
		ListEc2Instances()
	case "listLoadBalancers":
		ListLoadBalancers()
	case "ssh":
		if sshPem == "" {
			errUsage("A SSH PEM file must be specified")
		}
		if instanceId != "" {
			instance := GetInstanceForId(instanceId)
			Ssh(instance, sshPem, flag.Args())
		} else if instanceName != "" {
			instances := GetInstancesForName(instanceName)
			if (len(instances) == 1) {
				Ssh(instances[0], sshPem, flag.Args())
			} else {
				for i := 0; i < len(instances); i++ {
					fmt.Printf("[%s]\n", *instances[i].InstanceId)
					Ssh(instances[i], sshPem, flag.Args())
				}
			}
		} else {
			errUsage("Either instanceId or instanceName parameter has to be specified")
		}
	case "scp":
		if sshPem == "" {
			errUsage("A SSH PEM file must be specified")
		}
		if instanceId != "" {
			instance := GetInstanceForId(instanceId)
			Scp(instance, sshPem, flag.Args())
		} else if instanceName != "" {
			ips := GetInstancesForName(instanceName)
			if (len(ips) == 1) {
				fmt.Printf("[%s] ... ", *ips[0].InstanceId)
				Scp(ips[0], sshPem, flag.Args())
				fmt.Println("done")
			} else {
				for i := 0; i < len(ips); i++ {
					fmt.Printf("[%s] ... ", *ips[i].InstanceId)
					Scp(ips[i], sshPem, flag.Args())
					fmt.Println("done")
				}
			}
		} else {
			errUsage("Either instanceId or instanceName parameter has to be specified")
		}
	case "getEntity":
		if loadBalancer == "" {
			errUsage("loadBalancer must be specified")
		}
		if len(flag.Args()) != 1 {
			errUsage("Entity ID must be provided")
		}
		GetEntity(loadBalancer, flag.Args()[0])
	case "help":
		printCommandHelp()
	default:
		errUsage("Unknown command: " + command)
	}

}

<commit_msg>Added example for getEntity<commit_after>package main

import (
	"fmt"
	"flag"
	"os"
	"io/ioutil"
	"strings"
	"regexp"
	"sort"
)

var cluster, command, instanceId, instanceName, service, sshPem, output, credentialsFile, awsKey, awsSecretKey, version, loadBalancer string
var recursive, verbose, moreVerbose bool
var region = "eu-west-1"
var auth *Auth
var verboseLevel = 0

type Auth struct {
	key    string `toml:"aws_access_key_id"`
	secret string `toml:"aws_secret_access_key"`
}

func init() {
	flag.StringVar(&command, "command", "", "The command to use. The command 'help' displays available commands.")
	flag.StringVar(&cluster, "cluster", "", "Specify cluster to use")
	flag.StringVar(&instanceId, "instanceId", "", "Specify the EC2 instance")
	flag.StringVar(&instanceName, "instanceName", "", "Specify the EC2 instance(s) name")
	flag.StringVar(&service, "service", "", "Specify ECS service")
	flag.StringVar(&sshPem, "pemfile", "", "Specify PEM file for SSH access")
	flag.BoolVar(&recursive, "recursive", false, "Specify recursive operation")
	flag.StringVar(&output, "output", "", "Specify output directory")
	flag.StringVar(&credentialsFile, "credentials", "", "Specify credentials used for accessing AWS. Should be of format: .aws/credentials")
	flag.StringVar(&awsKey, "awsKey", "", "AWS key used for authentication. Overrides credentials file")
	flag.StringVar(&awsSecretKey, "awsSecretKey", "", "AWS secret key used for authentication, used in conjunction with 'awsKey'")
	flag.StringVar(&version, "version", "", "The version to use for docker image in the task definition")
	flag.StringVar(&loadBalancer, "loadBalancer", "", "Specifies the load balancer name to use")
	flag.BoolVar(&verbose, "v", false, "Making output more verbose, where applicable")
	flag.BoolVar(&moreVerbose, "vv", false, "Making output more verbose, where applicable")
}

func printCommandHelp() {
	var m = map[string]string{
		"help" : "Prints this help.",
		"listClusters" : "List available clusters.",
		"listServices" : "List available services. Needs -cluster flag.",
		"listTasks" : "List tasks for a service. Needs -cluster, -service flags.",
		"describeService" : "Describes the service. Needs -cluster, -service flags. Optionaly -v and -vv may be used.",
		"updateService" : "Stop/start all running tasks for the specified service. Needs -cluster, -service flags.",
		"releaseService" : "Creates a new release for the service. Neews -cluster, -service, -version flags.",
		"listEc2Instances" : "List available EC2 instances.",
		"listLoadBalancers" : "List available Load Balancers and their contained EC2 instances.",
		"getEntity" : "Gets an entity from the writer load balancer\n" +
			"                         -loadBalancer : The load balancer fronting the writer instances    (required)\n" +
			"                         {entityId}    : The ID of the entity to fetch    (required)\n" +
			"                         Example: -command getEntity -loadBalancer writer-loadbalancer a9fbd742-ea87-425d-ae86-045ab3ac91c1",
		"ssh" : "Executes a command over SSH for the specified service.\n" +
			"                         -instanceName : The aws instance(s) to use as source(s). Operation will occur on all instances with the specific name   (required if instanceId is not specified)\n" +
			"                         -instanceId   : The specific aws instance to use as source.   (required if instanceName is not specified)\n" +
			"                         -pemfile      : The SSH pem file used for authentication    (required)\n" +
			"                         {command}     : The command to execute (e.g. 'ls -l')   (required)\n" +
			"                         Example: -command ssh -instanceName writer -pemfile ~/.ssh/pem-files/im-dev tail -20 /var/log/writer/writer.log",
		"scp" : "Copies files from the specified instance(s). Needs -instanceName or -instanceId, -output and optionally -recursive flags.\n" +
			"                         -instanceName : The aws instance(s) to use as source(s). Operation will occur on all instances with the specific name   (required if instanceId is not specified)\n" +
			"                         -instanceId   : The specific aws instance to use as source.   (required if instanceName is not specified)\n" +
			"                         -pemfile      : The SSH pem file used for authentication    (required)\n" +
			"                         -output       : the target directory   (required)\n" +
			"                         -recursive    : copies from source recursively\n" +
			"                         Example: -command scp -instanceName writer -pemfile ~/.ssh/pem-files/im-dev -output Documents -recursive /var/log/writer",
	}

	k := sortKeys(m)

	for _, v := range k {
		fmt.Print(v);
		for j := 0; j < 20 - (len([]rune(v))); j++ {
			fmt.Print(" ")
		}
		fmt.Println(m[v])
	}

}

func sortKeys(m map[string]string) []string {
	keys := make([]string, 0, len(m))
	for key := range m {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	return keys
}

func errUsage(message string) {
	fmt.Println(message)
	os.Exit(1);
}

func errState(message string) {
	fmt.Println(message)
	os.Exit(2)
}

func errStatef(message string, a ...interface{}) {
	fmt.Printf(message, a)
	os.Exit(2)
}

func _getClusterArn() string {
	if (cluster == "") {
		errUsage("You must specify a cluster name with: -cluster");
	}

	arn := GetClusterArn(cluster)
	if (arn == "") {
		errUsage("Could not find cluster ARN for name: " + cluster);
	}

	return arn
}

func _getServiceArn() string {
	if (cluster == "") {
		errUsage("You must specify a cluster name with: -cluster");
	}

	if (service == "") {
		errUsage("You must specify a service name with: -service");
	}

	clusterArn := _getClusterArn();

	serviceArn := GetServiceArn(clusterArn, service)

	return serviceArn
}

func _getVersion() string {
	if (version == "") {
		errUsage("You must specify a version with: -version")
	}

	return version;
}

func getAwsCredentials(filepath string) (awsAccessKeyId, awsSecretKey string) {
	var awsAccessKeyIdRegexp = regexp.MustCompile("aws_access_key_id\\s*=\\s*(.*)")
	var awsSecretKeyRegexp = regexp.MustCompile("aws_secret_access_key\\s*=\\s*(.*)")

	file, err := ioutil.ReadFile(filepath)
	if err != nil {
		errUsage(err.Error())
	}

	matches := awsAccessKeyIdRegexp.FindStringSubmatch(string(file))
	if len(matches) > 1 {
		awsAccessKeyId = strings.TrimSpace(matches[1])
	} else {
		errUsage("Could not find aws access key id")
	}

	matches = awsSecretKeyRegexp.FindStringSubmatch(string(file))

	if len(matches) > 1 {
		awsSecretKey = strings.TrimSpace(matches[1])
	} else {
		errUsage("Could not find aws secret key")
	}

	return awsAccessKeyId, awsSecretKey
}

func main() {
	flag.Parse()

	if (verbose) {
		verboseLevel = 1
	}
	if (moreVerbose) {
		verboseLevel = 2
	}

	if (command == "") {
		flag.PrintDefaults();
		return
	}

	if (credentialsFile != "") {
		key, secret := getAwsCredentials(credentialsFile)
		auth = new(Auth)
		auth.key = key
		auth.secret = secret
	}

	if (awsKey != "" && awsSecretKey == "" ) {
		errUsage("Missing secretKey")
	}

	if (awsKey != "" && awsSecretKey != "" ) {
		auth = new(Auth)
		auth.key = awsKey
		auth.secret = awsSecretKey
	}

	switch command {
	case "listClusters":
		ListClusters();
	case "listServices":
		clusterArn := _getClusterArn()
		ListServices(clusterArn)
	case "listTasks":
		clusterArn := _getClusterArn()
		serviceArn := _getServiceArn()
		ListTasks(clusterArn, serviceArn)
	case "describeService":
		clusterArn := _getClusterArn()
		serviceArn := _getServiceArn()
		DescribeService(clusterArn, serviceArn)
	case "updateService":
		clusterArn := _getClusterArn()
		serviceArn := _getServiceArn()
		UpdateService(clusterArn, serviceArn)
	case "releaseService":
		clusterArn := _getClusterArn()
		serviceArn := _getServiceArn()
		version := _getVersion()
		ReleaseService(clusterArn, serviceArn, version)
	case "listEc2Instances":
		ListEc2Instances()
	case "listLoadBalancers":
		ListLoadBalancers()
	case "ssh":
		if sshPem == "" {
			errUsage("A SSH PEM file must be specified")
		}
		if instanceId != "" {
			instance := GetInstanceForId(instanceId)
			Ssh(instance, sshPem, flag.Args())
		} else if instanceName != "" {
			instances := GetInstancesForName(instanceName)
			if (len(instances) == 1) {
				Ssh(instances[0], sshPem, flag.Args())
			} else {
				for i := 0; i < len(instances); i++ {
					fmt.Printf("[%s]\n", *instances[i].InstanceId)
					Ssh(instances[i], sshPem, flag.Args())
				}
			}
		} else {
			errUsage("Either instanceId or instanceName parameter has to be specified")
		}
	case "scp":
		if sshPem == "" {
			errUsage("A SSH PEM file must be specified")
		}
		if instanceId != "" {
			instance := GetInstanceForId(instanceId)
			Scp(instance, sshPem, flag.Args())
		} else if instanceName != "" {
			ips := GetInstancesForName(instanceName)
			if (len(ips) == 1) {
				fmt.Printf("[%s] ... ", *ips[0].InstanceId)
				Scp(ips[0], sshPem, flag.Args())
				fmt.Println("done")
			} else {
				for i := 0; i < len(ips); i++ {
					fmt.Printf("[%s] ... ", *ips[i].InstanceId)
					Scp(ips[i], sshPem, flag.Args())
					fmt.Println("done")
				}
			}
		} else {
			errUsage("Either instanceId or instanceName parameter has to be specified")
		}
	case "getEntity":
		if loadBalancer == "" {
			errUsage("loadBalancer must be specified")
		}
		if len(flag.Args()) != 1 {
			errUsage("Entity ID must be provided")
		}
		GetEntity(loadBalancer, flag.Args()[0])
	case "help":
		printCommandHelp()
	default:
		errUsage("Unknown command: " + command)
	}

}

<|endoftext|>


<commit_before>package main

import "fmt"

var mapping = map[string]string{
	"ปา":   "pā",
	"ลิ":   "li",
	"ลี":   "lī",
	"ปาลิ": "pāli",
}

func main() {
	fmt.Println(mapping)
}
<commit_msg>update mapping<commit_after>package main

import "fmt"

var mapping = map[string]string{
	"ปา":   "pā",
	"ลิ":   "li",
	"ลี":   "lī",
	"ปาลิ": "pāli",
	"สาธุ": "sādhu",
}

func main() {
	fmt.Println(mapping)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/boltdb/bolt"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"
)

type RunningMachine struct {
	Id               uint64
	Name             string
	Path             string
	Input            string
	LastState        string
	NextState        string
	StatusMessage    string
	RunningStateCode bool
	NextStateRun     time.Time
}

type Scheduler struct {
	SchedulerLock   *sync.Mutex
	RunningMachines []RunningMachine
	Database        *bolt.DB
}

var globalScheduler Scheduler

func (s *Scheduler) GetRunningMachines() *[]RunningMachine {
	s.SchedulerLock.Lock()
	runningMachines := s.RunningMachines
	s.SchedulerLock.Unlock()

	return &runningMachines
}

func (s *Scheduler) GetMachineRun(id string) (*RunningMachine, error) {
	var machine RunningMachine

	returnErr := s.Database.View(func(tx *bolt.Tx) error {
		runsBucket := tx.Bucket([]byte("MachineRuns"))
		if runsBucket == nil {
			return fmt.Errorf("error getting database bucket")
		}

		machineJson := runsBucket.Get([]byte(id))
		if machineJson == nil {
			return fmt.Errorf("no running state machine with id %s found", id)
		}

		err := json.Unmarshal(machineJson, &machine)
		if err != nil {
			return fmt.Errorf("error deserializing machine run from persisted db: %s", err)
		}

		return nil
	})

	if returnErr != nil {
		return nil, returnErr
	} else {
		return &machine, nil
	}
}

func (s *Scheduler) UpdatePersistedMachine(machine *RunningMachine) (id uint64, returnErr error) {
	returnErr = s.Database.Update(func(tx *bolt.Tx) error {
		runningBucket := tx.Bucket([]byte("RunningMachines"))
		runsBucket := tx.Bucket([]byte("MachineRuns"))
		if runningBucket == nil || runsBucket == nil {
			return fmt.Errorf("error getting database bucket")
		}

		if machine.Id == 0 {
			var seqErr error
			id, seqErr = runsBucket.NextSequence()
			if seqErr != nil {
				return fmt.Errorf("error getting next value in MachineRuns sequence: %s", seqErr)
			}
			machine.Id = id
		} else {
			id = machine.Id
		}

		machineJson, err := json.Marshal(machine)
		if err != nil {
			return fmt.Errorf("error serializing machine as json for persisting: %s", err)
		}

		idKey := []byte(fmt.Sprintf("%d", id))
		err = runsBucket.Put(idKey, machineJson)
		if err != nil {
			return fmt.Errorf("error persisting machine run: %s", err)
		}

		err = runningBucket.Put(idKey, nil)
		if err != nil {
			return fmt.Errorf("error persisting machine run: %s", err)
		}

		return nil
	})

	return
}

func (s *Scheduler) ScheduleMachine(name string, path string, input string) (id uint64, returnErr error) {
	machine := RunningMachine{Id: 0, Name: name, Path: path, Input: input, NextState: "start", RunningStateCode: false, NextStateRun: time.Time{}}
	id, returnErr = s.UpdatePersistedMachine(&machine)
	if returnErr == nil {
		s.AddMachine(&machine)
	}

	return
}

func (s *Scheduler) AddMachine(machine *RunningMachine) {
	s.SchedulerLock.Lock()
	s.RunningMachines = append(s.RunningMachines, *machine)
	s.SchedulerLock.Unlock()
}

func (s *Scheduler) CancelMachineRun(id string) error {
	return s.Database.Update(func(tx *bolt.Tx) error {
		s.SchedulerLock.Lock()

		machineIdx := -1
		var machine RunningMachine
		var idx int
		for idx, machine = range s.RunningMachines {
			if fmt.Sprintf("%d", machine.Id) == id {
				machineIdx = idx
				break
			}
		}

		if machineIdx == -1 {
			s.SchedulerLock.Unlock()
			return fmt.Errorf("state machine run with id %s is not currently active", id)
		} else {
			// Delete from in-memory representation right away so that we can drop lock
			s.RunningMachines = append(s.RunningMachines[:machineIdx], s.RunningMachines[machineIdx+1:]...)
			s.SchedulerLock.Unlock()
		}

		runningBucket := tx.Bucket([]byte("RunningMachines"))
		runsBucket := tx.Bucket([]byte("MachineRuns"))
		if runningBucket == nil || runsBucket == nil {
			return fmt.Errorf("error getting database bucket")
		}

		err := runningBucket.Delete([]byte(id))
		if err != nil {
			return err
		}

		if machine.NextState != "stop" {
			machine.StatusMessage = "State machine run cancelled manually"
			machine.NextState = "stop"
			machine.RunningStateCode = false

			machineJson, jsonErr := json.Marshal(machine)
			if jsonErr != nil {
				return fmt.Errorf("error serializing machine as json for persisting: %s", err)
			}

			runsBucket.Put([]byte(id), machineJson)
		}

		return nil
	})
}

func (s *Scheduler) ExecuteState(machine *RunningMachine) {
	cmdPath := machine.Path + "/" + machine.NextState
	cmd := exec.Command(cmdPath)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	cmd.Stdin = strings.NewReader(machine.Input)
	err := cmd.Run()

	if err != nil {
		machine.StatusMessage = fmt.Sprintf("error executing state code at %s (will keep retrying): %s", cmdPath, err)
	} else {
		stderrStr := string(stderr.Bytes())
		stderrLines := strings.Split(stderrStr, "\n")
		if len(stderrLines) < 3 {
			machine.StatusMessage = fmt.Sprintf("state code at %s didn't return at least 3 lines correctly at stderr (will keep retrying), stderr was: %s",
				cmdPath, stderrStr)
		} else {
			machine.LastState = machine.NextState
			machine.NextState = strings.TrimSpace(stderrLines[0])

			numSeconds, intConvertError := strconv.Atoi(strings.TrimSpace(stderrLines[1]))
			if intConvertError != nil || numSeconds <= 0 {
				machine.NextStateRun = time.Time{}
			} else {
				machine.NextStateRun = time.Now().Add(time.Duration(numSeconds) * time.Second)
			}

			machine.Input = string(stdout.Bytes())
			machine.StatusMessage = strings.TrimSpace(stderrLines[2])
		}
	}

	machine.RunningStateCode = false

	s.UpdatePersistedMachine(machine)

	if machine.NextState == "stop" {
		s.CancelMachineRun(fmt.Sprintf("%d", machine.Id))
	}
}

func (s *Scheduler) HandleTick() {
	s.SchedulerLock.Lock()

	currentTime := time.Now()

	for idx, machine := range s.RunningMachines {
		if !machine.RunningStateCode && machine.NextState != "stop" && machine.NextStateRun.Before(currentTime) {
			var machinePtr *RunningMachine = &(s.RunningMachines[idx])
			machinePtr.RunningStateCode = true
			s.UpdatePersistedMachine(machinePtr)
			go s.ExecuteState(machinePtr)
		}
	}

	s.SchedulerLock.Unlock()
}

func (s *Scheduler) SchedulerTick(ticker *time.Ticker, quitChannel chan struct{}) {
	for {
		select {
		case <-ticker.C:
			s.HandleTick()
		case <-quitChannel:
			ticker.Stop()
			return
		}
	}
}

func (s *Scheduler) Init(db *bolt.DB) chan struct{} {
	s.SchedulerLock = &sync.Mutex{}
	s.Database = db
	s.RunningMachines = make([]RunningMachine, 0, 0)

	// Initialize database
	dbInitErr := db.Update(func(tx *bolt.Tx) error {
		runningBucket, err := tx.CreateBucketIfNotExists([]byte("RunningMachines"))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}

		var runsBucket *bolt.Bucket
		runsBucket, err = tx.CreateBucketIfNotExists([]byte("MachineRuns"))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}

		c := runningBucket.Cursor()

		for k, _ := c.First(); k != nil; k, _ = c.Next() {
			machineJson := runsBucket.Get(k)

			var machine RunningMachine
			err := json.Unmarshal(machineJson, &machine)
			if err != nil {
				return fmt.Errorf("error deserializing currently running machine from persisted db: %s", err)
			}

			s.AddMachine(&machine)
		}

		return nil
	})

	if dbInitErr != nil {
		fmt.Printf("error initializing database: %s\n", dbInitErr)
		os.Exit(1)
	}

	// Setup the timer that powers the scheduler
	ticker := time.NewTicker(1 * time.Second)
	stopSchedulerChannel := make(chan struct{})
	go s.SchedulerTick(ticker, stopSchedulerChannel)
	return stopSchedulerChannel
}
<commit_msg>Change so that in-memory representation of the running machines is an array of pointers<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/boltdb/bolt"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"
)

type RunningMachine struct {
	Id               uint64
	Name             string
	Path             string
	Input            string
	LastState        string
	NextState        string
	StatusMessage    string
	RunningStateCode bool
	NextStateRun     time.Time
}

type Scheduler struct {
	SchedulerLock   *sync.Mutex
	RunningMachines []*RunningMachine
	Database        *bolt.DB
}

var globalScheduler Scheduler

func (s *Scheduler) GetRunningMachines() *[]*RunningMachine {
	s.SchedulerLock.Lock()
	runningMachines := s.RunningMachines
	s.SchedulerLock.Unlock()

	return &runningMachines
}

func (s *Scheduler) GetMachineRun(id string) (*RunningMachine, error) {
	var machine RunningMachine

	returnErr := s.Database.View(func(tx *bolt.Tx) error {
		runsBucket := tx.Bucket([]byte("MachineRuns"))
		if runsBucket == nil {
			return fmt.Errorf("error getting database bucket")
		}

		machineJson := runsBucket.Get([]byte(id))
		if machineJson == nil {
			return fmt.Errorf("no running state machine with id %s found", id)
		}

		err := json.Unmarshal(machineJson, &machine)
		if err != nil {
			return fmt.Errorf("error deserializing machine run from persisted db: %s", err)
		}

		return nil
	})

	if returnErr != nil {
		return nil, returnErr
	} else {
		return &machine, nil
	}
}

func (s *Scheduler) UpdatePersistedMachine(machine *RunningMachine) (id uint64, returnErr error) {
	returnErr = s.Database.Update(func(tx *bolt.Tx) error {
		runningBucket := tx.Bucket([]byte("RunningMachines"))
		runsBucket := tx.Bucket([]byte("MachineRuns"))
		if runningBucket == nil || runsBucket == nil {
			return fmt.Errorf("error getting database bucket")
		}

		if machine.Id == 0 {
			var seqErr error
			id, seqErr = runsBucket.NextSequence()
			if seqErr != nil {
				return fmt.Errorf("error getting next value in MachineRuns sequence: %s", seqErr)
			}
			machine.Id = id
		} else {
			id = machine.Id
		}

		machineJson, err := json.Marshal(machine)
		if err != nil {
			return fmt.Errorf("error serializing machine as json for persisting: %s", err)
		}

		idKey := []byte(fmt.Sprintf("%d", id))
		err = runsBucket.Put(idKey, machineJson)
		if err != nil {
			return fmt.Errorf("error persisting machine run: %s", err)
		}

		err = runningBucket.Put(idKey, nil)
		if err != nil {
			return fmt.Errorf("error persisting machine run: %s", err)
		}

		return nil
	})

	return
}

func (s *Scheduler) ScheduleMachine(name string, path string, input string) (id uint64, returnErr error) {
	machine := RunningMachine{Id: 0, Name: name, Path: path, Input: input, NextState: "start", RunningStateCode: false, NextStateRun: time.Time{}}
	id, returnErr = s.UpdatePersistedMachine(&machine)
	if returnErr == nil {
		s.AddMachine(&machine)
	}

	return
}

func (s *Scheduler) AddMachine(machine *RunningMachine) {
	s.SchedulerLock.Lock()
	s.RunningMachines = append(s.RunningMachines, machine)
	s.SchedulerLock.Unlock()
}

func (s *Scheduler) CancelMachineRun(id string) error {
	return s.Database.Update(func(tx *bolt.Tx) error {
		s.SchedulerLock.Lock()

		machineIdx := -1
		var machine *RunningMachine
		var idx int
		for idx, machine = range s.RunningMachines {
			if fmt.Sprintf("%d", machine.Id) == id {
				machineIdx = idx
				break
			}
		}

		if machineIdx == -1 {
			s.SchedulerLock.Unlock()
			return fmt.Errorf("state machine run with id %s is not currently active", id)
		} else {
			// Delete from in-memory representation right away so that we can drop lock
			s.RunningMachines = append(s.RunningMachines[:machineIdx], s.RunningMachines[machineIdx+1:]...)
			s.SchedulerLock.Unlock()
		}

		runningBucket := tx.Bucket([]byte("RunningMachines"))
		runsBucket := tx.Bucket([]byte("MachineRuns"))
		if runningBucket == nil || runsBucket == nil {
			return fmt.Errorf("error getting database bucket")
		}

		err := runningBucket.Delete([]byte(id))
		if err != nil {
			return err
		}

		if machine.NextState != "stop" {
			machine.StatusMessage = "State machine run cancelled manually"
			machine.NextState = "stop"
			machine.RunningStateCode = false

			machineJson, jsonErr := json.Marshal(machine)
			if jsonErr != nil {
				return fmt.Errorf("error serializing machine as json for persisting: %s", err)
			}

			runsBucket.Put([]byte(id), machineJson)
		}

		return nil
	})
}

func (s *Scheduler) ExecuteState(machine *RunningMachine) {
	cmdPath := machine.Path + "/" + machine.NextState
	cmd := exec.Command(cmdPath)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	cmd.Stdin = strings.NewReader(machine.Input)
	err := cmd.Run()

	if err != nil {
		machine.StatusMessage = fmt.Sprintf("error executing state code at %s (will keep retrying): %s", cmdPath, err)
	} else {
		stderrStr := string(stderr.Bytes())
		stderrLines := strings.Split(stderrStr, "\n")
		if len(stderrLines) < 3 {
			machine.StatusMessage = fmt.Sprintf("state code at %s didn't return at least 3 lines correctly at stderr (will keep retrying), stderr was: %s",
				cmdPath, stderrStr)
		} else {
			machine.LastState = machine.NextState
			machine.NextState = strings.TrimSpace(stderrLines[0])

			numSeconds, intConvertError := strconv.Atoi(strings.TrimSpace(stderrLines[1]))
			if intConvertError != nil || numSeconds <= 0 {
				machine.NextStateRun = time.Time{}
			} else {
				machine.NextStateRun = time.Now().Add(time.Duration(numSeconds) * time.Second)
			}

			machine.Input = string(stdout.Bytes())
			machine.StatusMessage = strings.TrimSpace(stderrLines[2])
		}
	}

	machine.RunningStateCode = false

	s.UpdatePersistedMachine(machine)

	if machine.NextState == "stop" {
		s.CancelMachineRun(fmt.Sprintf("%d", machine.Id))
	}
}

func (s *Scheduler) HandleTick() {
	s.SchedulerLock.Lock()

	currentTime := time.Now()

	for idx, machine := range s.RunningMachines {
		if !machine.RunningStateCode && machine.NextState != "stop" && machine.NextStateRun.Before(currentTime) {
			var machinePtr *RunningMachine = s.RunningMachines[idx]
			machinePtr.RunningStateCode = true
			s.UpdatePersistedMachine(machinePtr)
			go s.ExecuteState(machinePtr)
		}
	}

	s.SchedulerLock.Unlock()
}

func (s *Scheduler) SchedulerTick(ticker *time.Ticker, quitChannel chan struct{}) {
	for {
		select {
		case <-ticker.C:
			s.HandleTick()
		case <-quitChannel:
			ticker.Stop()
			return
		}
	}
}

func (s *Scheduler) Init(db *bolt.DB) chan struct{} {
	s.SchedulerLock = &sync.Mutex{}
	s.Database = db
	s.RunningMachines = make([]*RunningMachine, 0, 0)

	// Initialize database
	dbInitErr := db.Update(func(tx *bolt.Tx) error {
		runningBucket, err := tx.CreateBucketIfNotExists([]byte("RunningMachines"))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}

		var runsBucket *bolt.Bucket
		runsBucket, err = tx.CreateBucketIfNotExists([]byte("MachineRuns"))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}

		c := runningBucket.Cursor()

		for k, _ := c.First(); k != nil; k, _ = c.Next() {
			machineJson := runsBucket.Get(k)

			var machine RunningMachine
			err := json.Unmarshal(machineJson, &machine)
			if err != nil {
				return fmt.Errorf("error deserializing currently running machine from persisted db: %s", err)
			}

			s.AddMachine(&machine)
		}

		return nil
	})

	if dbInitErr != nil {
		fmt.Printf("error initializing database: %s\n", dbInitErr)
		os.Exit(1)
	}

	// Setup the timer that powers the scheduler
	ticker := time.NewTicker(1 * time.Second)
	stopSchedulerChannel := make(chan struct{})
	go s.SchedulerTick(ticker, stopSchedulerChannel)
	return stopSchedulerChannel
}
<|endoftext|>


<commit_before>package json

import "github.com/gopherjs/gopherjs/js"

func Marshal(v interface{}) ([]byte, error) {
	return MarshalIndent(v, "", "")
}

func MarshalIndent(v interface{}, prefix, indent string) ([]byte, error) {
	s, err := MarshalIndentString(v, prefix, indent)
	return []byte(s), err
}

func MarshalString(v interface{}) (string, error) {
	return MarshalIndentString(v, "", "")
}

func MarshalIndentString(v interface{}, prefix, indent string) (json string, err error) {
	defer func() {
		if e := recover(); e != nil {
			if er, ok := e.(*js.Error); ok {
				err = er
			}
		}
	}()
	str := js.Global.Get("JSON").Call("stringify", js.InternalObject(jsonify).Invoke(js.InternalObject(v)), nil, indent)
	if len(prefix) > 0 {
		str = str.Call("replace", "\n", "\n"+prefix)
	}
	return str.String(), nil
}

const (
	boolKind       = 1
	intKind        = 2
	int8Kind       = 3
	int16Kind      = 4
	int32Kind      = 5
	int64Kind      = 6
	uintKind       = 7
	uint8Kind      = 8
	uint16Kind     = 9
	uint32Kind     = 10
	uint64Kind     = 11
	uintptrKind    = 12
	float32Kind    = 13
	float64Kind    = 14
	complex64Kind  = 15
	complex128Kind = 16
	arrayKind      = 17
	interfaceKind  = 20
	mapKind        = 21
	ptrKind        = 22
	sliceKind      = 23
	stringKind     = 24
	structKind     = 25
)

func unwrap(v, t *js.Object) (*js.Object, *js.Object) {
Loop:
	for {
		switch t.Get("kind").Int() {
		case ptrKind:
			if v.Get("$val") != t.Get("nil") {
				v = v.Call("$get")
				t = t.Get("elem")
			} else {
				break Loop
			}
		case interfaceKind:
			t = v.Get("constructor")
			v = v.Get("$val")
		default:
			break Loop
		}
	}
	return v, t
}

func wrap(v, t *js.Object) *js.Object {
	nt := js.Global.Get("Object").New()
	nt.Set("elem", t)
	nt.Set("kind", ptrKind)
	nv := js.Global.Get("Object").New()
	nv.Set("$val", v)
	nv.Set("constructor", nt)
	nv.Set("$get", js.MakeFunc(func(this *js.Object, _ []*js.Object) interface{} {
		return this.Get("$val")
	}))
	return nv
}

func jsonify(v *js.Object) *js.Object {
	t := v.Get("constructor")
	v, t = unwrap(v, t)
	switch t.Get("kind").Int() {
	case boolKind, intKind, int8Kind, int16Kind, int32Kind, uintKind, uint8Kind, uint16Kind, uint32Kind, uintptrKind, float32Kind, float64Kind, int64Kind, uint64Kind, complex64Kind, complex128Kind, ptrKind, stringKind:
		return js.Global.Call("$externalize", v, t)
	case arrayKind, sliceKind:
		l := v.Get("$length").Int()
		a := js.Global.Get("Array").New(l)
		for i := 0; i < l; i++ {
			a.SetIndex(i, js.InternalObject(jsonify).Invoke(v.Get("$array").Index(i)))
		}
		return a
	case mapKind:
		m := js.Global.Get("Object").New()
		keys := js.Global.Get("Object").Call("keys", v)
		len := keys.Length()
		for i := 0; i < len; i++ {
			val := v.Get(keys.Index(i).String())
			m.Set(val.Get("k").String(), js.InternalObject(jsonify).Invoke(val.Get("v")))
		}
		return m
	case structKind:
		s := js.Global.Get("Object").New()
		nextLevel := [][2]*js.Object{{v, t}}
		fieldsTodo := make(map[string]*js.Object)
		for len(nextLevel) > 0 {
			level := nextLevel
			//nextLevel = nextLevel[:0]
			nextLevel = make([][2]*js.Object, 0)
			for len(level) > 0 {
				v := level[0][0]
				t := level[0][1]
				level = level[1:]
				fields := t.Get("fields")
				fieldsLen := fields.Length()
				for i := 0; i < fieldsLen; i++ {
					f := fields.Index(i)
					if f.Get("pkg").Length() > 0 {
						continue
					}
					fName := f.Get("prop").String()
					name := f.Get("name").String()
					anon := false
					if name == "" {
						name = fName // filter $x?
						anon = true
					}
					n, o := parseTag(getJSONTag(f.Get("tag")))
					if n == "-" || (o.Contains("omitempty") && isEmpty(v.Get(fName), f.Get("typ"))) {
						continue
					}
					jName := name
					if n != "" {
						jName = n
					}
					if anon {
						nv, nt := unwrap(v.Get(fName), f.Get("typ"))
						if nv != nt.Get("nil") {
							nextLevel = append(nextLevel, [2]*js.Object{nv, nt})
						}
					} else if _, ok := fieldsTodo[jName]; ok {
						fieldsTodo[jName] = nil
					} else {
						val := v.Get(fName)
						vt := f.Get("typ")
						val, vt = unwrap(val, vt)
						if o.Contains("string") && stringable(vt) {
							s := js.Global.Get("Object").New()
							s.Set("kind", stringKind)
							if vt.Get("kind").Int() == stringKind {
								val = wrap(js.Global.Get("JSON").Call("stringify", val), s)
							} else {
								val = wrap(val.Call("toString"), s)
							}
						} else if val.Get("constructor").Get("kind") == js.Undefined {
							val = wrap(val, vt)
						}
						fieldsTodo[jName] = val
					}
				}
			}
			for name, f := range fieldsTodo {
				if f != nil {
					if isMarshaler(f) {
						tup := f.Call("MarshalJSON")
						if tup.Index(1) != js.Global.Get("$ifaceNil") {
							panic(tup.Index(1))
						}
						s.Set(name, js.Global.Get("JSON").Call("parse", js.Global.Call("$bytesToString", tup.Index(0)).String()))
					} else {
						s.Set(name, js.InternalObject(jsonify).Invoke(f))
					}
					fieldsTodo[name] = nil
				}
			}
		}
		return s
	}
	return js.Undefined
}

func isEmpty(v, t *js.Object) bool {
	switch t.Get("kind").Int() {
	case boolKind:
		return !v.Bool()
	case int8Kind, int16Kind, int32Kind, intKind, uint8Kind, uint16Kind, uint32Kind, uintKind:
		return v.Int() == 0
	case int64Kind, uint64Kind:
		return v.Int64() == 0
	case float32Kind, float64Kind:
		return v.Float() == 0
	case arrayKind, sliceKind, stringKind:
		return v.Length() == 0
	case ptrKind:
		return v == t.Get("nil")
	case interfaceKind:
		if v.Get("$val") == js.Undefined {
			return true
		}
		return isEmpty(v.Get("$val"), v.Get("constructor"))
	case mapKind:
		return js.Global.Get("Object").Call("keys", v).Length() == 0
	}
	return false
}

func stringable(t *js.Object) bool {
	switch t.Get("kind").Int() {
	case boolKind, int8Kind, int16Kind, int32Kind, intKind, uint8Kind, uint16Kind, uint32Kind, uintKind, int64Kind, uint64Kind, float32Kind, float64Kind, stringKind:
		return true
	}
	return false
}

func isMarshaler(v *js.Object) bool {
	_, t := unwrap(v, v.Get("constructor"))
	methods := t.Get("methods")
	if methods == js.Undefined {
		return false
	}
	methodNum := methods.Length()
	for i := 0; i < methodNum; i++ {
		method := methods.Index(i)
		if method.Get("name").String() == "MarshalJSON" {
			mt := method.Get("typ")
			return mt.Get("params").Length() == 0 && mt.Get("results").Length() == 2 && mt.Get("results").Index(0).Get("string").String() == "[]uint8" && mt.Get("results").Index(1).Get("string").String() == "error"
		}
	}
	return false
}
<commit_msg>renamed jsonify -> toObject<commit_after>package json

import "github.com/gopherjs/gopherjs/js"

func Marshal(v interface{}) ([]byte, error) {
	return MarshalIndent(v, "", "")
}

func MarshalIndent(v interface{}, prefix, indent string) ([]byte, error) {
	s, err := MarshalIndentString(v, prefix, indent)
	return []byte(s), err
}

func MarshalString(v interface{}) (string, error) {
	return MarshalIndentString(v, "", "")
}

func MarshalIndentString(v interface{}, prefix, indent string) (json string, err error) {
	defer func() {
		if e := recover(); e != nil {
			if er, ok := e.(*js.Error); ok {
				err = er
			}
		}
	}()
	str := js.Global.Get("JSON").Call("stringify", js.InternalObject(toObject).Invoke(js.InternalObject(v)), nil, indent)
	if len(prefix) > 0 {
		str = str.Call("replace", "\n", "\n"+prefix)
	}
	return str.String(), nil
}

const (
	boolKind       = 1
	intKind        = 2
	int8Kind       = 3
	int16Kind      = 4
	int32Kind      = 5
	int64Kind      = 6
	uintKind       = 7
	uint8Kind      = 8
	uint16Kind     = 9
	uint32Kind     = 10
	uint64Kind     = 11
	uintptrKind    = 12
	float32Kind    = 13
	float64Kind    = 14
	complex64Kind  = 15
	complex128Kind = 16
	arrayKind      = 17
	interfaceKind  = 20
	mapKind        = 21
	ptrKind        = 22
	sliceKind      = 23
	stringKind     = 24
	structKind     = 25
)

func unwrap(v, t *js.Object) (*js.Object, *js.Object) {
Loop:
	for {
		switch t.Get("kind").Int() {
		case ptrKind:
			if v.Get("$val") != t.Get("nil") {
				v = v.Call("$get")
				t = t.Get("elem")
			} else {
				break Loop
			}
		case interfaceKind:
			t = v.Get("constructor")
			v = v.Get("$val")
		default:
			break Loop
		}
	}
	return v, t
}

func wrap(v, t *js.Object) *js.Object {
	nt := js.Global.Get("Object").New()
	nt.Set("elem", t)
	nt.Set("kind", ptrKind)
	nv := js.Global.Get("Object").New()
	nv.Set("$val", v)
	nv.Set("constructor", nt)
	nv.Set("$get", js.MakeFunc(func(this *js.Object, _ []*js.Object) interface{} {
		return this.Get("$val")
	}))
	return nv
}

func toObject(v *js.Object) *js.Object {
	t := v.Get("constructor")
	v, t = unwrap(v, t)
	switch t.Get("kind").Int() {
	case boolKind, intKind, int8Kind, int16Kind, int32Kind, uintKind, uint8Kind, uint16Kind, uint32Kind, uintptrKind, float32Kind, float64Kind, int64Kind, uint64Kind, complex64Kind, complex128Kind, ptrKind, stringKind:
		return js.Global.Call("$externalize", v, t)
	case arrayKind, sliceKind:
		l := v.Get("$length").Int()
		a := js.Global.Get("Array").New(l)
		for i := 0; i < l; i++ {
			a.SetIndex(i, js.InternalObject(toObject).Invoke(v.Get("$array").Index(i)))
		}
		return a
	case mapKind:
		m := js.Global.Get("Object").New()
		keys := js.Global.Get("Object").Call("keys", v)
		len := keys.Length()
		for i := 0; i < len; i++ {
			val := v.Get(keys.Index(i).String())
			m.Set(val.Get("k").String(), js.InternalObject(toObject).Invoke(val.Get("v")))
		}
		return m
	case structKind:
		s := js.Global.Get("Object").New()
		nextLevel := [][2]*js.Object{{v, t}}
		fieldsTodo := make(map[string]*js.Object)
		for len(nextLevel) > 0 {
			level := nextLevel
			//nextLevel = nextLevel[:0]
			nextLevel = make([][2]*js.Object, 0)
			for len(level) > 0 {
				v := level[0][0]
				t := level[0][1]
				level = level[1:]
				fields := t.Get("fields")
				fieldsLen := fields.Length()
				for i := 0; i < fieldsLen; i++ {
					f := fields.Index(i)
					if f.Get("pkg").Length() > 0 {
						continue
					}
					fName := f.Get("prop").String()
					name := f.Get("name").String()
					anon := false
					if name == "" {
						name = fName // filter $x?
						anon = true
					}
					n, o := parseTag(getJSONTag(f.Get("tag")))
					if n == "-" || (o.Contains("omitempty") && isEmpty(v.Get(fName), f.Get("typ"))) {
						continue
					}
					jName := name
					if n != "" {
						jName = n
					}
					if anon {
						nv, nt := unwrap(v.Get(fName), f.Get("typ"))
						if nv != nt.Get("nil") {
							nextLevel = append(nextLevel, [2]*js.Object{nv, nt})
						}
					} else if _, ok := fieldsTodo[jName]; ok {
						fieldsTodo[jName] = nil
					} else {
						val := v.Get(fName)
						vt := f.Get("typ")
						val, vt = unwrap(val, vt)
						if o.Contains("string") && stringable(vt) {
							s := js.Global.Get("Object").New()
							s.Set("kind", stringKind)
							if vt.Get("kind").Int() == stringKind {
								val = wrap(js.Global.Get("JSON").Call("stringify", val), s)
							} else {
								val = wrap(val.Call("toString"), s)
							}
						} else if val.Get("constructor").Get("kind") == js.Undefined {
							val = wrap(val, vt)
						}
						fieldsTodo[jName] = val
					}
				}
			}
			for name, f := range fieldsTodo {
				if f != nil {
					if isMarshaler(f) {
						tup := f.Call("MarshalJSON")
						if tup.Index(1) != js.Global.Get("$ifaceNil") {
							panic(tup.Index(1))
						}
						s.Set(name, js.Global.Get("JSON").Call("parse", js.Global.Call("$bytesToString", tup.Index(0)).String()))
					} else {
						s.Set(name, js.InternalObject(toObject).Invoke(f))
					}
					fieldsTodo[name] = nil
				}
			}
		}
		return s
	}
	return js.Undefined
}

func isEmpty(v, t *js.Object) bool {
	switch t.Get("kind").Int() {
	case boolKind:
		return !v.Bool()
	case int8Kind, int16Kind, int32Kind, intKind, uint8Kind, uint16Kind, uint32Kind, uintKind:
		return v.Int() == 0
	case int64Kind, uint64Kind:
		return v.Int64() == 0
	case float32Kind, float64Kind:
		return v.Float() == 0
	case arrayKind, sliceKind, stringKind:
		return v.Length() == 0
	case ptrKind:
		return v == t.Get("nil")
	case interfaceKind:
		if v.Get("$val") == js.Undefined {
			return true
		}
		return isEmpty(v.Get("$val"), v.Get("constructor"))
	case mapKind:
		return js.Global.Get("Object").Call("keys", v).Length() == 0
	}
	return false
}

func stringable(t *js.Object) bool {
	switch t.Get("kind").Int() {
	case boolKind, int8Kind, int16Kind, int32Kind, intKind, uint8Kind, uint16Kind, uint32Kind, uintKind, int64Kind, uint64Kind, float32Kind, float64Kind, stringKind:
		return true
	}
	return false
}

func isMarshaler(v *js.Object) bool {
	_, t := unwrap(v, v.Get("constructor"))
	methods := t.Get("methods")
	if methods == js.Undefined {
		return false
	}
	methodNum := methods.Length()
	for i := 0; i < methodNum; i++ {
		method := methods.Index(i)
		if method.Get("name").String() == "MarshalJSON" {
			mt := method.Get("typ")
			return mt.Get("params").Length() == 0 && mt.Get("results").Length() == 2 && mt.Get("results").Index(0).Get("string").String() == "[]uint8" && mt.Get("results").Index(1).Get("string").String() == "error"
		}
	}
	return false
}
<|endoftext|>


<commit_before>package resourceviewer

import (
	"fmt"
	"strings"

	"github.com/heptio/developer-dash/internal/overview/link"

	"github.com/heptio/developer-dash/internal/log"
	"github.com/heptio/developer-dash/internal/overview/objectstatus"
	"github.com/heptio/developer-dash/internal/overview/objectvisitor"
	"github.com/heptio/developer-dash/internal/view/component"
	"github.com/pkg/errors"
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes/scheme"
)

// Collector collects objects to construct a resource viewer.
type Collector struct {
	edges  map[string][]string
	nodes  map[string]component.Node
	logger log.Logger

	// podStats counts pods in a replica set.
	podStats map[string]int
}

var _ objectvisitor.ObjectHandler = (*Collector)(nil)

// NewCollector creates an instance of Collector.
func NewCollector() *Collector {
	c := &Collector{
		podStats: make(map[string]int),
	}
	c.Reset()

	return c
}

// Reset resets a Collector's nodes and edges.
func (c *Collector) Reset() {
	c.edges = make(map[string][]string)
	c.nodes = make(map[string]component.Node)
}

// Process process an object by saving the object to a map.
func (c *Collector) Process(object objectvisitor.ClusterObject) error {
	var uid string
	var node component.Node
	var err error

	if c.isPod(object) {
		pod := &corev1.Pod{}
		if err := scheme.Scheme.Convert(object, pod, 0); err != nil {
			return errors.Wrap(err, "unable to convert object to pod")
		}

		if ownerReference := metav1.GetControllerOf(pod); ownerReference != nil {
			c.podStats[string(ownerReference.UID)]++
		}

		uid, node, err = c.createPodGroupNode(object)
	} else {
		uid, node, err = c.createObjectNode(object)
	}

	if err != nil {
		if isSkippedNode(err) {
			return nil
		}

		gvk := object.GetObjectKind().GroupVersionKind()
		accessor := meta.NewAccessor()
		name, err := accessor.Name(object)
		if err == nil {
			return errors.Wrapf(err, "processing unknown %s", gvk.String())
		}

		return errors.Wrapf(err, "processing %s %s", gvk.String(), name)
	}

	if _, ok := c.nodes[uid]; !ok {
		c.nodes[uid] = node
	}

	return nil
}

func (c *Collector) createPodGroupNode(object objectvisitor.ClusterObject) (string, component.Node, error) {
	pgd, err := c.podGroupDetails(object)
	if err != nil {
		return "", component.Node{}, errors.Wrap(err, "getting pod group id for pod")
	}

	objectKind := object.GetObjectKind()
	apiVersion, kind := objectKind.GroupVersionKind().ToAPIVersionAndKind()

	node := component.Node{
		Name:       pgd.name,
		APIVersion: apiVersion,
		Kind:       kind,
		// set all statuses to ok until we start checking for the real status.
		Status: component.NodeStatusOK,
	}

	return pgd.id, node, nil
}

type isSkipped interface {
	IsSkipped() bool
}

func isSkippedNode(err error) bool {
	sn, ok := err.(isSkipped)
	return ok && sn.IsSkipped()
}

type skipNode struct{}

func (e skipNode) IsSkipped() bool {
	return true
}

func (e skipNode) Error() string {
	return "skip node"
}

func (c *Collector) createObjectNode(object objectvisitor.ClusterObject) (string, component.Node, error) {
	objectKind := object.GetObjectKind()
	gvk := objectKind.GroupVersionKind()
	apiVersion, kind := gvk.ToAPIVersionAndKind()

	accessor := meta.NewAccessor()

	if (gvk.Group == "apps" || gvk.Group == "extensions") &&
		gvk.Kind == "ReplicaSet" {
		replicaSet := &appsv1.ReplicaSet{}
		if err := scheme.Scheme.Convert(object, replicaSet, nil); err != nil {
			return "", component.Node{}, errors.Wrap(err, "convert object to Replica Set")
		}

		replicas := replicaSet.Spec.Replicas
		if replicas == nil || *replicas < 1 {
			return "", component.Node{}, &skipNode{}
		}
	}

	uid, err := accessor.UID(object)
	if err != nil {
		return "", component.Node{}, err
	}

	name, err := accessor.Name(object)
	if err != nil {
		return "", component.Node{}, errors.New("unable to get object name")
	}

	objectStatus, err := objectstatus.Status(object)
	if err != nil {
		c.log().Errorf("error retrieving object status: %v", err)
		objectStatus.NodeStatus = component.NodeStatusOK
	}

	node := component.Node{
		Name:       name,
		APIVersion: apiVersion,
		Kind:       kind,
		Status:     objectStatus.NodeStatus,
		Details:    objectStatus.Details,
		Path:       link.ForObject(object, name),
	}

	return string(uid), node, nil
}

// AddChild adds children for an object to create edges. Pods are collated to a single object.
func (c *Collector) AddChild(parent objectvisitor.ClusterObject, children ...objectvisitor.ClusterObject) error {
	accessor := meta.NewAccessor()
	pid, err := accessor.UID(parent)
	if err != nil {
		return err
	}

	for _, child := range children {
		var cid string

		if c.isPod(child) {
			pgd, err := c.podGroupDetails(child)
			if err != nil {
				return errors.Wrap(err, "find pod group id for pod")
			}

			cid = pgd.id
		} else {
			id, err := accessor.UID(child)
			if err != nil {
				return err
			}

			cid = string(id)
		}

		if !listContains(c.edges[string(pid)], cid) {
			c.edges[string(pid)] = append(c.edges[string(pid)], cid)
		}
	}

	return nil
}

func (c *Collector) isPod(object objectvisitor.ClusterObject) bool {
	objectKind := object.GetObjectKind()
	gvk := objectKind.GroupVersionKind()

	return gvk.Group == "" &&
		gvk.Version == "v1" &&
		gvk.Kind == "Pod"
}

type podGroupDetails struct {
	id   string
	name string
}

func (c *Collector) podGroupDetails(object objectvisitor.ClusterObject) (podGroupDetails, error) {
	obj, err := meta.Accessor(object)
	if err != nil {
		return podGroupDetails{}, err
	}

	reference := metav1.GetControllerOf(obj)
	id := fmt.Sprintf("pods-%s", reference.UID)

	pgd := podGroupDetails{
		id:   id,
		name: fmt.Sprintf("%s pods", reference.Name),
	}

	return pgd, nil
}

func (c *Collector) ViewComponent() (component.ViewComponent, error) {
	rv := component.NewResourceViewer("Resource Viewer")

	var nodeIDs []string
	for nodeID, node := range c.nodes {
		if strings.HasPrefix(nodeID, "pods-") {
			ownerID := strings.TrimPrefix(nodeID, "pods-")
			node.Details = append(node.Details,
				component.NewText(fmt.Sprintf("Pod count: %d", c.podStats[ownerID])))
			c.nodes[nodeID] = node
		}

		rv.AddNode(nodeID, node)
		nodeIDs = append(nodeIDs, nodeID)
	}

	for nodeID, edges := range c.edges {
		for _, edgeID := range edges {
			if listContains(nodeIDs, edgeID) {
				rv.AddEdge(nodeID, edgeID, component.EdgeTypeExplicit)
			}
		}
	}

	return rv, nil
}

func (c *Collector) log() log.Logger {
	if c.logger != nil {
		return c.logger
	}

	return log.NopLogger()
}

func listContains(lst []string, s string) bool {
	for _, item := range lst {
		if item == s {
			return true
		}
	}
	return false
}
<commit_msg>pods don't have to belong to anything<commit_after>package resourceviewer

import (
	"fmt"
	"strings"

	"github.com/heptio/developer-dash/internal/overview/link"

	"github.com/heptio/developer-dash/internal/log"
	"github.com/heptio/developer-dash/internal/overview/objectstatus"
	"github.com/heptio/developer-dash/internal/overview/objectvisitor"
	"github.com/heptio/developer-dash/internal/view/component"
	"github.com/pkg/errors"
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes/scheme"
)

// Collector collects objects to construct a resource viewer.
type Collector struct {
	edges  map[string][]string
	nodes  map[string]component.Node
	logger log.Logger

	// podStats counts pods in a replica set.
	podStats map[string]int
}

var _ objectvisitor.ObjectHandler = (*Collector)(nil)

// NewCollector creates an instance of Collector.
func NewCollector() *Collector {
	c := &Collector{
		podStats: make(map[string]int),
	}
	c.Reset()

	return c
}

// Reset resets a Collector's nodes and edges.
func (c *Collector) Reset() {
	c.edges = make(map[string][]string)
	c.nodes = make(map[string]component.Node)
}

// Process process an object by saving the object to a map.
func (c *Collector) Process(object objectvisitor.ClusterObject) error {
	var uid string
	var node component.Node
	var err error

	if c.isPod(object) {
		pod := &corev1.Pod{}
		if err := scheme.Scheme.Convert(object, pod, 0); err != nil {
			return errors.Wrap(err, "unable to convert object to pod")
		}

		if ownerReference := metav1.GetControllerOf(pod); ownerReference != nil {
			c.podStats[string(ownerReference.UID)]++
		}

		uid, node, err = c.createPodGroupNode(object)
	} else {
		uid, node, err = c.createObjectNode(object)
	}

	if err != nil {
		if isSkippedNode(err) {
			return nil
		}

		gvk := object.GetObjectKind().GroupVersionKind()
		accessor := meta.NewAccessor()
		name, err := accessor.Name(object)
		if err == nil {
			return errors.Wrapf(err, "processing unknown %s", gvk.String())
		}

		return errors.Wrapf(err, "processing %s %s", gvk.String(), name)
	}

	if _, ok := c.nodes[uid]; !ok {
		c.nodes[uid] = node
	}

	return nil
}

func (c *Collector) createPodGroupNode(object objectvisitor.ClusterObject) (string, component.Node, error) {
	pgd, err := c.podGroupDetails(object)
	if err != nil {
		return "", component.Node{}, errors.Wrap(err, "getting pod group id for pod")
	}

	objectKind := object.GetObjectKind()
	apiVersion, kind := objectKind.GroupVersionKind().ToAPIVersionAndKind()

	node := component.Node{
		Name:       pgd.name,
		APIVersion: apiVersion,
		Kind:       kind,
		// set all statuses to ok until we start checking for the real status.
		Status: component.NodeStatusOK,
	}

	return pgd.id, node, nil
}

type isSkipped interface {
	IsSkipped() bool
}

func isSkippedNode(err error) bool {
	sn, ok := err.(isSkipped)
	return ok && sn.IsSkipped()
}

type skipNode struct{}

func (e skipNode) IsSkipped() bool {
	return true
}

func (e skipNode) Error() string {
	return "skip node"
}

func (c *Collector) createObjectNode(object objectvisitor.ClusterObject) (string, component.Node, error) {
	objectKind := object.GetObjectKind()
	gvk := objectKind.GroupVersionKind()
	apiVersion, kind := gvk.ToAPIVersionAndKind()

	accessor := meta.NewAccessor()

	if (gvk.Group == "apps" || gvk.Group == "extensions") &&
		gvk.Kind == "ReplicaSet" {
		replicaSet := &appsv1.ReplicaSet{}
		if err := scheme.Scheme.Convert(object, replicaSet, nil); err != nil {
			return "", component.Node{}, errors.Wrap(err, "convert object to Replica Set")
		}

		replicas := replicaSet.Spec.Replicas
		if replicas == nil || *replicas < 1 {
			return "", component.Node{}, &skipNode{}
		}
	}

	uid, err := accessor.UID(object)
	if err != nil {
		return "", component.Node{}, err
	}

	name, err := accessor.Name(object)
	if err != nil {
		return "", component.Node{}, errors.New("unable to get object name")
	}

	objectStatus, err := objectstatus.Status(object)
	if err != nil {
		c.log().Errorf("error retrieving object status: %v", err)
		objectStatus.NodeStatus = component.NodeStatusOK
	}

	node := component.Node{
		Name:       name,
		APIVersion: apiVersion,
		Kind:       kind,
		Status:     objectStatus.NodeStatus,
		Details:    objectStatus.Details,
		Path:       link.ForObject(object, name),
	}

	return string(uid), node, nil
}

// AddChild adds children for an object to create edges. Pods are collated to a single object.
func (c *Collector) AddChild(parent objectvisitor.ClusterObject, children ...objectvisitor.ClusterObject) error {
	accessor := meta.NewAccessor()
	pid, err := accessor.UID(parent)
	if err != nil {
		return err
	}

	for _, child := range children {
		var cid string

		if c.isPod(child) {
			pgd, err := c.podGroupDetails(child)
			if err != nil {
				return errors.Wrap(err, "find pod group id for pod")
			}

			cid = pgd.id
		} else {
			id, err := accessor.UID(child)
			if err != nil {
				return err
			}

			cid = string(id)
		}

		if !listContains(c.edges[string(pid)], cid) {
			c.edges[string(pid)] = append(c.edges[string(pid)], cid)
		}
	}

	return nil
}

func (c *Collector) isPod(object objectvisitor.ClusterObject) bool {
	objectKind := object.GetObjectKind()
	gvk := objectKind.GroupVersionKind()

	return gvk.Group == "" &&
		gvk.Version == "v1" &&
		gvk.Kind == "Pod"
}

type podGroupDetails struct {
	id   string
	name string
}

func (c *Collector) podGroupDetails(object objectvisitor.ClusterObject) (podGroupDetails, error) {
	obj, err := meta.Accessor(object)
	if err != nil {
		return podGroupDetails{}, err
	}

	reference := metav1.GetControllerOf(obj)
	if reference == nil {

		return podGroupDetails{
			id:   string(obj.GetUID()),
			name: obj.GetName(),
		}, nil
	}

	id := fmt.Sprintf("pods-%s", reference.UID)

	pgd := podGroupDetails{
		id:   id,
		name: fmt.Sprintf("%s pods", reference.Name),
	}

	return pgd, nil
}

func (c *Collector) ViewComponent() (component.ViewComponent, error) {
	rv := component.NewResourceViewer("Resource Viewer")

	var nodeIDs []string
	for nodeID, node := range c.nodes {
		if strings.HasPrefix(nodeID, "pods-") {
			ownerID := strings.TrimPrefix(nodeID, "pods-")
			node.Details = append(node.Details,
				component.NewText(fmt.Sprintf("Pod count: %d", c.podStats[ownerID])))
			c.nodes[nodeID] = node
		}

		rv.AddNode(nodeID, node)
		nodeIDs = append(nodeIDs, nodeID)
	}

	for nodeID, edges := range c.edges {
		for _, edgeID := range edges {
			if listContains(nodeIDs, edgeID) {
				rv.AddEdge(nodeID, edgeID, component.EdgeTypeExplicit)
			}
		}
	}

	return rv, nil
}

func (c *Collector) log() log.Logger {
	if c.logger != nil {
		return c.logger
	}

	return log.NopLogger()
}

func listContains(lst []string, s string) bool {
	for _, item := range lst {
		if item == s {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package generate

import (
	"bytes"
	"fmt"
	"go/format"
	"io/ioutil"
	"os"
	"path/filepath"
	"text/template"

	_ "github.com/wlMalk/gorator/driver/pgsql"
	"github.com/wlMalk/gorator/internal/util"

	"github.com/wlMalk/gorator/parser"
)

const (
	VERSION = "0.1"
)

var packageDescriptions = map[string]string{
	"database": `//
//
//`,
}

func getPackage(name string, config *parser.Config) interface{} {

	return struct {
		Name           string
		GoratorVersion string
		ConfigVersion  string
		Description    string
		Imports        []map[string]string
	}{
		Name:           name,
		GoratorVersion: config.GoratorVersion,
		ConfigVersion:  config.Version,
		Description:    packageDescriptions[name],
		Imports:        config.Imports[name],
	}

}

var tmpls *template.Template = template.New("")

var ormTmplsMap map[string]string = map[string]string{
	"database": "database",
	"orm":      "database/orm",
	"callback": "database/orm/internal/callback",
	"query":    "database/orm/query",
	"model":    "database/orm/model",
}
var ormDirs []string = []string{
	"database", "database/orm", "database/orm/internal/callback",
	"database/orm/query", "database/orm/model",
}

func init() {
	tmplsDir := util.GetFullPath(os.Getenv("GOPATH"), "src/github.com/wlMalk/gorator/templates/")
	tmplFiles, err := filepath.Glob(tmplsDir + "*/*.tmpl")
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(0)
	}
	tmpls, err = tmpls.Funcs(util.GetFuncsMap()).ParseFiles(tmplFiles...)
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(0)
	}
}

func Generate(path string, version string) error {
	path = util.GetPath(path)

	if version != "" {
		version = "." + version
	}

	configFiles, err := filepath.Glob(path + "/*.yml" + version)
	if err != nil {
		return fmt.Errorf("could not open config file")
	}

	if len(configFiles) == 0 {
		return fmt.Errorf("could not find any config file")
	}

	importPath := util.GetImportPath(path)
	var files [][]byte

	for _, f := range configFiles {
		file, err := ioutil.ReadFile(f)
		if err != nil {
			return fmt.Errorf("could not open config file")
		}
		files = append(files, file)
	}

	config, err := parser.Parse(importPath, files...)
	if err != nil {
		return err
	}

	return GenerateFrom(path, config)
}

func GenerateFrom(path string, config *parser.Config) error {
	// a, _ := json.Marshal(config)
	// fmt.Println(string(a))
	return generateORM(path, config)
}

func generateORM(path string, config *parser.Config) error {
	for _, d := range ormDirs {
		util.Mkdir(util.GetFullPath(path, d))
	}

	var w bytes.Buffer
	for t, dir := range ormTmplsMap {
		for _, db := range config.Databases {
			err := tmpls.ExecuteTemplate(&w, "heading", getPackage(t, config))
			if err != nil {
				return err
			}

			if db.Driver.Generate(t) {
				err = tmpls.ExecuteTemplate(&w, t, config)
				if err != nil {
					return err
				}
			}
			// driver headings
			err = db.Driver.Execute(&w, t, config)
			if err != nil {
				return err
			}
		}

		// fix white spaces

		b, err := format.Source(w.Bytes())
		if err != nil {
			return err
		}

		err = util.SaveFile(util.GetFullPath(path, dir+"/"+t+"_gen.go"), b)
		if err != nil {
			return err
		}

		w.Reset()
	}
	return nil
}
<commit_msg>Show feedback while generating<commit_after>package generate

import (
	"bytes"
	"fmt"
	"go/format"
	"io/ioutil"
	"os"
	"path/filepath"
	"text/template"

	_ "github.com/wlMalk/gorator/driver/pgsql"
	"github.com/wlMalk/gorator/internal/util"

	"github.com/wlMalk/gorator/parser"

	"github.com/fatih/color"
)

var magenta = color.New(color.FgMagenta, color.Underline).SprintFunc()
var okMsg = color.New(color.FgGreen, color.Bold).SprintFunc()("OK!")

const (
	VERSION = "0.1"
)

var packageDescriptions = map[string]string{
	"database": `//
//
//`,
}

func getPackage(name string, config *parser.Config) interface{} {

	return struct {
		Name           string
		GoratorVersion string
		ConfigVersion  string
		Description    string
		Imports        []map[string]string
	}{
		Name:           name,
		GoratorVersion: config.GoratorVersion,
		ConfigVersion:  config.Version,
		Description:    packageDescriptions[name],
		Imports:        config.Imports[name],
	}

}

var tmpls *template.Template = template.New("")

var ormTmplsMap map[string]string = map[string]string{
	"database": "database",
	"orm":      "database/orm",
	"callback": "database/orm/internal/callback",
	"query":    "database/orm/query",
	"model":    "database/orm/model",
}
var ormDirs []string = []string{
	"database", "database/orm", "database/orm/internal/callback",
	"database/orm/query", "database/orm/model",
}

func init() {
	tmplsDir := util.GetFullPath(os.Getenv("GOPATH"), "src/github.com/wlMalk/gorator/templates/")
	tmplFiles, err := filepath.Glob(tmplsDir + "*/*.tmpl")
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(0)
	}
	tmpls, err = tmpls.Funcs(util.GetFuncsMap()).ParseFiles(tmplFiles...)
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(0)
	}
}

func Generate(path string, version string) error {
	path = util.GetPath(path)

	if version != "" {
		version = "." + version
	}

	configFiles, err := filepath.Glob(path + "/*.yml" + version)
	if err != nil {
		return fmt.Errorf("could not open config file")
	}

	if len(configFiles) == 0 {
		return fmt.Errorf("could not find any config file")
	}

	importPath := util.GetImportPath(path)
	var files [][]byte

	for _, f := range configFiles {
		file, err := ioutil.ReadFile(f)
		if err != nil {
			return fmt.Errorf("could not open config file")
		}
		files = append(files, file)
	}

	fmt.Printf("parsing config file(s)...") // if show
	config, err := parser.Parse(importPath, files...)
	if err != nil {
		return err
	}
	fmt.Printf("\t%s\n", okMsg)

	return GenerateFrom(path, config)
}

func GenerateFrom(path string, config *parser.Config) error {
	// a, _ := json.Marshal(config)
	// fmt.Println(string(a))
	return generateORM(path, config)
}

func generateORM(path string, config *parser.Config) error {
	for _, d := range ormDirs {
		fmt.Printf("making %s folder...", magenta(d)) // if show
		err := util.Mkdir(util.GetFullPath(path, d))
		if err != nil {
			return err
		}
		fmt.Printf("\t%s\n", okMsg)
	}

	var w bytes.Buffer
	for t, dir := range ormTmplsMap {
		err := tmpls.ExecuteTemplate(&w, "heading", getPackage(t, config))
		if err != nil {
			return err
		}
		for _, db := range config.Databases {
			if t != "database" {
				for _, model := range db.Models {
					if db.Driver.Generate(t) {
						fmt.Printf("generating code for %s package - %s model...", magenta(t), magenta(model.Name)) // if show
						err = tmpls.ExecuteTemplate(&w, t, model)
						if err != nil {
							return err
						}
						fmt.Printf("\t%s\n", okMsg)
					} else {
						// driver headings
						fmt.Printf("generating %s driver code for %s package - %s model...", magenta(db.Driver.Name()), magenta(t), magenta(model.Name)) // if show
						err = db.Driver.Execute(&w, t, model)
						if err != nil {
							return err
						}
						fmt.Printf("\t%s\n", okMsg)
					}
				}
			} else {
				if db.Driver.Generate(t) {
					fmt.Printf("generating code for %s database...", magenta(db.Name)) // if show
					err = tmpls.ExecuteTemplate(&w, t, db)
					if err != nil {
						return err
					}
					fmt.Printf("\t%s\n", okMsg)
				}
				// driver headings
				// check if tmp is provided
				fmt.Printf("generating %s driver code for %s database...", magenta(db.Driver.Name()), magenta(db.Name)) // if show
				err = db.Driver.Execute(&w, t, db)
				if err != nil {
					return err
				}
				fmt.Printf("\t%s\n", okMsg)
			}
		}

		// fix white spaces

		fmt.Printf("formatting code for %s package...", magenta(t)) // if show
		b, err := format.Source(w.Bytes())
		if err != nil {
			return err
		}
		fmt.Printf("\t%s\n", okMsg)

		fmt.Printf("saving %s package file...", magenta(t)) // if show
		err = util.SaveFile(util.GetFullPath(path, dir+"/"+t+"_gen.go"), b)
		if err != nil {
			return err
		}
		fmt.Printf("\t%s\n", okMsg)

		w.Reset()
	}
	fmt.Printf("happy coding!\n") // if show
	return nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"flag"
	"fmt"
	"github.com/pkg/errors"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"

	usercorn "github.com/lunixbochs/usercorn/go"
	"github.com/lunixbochs/usercorn/go/debug"
	"github.com/lunixbochs/usercorn/go/models"
	"github.com/lunixbochs/usercorn/go/repl"
)

type strslice []string

func (s *strslice) String() string {
	return fmt.Sprintf("%v", *s)
}

func (s *strslice) Set(value string) error {
	*s = append(*s, value)
	return nil
}

// like go io.Copy(), but returns a channel to notify you upon completion
func copyNotify(dst io.Writer, src io.Reader) chan int {
	ret := make(chan int)
	go func() {
		io.Copy(dst, src)
		ret <- 1
	}()
	return ret
}

type UsercornCmd struct {
	Config *models.Config

	SetupFlags    func() error
	SetupUsercorn func() error
	MakeUsercorn  func(exe string) (models.Usercorn, error)
	RunUsercorn   func() error
	Teardown      func()

	NoExe, NoArgs bool

	Usercorn models.Usercorn
	Flags    *flag.FlagSet
}

func NewUsercornCmd() *UsercornCmd {
	fs := flag.NewFlagSet("cli", flag.ExitOnError)
	cmd := &UsercornCmd{Flags: fs}
	cmd.MakeUsercorn = func(exe string) (models.Usercorn, error) {
		// check permissions
		if stat, err := os.Stat(exe); err != nil {
			return nil, err
		} else if stat.Mode().Perm()&0111 == 0 {
			return nil, errors.Errorf("%s: permission denied (no execute bit)\n", exe)
		}
		return usercorn.NewUsercorn(exe, cmd.Config)
	}
	return cmd
}

type stackTracer interface {
	StackTrace() errors.StackTrace
}

func (c *UsercornCmd) PrintError(err error) {
	// print an error, and a stacktrace if available
	fmt.Fprintf(os.Stderr, "%s\n", strings.Repeat("-", 40))
	fmt.Fprintf(os.Stderr, "Error: %s\n", err)
	if err, ok := err.(stackTracer); ok {
		// parse full path and method name for each stack frame
		var frames [][]string
		for _, f := range err.StackTrace() {
			fullpath := ""
			fileline := fmt.Sprintf("%s:%d", f, f)
			method := fmt.Sprintf("%n", f)

			frame := fmt.Sprintf("%+s", f)
			tmp := strings.SplitN(frame, "\n", 3)
			if len(tmp) == 2 {
				pathsplit := strings.Split(tmp[0], "/")
				method = pathsplit[len(pathsplit)-1]
				fullpath = strings.TrimSpace(tmp[1])
			}
			frames = append(frames, []string{fullpath, fileline, method})
			if method == "main.main" {
				break
			}
		}
		// calculate column widths
		widths := make([]int, len(frames))
		for _, f := range frames {
			for i, s := range f {
				if len(s) > widths[i] {
					widths[i] = len(s)
				}
			}
		}
		// print pretty stacktrace
		for _, f := range frames {
			method := f[2]
			for i := 0; i < 2; i++ {
				if widths[i] > 0 {
					pad := strings.Repeat(" ", widths[i]-len(f[i]))
					fmt.Fprintf(os.Stderr, "%s%s | ", f[i], pad)
				}
			}
			fmt.Fprintf(os.Stderr, "%s()\n", method)
		}
	}
}

func (c *UsercornCmd) Run(argv, env []string) int {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	fs := c.Flags
	// tracing flags
	trace := fs.Bool("trace", false, "recommended tracing options: -loop 8 -strace -mtrace2 -etrace -rtrace -ftrace")
	strace := fs.Bool("strace", false, "trace syscalls")
	mtrace := fs.Bool("mtrace", false, "trace memory access (single)")
	mtrace2 := fs.Bool("mtrace2", false, "trace memory access (batched)")
	btrace := fs.Bool("btrace", false, "trace basic blocks")
	etrace := fs.Bool("etrace", false, "trace execution")
	rtrace := fs.Bool("rtrace", false, "trace register modification")
	ftrace := fs.Bool("ftrace", false, "trace source file:line")
	tracefile := fs.String("to", "", "binary trace output file")

	rewind := fs.Bool("rewind", false, "enable CPU rewind")
	ui := fs.Bool("ui", false, "display traced changes")

	match := fs.String("match", "", "trace from specific function(s) (func[,func...][+depth])")
	looproll := fs.Int("loop", 0, "collapse loop blocks of this depth")
	demangle := fs.Bool("demangle", false, "demangle symbols using c++filt")
	symfile := fs.Bool("symfile", false, "display symbols as sym@<mapped file>")
	disbytes := fs.Bool("disbytes", false, "show instruction bytes with disassembly")
	strsize := fs.Int("strsize", 30, "limited -strace'd strings to length (0 disables)")
	var src strslice
	fs.Var(&src, "src", "append source directory to search for -ftrace")
	// used for Usage grouping
	tnames := []string{
		"trace", "strace", "mtrace", "mtrace2", "btrace", "etrace", "rtrace", "ftrace",
		"match", "loop", "demangle", "symfile", "disbytes", "strsize",
	}

	inscount := fs.Bool("inscount", false, "print instruction count after execution")
	verbose := fs.Bool("v", false, "verbose output")
	prefix := fs.String("prefix", "", "library load prefix")
	base := fs.Uint64("base", 0, "force executable base address")
	ibase := fs.Uint64("ibase", 0, "force interpreter base address")
	skipinterp := fs.Bool("nointerp", false, "don't load binary's interpreter")
	native := fs.Bool("native", false, "[stub] use native syscall override (only works if host/guest arch/ABI matches)")
	stubsys := fs.Bool("stubsys", false, "stub missing syscalls")

	outfile := fs.String("o", "", "redirect debugging output to file (default stderr)")

	gdb := fs.Int("gdb", -1, "listen for gdb connection on localhost:<port>")

	startrepl := fs.Bool("repl", false, "experimental repl")

	cpuprofile := fs.String("cpuprofile", "", "write cpu profile to <file>")
	memprofile := fs.String("memprofile", "", "write mem profile to <file>")

	var envSet strslice
	var envUnset strslice
	fs.Var(&envSet, "set", "set environment var in the form name=value")
	fs.Var(&envUnset, "unset", "unset environment variable, '-' to unset all")

	fs.Usage = func() {
		usage := "Usage: %s [options]"
		if !c.NoExe {
			usage += " <exe>"
		}
		if !c.NoExe || !c.NoArgs {
			usage += " [args...]"
		}
		usage += "\n\nOptions:\n"
		fmt.Fprintf(os.Stderr, usage, os.Args[0])
		var flags []*flag.Flag
		var tflags []*flag.Flag
		fs.VisitAll(func(f *flag.Flag) {
			for _, name := range tnames {
				if name == f.Name {
					tflags = append(tflags, f)
					return
				}
			}
			flags = append(flags, f)
		})
		models.PrintFlags(flags)
		fmt.Fprintf(os.Stderr, "\nTrace Options:\n")
		models.PrintFlags(tflags)
		fmt.Fprintf(os.Stderr, "\nExample:\n  %s -trace -symfile bins/x86_64.linux.elf\n", os.Args[0])
	}
	if c.SetupFlags != nil {
		if err := c.SetupFlags(); err != nil {
			panic(err)
		}
	}
	fs.Parse(argv[1:])

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			panic(err)
		}
		pprof.StartCPUProfile(f)
	}

	var args []string
	if !c.NoExe {
		// make sure we were passed an executable
		args = fs.Args()
		if len(args) < 1 {
			fs.Usage()
			return 1
		}
	} else {
		args = []string{""}
	}

	// build configuration
	absPrefix := ""
	var err error
	if *prefix != "" {
		absPrefix, err = filepath.Abs(*prefix)
		if err != nil {
			panic(err)
		}
	}
	if *looproll == 0 && *trace {
		*looproll = 8
	}
	config := &models.Config{
		SymFile:      *symfile,
		StubSyscalls: *stubsys,

		Env:  env,
		Args: args,

		ForceBase:       *base,
		ForceInterpBase: *ibase,
		LoadPrefix:      absPrefix,
		NativeFallback:  *native,
		SkipInterp:      *skipinterp,
		Strsize:         *strsize,
		Verbose:         *verbose,

		Trace: models.TraceConfig{
			Tracefile:  *tracefile,
			Everything: *trace,

			Block: *btrace,
			Ins:   *etrace,
			Mem:   *mtrace,
			Reg:   *rtrace,
			Sys:   *strace,
		},
		Rewind: *rewind,
		UI:     *ui,

		// FIXME: these are UI tracing flags and now broken
		Demangle:      *demangle,
		DisBytes:      *disbytes,
		InsCount:      *inscount,
		LoopCollapse:  *looproll,
		SourcePaths:   src,
		TraceMemBatch: *mtrace2 || *trace,
		TraceSource:   *ftrace || *trace,
	}
	c.Config = config
	// FIXME: TraceMatch* is broken by trace changes
	if *match != "" {
		split := strings.SplitN(*match, "+", 2)
		if len(split) > 1 {
			if split[1] == "" {
				config.TraceMatchDepth = 999999
			} else {
				config.TraceMatchDepth, _ = strconv.Atoi(split[1])
			}
		}
		config.TraceMatch = strings.Split(split[0], ",")
	}

	if *outfile != "" {
		out, err := os.OpenFile(*outfile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
		if err != nil {
			panic(err)
		}
		config.Output = out
	}

	// merge environment with flags
	envSkip := make(map[string]bool)
	for _, v := range envSet {
		if strings.Contains(v, "=") {
			split := strings.SplitN(v, "=", 2)
			envSkip[split[0]] = true
		} else {
			fmt.Fprintf(os.Stderr, "warning: skipping invalid env set %#v\n", v)
		}
	}
	// unset specified env vars
	for _, v := range envUnset {
		if v == "-" {
			env = nil
			break
		}
		envSkip[v] = true
	}
	for _, v := range env {
		if strings.Contains(v, "=") {
			split := strings.SplitN(v, "=", 2)
			if _, ok := envSkip[split[0]]; !ok {
				envSet = append(envSet, v)
			}
		}
	}
	env = envSet

	corn, err := c.MakeUsercorn(args[0])
	if err != nil {
		c.PrintError(err)
		return 1
	}
	c.Usercorn = corn
	if c.SetupUsercorn != nil {
		if err := c.SetupUsercorn(); err != nil {
			c.PrintError(err)
			return 1
		}
	}
	teardown := func() {
		if *cpuprofile != "" {
			pprof.StopCPUProfile()
		}
		if *memprofile != "" {
			f, err := os.Create(*memprofile)
			if err != nil {
				fmt.Fprintf(os.Stderr, "could not write heap profile: %s\n", err)
			}
			pprof.WriteHeapProfile(f)
		}
		if c.Teardown != nil {
			c.Teardown()
		}
	}
	defer teardown()

	// start gdb server
	if *gdb > 0 {
		conn, err := debug.Accept("localhost", strconv.Itoa(*gdb))
		if err != nil {
			fmt.Fprintf(os.Stderr, "error accepting conn on port %d: %v\n", *gdb, err)
			return 1
		}
		go debug.NewGdbstub(corn).Run(conn)
	}

	if *startrepl {
		defer repl.Exit()
		if err := repl.Run(corn); err != nil {
			fmt.Fprintf(os.Stderr, "error starting repl: %v\n", err)
			return 1
		}
	}

	// start executable
	if c.RunUsercorn != nil {
		err = c.RunUsercorn()
	} else {
		err = corn.Run()
	}
	if err != nil {
		if e, ok := err.(models.ExitStatus); ok {
			return int(e)
		} else {
			c.PrintError(err)
			return 1
		}
	}
	return 0
}
<commit_msg>fix -set flag<commit_after>package cmd

import (
	"flag"
	"fmt"
	"github.com/pkg/errors"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"

	usercorn "github.com/lunixbochs/usercorn/go"
	"github.com/lunixbochs/usercorn/go/debug"
	"github.com/lunixbochs/usercorn/go/models"
	"github.com/lunixbochs/usercorn/go/repl"
)

type strslice []string

func (s *strslice) String() string {
	return fmt.Sprintf("%v", *s)
}

func (s *strslice) Set(value string) error {
	*s = append(*s, value)
	return nil
}

// like go io.Copy(), but returns a channel to notify you upon completion
func copyNotify(dst io.Writer, src io.Reader) chan int {
	ret := make(chan int)
	go func() {
		io.Copy(dst, src)
		ret <- 1
	}()
	return ret
}

type UsercornCmd struct {
	Config *models.Config

	SetupFlags    func() error
	SetupUsercorn func() error
	MakeUsercorn  func(exe string) (models.Usercorn, error)
	RunUsercorn   func() error
	Teardown      func()

	NoExe, NoArgs bool

	Usercorn models.Usercorn
	Flags    *flag.FlagSet
}

func NewUsercornCmd() *UsercornCmd {
	fs := flag.NewFlagSet("cli", flag.ExitOnError)
	cmd := &UsercornCmd{Flags: fs}
	cmd.MakeUsercorn = func(exe string) (models.Usercorn, error) {
		// check permissions
		if stat, err := os.Stat(exe); err != nil {
			return nil, err
		} else if stat.Mode().Perm()&0111 == 0 {
			return nil, errors.Errorf("%s: permission denied (no execute bit)\n", exe)
		}
		return usercorn.NewUsercorn(exe, cmd.Config)
	}
	return cmd
}

type stackTracer interface {
	StackTrace() errors.StackTrace
}

func (c *UsercornCmd) PrintError(err error) {
	// print an error, and a stacktrace if available
	fmt.Fprintf(os.Stderr, "%s\n", strings.Repeat("-", 40))
	fmt.Fprintf(os.Stderr, "Error: %s\n", err)
	if err, ok := err.(stackTracer); ok {
		// parse full path and method name for each stack frame
		var frames [][]string
		for _, f := range err.StackTrace() {
			fullpath := ""
			fileline := fmt.Sprintf("%s:%d", f, f)
			method := fmt.Sprintf("%n", f)

			frame := fmt.Sprintf("%+s", f)
			tmp := strings.SplitN(frame, "\n", 3)
			if len(tmp) == 2 {
				pathsplit := strings.Split(tmp[0], "/")
				method = pathsplit[len(pathsplit)-1]
				fullpath = strings.TrimSpace(tmp[1])
			}
			frames = append(frames, []string{fullpath, fileline, method})
			if method == "main.main" {
				break
			}
		}
		// calculate column widths
		widths := make([]int, len(frames))
		for _, f := range frames {
			for i, s := range f {
				if len(s) > widths[i] {
					widths[i] = len(s)
				}
			}
		}
		// print pretty stacktrace
		for _, f := range frames {
			method := f[2]
			for i := 0; i < 2; i++ {
				if widths[i] > 0 {
					pad := strings.Repeat(" ", widths[i]-len(f[i]))
					fmt.Fprintf(os.Stderr, "%s%s | ", f[i], pad)
				}
			}
			fmt.Fprintf(os.Stderr, "%s()\n", method)
		}
	}
}

func (c *UsercornCmd) Run(argv, env []string) int {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	fs := c.Flags
	// tracing flags
	trace := fs.Bool("trace", false, "recommended tracing options: -loop 8 -strace -mtrace2 -etrace -rtrace -ftrace")
	strace := fs.Bool("strace", false, "trace syscalls")
	mtrace := fs.Bool("mtrace", false, "trace memory access (single)")
	mtrace2 := fs.Bool("mtrace2", false, "trace memory access (batched)")
	btrace := fs.Bool("btrace", false, "trace basic blocks")
	etrace := fs.Bool("etrace", false, "trace execution")
	rtrace := fs.Bool("rtrace", false, "trace register modification")
	ftrace := fs.Bool("ftrace", false, "trace source file:line")
	tracefile := fs.String("to", "", "binary trace output file")

	rewind := fs.Bool("rewind", false, "enable CPU rewind")
	ui := fs.Bool("ui", false, "display traced changes")

	match := fs.String("match", "", "trace from specific function(s) (func[,func...][+depth])")
	looproll := fs.Int("loop", 0, "collapse loop blocks of this depth")
	demangle := fs.Bool("demangle", false, "demangle symbols using c++filt")
	symfile := fs.Bool("symfile", false, "display symbols as sym@<mapped file>")
	disbytes := fs.Bool("disbytes", false, "show instruction bytes with disassembly")
	strsize := fs.Int("strsize", 30, "limited -strace'd strings to length (0 disables)")
	var src strslice
	fs.Var(&src, "src", "append source directory to search for -ftrace")
	// used for Usage grouping
	tnames := []string{
		"trace", "strace", "mtrace", "mtrace2", "btrace", "etrace", "rtrace", "ftrace",
		"match", "loop", "demangle", "symfile", "disbytes", "strsize",
	}

	inscount := fs.Bool("inscount", false, "print instruction count after execution")
	verbose := fs.Bool("v", false, "verbose output")
	prefix := fs.String("prefix", "", "library load prefix")
	base := fs.Uint64("base", 0, "force executable base address")
	ibase := fs.Uint64("ibase", 0, "force interpreter base address")
	skipinterp := fs.Bool("nointerp", false, "don't load binary's interpreter")
	native := fs.Bool("native", false, "[stub] use native syscall override (only works if host/guest arch/ABI matches)")
	stubsys := fs.Bool("stubsys", false, "stub missing syscalls")

	outfile := fs.String("o", "", "redirect debugging output to file (default stderr)")

	gdb := fs.Int("gdb", -1, "listen for gdb connection on localhost:<port>")

	startrepl := fs.Bool("repl", false, "experimental repl")

	cpuprofile := fs.String("cpuprofile", "", "write cpu profile to <file>")
	memprofile := fs.String("memprofile", "", "write mem profile to <file>")

	var envSet strslice
	var envUnset strslice
	fs.Var(&envSet, "set", "set environment var in the form name=value")
	fs.Var(&envUnset, "unset", "unset environment variable, '-' to unset all")

	fs.Usage = func() {
		usage := "Usage: %s [options]"
		if !c.NoExe {
			usage += " <exe>"
		}
		if !c.NoExe || !c.NoArgs {
			usage += " [args...]"
		}
		usage += "\n\nOptions:\n"
		fmt.Fprintf(os.Stderr, usage, os.Args[0])
		var flags []*flag.Flag
		var tflags []*flag.Flag
		fs.VisitAll(func(f *flag.Flag) {
			for _, name := range tnames {
				if name == f.Name {
					tflags = append(tflags, f)
					return
				}
			}
			flags = append(flags, f)
		})
		models.PrintFlags(flags)
		fmt.Fprintf(os.Stderr, "\nTrace Options:\n")
		models.PrintFlags(tflags)
		fmt.Fprintf(os.Stderr, "\nExample:\n  %s -trace -symfile bins/x86_64.linux.elf\n", os.Args[0])
	}
	if c.SetupFlags != nil {
		if err := c.SetupFlags(); err != nil {
			panic(err)
		}
	}
	fs.Parse(argv[1:])

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			panic(err)
		}
		pprof.StartCPUProfile(f)
	}

	var args []string
	if !c.NoExe {
		// make sure we were passed an executable
		args = fs.Args()
		if len(args) < 1 {
			fs.Usage()
			return 1
		}
	} else {
		args = []string{""}
	}

	// build configuration
	absPrefix := ""
	var err error
	if *prefix != "" {
		absPrefix, err = filepath.Abs(*prefix)
		if err != nil {
			panic(err)
		}
	}
	if *looproll == 0 && *trace {
		*looproll = 8
	}
	config := &models.Config{
		SymFile:      *symfile,
		StubSyscalls: *stubsys,

		Env:  env,
		Args: args,

		ForceBase:       *base,
		ForceInterpBase: *ibase,
		LoadPrefix:      absPrefix,
		NativeFallback:  *native,
		SkipInterp:      *skipinterp,
		Strsize:         *strsize,
		Verbose:         *verbose,

		Trace: models.TraceConfig{
			Tracefile:  *tracefile,
			Everything: *trace,

			Block: *btrace,
			Ins:   *etrace,
			Mem:   *mtrace,
			Reg:   *rtrace,
			Sys:   *strace,
		},
		Rewind: *rewind,
		UI:     *ui,

		// FIXME: these are UI tracing flags and now broken
		Demangle:      *demangle,
		DisBytes:      *disbytes,
		InsCount:      *inscount,
		LoopCollapse:  *looproll,
		SourcePaths:   src,
		TraceMemBatch: *mtrace2 || *trace,
		TraceSource:   *ftrace || *trace,
	}
	c.Config = config
	// FIXME: TraceMatch* is broken by trace changes
	if *match != "" {
		split := strings.SplitN(*match, "+", 2)
		if len(split) > 1 {
			if split[1] == "" {
				config.TraceMatchDepth = 999999
			} else {
				config.TraceMatchDepth, _ = strconv.Atoi(split[1])
			}
		}
		config.TraceMatch = strings.Split(split[0], ",")
	}

	if *outfile != "" {
		out, err := os.OpenFile(*outfile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
		if err != nil {
			panic(err)
		}
		config.Output = out
	}

	// merge environment with flags
	envSkip := make(map[string]bool)
	for _, v := range envSet {
		if strings.Contains(v, "=") {
			split := strings.SplitN(v, "=", 2)
			envSkip[split[0]] = true
		} else {
			fmt.Fprintf(os.Stderr, "warning: skipping invalid env set %#v\n", v)
		}
	}
	// unset specified env vars
	for _, v := range envUnset {
		if v == "-" {
			env = nil
			break
		}
		envSkip[v] = true
	}
	for _, v := range env {
		if strings.Contains(v, "=") {
			split := strings.SplitN(v, "=", 2)
			if _, ok := envSkip[split[0]]; !ok {
				envSet = append(envSet, v)
			}
		}
	}
	config.Env = envSet

	corn, err := c.MakeUsercorn(args[0])
	if err != nil {
		c.PrintError(err)
		return 1
	}
	c.Usercorn = corn
	if c.SetupUsercorn != nil {
		if err := c.SetupUsercorn(); err != nil {
			c.PrintError(err)
			return 1
		}
	}
	teardown := func() {
		if *cpuprofile != "" {
			pprof.StopCPUProfile()
		}
		if *memprofile != "" {
			f, err := os.Create(*memprofile)
			if err != nil {
				fmt.Fprintf(os.Stderr, "could not write heap profile: %s\n", err)
			}
			pprof.WriteHeapProfile(f)
		}
		if c.Teardown != nil {
			c.Teardown()
		}
	}
	defer teardown()

	// start gdb server
	if *gdb > 0 {
		conn, err := debug.Accept("localhost", strconv.Itoa(*gdb))
		if err != nil {
			fmt.Fprintf(os.Stderr, "error accepting conn on port %d: %v\n", *gdb, err)
			return 1
		}
		go debug.NewGdbstub(corn).Run(conn)
	}

	if *startrepl {
		defer repl.Exit()
		if err := repl.Run(corn); err != nil {
			fmt.Fprintf(os.Stderr, "error starting repl: %v\n", err)
			return 1
		}
	}

	// start executable
	if c.RunUsercorn != nil {
		err = c.RunUsercorn()
	} else {
		err = corn.Run()
	}
	if err != nil {
		if e, ok := err.(models.ExitStatus); ok {
			return int(e)
		} else {
			c.PrintError(err)
			return 1
		}
	}
	return 0
}
<|endoftext|>


<commit_before>// Copyright (c) 2012 by Christoph Hack <christoph@tux21b.org>
// All rights reserved. Distributed under the Simplified BSD License.

package gocql

import (
	"bytes"
	"database/sql"
	"github.com/tux21b/gocql/uuid"
	"testing"
	"time"
)

func TestSimple(t *testing.T) {
	db, err := sql.Open("gocql", "localhost:8000 keyspace=system")
	if err != nil {
		t.Fatal(err)
	}

	rows, err := db.Query("SELECT keyspace_name FROM schema_keyspaces")
	if err != nil {
		t.Fatal(err)
	}

	for rows.Next() {
		var keyspace string
		if err := rows.Scan(&keyspace); err != nil {
			t.Fatal(err)
		}
	}
	if err != nil {
		t.Fatal(err)
	}
}

type Page struct {
	Title      string
	RevID      uuid.UUID
	Body       string
	Hits       int
	Protected  bool
	Modified   time.Time
	Attachment []byte
}

var pages = []*Page{
	&Page{"Frontpage", uuid.TimeUUID(), "Hello world!", 0, false,
		time.Date(2012, 8, 20, 10, 0, 0, 0, time.UTC), nil},
	&Page{"Frontpage", uuid.TimeUUID(), "Hello modified world!", 0, false,
		time.Date(2012, 8, 22, 10, 0, 0, 0, time.UTC), []byte("img data\x00")},
	&Page{"LoremIpsum", uuid.TimeUUID(), "Lorem ipsum dolor sit amet", 12,
		true, time.Date(2012, 8, 22, 10, 0, 8, 0, time.UTC), nil},
}

func TestWiki(t *testing.T) {
	db, err := sql.Open("gocql", "localhost:8000 compression=snappy")
	if err != nil {
		t.Fatal(err)
	}
	db.Exec("DROP KEYSPACE gocql_wiki")
	if _, err := db.Exec(`CREATE KEYSPACE gocql_wiki
        WITH strategy_class = 'SimpleStrategy'
        AND strategy_options:replication_factor = 1`); err != nil {
		t.Fatal(err)
	}
	if _, err := db.Exec("USE gocql_wiki"); err != nil {
		t.Fatal(err)
	}

	if _, err := db.Exec(`CREATE TABLE page (
        title varchar,
        revid timeuuid,
        body varchar,
        hits int,
        protected boolean,
        modified timestamp,
        attachment blob,
        PRIMARY KEY (title, revid)
        )`); err != nil {
		t.Fatal(err)
	}
	for _, p := range pages {
		if _, err := db.Exec(`INSERT INTO page (title, revid, body, hits,
            protected, modified, attachment) VALUES (?, ?, ?, ?, ?, ?, ?);`,
			p.Title, p.RevID, p.Body, p.Hits, p.Protected, p.Modified,
			p.Attachment); err != nil {
			t.Fatal(err)
		}
	}

	row := db.QueryRow(`SELECT count(*) FROM page`)
	var count int
	if err := row.Scan(&count); err != nil {
		t.Error(err)
	}
	if count != len(pages) {
		t.Fatalf("expected %d rows, got %d", len(pages), count)
	}

	for _, page := range pages {
		row := db.QueryRow(`SELECT title, revid, body, hits, protected,
            modified, attachment
            FROM page WHERE title = ? AND revid = ?`, page.Title, page.RevID)
		var p Page
		err := row.Scan(&p.Title, &p.RevID, &p.Body, &p.Hits, &p.Protected,
			&p.Modified, &p.Attachment)
		if err != nil {
			t.Fatal(err)
		}
		p.Modified = p.Modified.In(time.UTC)
		if page.Title != p.Title || page.RevID != p.RevID ||
			page.Body != p.Body || page.Modified != p.Modified ||
			page.Hits != p.Hits || page.Protected != p.Protected ||
			!bytes.Equal(page.Attachment, p.Attachment) {
			t.Errorf("expected %#v got %#v", *page, p)
		}
	}

	row = db.QueryRow(`SELECT title, revid, body, hits, protected,
        modified, attachment
        FROM page WHERE title = ? ORDER BY revid DESC`, "Frontpage")
	var p Page
	if err := row.Scan(&p.Title, &p.RevID, &p.Body, &p.Hits, &p.Protected,
		&p.Modified, &p.Attachment); err != nil {
		t.Error(err)
	}
	p.Modified = p.Modified.In(time.UTC)
	page := pages[1]
	if page.Title != p.Title || page.RevID != p.RevID ||
		page.Body != p.Body || page.Modified != p.Modified ||
		page.Hits != p.Hits || page.Protected != p.Protected ||
		!bytes.Equal(page.Attachment, p.Attachment) {
		t.Errorf("expected %#v got %#v", *page, p)
	}

}
<commit_msg>Fix CREATE KEYSPACE statement for Cassandra 1.2.0-beta1<commit_after>// Copyright (c) 2012 by Christoph Hack <christoph@tux21b.org>
// All rights reserved. Distributed under the Simplified BSD License.

package gocql

import (
	"bytes"
	"database/sql"
	"github.com/tux21b/gocql/uuid"
	"testing"
	"time"
)

func TestSimple(t *testing.T) {
	db, err := sql.Open("gocql", "localhost:8000 keyspace=system")
	if err != nil {
		t.Fatal(err)
	}

	rows, err := db.Query("SELECT keyspace_name FROM schema_keyspaces")
	if err != nil {
		t.Fatal(err)
	}

	for rows.Next() {
		var keyspace string
		if err := rows.Scan(&keyspace); err != nil {
			t.Fatal(err)
		}
	}
	if err != nil {
		t.Fatal(err)
	}
}

type Page struct {
	Title      string
	RevID      uuid.UUID
	Body       string
	Hits       int
	Protected  bool
	Modified   time.Time
	Attachment []byte
}

var pages = []*Page{
	&Page{"Frontpage", uuid.TimeUUID(), "Hello world!", 0, false,
		time.Date(2012, 8, 20, 10, 0, 0, 0, time.UTC), nil},
	&Page{"Frontpage", uuid.TimeUUID(), "Hello modified world!", 0, false,
		time.Date(2012, 8, 22, 10, 0, 0, 0, time.UTC), []byte("img data\x00")},
	&Page{"LoremIpsum", uuid.TimeUUID(), "Lorem ipsum dolor sit amet", 12,
		true, time.Date(2012, 8, 22, 10, 0, 8, 0, time.UTC), nil},
}

func TestWiki(t *testing.T) {
	db, err := sql.Open("gocql", "localhost:8000 compression=snappy")
	if err != nil {
		t.Fatal(err)
	}
	db.Exec("DROP KEYSPACE gocql_wiki")
	if _, err := db.Exec(`CREATE KEYSPACE gocql_wiki
	                      WITH replication = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 }`); err != nil {
		t.Fatal(err)
	}
	if _, err := db.Exec("USE gocql_wiki"); err != nil {
		t.Fatal(err)
	}

	if _, err := db.Exec(`CREATE TABLE page (
        title varchar,
        revid timeuuid,
        body varchar,
        hits int,
        protected boolean,
        modified timestamp,
        attachment blob,
        PRIMARY KEY (title, revid)
        )`); err != nil {
		t.Fatal(err)
	}
	for _, p := range pages {
		if _, err := db.Exec(`INSERT INTO page (title, revid, body, hits,
            protected, modified, attachment) VALUES (?, ?, ?, ?, ?, ?, ?);`,
			p.Title, p.RevID, p.Body, p.Hits, p.Protected, p.Modified,
			p.Attachment); err != nil {
			t.Fatal(err)
		}
	}

	row := db.QueryRow(`SELECT count(*) FROM page`)
	var count int
	if err := row.Scan(&count); err != nil {
		t.Error(err)
	}
	if count != len(pages) {
		t.Fatalf("expected %d rows, got %d", len(pages), count)
	}

	for _, page := range pages {
		row := db.QueryRow(`SELECT title, revid, body, hits, protected,
            modified, attachment
            FROM page WHERE title = ? AND revid = ?`, page.Title, page.RevID)
		var p Page
		err := row.Scan(&p.Title, &p.RevID, &p.Body, &p.Hits, &p.Protected,
			&p.Modified, &p.Attachment)
		if err != nil {
			t.Fatal(err)
		}
		p.Modified = p.Modified.In(time.UTC)
		if page.Title != p.Title || page.RevID != p.RevID ||
			page.Body != p.Body || page.Modified != p.Modified ||
			page.Hits != p.Hits || page.Protected != p.Protected ||
			!bytes.Equal(page.Attachment, p.Attachment) {
			t.Errorf("expected %#v got %#v", *page, p)
		}
	}

	row = db.QueryRow(`SELECT title, revid, body, hits, protected,
        modified, attachment
        FROM page WHERE title = ? ORDER BY revid DESC`, "Frontpage")
	var p Page
	if err := row.Scan(&p.Title, &p.RevID, &p.Body, &p.Hits, &p.Protected,
		&p.Modified, &p.Attachment); err != nil {
		t.Error(err)
	}
	p.Modified = p.Modified.In(time.UTC)
	page := pages[1]
	if page.Title != p.Title || page.RevID != p.RevID ||
		page.Body != p.Body || page.Modified != p.Modified ||
		page.Hits != p.Hits || page.Protected != p.Protected ||
		!bytes.Equal(page.Attachment, p.Attachment) {
		t.Errorf("expected %#v got %#v", *page, p)
	}

}
<|endoftext|>