Add parameter startup_sql (justwatchcom#22)

README.md

@@ -105,6 +105,11 @@ jobs:
   # each query will be executed on each connection
   connections:
   - 'postgres://postgres@localhost/postgres?sslmode=disable'
+  # startup_sql is an array of SQL statements
+  # each statements is executed once after connecting
+  startup_sql:
+  - 'SET lock_timeout = 1000'
+  - 'SET idle_in_transaction_session_timeout = 100'
   # queries is a map of Metric/Query mappings
   queries:
     # name is prefied with sql_ and used as the metric name

config.go

@@ -50,6 +50,7 @@ type Job struct {
 	Interval    time.Duration `yaml:"interval"`  // interval at which this job is run
 	Connections []string      `yaml:"connections"`
 	Queries     []*Query      `yaml:"queries"`
+	StartupSQL  []string      `yaml:"startup_sql"` // SQL executed on startup
 }
 
 type connection struct {

config.yml.dist

@@ -4,6 +4,9 @@ jobs:
   interval: '5m'
   connections:
   - 'postgres://postgres@localhost/postgres?sslmode=disable'
+  startup_sql:
+  - 'SET lock_timeout = 1000'
+  - 'SET idle_in_transaction_session_timeout = 100'
   queries:
   - name: "running_queries"
     help: "Number of running queries"
@@ -31,7 +34,7 @@ jobs:
     values:
       - "replication_lag"
     query:  |
-            WITH lag AS ( 
+            WITH lag AS (
             SELECT
             CASE
             WHEN pg_last_xlog_receive_location() = pg_last_xlog_replay_location() THEN 0

job.go

@@ -128,7 +128,7 @@ func (j *Job) runOnceConnection(conn *connection, done chan int) {
 	}()
 
 	// connect to DB if not connected already
-	if err := conn.connect(j.Interval); err != nil {
+	if err := conn.connect(j); err != nil {
 		level.Warn(j.log).Log("msg", "Failed to connect", "err", err)
 		return
 	}
@@ -173,7 +173,7 @@ func (j *Job) runOnce() error {
 	return nil
 }
 
-func (c *connection) connect(iv time.Duration) error {
+func (c *connection) connect(job *Job) error {
 	// already connected
 	if c.conn != nil {
 		return nil
@@ -192,7 +192,14 @@ func (c *connection) connect(iv time.Duration) error {
 	// be nice and don't use up too many connections for mere metrics
 	conn.SetMaxOpenConns(1)
 	conn.SetMaxIdleConns(1)
-	conn.SetConnMaxLifetime(iv * 2)
+	conn.SetConnMaxLifetime(job.Interval * 2)
+
+	// execute StartupSQL
+	for _, query := range job.StartupSQL {
+		level.Debug(job.log).Log("msg", "StartupSQL", "Query:", query)
+		conn.MustExec(query)
+	}
+
 	c.conn = conn
 	return nil
 }