diff --git a/README.md b/README.md index 5628e51..b0eb10e 100644 --- a/README.md +++ b/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 ``` diff --git a/fixtures/config.yml b/fixtures/config.yml index 490a080..f45e755 100644 --- a/fixtures/config.yml +++ b/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 diff --git a/pkg/heartbeat/configuration.go b/pkg/heartbeat/configuration.go index f3a6170..dca8756 100644 --- a/pkg/heartbeat/configuration.go +++ b/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"` } diff --git a/pkg/heartbeat/configuration_test.go b/pkg/heartbeat/configuration_test.go index 07a8327..b754782 100644 --- a/pkg/heartbeat/configuration_test.go +++ b/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, }, diff --git a/pkg/heartbeat/heartbeat_instance.go b/pkg/heartbeat/heartbeat_instance.go index 8521d41..b096604 100644 --- a/pkg/heartbeat/heartbeat_instance.go +++ b/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), } } diff --git a/pkg/heartbeat/heartbeat_instance_test.go b/pkg/heartbeat/heartbeat_instance_test.go index baeac9b..99ad058 100644 --- a/pkg/heartbeat/heartbeat_instance_test.go +++ b/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) { diff --git a/pkg/heartbeat/session.go b/pkg/heartbeat/session.go index 2cf504e..95452aa 100644 --- a/pkg/heartbeat/session.go +++ b/pkg/heartbeat/session.go @@ -4,6 +4,7 @@ package heartbeat type sessionHeartbeat struct { connection string url string + query string } // Session type @@ -11,6 +12,7 @@ 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, }, } } diff --git a/pkg/heartbeat/session_postgres.go b/pkg/heartbeat/session_postgres.go index 98123b5..70811d6 100644 --- a/pkg/heartbeat/session_postgres.go +++ b/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 diff --git a/pkg/heartbeat/session_postgres_test.go b/pkg/heartbeat/session_postgres_test.go index c9a0532..7474497 100644 --- a/pkg/heartbeat/session_postgres_test.go +++ b/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) { diff --git a/pkg/heartbeat/session_test.go b/pkg/heartbeat/session_test.go index d3b8ce2..dd544fa 100644 --- a/pkg/heartbeat/session_test.go +++ b/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()) +} diff --git a/pkg/heartbeat/test_helpers_test.go b/pkg/heartbeat/test_helpers_test.go index b0c401c..cc96a28 100644 --- a/pkg/heartbeat/test_helpers_test.go +++ b/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 diff --git a/pkg/heartbeat/test_mocks_test.go b/pkg/heartbeat/test_mocks_test.go index 728628d..35f6788 100644 --- a/pkg/heartbeat/test_mocks_test.go +++ b/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