Line data    Source code

       1              : use compute_api::responses::{InstalledExtension, InstalledExtensions};
       2              : use metrics::proto::MetricFamily;
       3              : use std::collections::HashMap;
       4              : use std::collections::HashSet;
       5              : 
       6              : use anyhow::Result;
       7              : use postgres::{Client, NoTls};
       8              : 
       9              : use metrics::core::Collector;
      10              : use metrics::{register_uint_gauge_vec, UIntGaugeVec};
      11              : use once_cell::sync::Lazy;
      12              : 
      13              : /// We don't reuse get_existing_dbs() just for code clarity
      14              : /// and to make database listing query here more explicit.
      15              : ///
      16              : /// Limit the number of databases to 500 to avoid excessive load.
      17            0 : fn list_dbs(client: &mut Client) -> Result<Vec<String>> {
      18              :     // `pg_database.datconnlimit = -2` means that the database is in the
      19              :     // invalid state
      20            0 :     let databases = client
      21            0 :         .query(
      22            0 :             "SELECT datname FROM pg_catalog.pg_database
      23            0 :                 WHERE datallowconn
      24            0 :                 AND datconnlimit <> - 2
      25            0 :                 LIMIT 500",
      26            0 :             &[],
      27            0 :         )?
      28            0 :         .iter()
      29            0 :         .map(|row| {
      30            0 :             let db: String = row.get("datname");
      31            0 :             db
      32            0 :         })
      33            0 :         .collect();
      34            0 : 
      35            0 :     Ok(databases)
      36            0 : }
      37              : 
      38              : /// Connect to every database (see list_dbs above) and get the list of installed extensions.
      39              : ///
      40              : /// Same extension can be installed in multiple databases with different versions,
      41              : /// we only keep the highest and lowest version across all databases.
      42            0 : pub fn get_installed_extensions(mut conf: postgres::config::Config) -> Result<InstalledExtensions> {
      43            0 :     conf.application_name("compute_ctl:get_installed_extensions");
      44            0 :     let mut client = conf.connect(NoTls)?;
      45              : 
      46            0 :     let databases: Vec<String> = list_dbs(&mut client)?;
      47              : 
      48            0 :     let mut extensions_map: HashMap<String, InstalledExtension> = HashMap::new();
      49            0 :     for db in databases.iter() {
      50            0 :         conf.dbname(db);
      51            0 :         let mut db_client = conf.connect(NoTls)?;
      52            0 :         let extensions: Vec<(String, String)> = db_client
      53            0 :             .query(
      54            0 :                 "SELECT extname, extversion FROM pg_catalog.pg_extension;",
      55            0 :                 &[],
      56            0 :             )?
      57            0 :             .iter()
      58            0 :             .map(|row| (row.get("extname"), row.get("extversion")))
      59            0 :             .collect();
      60              : 
      61            0 :         for (extname, v) in extensions.iter() {
      62            0 :             let version = v.to_string();
      63            0 : 
      64            0 :             // increment the number of databases where the version of extension is installed
      65            0 :             INSTALLED_EXTENSIONS
      66            0 :                 .with_label_values(&[extname, &version])
      67            0 :                 .inc();
      68            0 : 
      69            0 :             extensions_map
      70            0 :                 .entry(extname.to_string())
      71            0 :                 .and_modify(|e| {
      72            0 :                     e.versions.insert(version.clone());
      73            0 :                     // count the number of databases where the extension is installed
      74            0 :                     e.n_databases += 1;
      75            0 :                 })
      76            0 :                 .or_insert(InstalledExtension {
      77            0 :                     extname: extname.to_string(),
      78            0 :                     versions: HashSet::from([version.clone()]),
      79            0 :                     n_databases: 1,
      80            0 :                 });
      81            0 :         }
      82              :     }
      83              : 
      84            0 :     let res = InstalledExtensions {
      85            0 :         extensions: extensions_map.into_values().collect(),
      86            0 :     };
      87            0 : 
      88            0 :     Ok(res)
      89            0 : }
      90              : 
      91            0 : static INSTALLED_EXTENSIONS: Lazy<UIntGaugeVec> = Lazy::new(|| {
      92            0 :     register_uint_gauge_vec!(
      93            0 :         "compute_installed_extensions",
      94            0 :         "Number of databases where the version of extension is installed",
      95            0 :         &["extension_name", "version"]
      96            0 :     )
      97            0 :     .expect("failed to define a metric")
      98            0 : });
      99              : 
     100            0 : pub fn collect() -> Vec<MetricFamily> {
     101            0 :     INSTALLED_EXTENSIONS.collect()
     102            0 : }