LCOV - code coverage report
Current view: top level - control_plane/src - postgresql_conf.rs (source / functions) Coverage Total Hit
Test: f8d8f5b90fa487a9e82c42da223f012f5d4fece7.info Lines: 59.6 % 47 28
Test Date: 2024-09-19 20:36:02 Functions: 33.3 % 9 3

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

Generated by: LCOV version 2.1-beta