Feature/Ability to provide custom instance queries

README.md

@@ -19,6 +19,7 @@ Perfect for both small deployments and large-scale infrastructures, it seamlessl
 - [Requirements](#requirements)
 - [Installation](#installation)
 - [Usage](#usage)
+  - [Supported connections](#supported-connections)
   - [Configuring](#configuring)
   - [Starting server](#starting-server)
   - [Stopping server](#stopping-server)
@@ -62,10 +63,22 @@ import heartbeat "github.com/obstools/go-prometheus-heartbeat-exporter"
 
 ## Usage
 
+- [Supported connections](#supported-connections)
 - [Configuring](#configuring)
 - [Starting server](#starting-server)
 - [Stopping server](#stopping-server)
 
+### Supported connections
+
+Instances in the `heartbeat` configuration allow you to monitor various types of connections. Each instance can be configured with specific attributes such as `name`, `connection`, `url`, `query`, `interval`, and `timeout`. The `connection` attribute specifies the type of connection to be monitored, such as `postgres` for PostgreSQL, etc.
+By providing a `query`, you can define specific operations to be executed on the database or other infrastructure elements, enabling you to check not only the connection status but also the performance of specific queries. This flexibility allows for comprehensive monitoring of your services and ensures that you can quickly identify and respond to issues as they arise.
+
+> [!NOTE] Please make sure that your query is idempotent, as it can be executed multiple times during the `heartbeat` check.
+
+| Connection | Description | Query example |
+| --- | --- | --- |
+| `postgres` | Postgres database connection. If `query` is not provided, `heartbeat` will check if connection is established only. | `CREATE TABLE tmp (id SERIAL PRIMARY KEY); DROP TABLE tmp` |
+
 ### Configuring
 
 `heartbeat` configuration is available as YAML file. You can also use environment variable interpolation in your configuration. This allows you to set sensitive information, such as database URLs or API keys, as environment variables and reference them in your YAML configuration. For example, you can set `url: '${DB_URL}'` and ensure that the `DB_URL` environment variable is defined in your environment. Available configuration options are described below.
@@ -77,9 +90,10 @@ import heartbeat "github.com/obstools/go-prometheus-heartbeat-exporter"
 | `port` | required | Specifies Heartbeat Prometheus exporter server port number. | `port: 8080` |
 | `metrics_route` | required | Defines the route for Prometheus exporter metrics. | `metrics_route: '/metrics'` |
 | `instances` | required | List of instances to monitor. | `instances: [...]` |
-| `name` | required | Name of the instance. | `name: 'postgres_1'` |
+| `name` | required | Name of the instance. Should be unique. | `name: 'postgres_1'` |
 | `connection` | required | Type of connection to the instance. | `connection: 'postgres'` |
 | `url` | required | Connection URL for the instance. | `url: 'postgres://localhost:5432/heartbeat_test'` |
+| `query` | optional | Query to execute on the instance. | `query: 'CREATE TABLE tmp (id SERIAL PRIMARY KEY); DROP TABLE tmp'` |
 | `interval` | required | Check interval for the instance in seconds. | `interval: 3` |
 | `timeout` | required | Check timeout for the instance in seconds. | `timeout: 2` |
 
@@ -98,6 +112,7 @@ instances:
   - name: 'postgres_1'
     connection: 'postgres'
     url: 'postgres://localhost:5432/heartbeat_test'
+    query: 'CREATE TABLE tmp (id SERIAL PRIMARY KEY); DROP TABLE tmp'
     interval: 3
     timeout: 2 
 ```

fixtures/config.yml

@@ -8,5 +8,6 @@ instances:
   - name: 'postgres_1'
     connection: 'postgres'
     url: 'postgres://localhost:5432/heartbeat_test'
+    query: 'CREATE TABLE tmp (id SERIAL PRIMARY KEY); DROP TABLE tmp'
     interval: 3
     timeout: 2

pkg/heartbeat/configuration.go

@@ -13,6 +13,7 @@ type InstanceAttributes struct {
 	Name        string `yaml:"name"`
 	Connection  string `yaml:"connection"`
 	URL         string `yaml:"url"`
+	Query       string `yaml:"query"`
 	IntervalSec int    `yaml:"interval"`
 	TimeoutSec  int    `yaml:"timeout"`
 }

pkg/heartbeat/configuration_test.go

@@ -32,6 +32,7 @@ func TestLoadConfiguration(t *testing.T) {
 					Name:        "postgres_1",
 					Connection:  "postgres",
 					URL:         "postgres://localhost:5432/heartbeat_test",
+					Query:       "CREATE TABLE tmp (id SERIAL PRIMARY KEY); DROP TABLE tmp",
 					IntervalSec: 3,
 					TimeoutSec:  2,
 				},

pkg/heartbeat/heartbeat_instance.go

@@ -23,7 +23,7 @@ func newInstance(instanceAttrs *InstanceAttributes) *heartbeatInstance {
 		name:        instanceAttrs.Name,
 		intervalSec: instanceAttrs.IntervalSec,
 		timeoutSec:  instanceAttrs.TimeoutSec,
-		session:     newSession(instanceAttrs.Connection, instanceAttrs.URL),
+		session:     newSession(instanceAttrs.Connection, instanceAttrs.URL, instanceAttrs.Query),
 	}
 }
 

pkg/heartbeat/heartbeat_instance_test.go

@@ -13,11 +13,12 @@ import (
 func TestNewInstance(t *testing.T) {
 	name, url, intervalSec, timeoutSec := generateUniqueInstanceName(), "test_url", 1, 2
 	t.Run("when connection is postgres", func(t *testing.T) {
-		connection := "postgres"
+		connection, query := "postgres", "some_query"
 		instanceAttributes := &InstanceAttributes{
 			Connection:  connection,
 			Name:        name,
 			URL:         url,
+			Query:       query,
 			IntervalSec: intervalSec,
 			TimeoutSec:  timeoutSec,
 		}
@@ -30,6 +31,7 @@ func TestNewInstance(t *testing.T) {
 		assert.NotNil(t, instanceSession)
 		assert.Equal(t, connection, instanceSession.getConnection())
 		assert.Equal(t, url, instanceSession.getURL())
+		assert.Equal(t, query, instanceSession.getQuery())
 	})
 
 	t.Run("when connection is undefined", func(t *testing.T) {

pkg/heartbeat/session.go

@@ -4,13 +4,15 @@ package heartbeat
 type sessionHeartbeat struct {
 	connection string
 	url        string
+	query      string
 }
 
 // Session type
 type session interface {
 	run() error
 	getConnection() string
 	getURL() string
+	getQuery() string
 }
 
 // Session methods
@@ -25,14 +27,20 @@ func (session *sessionHeartbeat) getURL() string {
 	return session.url
 }
 
+// Returns session query
+func (session *sessionHeartbeat) getQuery() string {
+	return session.query
+}
+
 // New session builder. Returns new session interface based on the connection type
-func newSession(connection, url string) session {
+func newSession(connection, url, query string) session {
 	switch connection {
 	case connectionPostgres:
 		return &sessionPostgres{
 			sessionHeartbeat: &sessionHeartbeat{
 				connection: connection,
 				url:        url,
+				query:      query,
 			},
 		}
 	}

pkg/heartbeat/session_postgres.go

@@ -17,9 +17,10 @@ func (session *sessionPostgres) run() error {
 	}
 	defer db.Close()
 
-	err = db.Ping()
-	if err != nil {
-		return err
+	if session.query != "" {
+		_, err = db.Exec(session.query)
+	} else {
+		err = db.Ping()
 	}
 
 	return err

pkg/heartbeat/session_postgres_test.go

@@ -9,6 +9,7 @@ import (
 
 func TestSessionPostgresRun(t *testing.T) {
 	t.Run("returns nil if connection is established", func(t *testing.T) {
+		query := "CREATE TABLE tmp (id SERIAL PRIMARY KEY); DROP TABLE tmp"
 		if err := createPostgresDb(); err != nil {
 			t.Fatalf("Failed to create test database: %v", err)
 		}
@@ -18,11 +19,11 @@ func TestSessionPostgresRun(t *testing.T) {
 			}
 		}()
 
-		assert.Nil(t, createNewSession("postgres", composePostgresConnectionString()).run())
+		assert.Nil(t, createNewSession("postgres", composePostgresConnectionString(), query).run())
 	})
 
 	t.Run("returns error if ping fails", func(t *testing.T) {
-		assert.NotNil(t, createNewSession("postgres", "postgres://user:password@localhost:5432").run())
+		assert.NotNil(t, createNewSession("postgres", "postgres://user:password@localhost:5432", "").run())
 	})
 
 	t.Run("returns error if connection is not established", func(t *testing.T) {

pkg/heartbeat/session_test.go

@@ -8,15 +8,37 @@ import (
 
 func TestNewSession(t *testing.T) {
 	t.Run("returns new postgres session interface", func(t *testing.T) {
-		connection, url := "postgres", "some_url"
-		session := newSession(connection, url)
+		connection, url, query := "postgres", "some_url", "some_query"
+		session := newSession(connection, url, query)
 
 		assert.NotNil(t, session)
 		assert.Equal(t, connection, session.getConnection())
 		assert.Equal(t, url, session.getURL())
+		assert.Equal(t, query, session.getQuery())
 	})
 
 	t.Run("returns nil for undefined connection", func(t *testing.T) {
-		assert.Nil(t, newSession("undefined", "some_url"))
+		assert.Nil(t, newSession("undefined", "some_url", "some_query"))
 	})
 }
+
+func TestSessionGetConnection(t *testing.T) {
+	connection := "some_connection"
+	session := &sessionHeartbeat{connection: connection}
+
+	assert.Equal(t, connection, session.getConnection())
+}
+
+func TestSessionGetURL(t *testing.T) {
+	url := "some_url"
+	session := &sessionHeartbeat{url: url}
+
+	assert.Equal(t, url, session.getURL())
+}
+
+func TestSessionGetQuery(t *testing.T) {
+	query := "some_query"
+	session := &sessionHeartbeat{query: query}
+
+	assert.Equal(t, query, session.getQuery())
+}

pkg/heartbeat/test_helpers_test.go

@@ -59,8 +59,8 @@ func composePostgresConnectionString() string {
 }
 
 // Creates new session
-func createNewSession(connection, url string) session {
-	return newSession(connection, url)
+func createNewSession(connection, url, query string) session {
+	return newSession(connection, url, query)
 }
 
 // Creates new wait group

pkg/heartbeat/test_mocks_test.go

@@ -69,6 +69,11 @@ func (mock *sessionMock) getURL() string {
 	return args.String(0)
 }
 
+func (mock *sessionMock) getQuery() string {
+	args := mock.Called()
+	return args.String(0)
+}
+
 // prometheusCounterMock
 type prometheusCounterMock struct {
 	mock.Mock