<commit_before>package commands

import (
	"fmt"
	"github.com/kubeflow/arena/pkg/spark-operator/apis/sparkoperator.k8s.io/v1beta1"
	"github.com/kubeflow/arena/pkg/spark-operator/client/clientset/versioned"
	"github.com/kubeflow/arena/pkg/types"
	log "github.com/sirupsen/logrus"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"time"
)

// all spark jobs cache
var (
	allSparkJobs []v1beta1.SparkApplication
)

// spark application wrapper
type SparkJob struct {
	name        string
	sparkjob    v1beta1.SparkApplication
	trainerType string
	pods        []v1.Pod
	chiefPod    v1.Pod
}

func (sj *SparkJob) Name() string {
	return sj.name
}

// return driver pod
func (sj *SparkJob) ChiefPod() v1.Pod {
	return sj.chiefPod
}

// return trainerType: sparkjob
func (sj *SparkJob) Trainer() string {
	return sj.trainerType
}

// return pods from cache
func (sj *SparkJob) AllPods() []v1.Pod {
	return sj.pods
}

/*
				spark job driver state
	-------------------------------------------------------
	NewState              ApplicationStateType = ""
	SubmittedState        ApplicationStateType = "SUBMITTED"
	RunningState          ApplicationStateType = "RUNNING"
	CompletedState        ApplicationStateType = "COMPLETED"
	FailedState           ApplicationStateType = "FAILED"
	FailedSubmissionState ApplicationStateType = "SUBMISSION_FAILED"
	PendingRerunState     ApplicationStateType = "PENDING_RERUN"
	InvalidatingState     ApplicationStateType = "INVALIDATING"
	SucceedingState       ApplicationStateType = "SUCCEEDING"
	FailingState          ApplicationStateType = "FAILING"
	UnknownState          ApplicationStateType = "UNKNOWN"

				spark job executor state
	-------------------------------------------------------
	ExecutorPendingState   ExecutorState = "PENDING"
	ExecutorRunningState   ExecutorState = "RUNNING"
	ExecutorCompletedState ExecutorState = "COMPLETED"
	ExecutorFailedState    ExecutorState = "FAILED"
	ExecutorUnknownState   ExecutorState = "UNKNOWN"
*/
func (sj *SparkJob) GetStatus() (status string) {

	defer func() {
		if r := recover(); r != nil {
			fmt.Println("spark job may not complete,because of ", r)
		}
		return
	}()

	status = "UNKNOWN"

	// name is empty when the pod has not been scheduled
	if sj.sparkjob.Name == "" {
		return status
	}

	if sj.isSucceeded() {
		status = "SUCCEEDED"
	} else if sj.isFailed() {
		status = "FAILED"
	} else if sj.isPending() {
		status = "PENDING"
	} else if sj.isSubmitted() {
		status = "SUBMITTED"
	} else if sj.isRunning() {
		status = "RUNNING"
	} else {
		status = string(sj.sparkjob.Status.AppState.State)
	}

	return status
}

func (sj *SparkJob) isSucceeded() bool {
	return sj.sparkjob.Status.AppState.State == v1beta1.CompletedState
}

func (sj *SparkJob) isFailed() bool {
	return sj.sparkjob.Status.AppState.State == v1beta1.FailedState
}

func (sj *SparkJob) isPending() bool {
	return sj.sparkjob.Status.AppState.State == v1beta1.PendingRerunState
}

func (sj *SparkJob) isSubmitted() bool {
	return sj.sparkjob.Status.AppState.State == v1beta1.SubmittedState
}

func (sj *SparkJob) isRunning() bool {
	return sj.sparkjob.Status.AppState.State == v1beta1.RunningState
}

func (sj *SparkJob) StartTime() *metav1.Time {
	return &sj.sparkjob.CreationTimestamp
}

func (sj *SparkJob) Age() time.Duration {
	job := sj.sparkjob

	if job.CreationTimestamp.IsZero() {
		return 0
	}
	return metav1.Now().Sub(job.CreationTimestamp.Time)
}

// Get the Job Training Duration
func (sj *SparkJob) Duration() time.Duration {
	sparkjob := sj.sparkjob

	if sparkjob.CreationTimestamp.IsZero() {
		return 0
	}

	if sparkjob.Status.TerminationTime.IsZero() {
		return 0
	}

	//todo
	return sparkjob.Status.TerminationTime.Sub(sparkjob.CreationTimestamp.Time)
}

func (sj *SparkJob) GetJobDashboards(client *kubernetes.Clientset) ([]string, error) {
	urls := []string{}
	dashboardURL, err := dashboard(client, namespace, "kubernetes-dashboard")

	if err != nil {
		log.Debugf("Get dashboard failed due to %v", err)
		// retry for the existing customers, will be deprecated in the future
		dashboardURL, err = dashboard(client, arenaNamespace, "kubernetes-dashboard")
		if err != nil {
			log.Debugf("Get dashboard failed due to %v", err)
		}
	}

	if err != nil {
		log.Debugf("Get dashboard failed due to %v", err)
		// retry for the existing customers, will be deprecated in the future
		dashboardURL, err = dashboard(client, "kube-system", "kubernetes-dashboard")
		if err != nil {
			log.Debugf("Get dashboard failed due to %v", err)
		}
	}

	if dashboardURL == "" {
		return urls, fmt.Errorf("No LOGVIEWER Installed.")
	}

	if len(sj.chiefPod.Spec.Containers) == 0 {
		return urls, fmt.Errorf("spark driver pod is not ready!")
	}

	url := fmt.Sprintf("%s/#!/log/%s/%s/%s?namespace=%s\n",
		dashboardURL,
		sj.chiefPod.Namespace,
		sj.chiefPod.Name,
		sj.chiefPod.Spec.Containers[0].Name,
		sj.chiefPod.Namespace)

	urls = append(urls, url)

	return urls, nil
}

// spark job without gpu supported
func (sj *SparkJob) RequestedGPU() int64 {
	return 0
}

// spark job without gpu supported
func (sj *SparkJob) AllocatedGPU() int64 {
	return 0
}

// Get the hostIP of the driver Pod
func (sj *SparkJob) HostIPOfChief() (hostIP string) {
	hostIP = "N/A"
	if sj.GetStatus() == "RUNNING" {
		hostIP = sj.chiefPod.Status.HostIP
	}
	return hostIP
}

func (sj *SparkJob) Namespace() string {
	return sj.sparkjob.Namespace
}

func NewSparkJobTrainer(client *kubernetes.Clientset) Trainer {
	log.Debugf("Init Spark job trainer")
	sparkjobClient, err := initSparkJobClient()

	if err != nil {
		log.Debugf("unsupported spark job due to %v", err)
		return &SparkJobTrainer{
			trainerType: defaultSparkJobTrainingType,
			enabled:     false,
		}
	}
	// allPods have been cached, we do the same to allSparkJobs
	if useCache {
		ns := namespace
		if allNamespaces {
			ns = metav1.NamespaceAll
		}

		sparkJobList, err := sparkjobClient.SparkoperatorV1beta1().SparkApplications(ns).List(metav1.ListOptions{})
		if err != nil {
			log.Debugf("unsupported sparkJob due to %v", err)
			return &SparkJobTrainer{
				trainerType: defaultSparkJobTrainingType,
				enabled:     false,
			}
		}

		for _, sparkJob := range sparkJobList.Items {
			allSparkJobs = append(allSparkJobs, sparkJob)
		}
	}

	return &SparkJobTrainer{
		sparkjobClient: sparkjobClient,
		client:         client,
		trainerType:    defaultSparkJobTrainingType,
		enabled:        true,
	}
}

// init spark job client
func initSparkJobClient() (sparkjobClientset *versioned.Clientset, err error) {
	if restConfig == nil {
		restConfig, err = clientConfig.ClientConfig()
		if err != nil {
			log.Fatal(err)
			return nil, err
		}
	}

	sparkjobClientset = versioned.NewForConfigOrDie(restConfig)
	return sparkjobClientset, nil
}

// spark job trainer
type SparkJobTrainer struct {
	client         *kubernetes.Clientset
	sparkjobClient *versioned.Clientset
	trainerType    string
	enabled        bool
}

func (st *SparkJobTrainer) Type() string {
	return st.trainerType
}

func (st *SparkJobTrainer) IsSupported(name, ns string) bool {
	if !st.enabled {
		return false
	}

	isSpark := false

	if useCache {
		for _, job := range allSparkJobs {
			if st.isSparkJob(name, ns, job) {
				isSpark = true
				break
			}
		}
	} else {
		sparkJobList, err := st.sparkjobClient.SparkoperatorV1beta1().SparkApplications(ns).List(metav1.ListOptions{
			LabelSelector: fmt.Sprintf("release=%s", name),
		})

		if err != nil {
			log.Debugf("failed to search job %s in namespace %s due to %v", name, ns, err)
		}

		if len(sparkJobList.Items) > 0 {
			isSpark = true
		}
	}

	return isSpark
}

func (st *SparkJobTrainer) isSparkJob(name, ns string, job v1beta1.SparkApplication) bool {
	if val, ok := job.Labels["release"]; ok && (val == name) {
		log.Debugf("the sparkjob %s with labels %s", job.Name, val)
	} else {
		return false
	}

	if val, ok := job.Labels["app"]; ok && (val == "sparkjob") {
		log.Debugf("the sparkjob %s with labels %s is found.", job.Name, val)
	} else {
		return false
	}

	if job.Namespace != ns {
		return false
	}
	return true
}

func (st *SparkJobTrainer) GetTrainingJob(name, namespace string) (job TrainingJob, err error) {
	if len(allSparkJobs) > 0 {
		job, err = st.getTrainingJobFromCache(name, namespace)
	} else {
		job, err = st.getTrainingJob(name, namespace)
	}

	return job, err
}

func (st *SparkJobTrainer) getTrainingJobFromCache(name, namespace string) (job TrainingJob, err error) {
	var (
		sparkJob v1beta1.SparkApplication
	)

	for _, item := range allSparkJobs {
		if st.isSparkJob(name, namespace, item) {
			sparkJob = item
			break
		}
	}

	pods, chiefPod := getPodsOfSparkJob(name, st, allPods)

	return &SparkJob{
		chiefPod:    chiefPod,
		sparkjob:    sparkJob,
		pods:        pods,
		name:        name,
		trainerType: st.Type(),
	}, nil
}

func (st *SparkJobTrainer) getTrainingJob(name, namespace string) (job TrainingJob, err error) {
	var (
		sparkjob v1beta1.SparkApplication
	)

	sparkjobList, err := st.sparkjobClient.SparkoperatorV1beta1().SparkApplications(namespace).List(metav1.ListOptions{
		LabelSelector: fmt.Sprintf("release=%s", name),
	})
	if err != nil {
		return nil, err
	}
	if len(sparkjobList.Items) == 0 {
		return nil, nil
	} else {
		sparkjob = sparkjobList.Items[0]
	}

	podList, err := st.client.CoreV1().Pods(namespace).List(metav1.ListOptions{
		TypeMeta: metav1.TypeMeta{
			Kind:       "ListOptions",
			APIVersion: "v1",
		}, LabelSelector: fmt.Sprintf("release=%s", name),
	})

	if err != nil {
		return nil, err
	}

	pods, chiefPod := getPodsOfSparkJob(name, st, podList.Items)

	return &SparkJob{
		sparkjob:    sparkjob,
		chiefPod:    chiefPod,
		pods:        pods,
		name:        name,
		trainerType: st.Type(),
	}, nil
}

func (st *SparkJobTrainer) ListTrainingJobs() (jobs []TrainingJob, err error) {
	jobs = []TrainingJob{}
	jobInfos := []types.TrainingJobInfo{}
	for _, sparkJob := range allSparkJobs {
		jobInfo := types.TrainingJobInfo{}
		log.Debugf("find sparkjob %s in %s", sparkJob.Name, sparkJob.Namespace)
		if val, ok := sparkJob.Labels["release"]; ok && (sparkJob.Name == fmt.Sprintf("%s-%s", val, st.Type())) {
			log.Debugf("the sparkjob %s with labels %s found in List", sparkJob.Name, val)
			jobInfo.Name = val
		} else {
			jobInfo.Name = sparkJob.Name
		}

		jobInfo.Namespace = sparkJob.Namespace
		jobInfos = append(jobInfos, jobInfo)
	}
	log.Debugf("jobInfos %v", jobInfos)

	for _, jobInfo := range jobInfos {
		job, err := st.getTrainingJobFromCache(jobInfo.Name, jobInfo.Namespace)
		if err != nil {
			return jobs, err
		}
		jobs = append(jobs, job)
	}

	return jobs, nil
}

func (st *SparkJobTrainer) isSparkPod(name, ns string, item v1.Pod) bool {
	if val, ok := item.Labels["release"]; ok && (val == name) {
		log.Debugf("the sparkjob %s with labels %s", item.Name, val)
	} else {
		return false
	}

	if val, ok := item.Labels["app"]; ok && (val == "sparkjob") {
		log.Debugf("the sparkjob %s with labels %s is found.", item.Name, val)
	} else {
		return false
	}

	if item.Namespace != ns {
		return false
	}
	return true
}

func (st *SparkJobTrainer) isChiefPod(item v1.Pod) bool {
	if val, ok := item.Labels["spark-role"]; ok && (val == "driver") {
		log.Debugf("the sparkjob %s with labels %s", item.Name, val)
	} else {
		return false
	}
	return true
}

func getPodsOfSparkJob(name string, st *SparkJobTrainer, podList []v1.Pod) (pods []v1.Pod, chiefPod v1.Pod) {
	pods = []v1.Pod{}
	for _, item := range podList {
		if !st.isSparkPod(name, namespace, item) {
			continue
		}
		if st.isChiefPod(item) && item.CreationTimestamp.After(chiefPod.CreationTimestamp.Time) {
			// If there are some failed chiefPod, and the new chiefPod haven't started, set the latest failed pod as chief pod
			if chiefPod.Name != "" && item.Status.Phase == v1.PodPending {
				continue
			}
			chiefPod = item
		}

		// for non-job pod, add it into the pod list
		pods = append(pods, item)
		log.Debugf("add pod %v to pods", item)
	}
	return pods, chiefPod
}
<commit_msg>Fix the SIGSEGV on get spark job (#185)<commit_after>package commands

import (
	"fmt"
	"github.com/kubeflow/arena/pkg/spark-operator/apis/sparkoperator.k8s.io/v1beta1"
	"github.com/kubeflow/arena/pkg/spark-operator/client/clientset/versioned"
	"github.com/kubeflow/arena/pkg/types"
	log "github.com/sirupsen/logrus"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"time"
)

// all spark jobs cache
var (
	allSparkJobs []v1beta1.SparkApplication
)

// spark application wrapper
type SparkJob struct {
	name        string
	sparkjob    v1beta1.SparkApplication
	trainerType string
	pods        []v1.Pod
	chiefPod    v1.Pod
}

func (sj *SparkJob) Name() string {
	return sj.name
}

// return driver pod
func (sj *SparkJob) ChiefPod() v1.Pod {
	return sj.chiefPod
}

// return trainerType: sparkjob
func (sj *SparkJob) Trainer() string {
	return sj.trainerType
}

// return pods from cache
func (sj *SparkJob) AllPods() []v1.Pod {
	return sj.pods
}

/*
				spark job driver state
	-------------------------------------------------------
	NewState              ApplicationStateType = ""
	SubmittedState        ApplicationStateType = "SUBMITTED"
	RunningState          ApplicationStateType = "RUNNING"
	CompletedState        ApplicationStateType = "COMPLETED"
	FailedState           ApplicationStateType = "FAILED"
	FailedSubmissionState ApplicationStateType = "SUBMISSION_FAILED"
	PendingRerunState     ApplicationStateType = "PENDING_RERUN"
	InvalidatingState     ApplicationStateType = "INVALIDATING"
	SucceedingState       ApplicationStateType = "SUCCEEDING"
	FailingState          ApplicationStateType = "FAILING"
	UnknownState          ApplicationStateType = "UNKNOWN"

				spark job executor state
	-------------------------------------------------------
	ExecutorPendingState   ExecutorState = "PENDING"
	ExecutorRunningState   ExecutorState = "RUNNING"
	ExecutorCompletedState ExecutorState = "COMPLETED"
	ExecutorFailedState    ExecutorState = "FAILED"
	ExecutorUnknownState   ExecutorState = "UNKNOWN"
*/
func (sj *SparkJob) GetStatus() (status string) {

	defer func() {
		if r := recover(); r != nil {
			fmt.Println("spark job may not complete,because of ", r)
		}
		return
	}()

	status = "UNKNOWN"

	// name is empty when the pod has not been scheduled
	if sj.sparkjob.Name == "" {
		return status
	}

	if sj.isSucceeded() {
		status = "SUCCEEDED"
	} else if sj.isFailed() {
		status = "FAILED"
	} else if sj.isPending() {
		status = "PENDING"
	} else if sj.isSubmitted() {
		status = "SUBMITTED"
	} else if sj.isRunning() {
		status = "RUNNING"
	} else {
		status = string(sj.sparkjob.Status.AppState.State)
	}

	return status
}

func (sj *SparkJob) isSucceeded() bool {
	return sj.sparkjob.Status.AppState.State == v1beta1.CompletedState
}

func (sj *SparkJob) isFailed() bool {
	return sj.sparkjob.Status.AppState.State == v1beta1.FailedState
}

func (sj *SparkJob) isPending() bool {
	return sj.sparkjob.Status.AppState.State == v1beta1.PendingRerunState
}

func (sj *SparkJob) isSubmitted() bool {
	return sj.sparkjob.Status.AppState.State == v1beta1.SubmittedState
}

func (sj *SparkJob) isRunning() bool {
	return sj.sparkjob.Status.AppState.State == v1beta1.RunningState
}

func (sj *SparkJob) StartTime() *metav1.Time {
	return &sj.sparkjob.CreationTimestamp
}

func (sj *SparkJob) Age() time.Duration {
	job := sj.sparkjob

	if job.CreationTimestamp.IsZero() {
		return 0
	}
	return metav1.Now().Sub(job.CreationTimestamp.Time)
}

// Get the Job Training Duration
func (sj *SparkJob) Duration() time.Duration {
	sparkjob := sj.sparkjob

	if sparkjob.CreationTimestamp.IsZero() {
		return 0
	}

	if sparkjob.Status.TerminationTime.IsZero() {
		return 0
	}

	//todo
	return sparkjob.Status.TerminationTime.Sub(sparkjob.CreationTimestamp.Time)
}

func (sj *SparkJob) GetJobDashboards(client *kubernetes.Clientset) ([]string, error) {
	urls := []string{}
	dashboardURL, err := dashboard(client, namespace, "kubernetes-dashboard")

	if err != nil {
		log.Debugf("Get dashboard failed due to %v", err)
		// retry for the existing customers, will be deprecated in the future
		dashboardURL, err = dashboard(client, arenaNamespace, "kubernetes-dashboard")
		if err != nil {
			log.Debugf("Get dashboard failed due to %v", err)
		}
	}

	if err != nil {
		log.Debugf("Get dashboard failed due to %v", err)
		// retry for the existing customers, will be deprecated in the future
		dashboardURL, err = dashboard(client, "kube-system", "kubernetes-dashboard")
		if err != nil {
			log.Debugf("Get dashboard failed due to %v", err)
		}
	}

	if dashboardURL == "" {
		return urls, fmt.Errorf("No LOGVIEWER Installed.")
	}

	if len(sj.chiefPod.Spec.Containers) == 0 {
		return urls, fmt.Errorf("spark driver pod is not ready!")
	}

	url := fmt.Sprintf("%s/#!/log/%s/%s/%s?namespace=%s\n",
		dashboardURL,
		sj.chiefPod.Namespace,
		sj.chiefPod.Name,
		sj.chiefPod.Spec.Containers[0].Name,
		sj.chiefPod.Namespace)

	urls = append(urls, url)

	return urls, nil
}

// spark job without gpu supported
func (sj *SparkJob) RequestedGPU() int64 {
	return 0
}

// spark job without gpu supported
func (sj *SparkJob) AllocatedGPU() int64 {
	return 0
}

// Get the hostIP of the driver Pod
func (sj *SparkJob) HostIPOfChief() (hostIP string) {
	hostIP = "N/A"
	if sj.GetStatus() == "RUNNING" {
		hostIP = sj.chiefPod.Status.HostIP
	}
	return hostIP
}

func (sj *SparkJob) Namespace() string {
	return sj.sparkjob.Namespace
}

func NewSparkJobTrainer(client *kubernetes.Clientset) Trainer {
	log.Debugf("Init Spark job trainer")
	sparkjobClient, err := initSparkJobClient()

	if err != nil {
		log.Debugf("unsupported spark job due to %v", err)
		return &SparkJobTrainer{
			trainerType: defaultSparkJobTrainingType,
			enabled:     false,
		}
	}
	// allPods have been cached, we do the same to allSparkJobs
	if useCache {
		ns := namespace
		if allNamespaces {
			ns = metav1.NamespaceAll
		}

		sparkJobList, err := sparkjobClient.SparkoperatorV1beta1().SparkApplications(ns).List(metav1.ListOptions{})
		if err != nil {
			log.Debugf("unsupported sparkJob due to %v", err)
			return &SparkJobTrainer{
				trainerType: defaultSparkJobTrainingType,
				enabled:     false,
			}
		}

		for _, sparkJob := range sparkJobList.Items {
			allSparkJobs = append(allSparkJobs, sparkJob)
		}
	}

	return &SparkJobTrainer{
		sparkjobClient: sparkjobClient,
		client:         client,
		trainerType:    defaultSparkJobTrainingType,
		enabled:        true,
	}
}

// init spark job client
func initSparkJobClient() (sparkjobClientset *versioned.Clientset, err error) {
	if restConfig == nil {
		restConfig, err = clientConfig.ClientConfig()
		if err != nil {
			log.Fatal(err)
			return nil, err
		}
	}

	sparkjobClientset = versioned.NewForConfigOrDie(restConfig)
	return sparkjobClientset, nil
}

// spark job trainer
type SparkJobTrainer struct {
	client         *kubernetes.Clientset
	sparkjobClient *versioned.Clientset
	trainerType    string
	enabled        bool
}

func (st *SparkJobTrainer) Type() string {
	return st.trainerType
}

func (st *SparkJobTrainer) IsSupported(name, ns string) bool {
	if !st.enabled {
		return false
	}

	isSpark := false

	if useCache {
		for _, job := range allSparkJobs {
			if st.isSparkJob(name, ns, job) {
				isSpark = true
				break
			}
		}
	} else {
		sparkJobList, err := st.sparkjobClient.SparkoperatorV1beta1().SparkApplications(ns).List(metav1.ListOptions{
			LabelSelector: fmt.Sprintf("release=%s", name),
		})

		if err != nil {
			log.Debugf("failed to search job %s in namespace %s due to %v", name, ns, err)
		}

		if len(sparkJobList.Items) > 0 {
			isSpark = true
		}
	}

	return isSpark
}

func (st *SparkJobTrainer) isSparkJob(name, ns string, job v1beta1.SparkApplication) bool {
	if val, ok := job.Labels["release"]; ok && (val == name) {
		log.Debugf("the sparkjob %s with labels %s", job.Name, val)
	} else {
		return false
	}

	if val, ok := job.Labels["app"]; ok && (val == "sparkjob") {
		log.Debugf("the sparkjob %s with labels %s is found.", job.Name, val)
	} else {
		return false
	}

	if job.Namespace != ns {
		return false
	}
	return true
}

func (st *SparkJobTrainer) GetTrainingJob(name, namespace string) (job TrainingJob, err error) {
	if len(allSparkJobs) > 0 {
		job, err = st.getTrainingJobFromCache(name, namespace)
	} else {
		job, err = st.getTrainingJob(name, namespace)
	}

	return job, err
}

func (st *SparkJobTrainer) getTrainingJobFromCache(name, namespace string) (job TrainingJob, err error) {
	var (
		sparkJob v1beta1.SparkApplication
	)

	for _, item := range allSparkJobs {
		if st.isSparkJob(name, namespace, item) {
			sparkJob = item
			break
		}
	}

	pods, chiefPod := getPodsOfSparkJob(name, st, allPods)

	return &SparkJob{
		chiefPod:    chiefPod,
		sparkjob:    sparkJob,
		pods:        pods,
		name:        name,
		trainerType: st.Type(),
	}, nil
}

func (st *SparkJobTrainer) getTrainingJob(name, namespace string) (job TrainingJob, err error) {
	var (
		sparkjob v1beta1.SparkApplication
	)

	sparkjobList, err := st.sparkjobClient.SparkoperatorV1beta1().SparkApplications(namespace).List(metav1.ListOptions{
		LabelSelector: fmt.Sprintf("release=%s", name),
	})
	if err != nil {
		return nil, err
	}
	if len(sparkjobList.Items) == 0 {
		return nil, fmt.Errorf("Failed to find the job for %s", name)
	} else {
		sparkjob = sparkjobList.Items[0]
	}

	podList, err := st.client.CoreV1().Pods(namespace).List(metav1.ListOptions{
		TypeMeta: metav1.TypeMeta{
			Kind:       "ListOptions",
			APIVersion: "v1",
		}, LabelSelector: fmt.Sprintf("release=%s", name),
	})

	if err != nil {
		return nil, err
	}

	pods, chiefPod := getPodsOfSparkJob(name, st, podList.Items)

	return &SparkJob{
		sparkjob:    sparkjob,
		chiefPod:    chiefPod,
		pods:        pods,
		name:        name,
		trainerType: st.Type(),
	}, nil
}

func (st *SparkJobTrainer) ListTrainingJobs() (jobs []TrainingJob, err error) {
	jobs = []TrainingJob{}
	jobInfos := []types.TrainingJobInfo{}
	for _, sparkJob := range allSparkJobs {
		jobInfo := types.TrainingJobInfo{}
		log.Debugf("find sparkjob %s in %s", sparkJob.Name, sparkJob.Namespace)
		if val, ok := sparkJob.Labels["release"]; ok && (sparkJob.Name == fmt.Sprintf("%s-%s", val, st.Type())) {
			log.Debugf("the sparkjob %s with labels %s found in List", sparkJob.Name, val)
			jobInfo.Name = val
		} else {
			jobInfo.Name = sparkJob.Name
		}

		jobInfo.Namespace = sparkJob.Namespace
		jobInfos = append(jobInfos, jobInfo)
	}
	log.Debugf("jobInfos %v", jobInfos)

	for _, jobInfo := range jobInfos {
		job, err := st.getTrainingJobFromCache(jobInfo.Name, jobInfo.Namespace)
		if err != nil {
			return jobs, err
		}
		jobs = append(jobs, job)
	}

	return jobs, nil
}

func (st *SparkJobTrainer) isSparkPod(name, ns string, item v1.Pod) bool {
	if val, ok := item.Labels["release"]; ok && (val == name) {
		log.Debugf("the sparkjob %s with labels %s", item.Name, val)
	} else {
		return false
	}

	if val, ok := item.Labels["app"]; ok && (val == "sparkjob") {
		log.Debugf("the sparkjob %s with labels %s is found.", item.Name, val)
	} else {
		return false
	}

	if item.Namespace != ns {
		return false
	}
	return true
}

func (st *SparkJobTrainer) isChiefPod(item v1.Pod) bool {
	if val, ok := item.Labels["spark-role"]; ok && (val == "driver") {
		log.Debugf("the sparkjob %s with labels %s", item.Name, val)
	} else {
		return false
	}
	return true
}

func getPodsOfSparkJob(name string, st *SparkJobTrainer, podList []v1.Pod) (pods []v1.Pod, chiefPod v1.Pod) {
	pods = []v1.Pod{}
	for _, item := range podList {
		if !st.isSparkPod(name, namespace, item) {
			continue
		}
		if st.isChiefPod(item) && item.CreationTimestamp.After(chiefPod.CreationTimestamp.Time) {
			// If there are some failed chiefPod, and the new chiefPod haven't started, set the latest failed pod as chief pod
			if chiefPod.Name != "" && item.Status.Phase == v1.PodPending {
				continue
			}
			chiefPod = item
		}

		// for non-job pod, add it into the pod list
		pods = append(pods, item)
		log.Debugf("add pod %v to pods", item)
	}
	return pods, chiefPod
}
<|endoftext|>


<commit_before>package gorm_test

import (
	"fmt"
	"testing"
	"time"
)

func runMigration() {
	if err := DB.DropTableIfExists(&User{}).Error; err != nil {
		fmt.Printf("Got error when try to delete table users, %+v\n", err)
	}

	for _, table := range []string{"animals", "user_languages"} {
		DB.Exec(fmt.Sprintf("drop table %v;", table))
	}

	values := []interface{}{&Product{}, &Email{}, &Address{}, &CreditCard{}, &Company{}, &Role{}, &Language{}, &HNPost{}, &EngadgetPost{}, &Animal{}, &User{}, &JoinTable{}}
	for _, value := range values {
		DB.DropTable(value)
	}

	if err := DB.AutoMigrate(values...).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
}

func TestIndexes(t *testing.T) {
	if err := DB.Model(&Email{}).AddIndex("idx_email_email", "email").Error; err != nil {
		t.Errorf("Got error when tried to create index: %+v", err)
	}

	scope := DB.NewScope(&Email{})
	if !scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email") {
		t.Errorf("Email should have index idx_email_email")
	}

	if err := DB.Model(&Email{}).RemoveIndex("idx_email_email").Error; err != nil {
		t.Errorf("Got error when tried to remove index: %+v", err)
	}

	if scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email") {
		t.Errorf("Email's index idx_email_email should be deleted")
	}

	if err := DB.Model(&Email{}).AddIndex("idx_email_email_and_user_id", "user_id", "email").Error; err != nil {
		t.Errorf("Got error when tried to create index: %+v", err)
	}

	if !scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email_and_user_id") {
		t.Errorf("Email should have index idx_email_email_and_user_id")
	}

	if err := DB.Model(&Email{}).RemoveIndex("idx_email_email_and_user_id").Error; err != nil {
		t.Errorf("Got error when tried to remove index: %+v", err)
	}

	if scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email_and_user_id") {
		t.Errorf("Email's index idx_email_email_and_user_id should be deleted")
	}

	if err := DB.Model(&Email{}).AddUniqueIndex("idx_email_email_and_user_id", "user_id", "email").Error; err != nil {
		t.Errorf("Got error when tried to create index: %+v", err)
	}

	if !scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email_and_user_id") {
		t.Errorf("Email should have index idx_email_email_and_user_id")
	}

	if DB.Save(&User{Name: "unique_indexes", Emails: []Email{{Email: "user1@example.comiii"}, {Email: "user1@example.com"}, {Email: "user1@example.com"}}}).Error == nil {
		t.Errorf("Should get to create duplicate record when having unique index")
	}

	if err := DB.Model(&Email{}).RemoveIndex("idx_email_email_and_user_id").Error; err != nil {
		t.Errorf("Got error when tried to remove index: %+v", err)
	}

	if scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email_and_user_id") {
		t.Errorf("Email's index idx_email_email_and_user_id should be deleted")
	}

	if DB.Save(&User{Name: "unique_indexes", Emails: []Email{{Email: "user1@example.com"}, {Email: "user1@example.com"}}}).Error != nil {
		t.Errorf("Should be able to create duplicated emails after remove unique index")
	}
}

type BigEmail struct {
	Id           int64
	UserId       int64
	Email        string    `sql:"index:idx_email_agent"`
	UserAgent    string    `sql:"index:idx_email_agent"`
	RegisteredAt time.Time `sql:"unique_index"`
	CreatedAt    time.Time
	UpdatedAt    time.Time
}

func (b BigEmail) TableName() string {
	return "emails"
}

func TestAutoMigration(t *testing.T) {
	DB.AutoMigrate(&Address{})
	if err := DB.Table("emails").AutoMigrate(&BigEmail{}).Error; err != nil {
		t.Errorf("Auto Migrate should not raise any error")
	}

	DB.Save(&BigEmail{Email: "jinzhu@example.org", UserAgent: "pc", RegisteredAt: time.Now()})

	scope := DB.NewScope(&BigEmail{})
	if !scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_agent") {
		t.Errorf("Failed to create index")
	}

	if !scope.Dialect().HasIndex(scope, scope.TableName(), "uix_emails_registered_at") {
		t.Errorf("Failed to create index")
	}

	var bigemail BigEmail
	DB.First(&bigemail, "user_agent = ?", "pc")
	if bigemail.Email != "jinzhu@example.org" || bigemail.UserAgent != "pc" || bigemail.RegisteredAt.IsZero() {
		t.Error("Big Emails should be saved and fetched correctly")
	}
}
<commit_msg>Fix tests on CI<commit_after>package gorm_test

import (
	"fmt"
	"testing"
	"time"
)

func runMigration() {
	if err := DB.DropTableIfExists(&User{}).Error; err != nil {
		fmt.Printf("Got error when try to delete table users, %+v\n", err)
	}

	for _, table := range []string{"animals", "user_languages"} {
		DB.Exec(fmt.Sprintf("drop table %v;", table))
	}

	values := []interface{}{&Product{}, &Email{}, &Address{}, &CreditCard{}, &Company{}, &Role{}, &Language{}, &HNPost{}, &EngadgetPost{}, &Animal{}, &User{}, &JoinTable{}, &Post{}, &Category{}, &Comment{}}
	for _, value := range values {
		DB.DropTable(value)
	}

	if err := DB.AutoMigrate(values...).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
}

func TestIndexes(t *testing.T) {
	if err := DB.Model(&Email{}).AddIndex("idx_email_email", "email").Error; err != nil {
		t.Errorf("Got error when tried to create index: %+v", err)
	}

	scope := DB.NewScope(&Email{})
	if !scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email") {
		t.Errorf("Email should have index idx_email_email")
	}

	if err := DB.Model(&Email{}).RemoveIndex("idx_email_email").Error; err != nil {
		t.Errorf("Got error when tried to remove index: %+v", err)
	}

	if scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email") {
		t.Errorf("Email's index idx_email_email should be deleted")
	}

	if err := DB.Model(&Email{}).AddIndex("idx_email_email_and_user_id", "user_id", "email").Error; err != nil {
		t.Errorf("Got error when tried to create index: %+v", err)
	}

	if !scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email_and_user_id") {
		t.Errorf("Email should have index idx_email_email_and_user_id")
	}

	if err := DB.Model(&Email{}).RemoveIndex("idx_email_email_and_user_id").Error; err != nil {
		t.Errorf("Got error when tried to remove index: %+v", err)
	}

	if scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email_and_user_id") {
		t.Errorf("Email's index idx_email_email_and_user_id should be deleted")
	}

	if err := DB.Model(&Email{}).AddUniqueIndex("idx_email_email_and_user_id", "user_id", "email").Error; err != nil {
		t.Errorf("Got error when tried to create index: %+v", err)
	}

	if !scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email_and_user_id") {
		t.Errorf("Email should have index idx_email_email_and_user_id")
	}

	if DB.Save(&User{Name: "unique_indexes", Emails: []Email{{Email: "user1@example.comiii"}, {Email: "user1@example.com"}, {Email: "user1@example.com"}}}).Error == nil {
		t.Errorf("Should get to create duplicate record when having unique index")
	}

	if err := DB.Model(&Email{}).RemoveIndex("idx_email_email_and_user_id").Error; err != nil {
		t.Errorf("Got error when tried to remove index: %+v", err)
	}

	if scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_email_and_user_id") {
		t.Errorf("Email's index idx_email_email_and_user_id should be deleted")
	}

	if DB.Save(&User{Name: "unique_indexes", Emails: []Email{{Email: "user1@example.com"}, {Email: "user1@example.com"}}}).Error != nil {
		t.Errorf("Should be able to create duplicated emails after remove unique index")
	}
}

type BigEmail struct {
	Id           int64
	UserId       int64
	Email        string    `sql:"index:idx_email_agent"`
	UserAgent    string    `sql:"index:idx_email_agent"`
	RegisteredAt time.Time `sql:"unique_index"`
	CreatedAt    time.Time
	UpdatedAt    time.Time
}

func (b BigEmail) TableName() string {
	return "emails"
}

func TestAutoMigration(t *testing.T) {
	DB.AutoMigrate(&Address{})
	if err := DB.Table("emails").AutoMigrate(&BigEmail{}).Error; err != nil {
		t.Errorf("Auto Migrate should not raise any error")
	}

	DB.Save(&BigEmail{Email: "jinzhu@example.org", UserAgent: "pc", RegisteredAt: time.Now()})

	scope := DB.NewScope(&BigEmail{})
	if !scope.Dialect().HasIndex(scope, scope.TableName(), "idx_email_agent") {
		t.Errorf("Failed to create index")
	}

	if !scope.Dialect().HasIndex(scope, scope.TableName(), "uix_emails_registered_at") {
		t.Errorf("Failed to create index")
	}

	var bigemail BigEmail
	DB.First(&bigemail, "user_agent = ?", "pc")
	if bigemail.Email != "jinzhu@example.org" || bigemail.UserAgent != "pc" || bigemail.RegisteredAt.IsZero() {
		t.Error("Big Emails should be saved and fetched correctly")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"database/sql"
	ggproto "github.com/gogo/protobuf/proto"
	sqlite3 "github.com/mattn/go-sqlite3"
	mcq "github.com/mediachain/concat/mc/query"
	pb "github.com/mediachain/concat/proto"
	"os"
	"path"
)

type SQLDB struct {
	db                 *sql.DB
	insertStmtData     *sql.Stmt
	insertStmtEnvelope *sql.Stmt
	selectStmtData     *sql.Stmt
	deleteStmtData     *sql.Stmt
	deleteStmtEnvelope *sql.Stmt
}

func (sdb *SQLDB) Put(stmt *pb.Statement) error {
	bytes, err := ggproto.Marshal(stmt)
	if err != nil {
		return err
	}

	tx, err := sdb.db.Begin()
	if err != nil {
		return err
	}

	xstmt := tx.Stmt(sdb.insertStmtData)
	_, err = xstmt.Exec(stmt.Id, bytes)
	if err != nil {
		tx.Rollback()
		return err
	}

	xstmt = tx.Stmt(sdb.insertStmtEnvelope)
	// XXX source = publisher only for simple statements
	_, err = xstmt.Exec(stmt.Id, stmt.Namespace, stmt.Publisher, stmt.Publisher, stmt.Timestamp)
	if err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit()
}

func (sdb *SQLDB) PutBatch(stmts []*pb.Statement) error {
	tx, err := sdb.db.Begin()
	if err != nil {
		return err
	}

	insertStmtData := tx.Stmt(sdb.insertStmtData)
	insertStmtEnvelope := tx.Stmt(sdb.insertStmtEnvelope)

	for _, stmt := range stmts {
		bytes, err := ggproto.Marshal(stmt)
		if err != nil {
			tx.Rollback()
			return err
		}

		_, err = insertStmtData.Exec(stmt.Id, bytes)
		if err != nil {
			tx.Rollback()
			return err
		}

		// XXX source = publisher only for simple statements
		_, err = insertStmtEnvelope.Exec(stmt.Id, stmt.Namespace, stmt.Publisher, stmt.Publisher, stmt.Timestamp)
		if err != nil {
			tx.Rollback()
			return err
		}
	}

	return tx.Commit()
}

func (sdb *SQLDB) Get(id string) (*pb.Statement, error) {
	row := sdb.selectStmtData.QueryRow(id)

	var bytes []byte
	err := row.Scan(&bytes)
	if err != nil {
		if err == sql.ErrNoRows {
			err = UnknownStatement
		}

		return nil, err
	}

	stmt := new(pb.Statement)
	err = ggproto.Unmarshal(bytes, stmt)
	if err != nil {
		return nil, err
	}

	return stmt, nil
}

func (sdb *SQLDB) Query(q *mcq.Query) ([]interface{}, error) {
	sq, rsel, err := mcq.CompileQuery(q)
	if err != nil {
		return nil, err
	}

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

	res := make([]interface{}, 0)
	for rows.Next() {
		obj, err := rsel.Scan(rows)
		if err != nil {
			return nil, err
		}
		res = append(res, obj)
	}

	return res, nil
}

func (sdb *SQLDB) QueryStream(ctx context.Context, q *mcq.Query) (<-chan interface{}, error) {
	sq, rsel, err := mcq.CompileQuery(q)
	if err != nil {
		return nil, err
	}

	rows, err := sdb.db.Query(sq)
	if err != nil {
		return nil, err
	}

	ch := make(chan interface{})
	go func() {
		defer close(ch)
		defer rows.Close()

		for rows.Next() {
			obj, err := rsel.Scan(rows)
			if err != nil {
				sendStreamError(ctx, ch, err.Error())
				return
			}

			select {
			case ch <- obj:
				continue
			case <-ctx.Done():
				return
			}
		}
	}()

	return ch, nil
}

func (sdb *SQLDB) QueryOne(q *mcq.Query) (interface{}, error) {
	sq, rsel, err := mcq.CompileQuery(q)
	if err != nil {
		return nil, err
	}

	row := sdb.db.QueryRow(sq)
	res, err := rsel.Scan(row)
	if err != nil {
		if err == sql.ErrNoRows {
			err = NoResult
		}
		return nil, err
	}

	return res, nil
}

func (sdb *SQLDB) Delete(q *mcq.Query) (count int, err error) {
	if q.Op != mcq.OpDelete {
		return 0, BadQuery
	}

	// Delete collects the target ids and deletes in batches, to avoid
	// excessive buffer memory when deleting large sets.
	// It cannot use the natural streaming query solution to delete in a single
	// tx because it deadlocks when connection pooling is disabled.
	// Partial deletes are possible because of an error in some batch,
	// which will result in both count > 0 and the error being returned.

	const batch = 8192 // 1MB worth of ids
	q = q.WithLimit(batch)

loop:
	for {
		xcount := 0

		res, err := sdb.Query(q)
		if err != nil || len(res) == 0 {
			break
		}

		tx, err := sdb.db.Begin()
		if err != nil {
			break
		}

		delData := tx.Stmt(sdb.deleteStmtData)
		delEnvelope := tx.Stmt(sdb.deleteStmtEnvelope)

		for _, id := range res {
			_, err = delData.Exec(id)
			if err != nil {
				tx.Rollback()
				break loop
			}

			_, err = delEnvelope.Exec(id)
			if err != nil {
				tx.Rollback()
				break loop
			}

			xcount += 1
		}

		err = tx.Commit()
		if err != nil {
			break
		}

		count += xcount

		if xcount < batch {
			break
		}
	}

	return count, err
}

func (sdb *SQLDB) Close() error {
	return sdb.db.Close()
}

func (sdb *SQLDB) createTables() error {
	_, err := sdb.db.Exec("CREATE TABLE Statement (id VARCHAR(128) PRIMARY KEY, data VARBINARY)")
	if err != nil {
		return err
	}

	_, err = sdb.db.Exec("CREATE TABLE Envelope (counter INTEGER PRIMARY KEY AUTOINCREMENT, id VARCHAR(128), namespace VARCHAR, publisher VARCHAR, source VARCHAR, timestamp INTEGER)")
	if err != nil {
		return err
	}

	_, err = sdb.db.Exec("CREATE UNIQUE INDEX EnvelopeId ON Envelope (id)")
	if err != nil {
		return err
	}

	_, err = sdb.db.Exec("CREATE INDEX EnvelopeNS ON Envelope (namespace)")
	return err
}

func (sdb *SQLDB) prepareStatements() error {
	stmt, err := sdb.db.Prepare("INSERT INTO Statement VALUES (?, ?)")
	if err != nil {
		return err
	}
	sdb.insertStmtData = stmt

	stmt, err = sdb.db.Prepare("INSERT INTO Envelope VALUES (NULL, ?, ?, ?, ?, ?)")
	if err != nil {
		return err
	}
	sdb.insertStmtEnvelope = stmt

	stmt, err = sdb.db.Prepare("SELECT data FROM Statement WHERE id = ?")
	if err != nil {
		return err
	}
	sdb.selectStmtData = stmt

	stmt, err = sdb.db.Prepare("DELETE FROM Statement WHERE id = ?")
	if err != nil {
		return err
	}
	sdb.deleteStmtData = stmt

	stmt, err = sdb.db.Prepare("DELETE FROM Envelope WHERE id = ?")
	if err != nil {
		return err
	}
	sdb.deleteStmtEnvelope = stmt

	return nil
}

// SQLite backend
type SQLiteDB struct {
	SQLDB
}

func (sdb *SQLiteDB) Open(home string) error {
	var dbpath string
	var mktables bool

	if home == ":memory:" { // allow testing
		dbpath = home
		mktables = true
	} else {
		dbdir := path.Join(home, "stmt")
		err := os.MkdirAll(dbdir, 0755)
		if err != nil {
			return err
		}

		dbpath = path.Join(dbdir, "stmt.db")
		_, err = os.Stat(dbpath)
		switch {
		case os.IsNotExist(err):
			mktables = true
		case err != nil:
			return err
		}
	}

	err := sdb.openDB(dbpath)
	if err != nil {
		return err
	}

	if mktables {
		err = sdb.createTables()
		if err != nil {
			return err
		}

		err = sdb.tuneDB()
		if err != nil {
			return err
		}
	}

	sdb.configPool()

	return sdb.prepareStatements()
}

func (sdb *SQLiteDB) openDB(dbpath string) error {
	db, err := sql.Open("sqlite3", dbpath)
	if err != nil {
		return err
	}

	sdb.db = db
	return nil
}

func (sdb *SQLiteDB) tuneDB() error {
	_, err := sdb.db.Exec("PRAGMA journal_mode=WAL")
	return err
}

func (sdb *SQLiteDB) configPool() {
	// disable connection pooling as lock contention totally kills
	// concurrent write performance
	sdb.db.SetMaxOpenConns(1)
}

func (sdb *SQLiteDB) Merge(stmt *pb.Statement) (bool, error) {
	err := sdb.Put(stmt)
	if err != nil {
		xerr, ok := err.(sqlite3.Error)
		if ok && xerr.ExtendedCode == sqlite3.ErrConstraintPrimaryKey {
			return false, nil
		}
		return false, err
	}
	return true, nil
}
<commit_msg>mcnode/db: marshall statements outside the tx in PutBatch<commit_after>package main

import (
	"context"
	"database/sql"
	ggproto "github.com/gogo/protobuf/proto"
	sqlite3 "github.com/mattn/go-sqlite3"
	mcq "github.com/mediachain/concat/mc/query"
	pb "github.com/mediachain/concat/proto"
	"os"
	"path"
)

type SQLDB struct {
	db                 *sql.DB
	insertStmtData     *sql.Stmt
	insertStmtEnvelope *sql.Stmt
	selectStmtData     *sql.Stmt
	deleteStmtData     *sql.Stmt
	deleteStmtEnvelope *sql.Stmt
}

func (sdb *SQLDB) Put(stmt *pb.Statement) error {
	bytes, err := ggproto.Marshal(stmt)
	if err != nil {
		return err
	}

	tx, err := sdb.db.Begin()
	if err != nil {
		return err
	}

	xstmt := tx.Stmt(sdb.insertStmtData)
	_, err = xstmt.Exec(stmt.Id, bytes)
	if err != nil {
		tx.Rollback()
		return err
	}

	xstmt = tx.Stmt(sdb.insertStmtEnvelope)
	// XXX source = publisher only for simple statements
	_, err = xstmt.Exec(stmt.Id, stmt.Namespace, stmt.Publisher, stmt.Publisher, stmt.Timestamp)
	if err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit()
}

func (sdb *SQLDB) PutBatch(stmts []*pb.Statement) error {
	// marshal statements before starting the tx, minimize critical section
	data := make([][]byte, len(stmts))
	for x, stmt := range stmts {
		bytes, err := ggproto.Marshal(stmt)
		if err != nil {
			return err
		}
		data[x] = bytes
	}

	tx, err := sdb.db.Begin()
	if err != nil {
		return err
	}

	insertStmtData := tx.Stmt(sdb.insertStmtData)
	insertStmtEnvelope := tx.Stmt(sdb.insertStmtEnvelope)

	for x, stmt := range stmts {
		_, err = insertStmtData.Exec(stmt.Id, data[x])
		if err != nil {
			tx.Rollback()
			return err
		}

		// XXX source = publisher only for simple statements
		_, err = insertStmtEnvelope.Exec(stmt.Id, stmt.Namespace, stmt.Publisher, stmt.Publisher, stmt.Timestamp)
		if err != nil {
			tx.Rollback()
			return err
		}
	}

	return tx.Commit()
}

func (sdb *SQLDB) Get(id string) (*pb.Statement, error) {
	row := sdb.selectStmtData.QueryRow(id)

	var bytes []byte
	err := row.Scan(&bytes)
	if err != nil {
		if err == sql.ErrNoRows {
			err = UnknownStatement
		}

		return nil, err
	}

	stmt := new(pb.Statement)
	err = ggproto.Unmarshal(bytes, stmt)
	if err != nil {
		return nil, err
	}

	return stmt, nil
}

func (sdb *SQLDB) Query(q *mcq.Query) ([]interface{}, error) {
	sq, rsel, err := mcq.CompileQuery(q)
	if err != nil {
		return nil, err
	}

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

	res := make([]interface{}, 0)
	for rows.Next() {
		obj, err := rsel.Scan(rows)
		if err != nil {
			return nil, err
		}
		res = append(res, obj)
	}

	return res, nil
}

func (sdb *SQLDB) QueryStream(ctx context.Context, q *mcq.Query) (<-chan interface{}, error) {
	sq, rsel, err := mcq.CompileQuery(q)
	if err != nil {
		return nil, err
	}

	rows, err := sdb.db.Query(sq)
	if err != nil {
		return nil, err
	}

	ch := make(chan interface{})
	go func() {
		defer close(ch)
		defer rows.Close()

		for rows.Next() {
			obj, err := rsel.Scan(rows)
			if err != nil {
				sendStreamError(ctx, ch, err.Error())
				return
			}

			select {
			case ch <- obj:
				continue
			case <-ctx.Done():
				return
			}
		}
	}()

	return ch, nil
}

func (sdb *SQLDB) QueryOne(q *mcq.Query) (interface{}, error) {
	sq, rsel, err := mcq.CompileQuery(q)
	if err != nil {
		return nil, err
	}

	row := sdb.db.QueryRow(sq)
	res, err := rsel.Scan(row)
	if err != nil {
		if err == sql.ErrNoRows {
			err = NoResult
		}
		return nil, err
	}

	return res, nil
}

func (sdb *SQLDB) Delete(q *mcq.Query) (count int, err error) {
	if q.Op != mcq.OpDelete {
		return 0, BadQuery
	}

	// Delete collects the target ids and deletes in batches, to avoid
	// excessive buffer memory when deleting large sets.
	// It cannot use the natural streaming query solution to delete in a single
	// tx because it deadlocks when connection pooling is disabled.
	// Partial deletes are possible because of an error in some batch,
	// which will result in both count > 0 and the error being returned.

	const batch = 8192 // 1MB worth of ids
	q = q.WithLimit(batch)

loop:
	for {
		xcount := 0

		res, err := sdb.Query(q)
		if err != nil || len(res) == 0 {
			break
		}

		tx, err := sdb.db.Begin()
		if err != nil {
			break
		}

		delData := tx.Stmt(sdb.deleteStmtData)
		delEnvelope := tx.Stmt(sdb.deleteStmtEnvelope)

		for _, id := range res {
			_, err = delData.Exec(id)
			if err != nil {
				tx.Rollback()
				break loop
			}

			_, err = delEnvelope.Exec(id)
			if err != nil {
				tx.Rollback()
				break loop
			}

			xcount += 1
		}

		err = tx.Commit()
		if err != nil {
			break
		}

		count += xcount

		if xcount < batch {
			break
		}
	}

	return count, err
}

func (sdb *SQLDB) Close() error {
	return sdb.db.Close()
}

func (sdb *SQLDB) createTables() error {
	_, err := sdb.db.Exec("CREATE TABLE Statement (id VARCHAR(128) PRIMARY KEY, data VARBINARY)")
	if err != nil {
		return err
	}

	_, err = sdb.db.Exec("CREATE TABLE Envelope (counter INTEGER PRIMARY KEY AUTOINCREMENT, id VARCHAR(128), namespace VARCHAR, publisher VARCHAR, source VARCHAR, timestamp INTEGER)")
	if err != nil {
		return err
	}

	_, err = sdb.db.Exec("CREATE UNIQUE INDEX EnvelopeId ON Envelope (id)")
	if err != nil {
		return err
	}

	_, err = sdb.db.Exec("CREATE INDEX EnvelopeNS ON Envelope (namespace)")
	return err
}

func (sdb *SQLDB) prepareStatements() error {
	stmt, err := sdb.db.Prepare("INSERT INTO Statement VALUES (?, ?)")
	if err != nil {
		return err
	}
	sdb.insertStmtData = stmt

	stmt, err = sdb.db.Prepare("INSERT INTO Envelope VALUES (NULL, ?, ?, ?, ?, ?)")
	if err != nil {
		return err
	}
	sdb.insertStmtEnvelope = stmt

	stmt, err = sdb.db.Prepare("SELECT data FROM Statement WHERE id = ?")
	if err != nil {
		return err
	}
	sdb.selectStmtData = stmt

	stmt, err = sdb.db.Prepare("DELETE FROM Statement WHERE id = ?")
	if err != nil {
		return err
	}
	sdb.deleteStmtData = stmt

	stmt, err = sdb.db.Prepare("DELETE FROM Envelope WHERE id = ?")
	if err != nil {
		return err
	}
	sdb.deleteStmtEnvelope = stmt

	return nil
}

// SQLite backend
type SQLiteDB struct {
	SQLDB
}

func (sdb *SQLiteDB) Open(home string) error {
	var dbpath string
	var mktables bool

	if home == ":memory:" { // allow testing
		dbpath = home
		mktables = true
	} else {
		dbdir := path.Join(home, "stmt")
		err := os.MkdirAll(dbdir, 0755)
		if err != nil {
			return err
		}

		dbpath = path.Join(dbdir, "stmt.db")
		_, err = os.Stat(dbpath)
		switch {
		case os.IsNotExist(err):
			mktables = true
		case err != nil:
			return err
		}
	}

	err := sdb.openDB(dbpath)
	if err != nil {
		return err
	}

	if mktables {
		err = sdb.createTables()
		if err != nil {
			return err
		}

		err = sdb.tuneDB()
		if err != nil {
			return err
		}
	}

	sdb.configPool()

	return sdb.prepareStatements()
}

func (sdb *SQLiteDB) openDB(dbpath string) error {
	db, err := sql.Open("sqlite3", dbpath)
	if err != nil {
		return err
	}

	sdb.db = db
	return nil
}

func (sdb *SQLiteDB) tuneDB() error {
	_, err := sdb.db.Exec("PRAGMA journal_mode=WAL")
	return err
}

func (sdb *SQLiteDB) configPool() {
	// disable connection pooling as lock contention totally kills
	// concurrent write performance
	sdb.db.SetMaxOpenConns(1)
}

func (sdb *SQLiteDB) Merge(stmt *pb.Statement) (bool, error) {
	err := sdb.Put(stmt)
	if err != nil {
		xerr, ok := err.(sqlite3.Error)
		if ok && xerr.ExtendedCode == sqlite3.ErrConstraintPrimaryKey {
			return false, nil
		}
		return false, err
	}
	return true, nil
}
<|endoftext|>


<commit_before>package errors

import (
	"fmt"

	"github.com/mondough/typhon/errors/stack"
)

type Error struct {
	Code           int
	Message        string
	PublicContext  map[string]string
	PrivateContext map[string]string
	Stack          stack.Stack
}

// Generic error codes. Each of these has their own constructor for convenience.
// You can use any integer as a code, just use the `New` method.
const (
	ErrUnknown         = 0
	ErrInternalService = 1
	ErrBadRequest      = 2
	ErrBadResponse     = 3
	ErrForbidden       = 4
	ErrUnauthorized    = 5
	ErrNotFound        = 6
	ErrTimeout         = 7
)

// Error returns a string message of the error
// This means the Error implements the error interface
func (p *Error) Error() string {
	if p == nil {
		return ""
	}
	return p.Message
}

// StackString formats the stack as a beautiful string with newlines
func (p *Error) StackString() string {
	stackStr := ""
	for _, frame := range p.Stack {
		stackStr = fmt.Sprintf("%s\n  %s:%d in %s", stackStr, frame.Filename, frame.Line, frame.Method)
	}
	return stackStr
}

// VerboseString returns the error message, stack trace and contexts
func (p *Error) VerboseString() string {
	return fmt.Sprintf("%s\nPublic context: %+v\nPrivate context: %+v\n%s", p.Error(), p.PublicContext, p.PrivateContext, p.StackString())
}

// New creates a new error for you. Use this if you want to pass along a custom error code.
// Otherwise use the handy shorthand factories below
func New(code int, message string, context ...map[string]string) *Error {
	return errorFactory(code, message, context...)
}

// Wrap takes any error interface and wraps it into an Error.
// This is useful because an Error contains lots of useful goodies, like the stacktrace of the error.
// NOTE: If `err` is already an `Error` the passed contexts will be ignored
func Wrap(err error, context ...map[string]string) *Error {
	if err == nil {
		return nil
	}
	switch err := err.(type) {
	case *Error:
		return err
	default:
		return errorFactory(ErrInternalService, err.Error(), context...)
	}
}

// InternalService creates a new error to represent an internal service error.
// Only use internal service error if we know very little about the error. Most
// internal service errors will come from `Wrap`ing a vanilla `error` interface
func InternalService(message string, context ...map[string]string) *Error {
	return errorFactory(ErrInternalService, message, context...)
}

// BadRequest creates a new error to represent an error caused by the client sending
// an invalid request. This is non-retryable unless the request is modified.
func BadRequest(message string, context ...map[string]string) *Error {
	return errorFactory(ErrBadRequest, message, context...)
}

// BadResponse creates a new error representing a failure to response with a valid response
// Examples of this would be a handler returning an invalid message format
func BadResponse(message string, context ...map[string]string) *Error {
	return errorFactory(ErrBadResponse, message, context...)
}

// Timeout creates a new error representing a timeout from client to server
func Timeout(message string, context ...map[string]string) *Error {
	return errorFactory(ErrTimeout, message, context...)
}

// NotFound creates a new error representing a resource that cannot be found. In some
// cases this is not an error, and would be better represented by a zero length slice of elements
func NotFound(message string, context ...map[string]string) *Error {
	return errorFactory(ErrNotFound, message, context...)
}

// Forbidden creates a new error representing a resource that cannot be accessed with
// the current authorisation credentials. The user may need authorising, or if authorised,
// may not be permitted to perform this action
func Forbidden(message string, context ...map[string]string) *Error {
	return errorFactory(ErrForbidden, message, context...)
}

// Unauthorized creates a new error indicating that authentication is required,
// but has either failed or not been provided.
func Unauthorized(message string, context ...map[string]string) *Error {
	return errorFactory(ErrUnauthorized, message, context...)
}

// errorConstructor returns a `*Error` with the specified code, message and context. The main work
// consists of managing the map[string]string's at the end of the arguments list.
// In practice we only ever pass in two: the first one is public and will be sent to the client,
// the second one is private and can contain internal information that is useful for debugging
func errorFactory(code int, message string, context ...map[string]string) *Error {
	err := &Error{
		Code:           code,
		Message:        message,
		PrivateContext: map[string]string{},
		PublicContext:  map[string]string{},
	}
	// The first context map is the PublicContext
	if len(context) > 0 && context[0] != nil {
		err.PublicContext = context[0]
	}
	// The second context map is the privateContext
	if len(context) > 1 && context[1] != nil {
		err.PrivateContext = context[1]
	}

	// Build stack and skip first three lines:
	//  - stack.go BuildStack()
	//  - errors.go errorFactory()
	//  - errors.go public constructor method
	err.Stack = stack.BuildStack(3)

	// ... ignore all remaining map[string]string
	return err
}
<commit_msg>Add json name to Error struct<commit_after>package errors

import (
	"fmt"

	"github.com/mondough/typhon/errors/stack"
)

type Error struct {
	Code           int               `json:"code"`
	Message        string            `json:"message"`
	PublicContext  map[string]string `json:"public_context"`
	PrivateContext map[string]string `json:"private_context"`
	Stack          stack.Stack       `json:"stack"`
}

// Generic error codes. Each of these has their own constructor for convenience.
// You can use any integer as a code, just use the `New` method.
const (
	ErrUnknown         = 0
	ErrInternalService = 1
	ErrBadRequest      = 2
	ErrBadResponse     = 3
	ErrForbidden       = 4
	ErrUnauthorized    = 5
	ErrNotFound        = 6
	ErrTimeout         = 7
)

// Error returns a string message of the error
// This means the Error implements the error interface
func (p *Error) Error() string {
	if p == nil {
		return ""
	}
	return p.Message
}

// StackString formats the stack as a beautiful string with newlines
func (p *Error) StackString() string {
	stackStr := ""
	for _, frame := range p.Stack {
		stackStr = fmt.Sprintf("%s\n  %s:%d in %s", stackStr, frame.Filename, frame.Line, frame.Method)
	}
	return stackStr
}

// VerboseString returns the error message, stack trace and contexts
func (p *Error) VerboseString() string {
	return fmt.Sprintf("%s\nPublic context: %+v\nPrivate context: %+v\n%s", p.Error(), p.PublicContext, p.PrivateContext, p.StackString())
}

// New creates a new error for you. Use this if you want to pass along a custom error code.
// Otherwise use the handy shorthand factories below
func New(code int, message string, context ...map[string]string) *Error {
	return errorFactory(code, message, context...)
}

// Wrap takes any error interface and wraps it into an Error.
// This is useful because an Error contains lots of useful goodies, like the stacktrace of the error.
// NOTE: If `err` is already an `Error` the passed contexts will be ignored
func Wrap(err error, context ...map[string]string) *Error {
	if err == nil {
		return nil
	}
	switch err := err.(type) {
	case *Error:
		return err
	default:
		return errorFactory(ErrInternalService, err.Error(), context...)
	}
}

// InternalService creates a new error to represent an internal service error.
// Only use internal service error if we know very little about the error. Most
// internal service errors will come from `Wrap`ing a vanilla `error` interface
func InternalService(message string, context ...map[string]string) *Error {
	return errorFactory(ErrInternalService, message, context...)
}

// BadRequest creates a new error to represent an error caused by the client sending
// an invalid request. This is non-retryable unless the request is modified.
func BadRequest(message string, context ...map[string]string) *Error {
	return errorFactory(ErrBadRequest, message, context...)
}

// BadResponse creates a new error representing a failure to response with a valid response
// Examples of this would be a handler returning an invalid message format
func BadResponse(message string, context ...map[string]string) *Error {
	return errorFactory(ErrBadResponse, message, context...)
}

// Timeout creates a new error representing a timeout from client to server
func Timeout(message string, context ...map[string]string) *Error {
	return errorFactory(ErrTimeout, message, context...)
}

// NotFound creates a new error representing a resource that cannot be found. In some
// cases this is not an error, and would be better represented by a zero length slice of elements
func NotFound(message string, context ...map[string]string) *Error {
	return errorFactory(ErrNotFound, message, context...)
}

// Forbidden creates a new error representing a resource that cannot be accessed with
// the current authorisation credentials. The user may need authorising, or if authorised,
// may not be permitted to perform this action
func Forbidden(message string, context ...map[string]string) *Error {
	return errorFactory(ErrForbidden, message, context...)
}

// Unauthorized creates a new error indicating that authentication is required,
// but has either failed or not been provided.
func Unauthorized(message string, context ...map[string]string) *Error {
	return errorFactory(ErrUnauthorized, message, context...)
}

// errorConstructor returns a `*Error` with the specified code, message and context. The main work
// consists of managing the map[string]string's at the end of the arguments list.
// In practice we only ever pass in two: the first one is public and will be sent to the client,
// the second one is private and can contain internal information that is useful for debugging
func errorFactory(code int, message string, context ...map[string]string) *Error {
	err := &Error{
		Code:           code,
		Message:        message,
		PrivateContext: map[string]string{},
		PublicContext:  map[string]string{},
	}
	// The first context map is the PublicContext
	if len(context) > 0 && context[0] != nil {
		err.PublicContext = context[0]
	}
	// The second context map is the privateContext
	if len(context) > 1 && context[1] != nil {
		err.PrivateContext = context[1]
	}

	// Build stack and skip first three lines:
	//  - stack.go BuildStack()
	//  - errors.go errorFactory()
	//  - errors.go public constructor method
	err.Stack = stack.BuildStack(3)

	// ... ignore all remaining map[string]string
	return err
}
<|endoftext|>


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

// Package errors 公用的错误信息
package errors

import (
	"golang.org/x/text/message"

	"github.com/caixw/apidoc/internal/locale"
)

// 错误分类
const (
	SyntaxError int8 = iota + 1
	SyntaxWarn
	Other
)

// LocaleError 本地化的错误信息
type LocaleError struct {
	// 保存着错误内容的本地化信息。
	// 仅在返回错误信息时，才会转换成本地化内容。
	MessageKey  message.Reference
	MessageArgs []interface{}
}

// Error 错误信息
type Error struct {
	LocaleError
	Type  int8
	File  string
	Line  int
	Field string
}

// Message 输出的错误信息。
//
// 仅是错误信息，但是不会包含行号等内容。
func (err *Error) Message() string {
	return locale.Sprintf(err.MessageKey, err.MessageArgs...)
}

func (err *Error) Error() string {
	// ErrMessage = "错误信息 %s 位次于 %s:%d 的 %s",
	return locale.Sprintf(locale.ErrMessage, err.Message(), err.File, err.Line, err.Field)
}
<commit_msg>[errors] 为 LocaleError 实现 error 接口<commit_after>// Copyright 2018 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

// Package errors 公用的错误信息
package errors

import (
	"golang.org/x/text/message"

	"github.com/caixw/apidoc/internal/locale"
)

// 错误分类
const (
	SyntaxError int8 = iota + 1
	SyntaxWarn
	Other
)

// LocaleError 本地化的错误信息
type LocaleError struct {
	// 保存着错误内容的本地化信息。
	// 仅在返回错误信息时，才会转换成本地化内容。
	MessageKey  message.Reference
	MessageArgs []interface{}
}

// Error 错误信息
type Error struct {
	LocaleError
	Type  int8
	File  string
	Line  int
	Field string
}

func (err *LocaleError) Error() string {
	return locale.Sprintf(err.MessageKey, err.MessageArgs...)
}

func (err *Error) Error() string {
	msg := err.LocaleError.Error()
	return locale.Sprintf(locale.ErrMessage, msg, err.File, err.Line, err.Field)
}
<|endoftext|>


<commit_before>package class

const (
    //Array Type  atype
    AT_BOOLEAN   = 4
    AT_CHAR      = 5
    AT_FLOAT     = 6
    AT_DOUBLE    = 7
    AT_BYTE      = 8
    AT_SHORT     = 9
    AT_INT       = 10
    AT_LONG      = 11
)

func NewPrimitiveArray(atype uint8, count int32) (*Obj) {
    switch atype {
    case AT_BOOLEAN: return &Obj{make([]uint8, count)}
    case AT_CHAR:    return &Obj{make([]uint16, count)}
    case AT_FLOAT:   return &Obj{make([]float32, count)}
    case AT_DOUBLE:  return &Obj{make([]float64, count)}
    case AT_BYTE:    return &Obj{make([]uint8, count)}
    case AT_SHORT:   return &Obj{make([]int16, count)}
    case AT_INT:     return &Obj{make([]int32, count)}
    case AT_LONG:    return &Obj{make([]int64, count)}
    default: panic("BAD atype!") // todo
    }
}

func NewRefArray(count int32) (*Obj) {
    arr := make([]*Obj, count)
    return &Obj{arr}
}

func ArrayLength(arr *Obj) (int32) {
    // todo
    return 0
}
<commit_msg>arraylength..<commit_after>package class

const (
    //Array Type  atype
    AT_BOOLEAN   = 4
    AT_CHAR      = 5
    AT_FLOAT     = 6
    AT_DOUBLE    = 7
    AT_BYTE      = 8
    AT_SHORT     = 9
    AT_INT       = 10
    AT_LONG      = 11
)

func NewPrimitiveArray(atype uint8, count int32) (*Obj) {
    switch atype {
    case AT_BOOLEAN: return &Obj{make([]int8, count)}
    case AT_CHAR:    return &Obj{make([]uint16, count)}
    case AT_FLOAT:   return &Obj{make([]float32, count)}
    case AT_DOUBLE:  return &Obj{make([]float64, count)}
    case AT_BYTE:    return &Obj{make([]int8, count)}
    case AT_SHORT:   return &Obj{make([]int16, count)}
    case AT_INT:     return &Obj{make([]int32, count)}
    case AT_LONG:    return &Obj{make([]int64, count)}
    default: panic("BAD atype!") // todo
    }
}

func NewRefArray(count int32) (*Obj) {
    arr := make([]*Obj, count)
    return &Obj{arr}
}

func ArrayLength(arr *Obj) (int32) {
    switch arr.fields.(type) {
        case []int8: return int32(len(arr.fields.([]int8)))
        case []int16: return int32(len(arr.fields.([]int16)))
        case []int32: return int32(len(arr.fields.([]int32)))
        case []int64: return int32(len(arr.fields.([]int64)))
        case []uint16: return int32(len(arr.fields.([]uint16)))
        case []float32: return int32(len(arr.fields.([]float32)))
        case []float64: return int32(len(arr.fields.([]float64)))
        case []*Obj: return int32(len(arr.fields.([]*Obj)))
        default: panic("Not array!") // todo
    }
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime/pprof"
	"sort"
	"text/tabwriter"
)

const VERSION = `0.1.2`

var languages = []Language{
	Language{"Thrift", mExt(".thrift"), cComments},

	Language{"C", mExt(".c", ".h"), cComments},
	Language{"C++", mExt(".cc", ".cpp", ".cxx", ".hh", ".hpp", ".hxx"), cComments},
	Language{"Go", mExt(".go"), cComments},
	Language{"Scala", mExt(".scala"), cComments},
	Language{"Java", mExt(".java"), cComments},

	Language{"YACC", mExt(".y"), cComments},
	Language{"Lex", mExt(".l"), cComments},

	Language{"SQL", mExt(".sql"), sqlComments},

	Language{"Haskell", mExt(".hs", ".lhs"), hsComments},

	Language{"Perl", mExt(".pl", ".pm"), shComments},
	Language{"PHP", mExt(".php"), cComments},

	Language{"Shell", mExt(".sh"), shComments},
	Language{"Bash", mExt(".bash"), shComments},

	Language{"Ruby", mExt(".rb"), shComments},
	Language{"Python", mExt(".py"), pyComments},
	Language{"Assembly", mExt(".asm", ".s"), semiComments},
	Language{"Lisp", mExt(".lsp", ".lisp"), semiComments},
	Language{"Scheme", mExt(".scm", ".scheme"), semiComments},

	Language{"Make", mName("makefile", "Makefile", "MAKEFILE"), shComments},
	Language{"CMake", mName("CMakeLists.txt"), shComments},
	Language{"Jam", mName("Jamfile", "Jamrules"), shComments},

	Language{"Markdown", mExt(".md"), noComments},

	Language{"HAML", mExt(".haml"), noComments},
	Language{"SASS", mExt(".sass"), cssComments},
	Language{"SCSS", mExt(".scss"), cssComments},

	Language{"HTML", mExt(".htm", ".html", ".xhtml"), xmlComments},
	Language{"XML", mExt(".xml"), xmlComments},
	Language{"CSS", mExt(".css"), cssComments},
	Language{"JavaScript", mExt(".js"), cComments},
}

type Commenter struct {
	LineComment  string
	StartComment string
	EndComment   string
	Nesting      bool
}

var (
	noComments   = Commenter{"\000", "\000", "\000", false}
	xmlComments  = Commenter{"\000", `<!--`, `-->`, false}
	cComments    = Commenter{`//`, `/*`, `*/`, false}
	cssComments  = Commenter{"\000", `/*`, `*/`, false}
	shComments   = Commenter{`#`, "\000", "\000", false}
	semiComments = Commenter{`;`, "\000", "\000", false}
	hsComments   = Commenter{`--`, `{-`, `-}`, true}
	sqlComments  = Commenter{`--`, "\000", "\000", false}
	pyComments   = Commenter{`#`, `"""`, `"""`, false}
)

type Language struct {
	Namer
	Matcher
	Commenter
}

// TODO work properly with unicode
func (l Language) Update(c []byte, s *Stats) {
	s.FileCount++

	inComment := 0 // this is an int for nesting
	inLComment := false
	blank := true
	lc := []byte(l.LineComment)
	sc := []byte(l.StartComment)
	ec := []byte(l.EndComment)
	lp, sp, ep := 0, 0, 0

	for _, b := range c {
		if inComment == 0 && b == lc[lp] {
			lp++
			if lp == len(lc) {
				inLComment = true
				lp = 0
			}
		} else {
			lp = 0
		}
		if !inLComment && b == sc[sp] {
			sp++
			if sp == len(sc) {
				inComment++
				if inComment > 1 && !l.Nesting {
					inComment = 1
				}
				sp = 0
			}
		} else {
			sp = 0
		}
		if !inLComment && inComment > 0 && b == ec[ep] {
			ep++
			if ep == len(ec) {
				if inComment > 0 {
					inComment--
				}
				ep = 0
			}
		} else {
			ep = 0
		}

		if b != byte(' ') && b != byte('\t') && b != byte('\n') {
			blank = false
		}

		// BUG(srl): lines with comment don't count towards code
		// Note that lines with both code and comment count towards
		// each, but are not counted twice in the total.
		if b == byte('\n') {
			s.TotalLines++
			if inComment > 0 || inLComment {
				inLComment = false
				s.CommentLines++
			} else if blank {
				s.BlankLines++
			} else {
				s.CodeLines++
			}
			blank = true
			continue
		}
	}
}

type Namer string

func (l Namer) Name() string { return string(l) }

type Matcher func(string) bool

func (m Matcher) Match(fname string) bool { return m(fname) }

func mExt(exts ...string) Matcher {
	return func(fname string) bool {
		for _, ext := range exts {
			if ext == path.Ext(fname) {
				return true
			}
		}
		return false
	}
}

func mName(names ...string) Matcher {
	return func(fname string) bool {
		for _, name := range names {
			if name == path.Base(fname) {
				return true
			}
		}
		return false
	}
}

type Stats struct {
	FileCount    int
	TotalLines   int
	CodeLines    int
	BlankLines   int
	CommentLines int
}

var info = map[string]*Stats{}

func handleFileLang(fname string, l Language) {
	i, ok := info[l.Name()]
	if !ok {
		i = &Stats{}
		info[l.Name()] = i
	}
	c, err := ioutil.ReadFile(fname)
	if err != nil {
		fmt.Fprintf(os.Stderr, "  ! %s\n", fname)
		return
	}
	l.Update(c, i)
}

func handleFile(fname string) {
	for _, lang := range languages {
		if lang.Match(fname) {
			handleFileLang(fname, lang)
			return
		}
	}
	// TODO No recognized extension - check for hashbang
}

var files []string

func add(n string) {
	fi, err := os.Stat(n)
	if err != nil {
		goto invalid
	}
	if fi.IsDir() {
		fs, err := ioutil.ReadDir(n)
		if err != nil {
			goto invalid
		}
		for _, f := range fs {
			if f.Name()[0] != '.' {
				add(path.Join(n, f.Name()))
			}
		}
		return
	}
	if fi.Mode()&os.ModeType == 0 {
		files = append(files, n)
		return
	}

	println(fi.Mode())

invalid:
	fmt.Fprintf(os.Stderr, "  ! %s\n", n)
}

type LData []LResult

func (d LData) Len() int { return len(d) }

func (d LData) Less(i, j int) bool {
	if d[i].CodeLines == d[j].CodeLines {
		return d[i].Name > d[j].Name
	}
	return d[i].CodeLines > d[j].CodeLines
}

func (d LData) Swap(i, j int) {
	d[i], d[j] = d[j], d[i]
}

type LResult struct {
	Name         string
	FileCount    int
	CodeLines    int
	CommentLines int
	BlankLines   int
	TotalLines   int
}

func (r *LResult) Add(a LResult) {
	r.FileCount += a.FileCount
	r.CodeLines += a.CodeLines
	r.CommentLines += a.CommentLines
	r.BlankLines += a.BlankLines
	r.TotalLines += a.TotalLines
}

func printJSON() {
	bs, err := json.MarshalIndent(info, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(bs))
}

func printInfo() {
	w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', tabwriter.AlignRight)
	fmt.Fprintln(w, "Language\tFiles\tCode\tComment\tBlank\tTotal\t")
	d := LData([]LResult{})
	total := &LResult{}
	total.Name = "Total"
	for n, i := range info {
		r := LResult{n, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines}
		d = append(d, r)
		total.Add(r)
	}
	d = append(d, *total)
	sort.Sort(d)
	for _, i := range d {
		fmt.Fprintf(w, "%s\t%d\t%d\t%d\t%d\t%d\t\n", i.Name, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines)
	}

	w.Flush()
}

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	useJson    = flag.Bool("json", false, "JSON-format output")
	version    = flag.Bool("V", false, "display version info and exit")
)

func main() {
	flag.Parse()
	if *version {
		fmt.Printf("sloc %s\n", VERSION)
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %s\n", err.Error())
			return
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	args := flag.Args()
	if len(args) == 0 {
		args = append(args, `.`)
	}

	for _, n := range args {
		add(n)
	}

	for _, f := range files {
		handleFile(f)
	}

	if *useJson {
		printJSON()
	} else {
		printInfo()
	}
}
<commit_msg>Adding more TODO<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime/pprof"
	"sort"
	"text/tabwriter"
)

const VERSION = `0.1.2`

var languages = []Language{
	Language{"Thrift", mExt(".thrift"), cComments},

	Language{"C", mExt(".c", ".h"), cComments},
	Language{"C++", mExt(".cc", ".cpp", ".cxx", ".hh", ".hpp", ".hxx"), cComments},
	Language{"Go", mExt(".go"), cComments},
	Language{"Scala", mExt(".scala"), cComments},
	Language{"Java", mExt(".java"), cComments},

	Language{"YACC", mExt(".y"), cComments},
	Language{"Lex", mExt(".l"), cComments},

	Language{"SQL", mExt(".sql"), sqlComments},

	Language{"Haskell", mExt(".hs", ".lhs"), hsComments},

	Language{"Perl", mExt(".pl", ".pm"), perlComments},
	Language{"PHP", mExt(".php"), cComments},

	Language{"Shell", mExt(".sh"), shComments},
	Language{"Bash", mExt(".bash"), shComments},

	Language{"Ruby", mExt(".rb"), shComments},
	Language{"Python", mExt(".py"), pyComments},
	Language{"Assembly", mExt(".asm", ".s"), semiComments},
	Language{"Lisp", mExt(".lsp", ".lisp"), semiComments},
	Language{"Scheme", mExt(".scm", ".scheme"), semiComments},

	Language{"Make", mName("makefile", "Makefile", "MAKEFILE"), shComments},
	Language{"CMake", mName("CMakeLists.txt"), shComments},
	Language{"Jam", mName("Jamfile", "Jamrules"), shComments},

	Language{"Markdown", mExt(".md"), noComments},

	Language{"HAML", mExt(".haml"), noComments},
	Language{"SASS", mExt(".sass"), cssComments},
	Language{"SCSS", mExt(".scss"), cssComments},

	Language{"HTML", mExt(".htm", ".html", ".xhtml"), xmlComments},
	Language{"XML", mExt(".xml"), xmlComments},
	Language{"CSS", mExt(".css"), cssComments},
	Language{"JavaScript", mExt(".js"), cComments},
}

type Commenter struct {
	LineComment  string
	StartComment string
	EndComment   string
	Nesting      bool
}

var (
	noComments   = Commenter{"\000", "\000", "\000", false}
	xmlComments  = Commenter{"\000", `<!--`, `-->`, false}
	cComments    = Commenter{`//`, `/*`, `*/`, false}
	cssComments  = Commenter{"\000", `/*`, `*/`, false}
	shComments   = Commenter{`#`, "\000", "\000", false}
	semiComments = Commenter{`;`, "\000", "\000", false}
	hsComments   = Commenter{`--`, `{-`, `-}`, true}
	sqlComments  = Commenter{`--`, "\000", "\000", false}
	pyComments   = Commenter{`#`, `"""`, `"""`, false}

	// TODO support POD and __END__
	perlComments = Commenter{`#`, "\000", "\000", false}
)

type Language struct {
	Namer
	Matcher
	Commenter
}

// TODO work properly with unicode
func (l Language) Update(c []byte, s *Stats) {
	s.FileCount++

	inComment := 0 // this is an int for nesting
	inLComment := false
	blank := true
	lc := []byte(l.LineComment)
	sc := []byte(l.StartComment)
	ec := []byte(l.EndComment)
	lp, sp, ep := 0, 0, 0

	for _, b := range c {
		if inComment == 0 && b == lc[lp] {
			lp++
			if lp == len(lc) {
				inLComment = true
				lp = 0
			}
		} else {
			lp = 0
		}
		if !inLComment && b == sc[sp] {
			sp++
			if sp == len(sc) {
				inComment++
				if inComment > 1 && !l.Nesting {
					inComment = 1
				}
				sp = 0
			}
		} else {
			sp = 0
		}
		if !inLComment && inComment > 0 && b == ec[ep] {
			ep++
			if ep == len(ec) {
				if inComment > 0 {
					inComment--
				}
				ep = 0
			}
		} else {
			ep = 0
		}

		if b != byte(' ') && b != byte('\t') && b != byte('\n') {
			blank = false
		}

		// BUG(srl): lines with comment don't count towards code
		// Note that lines with both code and comment count towards
		// each, but are not counted twice in the total.
		if b == byte('\n') {
			s.TotalLines++
			if inComment > 0 || inLComment {
				inLComment = false
				s.CommentLines++
			} else if blank {
				s.BlankLines++
			} else {
				s.CodeLines++
			}
			blank = true
			continue
		}
	}
}

type Namer string

func (l Namer) Name() string { return string(l) }

type Matcher func(string) bool

func (m Matcher) Match(fname string) bool { return m(fname) }

func mExt(exts ...string) Matcher {
	return func(fname string) bool {
		for _, ext := range exts {
			if ext == path.Ext(fname) {
				return true
			}
		}
		return false
	}
}

func mName(names ...string) Matcher {
	return func(fname string) bool {
		for _, name := range names {
			if name == path.Base(fname) {
				return true
			}
		}
		return false
	}
}

type Stats struct {
	FileCount    int
	TotalLines   int
	CodeLines    int
	BlankLines   int
	CommentLines int
}

var info = map[string]*Stats{}

func handleFileLang(fname string, l Language) {
	i, ok := info[l.Name()]
	if !ok {
		i = &Stats{}
		info[l.Name()] = i
	}
	c, err := ioutil.ReadFile(fname)
	if err != nil {
		fmt.Fprintf(os.Stderr, "  ! %s\n", fname)
		return
	}
	l.Update(c, i)
}

func handleFile(fname string) {
	for _, lang := range languages {
		if lang.Match(fname) {
			handleFileLang(fname, lang)
			return
		}
	}
	// TODO No recognized extension - check for hashbang
}

var files []string

func add(n string) {
	fi, err := os.Stat(n)
	if err != nil {
		goto invalid
	}
	if fi.IsDir() {
		fs, err := ioutil.ReadDir(n)
		if err != nil {
			goto invalid
		}
		for _, f := range fs {
			if f.Name()[0] != '.' {
				add(path.Join(n, f.Name()))
			}
		}
		return
	}
	if fi.Mode()&os.ModeType == 0 {
		files = append(files, n)
		return
	}

	println(fi.Mode())

invalid:
	fmt.Fprintf(os.Stderr, "  ! %s\n", n)
}

type LData []LResult

func (d LData) Len() int { return len(d) }

func (d LData) Less(i, j int) bool {
	if d[i].CodeLines == d[j].CodeLines {
		return d[i].Name > d[j].Name
	}
	return d[i].CodeLines > d[j].CodeLines
}

func (d LData) Swap(i, j int) {
	d[i], d[j] = d[j], d[i]
}

type LResult struct {
	Name         string
	FileCount    int
	CodeLines    int
	CommentLines int
	BlankLines   int
	TotalLines   int
}

func (r *LResult) Add(a LResult) {
	r.FileCount += a.FileCount
	r.CodeLines += a.CodeLines
	r.CommentLines += a.CommentLines
	r.BlankLines += a.BlankLines
	r.TotalLines += a.TotalLines
}

func printJSON() {
	bs, err := json.MarshalIndent(info, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(bs))
}

func printInfo() {
	w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', tabwriter.AlignRight)
	fmt.Fprintln(w, "Language\tFiles\tCode\tComment\tBlank\tTotal\t")
	d := LData([]LResult{})
	total := &LResult{}
	total.Name = "Total"
	for n, i := range info {
		r := LResult{n, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines}
		d = append(d, r)
		total.Add(r)
	}
	d = append(d, *total)
	sort.Sort(d)
	for _, i := range d {
		fmt.Fprintf(w, "%s\t%d\t%d\t%d\t%d\t%d\t\n", i.Name, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines)
	}

	w.Flush()
}

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	useJson    = flag.Bool("json", false, "JSON-format output")
	version    = flag.Bool("V", false, "display version info and exit")
)

func main() {
	flag.Parse()
	if *version {
		fmt.Printf("sloc %s\n", VERSION)
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %s\n", err.Error())
			return
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	args := flag.Args()
	if len(args) == 0 {
		args = append(args, `.`)
	}

	for _, n := range args {
		add(n)
	}

	for _, f := range files {
		handleFile(f)
	}

	if *useJson {
		printJSON()
	} else {
		printInfo()
	}
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com

// Generate scribe policy checks using data in the Ubuntu CVE tracker
// repository.
package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"regexp"
	"github.com/mozilla/scribe"
	"strings"
	"unicode"
)

type distPatchInfo map[string]string

type cveEntry struct {
	cveID  string
	pkgMap map[string]distPatchInfo
}

type releaseInformation struct {
	identifier string
	lsbmatch   string
	defid      string
}

var releaseList = []releaseInformation{
	{"utopic", "14.10", ""},
	{"vivid", "15.04", ""},
	{"trusty", "14.04", ""},
	{"precise", "12.04", ""},
	{"lucid", "10.04", ""},
}

var entries []cveEntry
var matchFilter *regexp.Regexp

// The hackTranslate* functions perform conversion on name and version
// strings if the package name is either "linux" or begins with "linux-". If
// this is the case, linux in the package name is translated to be
// "linux-image-generic". Additionally, the upload revision element if
// present in the version string is removed. This allows a comparison of
// the version to occur against a known package name (linux-image-generic),
// instead of requiring resolution of whatever the most recent installed
// kernel image package is which will be named linux-image-<ver>-generic.
func hackTranslateName(pkgname string) string {
	if pkgname != "linux" && !strings.HasPrefix(pkgname, "linux-") {
		return pkgname
	}
	return strings.Replace(pkgname, "linux", "linux-image-generic", 1)
}

func hackTranslateVersion(pkgname string, ver string) string {
	if !strings.HasPrefix(pkgname, "linux-image-generic") {
		return ver
	}

	ffunc := func(c rune) bool {
		if c == '-' {
			return true
		}
		return false
	}
	f := strings.FieldsFunc(ver, ffunc)
	if len(f) != 2 {
		return ver
	}
	ret := f[1]
	idx := strings.Index(ret, ".")
	if idx == -1 {
		return ver
	}
	idx2 := idx + 1
	for _, c := range ret[idx2:] {
		if !unicode.IsDigit(c) {
			break
		}
		idx2++
	}
	ret = f[0] + "." + ret[:idx] + ret[idx2:]

	return ret
}

func addReleaseDefinition(o *scribe.Document, rinfo *releaseInformation) {
	identifier := fmt.Sprintf("reldef-%v", rinfo.identifier)
	rinfo.defid = identifier

	obj := scribe.Object{}
	obj.Object = identifier + "-object"
	obj.FileContent.Path = "/etc"
	obj.FileContent.File = "^lsb-release$"
	obj.FileContent.Expression = "DISTRIB_RELEASE=(\\d{1,2}\\.\\d{1,2})"

	test := scribe.Test{}
	test.TestID = identifier + "-test"
	test.Object = obj.Object
	test.EMatch.Value = rinfo.lsbmatch

	o.Tests = append(o.Tests, test)
	o.Objects = append(o.Objects, obj)
}

func addReleaseDefinitions(o *scribe.Document) {
	for x := range releaseList {
		addReleaseDefinition(o, &releaseList[x])
	}
}

func getReleaseDefinition(dist string) string {
	for _, x := range releaseList {
		if dist == x.identifier {
			return x.defid
		}
	}
	return ""
}

func parseEntryFile(fpath string) (ret cveEntry) {
	const (
		_ = iota
		INNER_NONE
		INNER_PATCH
	)
	fd, err := os.Open(fpath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	defer func() {
		fd.Close()
	}()

	scanner := bufio.NewScanner(fd)
	parserMode := INNER_NONE
	curPkgName := ""
	ret.cveID = path.Base(fpath)
	ret.pkgMap = make(map[string]distPatchInfo)
	for scanner.Scan() {
		tokens := strings.Fields(scanner.Text())
		if len(tokens) == 0 {
			parserMode = INNER_NONE
			curPkgName = ""
			continue
		}

		if strings.HasPrefix(tokens[0], "Patches_") {
			parserMode = INNER_PATCH
			curPkgName = strings.TrimPrefix(tokens[0], "Patches_")
			curPkgName = strings.TrimRight(curPkgName, ":")
			curPkgName = hackTranslateName(curPkgName)
			ret.pkgMap[curPkgName] = make(map[string]string)
			continue
		}

		if parserMode == INNER_PATCH {
			if len(tokens) < 2 {
				continue
			}
			idx := strings.Index(tokens[0], "_")
			if idx == -1 {
				continue
			}
			distname := tokens[0][:idx]
			if tokens[1] == "released" && len(tokens) > 2 {
				patchver := tokens[2]
				patchver = strings.Trim(patchver, "()")
				patchver = hackTranslateVersion(curPkgName, patchver)
				ret.pkgMap[curPkgName][distname] = patchver
			}
		}
	}
	if err = scanner.Err(); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	return
}

func loadEntries(dirpath string) {
	dirents, err := ioutil.ReadDir(dirpath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	for _, i := range dirents {
		if !strings.HasPrefix(i.Name(), "CVE-") {
			continue
		}
		if matchFilter != nil {
			if !matchFilter.MatchString(i.Name()) {
				continue
			}
		}
		fname := path.Join(dirpath, i.Name())
		entries = append(entries, parseEntryFile(fname))
	}
}

func addDefinition(o *scribe.Document, prefix string, pkgname string, dist string, cve cveEntry) {
	// Don't create a definition for anything that is not in our release
	// list.
	reldefid := getReleaseDefinition(dist)
	if reldefid == "" {
		return
	}

	// Create an object definition for the package
	objid := fmt.Sprintf("%v-object", prefix)
	obj := scribe.Object{}
	obj.Object = objid
	obj.Package.Name = pkgname

	// Create a test
	testid := fmt.Sprintf("%v-test", prefix)
	test := scribe.Test{}
	test.TestID = testid
	test.Object = obj.Object
	test.EVR.Value = cve.pkgMap[pkgname][dist]
	test.EVR.Operation = "<"
	disttestref := fmt.Sprintf("reldef-%v-test", dist)
	test.If = append(test.If, disttestref)

	o.Tests = append(o.Tests, test)
	o.Objects = append(o.Objects, obj)
}

func processEntries() {
	root := scribe.Document{}

	addReleaseDefinitions(&root)

	for i, ent := range entries {
		for x := range ent.pkgMap {
			for y := range ent.pkgMap[x] {
				prefix := fmt.Sprintf("ubuntu-%v-%v-%v", i, x, y)
				addDefinition(&root, prefix, x, y, ent)
			}
		}
	}

	buf, err := json.MarshalIndent(&root, "", "    ")
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}
	fmt.Fprintf(os.Stdout, "%v\n", string(buf))
}

func main() {
	var fMatch string

	flag.StringVar(&fMatch, "i", "", "filter regexp")
	flag.Parse()
	args := flag.Args()
	if len(args) < 1 {
		fmt.Fprintf(os.Stderr, "specify path to ubuntu-cve-tracker directory\n")
		os.Exit(1)
	}

	fm, err := regexp.Compile(fMatch)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	matchFilter = fm

	entries = make([]cveEntry, 0)

	procdir := path.Join(args[0], "active")
	loadEntries(procdir)
	procdir = path.Join(args[0], "retired")
	loadEntries(procdir)
	processEntries()
}
<commit_msg>remove ubuntu-cve-tracker<commit_after><|endoftext|>