<commit_before>package state

import (
	"database/sql/driver"
	"encoding/json"
	"fmt"
	"github.com/stitchfix/flotilla-os/clients/metrics"

	"github.com/jmoiron/sqlx"

	// Pull in postgres specific drivers
	"database/sql"
	"math"
	"strings"
	"time"

	"github.com/lib/pq"
	_ "github.com/lib/pq"
	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/exceptions"
	"go.uber.org/multierr"
	sqltrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql"
	sqlxtrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/jmoiron/sqlx"
)

//
// SQLStateManager uses postgresql to manage state
//
type SQLStateManager struct {
	db         *sqlx.DB
	readonlyDB *sqlx.DB
}

func (sm *SQLStateManager) ListFailingNodes() (NodeList, error) {
	var err error
	var nodeList NodeList

	err = sm.readonlyDB.Select(&nodeList, ListFailingNodesSQL)

	if err != nil {
		if err == sql.ErrNoRows {
			return nodeList, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Error fetching node list")}
		} else {
			return nodeList, errors.Wrapf(err, "Error fetching node list")
		}
	}
	return nodeList, err
}

func (sm *SQLStateManager) GetPodReAttemptRate() (float32, error) {
	var err error
	attemptRate := float32(1.0)
	err = sm.readonlyDB.Get(&attemptRate, PodReAttemptRate)

	if err != nil {
		if err == sql.ErrNoRows {
			return attemptRate, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Error fetching attempt rate")}
		} else {
			return attemptRate, errors.Wrapf(err, "Error fetching attempt rate")
		}
	}
	return attemptRate, err
}

func (sm *SQLStateManager) GetNodeLifecycle(executableID string, commandHash string) (string, error) {
	var err error
	nodeType := "spot"
	err = sm.readonlyDB.Get(&nodeType, TaskResourcesExecutorNodeLifecycleSQL, executableID, commandHash)

	if err != nil {
		if err == sql.ErrNoRows {
			return nodeType, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Error fetching node type")}
		} else {
			return nodeType, errors.Wrapf(err, "Error fetching node type")
		}
	}
	return nodeType, err
}

func (sm *SQLStateManager) GetTaskHistoricalRuntime(executableID string, runID string) (float32, error) {
	var err error
	minutes := float32(1.0)
	err = sm.readonlyDB.Get(&minutes, TaskExecutionRuntimeCommandSQL, executableID, runID)

	if err != nil {
		if err == sql.ErrNoRows {
			return minutes, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Error fetching TaskRuntime rate")}
		} else {
			return minutes, errors.Wrapf(err, "Error fetching attempt rate")
		}
	}
	return minutes, err
}

func (sm *SQLStateManager) EstimateRunResources(executableID string, runID string) (TaskResources, error) {
	var err error
	var taskResources TaskResources

	err = sm.readonlyDB.Get(&taskResources, TaskResourcesSelectCommandSQL, executableID, runID)

	if err != nil {
		if err == sql.ErrNoRows {
			return taskResources, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Resource usage with executable %s not found", executableID)}
		} else {
			return taskResources, errors.Wrapf(err, "issue getting resources with executable [%s]", executableID)
		}
	}
	return taskResources, err
}

func (sm *SQLStateManager) EstimateExecutorCount(executableID string, commandHash string) (int64, error) {
	var err error
	executorCount := int64(25)
	err = sm.readonlyDB.Get(&executorCount, TaskResourcesExecutorCountSQL, executableID, commandHash)

	if err != nil {
		if err == sql.ErrNoRows {
			return executorCount, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Resource usage with executable %s not found", executableID)}
		} else {
			return executorCount, errors.Wrapf(err, "issue getting resources with executable [%s]", executableID)
		}
	}
	return executorCount, err
}

func (sm *SQLStateManager) ExecutorOOM(executableID string, commandHash string) (bool, error) {
	var err error
	executorOOM := false
	err = sm.readonlyDB.Get(&executorOOM, TaskResourcesExecutorOOMSQL, executableID, commandHash)

	if err != nil {
		if err == sql.ErrNoRows {
			return executorOOM, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Resource oom for executable %s not found", executableID)}
		} else {
			return executorOOM, errors.Wrapf(err, "issue getting resources with executable [%s]", executableID)
		}
	}
	return executorOOM, err
}

func (sm *SQLStateManager) DriverOOM(executableID string, commandHash string) (bool, error) {
	var err error
	driverOOM := false
	err = sm.readonlyDB.Get(&driverOOM, TaskResourcesDriverOOMSQL, executableID, commandHash)

	if err != nil {
		if err == sql.ErrNoRows {
			return driverOOM, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Resource oom for driver %s not found", executableID)}
		} else {
			return driverOOM, errors.Wrapf(err, "issue getting resources with executable [%s]", executableID)
		}
	}
	return driverOOM, err
}

//
// Name is the name of the state manager - matches value in configuration
//
func (sm *SQLStateManager) Name() string {
	return "postgres"
}

//
// likeFields are the set of fields
// that are filtered using a `like` clause
//
var likeFields = map[string]bool{
	"image":       true,
	"alias":       true,
	"group_name":  true,
	"command":     true,
	"text":        true,
	"exit_reason": true,
}

//
// Initialize creates tables if they do not exist
//
func (sm *SQLStateManager) Initialize(conf config.Config) error {
	dburl := conf.GetString("database_url")
	readonlyDbUrl := conf.GetString("readonly_database_url")

	createSchema := conf.GetBool("create_database_schema")
	sqltrace.Register("postgres", &pq.Driver{})
	var err error
	if sm.db, err = sqlxtrace.Open("postgres", dburl); err != nil {
		return errors.Wrap(err, "unable to open postgres db")
	}

	if sm.readonlyDB, err = sqlxtrace.Open("postgres", readonlyDbUrl); err != nil {
		return errors.Wrap(err, "unable to open readonly postgres db")
	}

	if conf.IsSet("database_max_idle_connections") {
		sm.db.SetMaxIdleConns(conf.GetInt("database_max_idle_connections"))
		sm.readonlyDB.SetMaxIdleConns(conf.GetInt("database_max_idle_connections"))
	}

	if createSchema {
		// Since this happens at initialization we<commit_msg>update tracing tag<commit_after>package state

import (
	"database/sql/driver"
	"encoding/json"
	"fmt"
	"github.com/stitchfix/flotilla-os/clients/metrics"

	"github.com/jmoiron/sqlx"

	// Pull in postgres specific drivers
	"database/sql"
	"math"
	"strings"
	"time"

	"github.com/lib/pq"
	_ "github.com/lib/pq"
	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/exceptions"
	"go.uber.org/multierr"
	sqltrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql"
	sqlxtrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/jmoiron/sqlx"
)

//
// SQLStateManager uses postgresql to manage state
//
type SQLStateManager struct {
	db         *sqlx.DB
	readonlyDB *sqlx.DB
}

func (sm *SQLStateManager) ListFailingNodes() (NodeList, error) {
	var err error
	var nodeList NodeList

	err = sm.readonlyDB.Select(&nodeList, ListFailingNodesSQL)

	if err != nil {
		if err == sql.ErrNoRows {
			return nodeList, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Error fetching node list")}
		} else {
			return nodeList, errors.Wrapf(err, "Error fetching node list")
		}
	}
	return nodeList, err
}

func (sm *SQLStateManager) GetPodReAttemptRate() (float32, error) {
	var err error
	attemptRate := float32(1.0)
	err = sm.readonlyDB.Get(&attemptRate, PodReAttemptRate)

	if err != nil {
		if err == sql.ErrNoRows {
			return attemptRate, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Error fetching attempt rate")}
		} else {
			return attemptRate, errors.Wrapf(err, "Error fetching attempt rate")
		}
	}
	return attemptRate, err
}

func (sm *SQLStateManager) GetNodeLifecycle(executableID string, commandHash string) (string, error) {
	var err error
	nodeType := "spot"
	err = sm.readonlyDB.Get(&nodeType, TaskResourcesExecutorNodeLifecycleSQL, executableID, commandHash)

	if err != nil {
		if err == sql.ErrNoRows {
			return nodeType, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Error fetching node type")}
		} else {
			return nodeType, errors.Wrapf(err, "Error fetching node type")
		}
	}
	return nodeType, err
}

func (sm *SQLStateManager) GetTaskHistoricalRuntime(executableID string, runID string) (float32, error) {
	var err error
	minutes := float32(1.0)
	err = sm.readonlyDB.Get(&minutes, TaskExecutionRuntimeCommandSQL, executableID, runID)

	if err != nil {
		if err == sql.ErrNoRows {
			return minutes, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Error fetching TaskRuntime rate")}
		} else {
			return minutes, errors.Wrapf(err, "Error fetching attempt rate")
		}
	}
	return minutes, err
}

func (sm *SQLStateManager) EstimateRunResources(executableID string, runID string) (TaskResources, error) {
	var err error
	var taskResources TaskResources

	err = sm.readonlyDB.Get(&taskResources, TaskResourcesSelectCommandSQL, executableID, runID)

	if err != nil {
		if err == sql.ErrNoRows {
			return taskResources, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Resource usage with executable %s not found", executableID)}
		} else {
			return taskResources, errors.Wrapf(err, "issue getting resources with executable [%s]", executableID)
		}
	}
	return taskResources, err
}

func (sm *SQLStateManager) EstimateExecutorCount(executableID string, commandHash string) (int64, error) {
	var err error
	executorCount := int64(25)
	err = sm.readonlyDB.Get(&executorCount, TaskResourcesExecutorCountSQL, executableID, commandHash)

	if err != nil {
		if err == sql.ErrNoRows {
			return executorCount, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Resource usage with executable %s not found", executableID)}
		} else {
			return executorCount, errors.Wrapf(err, "issue getting resources with executable [%s]", executableID)
		}
	}
	return executorCount, err
}

func (sm *SQLStateManager) ExecutorOOM(executableID string, commandHash string) (bool, error) {
	var err error
	executorOOM := false
	err = sm.readonlyDB.Get(&executorOOM, TaskResourcesExecutorOOMSQL, executableID, commandHash)

	if err != nil {
		if err == sql.ErrNoRows {
			return executorOOM, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Resource oom for executable %s not found", executableID)}
		} else {
			return executorOOM, errors.Wrapf(err, "issue getting resources with executable [%s]", executableID)
		}
	}
	return executorOOM, err
}

func (sm *SQLStateManager) DriverOOM(executableID string, commandHash string) (bool, error) {
	var err error
	driverOOM := false
	err = sm.readonlyDB.Get(&driverOOM, TaskResourcesDriverOOMSQL, executableID, commandHash)

	if err != nil {
		if err == sql.ErrNoRows {
			return driverOOM, exceptions.MissingResource{
				ErrorString: fmt.Sprintf("Resource oom for driver %s not found", executableID)}
		} else {
			return driverOOM, errors.Wrapf(err, "issue getting resources with executable [%s]", executableID)
		}
	}
	return driverOOM, err
}

//
// Name is the name of the state manager - matches value in configuration
//
func (sm *SQLStateManager) Name() string {
	return "postgres"
}

//
// likeFields are the set of fields
// that are filtered using a `like` clause
//
var likeFields = map[string]bool{
	"image":       true,
	"alias":       true,
	"group_name":  true,
	"command":     true,
	"text":        true,
	"exit_reason": true,
}

//
// Initialize creates tables if they do not exist
//
func (sm *SQLStateManager) Initialize(conf config.Config) error {
	dburl := conf.GetString("database_url")
	readonlyDbUrl := conf.GetString("readonly_database_url")

	createSchema := conf.GetBool("create_database_schema")
	sqltrace.Register("postgres", &pq.Driver{}, sqltrace.WithServiceName("flotilla"))
	var err error
	if sm.db, err = sqlxtrace.Open("postgres", dburl); err != nil {
		return errors.Wrap(err, "unable to open postgres db")
	}

	sqltrace.Register("postgres", &pq.Driver{}, sqltrace.WithServiceName("flotilla"))
	if sm.readonlyDB, err = sqlxtrace.Open("postgres", readonlyDbUrl); err != nil {
		return errors.Wrap(err, "unable to open readonly postgres db")
	}

	if conf.IsSet("database_max_idle_connections") {
		sm.db.SetMaxIdleConns(conf.GetInt("database_max_idle_connections"))
		sm.readonlyDB.SetMaxIdleConns(conf.GetInt("database_max_idle_connections"))
	}

	if createSchema {
		// Since this happens at initialization we
{'fraction_non_alphanumeric': 0.11520514883346741, 'fraction_numerical': 0.003218020917135961, 'mean_word_length': 4.619801084990958, 'pattern_counts': {'":': 12, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7359625', 'n_tokens_mistral': 4325, 'n_tokens_neox': 4017, 'n_words': 1055}