Initial pgterminate code

base/config.go

@@ -0,0 +1,97 @@
+package base
+
+import (
+	"fmt"
+	"gopkg.in/yaml.v2"
+	"io/ioutil"
+	"log"
+	"path/filepath"
+	"strings"
+	"sync"
+)
+
+// AppName exposes application name to config module
+var AppName string
+
+// Config receives configuration options
+type Config struct {
+	mutex          sync.Mutex
+	File           string
+	Host           string  `yaml:"host"`
+	Port           int     `yaml:"port"`
+	User           string  `yaml:"user"`
+	Password       string  `yaml:"password"`
+	Database       string  `yaml:"database"`
+	Interval       float64 `yaml:"interval"`
+	ConnectTimeout int     `yaml:"connect-timeout"`
+	IdleTimeout    float64 `yaml:"idle-timeout"`
+	ActiveTimeout  float64 `yaml:"active-timeout"`
+	LogFile        string  `yaml:"log-file"`
+	PidFile        string  `yaml:"pid-file"`
+}
+
+func init() {
+	AppName = "pgterminate"
+}
+
+// NewConfig creates a Config object
+func NewConfig() *Config {
+	return &Config{}
+}
+
+// Read loads options from a configuration file to Config
+func (c *Config) Read(file string) error {
+	file, err := filepath.Abs(file)
+	if err != nil {
+		return err
+	}
+
+	yamlFile, err := ioutil.ReadFile(file)
+	if err != nil {
+		return err
+	}
+
+	err = yaml.Unmarshal(yamlFile, &c)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Reload reads from file and update configuration
+func (c *Config) Reload() {
+	log.Println("Reloading configuration")
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+	if c.File != "" {
+		c.Read(c.File)
+	}
+}
+
+// Dsn formats a connection string based on Config
+func (c *Config) Dsn() string {
+	var parameters []string
+	if c.Host != "" {
+		parameters = append(parameters, fmt.Sprintf("host=%s", c.Host))
+	}
+	if c.Port != 0 {
+		parameters = append(parameters, fmt.Sprintf("port=%d", c.Port))
+	}
+	if c.User != "" {
+		parameters = append(parameters, fmt.Sprintf("user=%s", c.User))
+	}
+	if c.Password != "" {
+		parameters = append(parameters, fmt.Sprintf("password=%s", c.Password))
+	}
+	if c.Database != "" {
+		parameters = append(parameters, fmt.Sprintf("database=%s", c.Database))
+	}
+	if c.ConnectTimeout != 0 {
+		parameters = append(parameters, fmt.Sprintf("connect_timeout=%d", c.ConnectTimeout))
+	}
+	if AppName != "" {
+		parameters = append(parameters, fmt.Sprintf("application_name=%s", AppName))
+	}
+	return strings.Join(parameters, " ")
+}

base/context.go

@@ -0,0 +1,17 @@
+package base
+
+// Context stores dynamic values like channels and exposes configuration
+type Context struct {
+	Sessions chan Session
+	Done     chan bool
+	Config   *Config
+}
+
+// NewContext instanciates a Context
+func NewContext(config *Config, sessions chan Session, done chan bool) *Context {
+	return &Context{
+		Config:   config,
+		Sessions: sessions,
+		Done:     done,
+	}
+}

base/db.go

@@ -0,0 +1,76 @@
+package base
+
+import (
+	"database/sql"
+	"github.com/lib/pq"
+	"strconv"
+)
+
+const (
+	maxQueryLength = 1000
+)
+
+// Db centralizes connection to the database
+type Db struct {
+	dsn  string
+	conn *sql.DB
+}
+
+// NewDb creates a Db object
+func NewDb(dsn string) *Db {
+	return &Db{
+		dsn: dsn,
+	}
+}
+
+// Connect connects to the instance and ping it to ensure connection is working
+func (db *Db) Connect() {
+	conn, err := sql.Open("postgres", db.dsn)
+	Panic(err)
+
+	err = conn.Ping()
+	Panic(err)
+
+	db.conn = conn
+}
+
+// Disconnect ends connection cleanly
+func (db *Db) Disconnect() {
+	err := db.conn.Close()
+	Panic(err)
+}
+
+// Sessions connects to the database and returns current sessions
+func (db *Db) Sessions() (sessions []Session) {
+	query := `select pid as pid, usename as user, datname as db, host(client_addr)::text || ':' || client_port::text as client, state as state, substring(query from 1 for ` + strconv.Itoa(maxQueryLength) + `) as query, coalesce(extract(epoch from now() - backend_start), 0) as "backendDuration", coalesce(extract(epoch from now() - xact_start), 0) as "xactDuration", coalesce(extract(epoch from now() - query_start), 0) as "queryDuration" from pg_catalog.pg_stat_activity where pid <> pg_backend_pid();`
+	rows, err := db.conn.Query(query)
+	Panic(err)
+	defer rows.Close()
+
+	for rows.Next() {
+		var pid sql.NullInt64
+		var user, db, client, state, query sql.NullString
+		var backendDuration, xactDuration, queryDuration float64
+		err := rows.Scan(&pid, &user, &db, &client, &state, &query, &backendDuration, &xactDuration, &queryDuration)
+		Panic(err)
+
+		if pid.Valid && user.Valid && db.Valid && client.Valid && state.Valid && query.Valid {
+			sessions = append(sessions, NewSession(pid.Int64, user.String, db.String, client.String, state.String, query.String, backendDuration, xactDuration, queryDuration))
+		}
+	}
+
+	return sessions
+}
+
+// TerminateSessions terminates a list of sessions
+func (db *Db) TerminateSessions(sessions []Session) {
+	var pids []int64
+	for _, session := range sessions {
+		pids = append(pids, session.Pid)
+	}
+	if len(pids) > 0 {
+		query := `select pg_terminate_backend(pid) from pg_stat_activity where pid = any($1);`
+		_, err := db.conn.Exec(query, pq.Array(pids))
+		Panic(err)
+	}
+}

base/session.go

@@ -0,0 +1,67 @@
+package base
+
+import (
+	"fmt"
+	"strings"
+)
+
+// Session represents a PostgreSQL backend
+type Session struct {
+	Pid             int64
+	User            string
+	Db              string
+	Client          string
+	State           string
+	Query           string
+	BackendDuration float64
+	XactDuration    float64
+	QueryDuration   float64
+}
+
+// NewSession instanciates a Session
+func NewSession(pid int64, user string, db string, client string, state string, query string, backendDuration float64, xactDuration float64, queryDuration float64) Session {
+	return Session{
+		Pid:             pid,
+		User:            user,
+		Db:              db,
+		Client:          client,
+		State:           state,
+		Query:           query,
+		BackendDuration: backendDuration,
+		XactDuration:    xactDuration,
+		QueryDuration:   queryDuration,
+	}
+}
+
+// String represents a Session as a string
+func (s Session) String() string {
+	var output []string
+	if s.Pid != 0 {
+		output = append(output, fmt.Sprintf("pid=%d", s.Pid))
+	}
+	if s.User != "" {
+		output = append(output, fmt.Sprintf("user=%s", s.User))
+	}
+	if s.Db != "" {
+		output = append(output, fmt.Sprintf("db=%s", s.Db))
+	}
+	if s.Client != "" {
+		output = append(output, fmt.Sprintf("client=%s", s.Client))
+	}
+	if s.State != "" {
+		output = append(output, fmt.Sprintf("state=%s", s.State))
+	}
+	if s.BackendDuration != 0 {
+		output = append(output, fmt.Sprintf("backend_duration=%f", s.BackendDuration))
+	}
+	if s.XactDuration != 0 {
+		output = append(output, fmt.Sprintf("xact_duration=%f", s.XactDuration))
+	}
+	if s.QueryDuration != 0 {
+		output = append(output, fmt.Sprintf("query_duration=%f", s.QueryDuration))
+	}
+	if s.Query != "" {
+		output = append(output, fmt.Sprintf("query=%s", s.Query))
+	}
+	return strings.Join(output, " ")
+}

base/utils.go

@@ -0,0 +1,12 @@
+package base
+
+import (
+	"log"
+)
+
+// Panic prints a non-nil error and terminates the program
+func Panic(err error) {
+	if err != nil {
+		log.Fatalln(err)
+	}
+}