LCOV - code coverage report
Current view: top level - control_plane/src - postgresql_conf.rs (source / functions) Coverage Total Hit
Test: 1b0a6a0c05cee5a7de360813c8034804e105ce1c.info Lines: 59.6 % 47 28
Test Date: 2025-03-12 00:01:28 Functions: 33.3 % 9 3

            Line data    Source code
       1              : use std::collections::HashMap;
       2              : use std::fmt;
       3              : 
       4              : ///
       5              : /// Module for parsing postgresql.conf file.
       6              : ///
       7              : /// NOTE: This doesn't implement the full, correct postgresql.conf syntax. Just
       8              : /// enough to extract a few settings we need in Neon, assuming you don't do
       9              : /// funny stuff like include-directives or funny escaping.
      10              : use once_cell::sync::Lazy;
      11              : use regex::Regex;
      12              : 
      13              : /// In-memory representation of a postgresql.conf file
      14              : #[derive(Default, Debug)]
      15              : pub struct PostgresConf {
      16              :     lines: Vec<String>,
      17              :     hash: HashMap<String, String>,
      18              : }
      19              : 
      20              : impl PostgresConf {
      21            0 :     pub fn new() -> PostgresConf {
      22            0 :         PostgresConf::default()
      23            0 :     }
      24              : 
      25              :     /// Return the current value of 'option'
      26            0 :     pub fn get(&self, option: &str) -> Option<&str> {
      27            0 :         self.hash.get(option).map(|x| x.as_ref())
      28            0 :     }
      29              : 
      30              :     ///
      31              :     /// Note: if you call this multiple times for the same option, the config
      32              :     /// file will a line for each call. It would be nice to have a function
      33              :     /// to change an existing line, but that's a TODO.
      34              :     ///
      35            0 :     pub fn append(&mut self, option: &str, value: &str) {
      36            0 :         self.lines
      37            0 :             .push(format!("{}={}\n", option, escape_str(value)));
      38            0 :         self.hash.insert(option.to_string(), value.to_string());
      39            0 :     }
      40              : 
      41              :     /// Append an arbitrary non-setting line to the config file
      42            0 :     pub fn append_line(&mut self, line: &str) {
      43            0 :         self.lines.push(line.to_string());
      44            0 :     }
      45              : }
      46              : 
      47              : impl fmt::Display for PostgresConf {
      48              :     /// Return the whole configuration file as a string
      49            0 :     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
      50            0 :         for line in self.lines.iter() {
      51            0 :             f.write_str(line)?;
      52              :         }
      53            0 :         Ok(())
      54            0 :     }
      55              : }
      56              : 
      57              : /// Escape a value for putting in postgresql.conf.
      58           14 : fn escape_str(s: &str) -> String {
      59              :     // If the string doesn't contain anything that needs quoting or escaping, return it
      60              :     // as it is.
      61              :     //
      62              :     // The first part of the regex, before the '|', matches the INTEGER rule in the
      63              :     // PostgreSQL flex grammar (guc-file.l). It matches plain integers like "123" and
      64              :     // "-123", and also accepts units like "10MB". The second part of the regex matches
      65              :     // the UNQUOTED_STRING rule, and accepts strings that contain a single word, beginning
      66              :     // with a letter. That covers words like "off" or "posix". Everything else is quoted.
      67              :     //
      68              :     // This regex is a bit more conservative than the rules in guc-file.l, so we quote some
      69              :     // strings that PostgreSQL would accept without quoting, but that's OK.
      70              : 
      71              :     static UNQUOTED_RE: Lazy<Regex> =
      72            1 :         Lazy::new(|| Regex::new(r"(^[-+]?[0-9]+[a-zA-Z]*$)|(^[a-zA-Z][a-zA-Z0-9]*$)").unwrap());
      73              : 
      74           14 :     if UNQUOTED_RE.is_match(s) {
      75            8 :         s.to_string()
      76              :     } else {
      77              :         // Otherwise escape and quote it
      78            6 :         let s = s
      79            6 :             .replace('\\', "\\\\")
      80            6 :             .replace('\n', "\\n")
      81            6 :             .replace('\'', "''");
      82            6 : 
      83            6 :         "\'".to_owned() + &s + "\'"
      84              :     }
      85           14 : }
      86              : 
      87              : #[test]
      88            1 : fn test_postgresql_conf_escapes() -> anyhow::Result<()> {
      89            1 :     assert_eq!(escape_str("foo bar"), "'foo bar'");
      90              :     // these don't need to be quoted
      91            1 :     assert_eq!(escape_str("foo"), "foo");
      92            1 :     assert_eq!(escape_str("123"), "123");
      93            1 :     assert_eq!(escape_str("+123"), "+123");
      94            1 :     assert_eq!(escape_str("-10"), "-10");
      95            1 :     assert_eq!(escape_str("1foo"), "1foo");
      96            1 :     assert_eq!(escape_str("foo1"), "foo1");
      97            1 :     assert_eq!(escape_str("10MB"), "10MB");
      98            1 :     assert_eq!(escape_str("-10kB"), "-10kB");
      99              : 
     100              :     // these need quoting and/or escaping
     101            1 :     assert_eq!(escape_str("foo bar"), "'foo bar'");
     102            1 :     assert_eq!(escape_str("fo'o"), "'fo''o'");
     103            1 :     assert_eq!(escape_str("fo\no"), "'fo\\no'");
     104            1 :     assert_eq!(escape_str("fo\\o"), "'fo\\\\o'");
     105            1 :     assert_eq!(escape_str("10 cats"), "'10 cats'");
     106              : 
     107            1 :     Ok(())
     108            1 : }
        

Generated by: LCOV version 2.1-beta