chore(coprocessor): append default application-name to the db_url

Author: goshawk-3

coprocessor/fhevm-engine/fhevm-engine-common/src/utils.rs

@@ -114,32 +114,70 @@ impl Default for HeartBeat {
         Self::new()
     }
 }
-
+/// Simple wrapper around Database URL string to provide
+/// url constraints and masking functionality.
 #[derive(Clone)]
 pub struct DatabaseURL(String);
 
 impl From<&str> for DatabaseURL {
     fn from(s: &str) -> Self {
-        Self(s.to_owned())
+        let url: Self = Self(s.to_owned());
+        url.parse().expect("DatabaseURL is valid");
+        url
     }
 }
 impl From<String> for DatabaseURL {
     fn from(s: String) -> Self {
-        Self(s)
+        let url: Self = Self(s);
+        url.parse().expect("DatabaseURL is valid");
+        url
     }
 }
 
 impl Default for DatabaseURL {
     fn default() -> Self {
-        std::env::var("DATABASE_URL")
-            .unwrap_or("postgres://postgres:postgres@localhost:5432/coprocessor".to_owned())
-            .into()
+        let url = std::env::var("DATABASE_URL")
+            .unwrap_or("postgres://postgres:postgres@localhost:5432/coprocessor".to_owned());
+
+        let app_name = Self::default_app_name();
+        Self::new_with_app_name(&url, &app_name)
     }
 }
 
 impl DatabaseURL {
+    /// Create a new DatabaseURL, appending application_name if not present
+    /// If the base URL already contains an application_name, it will be preserved.
+    ///
+    /// application_name is useful for identifying the source of DB conns
+    pub fn new_with_app_name(base: &str, app_name: &str) -> Self {
+        // Append application_name if not present
+        let mut url = base.to_owned();
+        if !url.contains("application_name=") {
+            if url.contains('?') {
+                url.push_str(&format!("&application_name={}", app_name));
+            } else {
+                url.push_str(&format!("?application_name={}", app_name));
+            }
+        }
+        let url: Self = Self(url);
+        let _ = url.parse().expect("DatabaseURL is valid");
+        url
+    }
+
+    /// Get default app name from the executable name
+    fn default_app_name() -> String {
+        std::env::args()
+            .next()
+            .and_then(|path| {
+                std::path::Path::new(&path)
+                    .file_name()
+                    .map(|s| s.to_string_lossy().into_owned())
+            })
+            .unwrap_or_else(|| "unknown".to_string())
+    }
+
     pub fn as_str(&self) -> &str {
-        &self.0
+        self.0.as_str()
     }
 
     pub fn into_inner(self) -> String {
@@ -148,15 +186,20 @@ impl DatabaseURL {
 
     fn mask_password(options: &PgConnectOptions) -> String {
         let new_url = format!(
-            "postgres://{}:{}@{}:{}/{}",
+            "postgres://{}:{}@{}:{}/{}?application_name={}",
             options.get_username(),
             "*****",
             options.get_host(),
             options.get_port(),
-            options.get_database().unwrap_or_default()
+            options.get_database().unwrap_or_default(),
+            options.get_application_name().unwrap_or_default()
         );
         new_url
     }
+
+    pub fn parse(&self) -> Result<PgConnectOptions, sqlx::Error> {
+        PgConnectOptions::from_str(self.as_str())
+    }
 }
 
 impl fmt::Display for DatabaseURL {