Line data    Source code

       1              : use compute_api::responses::{InstalledExtension, InstalledExtensions};
       2              : use std::collections::HashMap;
       3              : use std::collections::HashSet;
       4              : use tracing::info;
       5              : use url::Url;
       6              : 
       7              : use anyhow::Result;
       8              : use postgres::{Client, NoTls};
       9              : use tokio::task;
      10              : 
      11              : /// We don't reuse get_existing_dbs() just for code clarity
      12              : /// and to make database listing query here more explicit.
      13              : ///
      14              : /// Limit the number of databases to 500 to avoid excessive load.
      15            0 : fn list_dbs(client: &mut Client) -> Result<Vec<String>> {
      16              :     // `pg_database.datconnlimit = -2` means that the database is in the
      17              :     // invalid state
      18            0 :     let databases = client
      19            0 :         .query(
      20            0 :             "SELECT datname FROM pg_catalog.pg_database
      21            0 :                 WHERE datallowconn
      22            0 :                 AND datconnlimit <> - 2
      23            0 :                 LIMIT 500",
      24            0 :             &[],
      25            0 :         )?
      26            0 :         .iter()
      27            0 :         .map(|row| {
      28            0 :             let db: String = row.get("datname");
      29            0 :             db
      30            0 :         })
      31            0 :         .collect();
      32            0 : 
      33            0 :     Ok(databases)
      34            0 : }
      35              : 
      36              : /// Connect to every database (see list_dbs above) and get the list of installed extensions.
      37              : ///
      38              : /// Same extension can be installed in multiple databases with different versions,
      39              : /// we only keep the highest and lowest version across all databases.
      40            0 : pub async fn get_installed_extensions(connstr: Url) -> Result<InstalledExtensions> {
      41            0 :     let mut connstr = connstr.clone();
      42            0 : 
      43            0 :     task::spawn_blocking(move || {
      44            0 :         let mut client = Client::connect(connstr.as_str(), NoTls)?;
      45            0 :         let databases: Vec<String> = list_dbs(&mut client)?;
      46              : 
      47            0 :         let mut extensions_map: HashMap<String, InstalledExtension> = HashMap::new();
      48            0 :         for db in databases.iter() {
      49            0 :             connstr.set_path(db);
      50            0 :             let mut db_client = Client::connect(connstr.as_str(), NoTls)?;
      51            0 :             let extensions: Vec<(String, String)> = db_client
      52            0 :                 .query(
      53            0 :                     "SELECT extname, extversion FROM pg_catalog.pg_extension;",
      54            0 :                     &[],
      55            0 :                 )?
      56            0 :                 .iter()
      57            0 :                 .map(|row| (row.get("extname"), row.get("extversion")))
      58            0 :                 .collect();
      59              : 
      60            0 :             for (extname, v) in extensions.iter() {
      61            0 :                 let version = v.to_string();
      62            0 :                 extensions_map
      63            0 :                     .entry(extname.to_string())
      64            0 :                     .and_modify(|e| {
      65            0 :                         e.versions.insert(version.clone());
      66            0 :                         // count the number of databases where the extension is installed
      67            0 :                         e.n_databases += 1;
      68            0 :                     })
      69            0 :                     .or_insert(InstalledExtension {
      70            0 :                         extname: extname.to_string(),
      71            0 :                         versions: HashSet::from([version.clone()]),
      72            0 :                         n_databases: 1,
      73            0 :                     });
      74            0 :             }
      75              :         }
      76              : 
      77            0 :         Ok(InstalledExtensions {
      78            0 :             extensions: extensions_map.values().cloned().collect(),
      79            0 :         })
      80            0 :     })
      81            0 :     .await?
      82            0 : }
      83              : 
      84              : // Gather info about installed extensions
      85            0 : pub fn get_installed_extensions_sync(connstr: Url) -> Result<()> {
      86            0 :     let rt = tokio::runtime::Builder::new_current_thread()
      87            0 :         .enable_all()
      88            0 :         .build()
      89            0 :         .expect("failed to create runtime");
      90            0 :     let result = rt
      91            0 :         .block_on(crate::installed_extensions::get_installed_extensions(
      92            0 :             connstr,
      93            0 :         ))
      94            0 :         .expect("failed to get installed extensions");
      95            0 : 
      96            0 :     info!(
      97            0 :         "[NEON_EXT_STAT] {}",
      98            0 :         serde_json::to_string(&result).expect("failed to serialize extensions list")
      99              :     );
     100              : 
     101            0 :     Ok(())
     102            0 : }