controller/
postgres_exporter.rs

1use crate::{
2    apis::coredb_types::CoreDB,
3    configmap::{create_configmap_ifnotexist, set_configmap},
4    defaults, Context, Error,
5};
6use kube::Client;
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9use std::{collections::BTreeMap, sync::Arc};
10use tracing::debug;
11
12pub const QUERIES_YAML: &str = "queries.yaml";
13pub const EXPORTER_VOLUME: &str = "postgres-exporter";
14pub const EXPORTER_CONFIGMAP: &str = "postgres-exporter";
15
16#[derive(Deserialize, Serialize, Clone, Debug, JsonSchema, Default)]
17#[allow(non_snake_case)]
18pub struct PostgresMetrics {
19    #[serde(default = "defaults::default_postgres_exporter_image")]
20    pub image: String,
21    #[serde(default = "defaults::default_postgres_exporter_enabled")]
22    pub enabled: bool,
23
24    #[schemars(schema_with = "preserve_arbitrary")]
25    pub queries: Option<QueryConfig>,
26}
27
28#[derive(Clone, Debug, JsonSchema, PartialEq, Serialize, Deserialize)]
29pub struct Metric {
30    pub usage: Usage,
31    pub description: String,
32}
33
34#[derive(Clone, Debug, JsonSchema, PartialEq, Serialize, Deserialize)]
35pub struct Metrics {
36    #[serde(flatten)]
37    pub metrics: BTreeMap<String, Metric>,
38}
39
40#[derive(Clone, Debug, JsonSchema, PartialEq, Serialize, Deserialize)]
41pub struct QueryItem {
42    pub query: String,
43    pub master: bool,
44    pub metrics: Vec<Metrics>,
45}
46
47#[derive(Clone, Debug, JsonSchema, PartialEq, Serialize, Deserialize)]
48pub struct QueryConfig {
49    #[serde(flatten)]
50    pub queries: BTreeMap<String, QueryItem>,
51}
52
53// source: https://github.com/kube-rs/kube/issues/844
54fn preserve_arbitrary(_gen: &mut schemars::gen::SchemaGenerator) -> schemars::schema::Schema {
55    let mut obj = schemars::schema::SchemaObject::default();
56    obj.extensions
57        .insert("x-kubernetes-preserve-unknown-fields".into(), true.into());
58    schemars::schema::Schema::Object(obj)
59}
60
61use std::str::FromStr;
62
63#[derive(Clone, Debug, JsonSchema, Serialize, Deserialize, PartialEq)]
64#[serde(rename_all = "UPPERCASE")]
65pub enum Usage {
66    Counter,
67    Gauge,
68    Histogram,
69    Label,
70}
71
72impl FromStr for Usage {
73    type Err = ();
74
75    fn from_str(input: &str) -> Result<Usage, Self::Err> {
76        match input {
77            "COUNTER" => Ok(Usage::Counter),
78            "GAUGE" => Ok(Usage::Gauge),
79            "HISTOGRAM" => Ok(Usage::Histogram),
80            "LABEL" => Ok(Usage::Label),
81            _ => Err(()),
82        }
83    }
84}
85
86pub async fn create_postgres_exporter_role(cdb: &CoreDB, ctx: Arc<Context>) -> Result<(), Error> {
87    let client = ctx.client.clone();
88    if !(cdb.spec.postgresExporterEnabled) {
89        return Ok(());
90    }
91    debug!(
92        "Creating postgres_exporter role for database {} in namespace {}",
93        cdb.metadata.name.clone().unwrap(),
94        cdb.metadata.namespace.clone().unwrap()
95    );
96    // https://github.com/prometheus-community/postgres_exporter#running-as-non-superuser
97    let _ = cdb
98        .psql(
99            "
100            CREATE OR REPLACE FUNCTION __tmp_create_user() returns void as $$
101            BEGIN
102              IF NOT EXISTS (
103                      SELECT
104                      FROM   pg_catalog.pg_user
105                      WHERE  usename = 'postgres_exporter') THEN
106                CREATE USER postgres_exporter;
107              END IF;
108            END;
109            $$ language plpgsql;
110
111            SELECT __tmp_create_user();
112            DROP FUNCTION __tmp_create_user();
113
114            ALTER USER postgres_exporter SET SEARCH_PATH TO postgres_exporter,pg_catalog;
115            GRANT CONNECT ON DATABASE postgres TO postgres_exporter;
116            GRANT pg_monitor to postgres_exporter;
117            GRANT pg_read_all_stats to postgres_exporter;
118            "
119            .to_string(),
120            "postgres".to_owned(),
121            client.clone(),
122        )
123        .await?;
124    Ok(())
125}
126
127
128pub async fn reconcile_prom_configmap(cdb: &CoreDB, client: Client, ns: &str) -> Result<(), Error> {
129    create_configmap_ifnotexist(client.clone(), ns, EXPORTER_CONFIGMAP).await?;
130    // set custom pg-prom metrics in configmap values if they are specified
131    match cdb.spec.metrics.clone().and_then(|m| m.queries) {
132        Some(queries) => {
133            let qdata = serde_yaml::to_string(&queries).unwrap();
134            let d: BTreeMap<String, String> = BTreeMap::from([(QUERIES_YAML.to_string(), qdata)]);
135            set_configmap(client.clone(), ns, EXPORTER_CONFIGMAP, d).await?
136        }
137        None => {
138            debug!("No queries specified in CoreDB spec");
139        }
140    }
141    Ok(())
142}
143
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    use serde_yaml;
149
150    #[test]
151    fn query_deserialize_serialize() {
152        // query data received as json. map to struct.
153        // serialize struct to yaml
154        let incoming_data = serde_json::json!(
155            {
156                "pg_postmaster": {
157                  "query": "SELECT pg_postmaster_start_time as start_time_seconds from pg_postmaster_start_time()",
158                  "master": true,
159                  "metrics": [
160                    {
161                      "start_time_seconds": {
162                        "usage": "GAUGE",
163                        "description": "Time at which postmaster started"
164                      }
165                    }
166                  ]
167                },
168                "extensions": {
169                  "query": "select count(*) as num_ext from pg_available_extensions",
170                  "master": true,
171                  "metrics": [
172                    {
173                      "num_ext": {
174                        "usage": "GAUGE",
175                        "description": "Num extensions"
176                      }
177                    }
178                  ]
179                }
180              }
181        );
182
183        let query_config: QueryConfig = serde_json::from_value(incoming_data).expect("failed to deserialize");
184
185        assert!(query_config.queries.contains_key("pg_postmaster"));
186        assert!(query_config.queries.contains_key("extensions"));
187
188        let pg_postmaster = query_config.queries.get("pg_postmaster").unwrap();
189        assert_eq!(
190            pg_postmaster.query,
191            "SELECT pg_postmaster_start_time as start_time_seconds from pg_postmaster_start_time()"
192        );
193        assert!(pg_postmaster.master);
194        assert!(pg_postmaster.metrics[0]
195            .metrics
196            .contains_key("start_time_seconds"));
197
198        let start_time_seconds_metric = pg_postmaster.metrics[0]
199            .metrics
200            .get("start_time_seconds")
201            .unwrap();
202        assert_eq!(
203            start_time_seconds_metric.description,
204            "Time at which postmaster started"
205        );
206
207        let extensions = query_config
208            .queries
209            .get("extensions")
210            .expect("extensions not found");
211        assert_eq!(
212            extensions.query,
213            "select count(*) as num_ext from pg_available_extensions"
214        );
215        assert!(extensions.master);
216        assert!(extensions.metrics[0].metrics.contains_key("num_ext"));
217
218        // yaml to yaml
219
220        let yaml = serde_yaml::to_string(&query_config).expect("failed to serialize to yaml");
221
222        let data = r#"extensions:
223  query: select count(*) as num_ext from pg_available_extensions
224  master: true
225  metrics:
226  - num_ext:
227      usage: GAUGE
228      description: Num extensions
229pg_postmaster:
230  query: SELECT pg_postmaster_start_time as start_time_seconds from pg_postmaster_start_time()
231  master: true
232  metrics:
233  - start_time_seconds:
234      usage: GAUGE
235      description: Time at which postmaster started
236"#;
237        // formmatted correctly as yaml (for configmap)
238        assert_eq!(yaml, data);
239    }
240}
controller/postgres_exporter.rs

controller/
postgres_exporter.rs
