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