Initial pgterminate code

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)
+	}
+}