From 3371c01e0e3d2185937666537842137310dbf7ef Mon Sep 17 00:00:00 2001
From: Mohammad Dashti <mdashti@gmail.com>
Date: Tue, 3 Oct 2023 13:13:21 -0700
Subject: [PATCH] Added a `Plugin` trait (#536)

* Improved logging

* Improved logging for more `Address` usages

* Fixed lint issues.

* Reverted the `Address` logging changes.

* Applied the PR comment by @levkk.

* Applied the PR comment by @levkk.

* Applied the PR comment by @levkk.

* Applied the PR comment by @levkk.
---
 src/config.rs | 43 +++++++++++++++++++++++++++++++++++++++----
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 0404abc..90b4beb 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -867,15 +867,26 @@ pub struct Plugins {
     pub prewarmer: Option<Prewarmer>,
 }
 
+pub trait Plugin {
+    fn is_enabled(&self) -> bool;
+}
+
 impl std::fmt::Display for Plugins {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        fn is_enabled<T: Plugin>(arg: Option<&T>) -> bool {
+            if let Some(ref arg) = arg {
+                arg.is_enabled()
+            } else {
+                false
+            }
+        }
         write!(
             f,
             "interceptor: {}, table_access: {}, query_logger: {}, prewarmer: {}",
-            self.intercept.is_some(),
-            self.table_access.is_some(),
-            self.query_logger.is_some(),
-            self.prewarmer.is_some(),
+            is_enabled(self.intercept.as_ref()),
+            is_enabled(self.table_access.as_ref()),
+            is_enabled(self.query_logger.as_ref()),
+            is_enabled(self.prewarmer.as_ref()),
         )
     }
 }
@@ -886,23 +897,47 @@ pub struct Intercept {
     pub queries: BTreeMap<String, Query>,
 }
 
+impl Plugin for Intercept {
+    fn is_enabled(&self) -> bool {
+        self.enabled
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Hash, Eq)]
 pub struct TableAccess {
     pub enabled: bool,
     pub tables: Vec<String>,
 }
 
+impl Plugin for TableAccess {
+    fn is_enabled(&self) -> bool {
+        self.enabled
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Hash, Eq)]
 pub struct QueryLogger {
     pub enabled: bool,
 }
 
+impl Plugin for QueryLogger {
+    fn is_enabled(&self) -> bool {
+        self.enabled
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Hash, Eq)]
 pub struct Prewarmer {
     pub enabled: bool,
     pub queries: Vec<String>,
 }
 
+impl Plugin for Prewarmer {
+    fn is_enabled(&self) -> bool {
+        self.enabled
+    }
+}
+
 impl Intercept {
     pub fn substitute(&mut self, db: &str, user: &str) {
         for (_, query) in self.queries.iter_mut() {