Modular Architecture Refactoring: Improve Maintainability and Test Organization

examples/two-tables/src/t2.rs

@@ -3,7 +3,6 @@ use osquery_rust_ng::plugin::{
     ColumnDef, ColumnOptions, ColumnType, DeleteResult, InsertResult, Table, UpdateResult,
 };
 use osquery_rust_ng::{ExtensionPluginRequest, ExtensionResponse, ExtensionStatus};
-use serde_json::Value;
 use std::collections::BTreeMap;
 
 pub struct Table2 {}
@@ -26,7 +25,7 @@ impl Table for Table2 {
         ]
     }
 
-    fn generate(&self, _req: ExtensionPluginRequest) -> ExtensionResponse {
+    fn generate(&mut self, _req: ExtensionPluginRequest) -> ExtensionResponse {
         let resp = BTreeMap::from([
             ("top".to_string(), "top".to_string()),
             ("bottom".to_string(), "bottom".to_string()),
@@ -35,16 +34,16 @@ impl Table for Table2 {
         ExtensionResponse::new(ExtensionStatus::default(), vec![resp])
     }
 
-    fn update(&mut self, _rowid: u64, _row: &Value) -> UpdateResult {
-        UpdateResult::Constraint
+    fn update(&mut self, _rowid: String, _row: serde_json::Value) -> UpdateResult {
+        UpdateResult::Error("Table t2 is read-only".to_string())
     }
 
-    fn delete(&mut self, _rowid: u64) -> DeleteResult {
-        DeleteResult::Err("Not yet implemented".to_string())
+    fn delete(&mut self, _rowid: String) -> DeleteResult {
+        DeleteResult::Error("Table t2 is read-only".to_string())
     }
 
-    fn insert(&mut self, _auto_rowid: bool, _row: &Value) -> InsertResult {
-        InsertResult::Constraint
+    fn insert(&mut self, _row: serde_json::Value) -> InsertResult {
+        InsertResult::Error("Table t2 is read-only".to_string())
     }
 
     fn shutdown(&self) {

examples/writeable-table/src/main.rs

@@ -6,12 +6,12 @@ use log::info;
 use osquery_rust_ng::plugin::{ColumnDef, ColumnOptions, ColumnType, Plugin, Table};
 use osquery_rust_ng::plugin::{DeleteResult, InsertResult, UpdateResult};
 use osquery_rust_ng::{ExtensionPluginRequest, ExtensionResponse, ExtensionStatus, Server};
-use serde_json::Value;
 use std::collections::BTreeMap;
 use std::io::{Error, ErrorKind};
 
 struct WriteableTable {
-    items: BTreeMap<u64, (String, String)>,
+    items: BTreeMap<String, (String, String)>,
+    next_id: u64,
 }
 
 impl WriteableTable {
@@ -20,8 +20,9 @@ impl WriteableTable {
             items: vec!["foo".to_string(), "bar".to_string(), "baz".to_string()]
                 .into_iter()
                 .enumerate()
-                .map(|(idx, item)| (idx as u64, (item.clone(), item.clone())))
+                .map(|(idx, item)| (idx.to_string(), (item.clone(), item.clone())))
                 .collect(),
+            next_id: 3,
         }
     }
 }
@@ -43,7 +44,7 @@ impl Table for WriteableTable {
         ]
     }
 
-    fn generate(&self, _req: ExtensionPluginRequest) -> ExtensionResponse {
+    fn generate(&mut self, _req: ExtensionPluginRequest) -> ExtensionResponse {
         let resp = self
             .items
             .iter()
@@ -59,76 +60,69 @@ impl Table for WriteableTable {
         ExtensionResponse::new(ExtensionStatus::default(), resp)
     }
 
-    fn update(&mut self, rowid: u64, row: &Value) -> UpdateResult {
+    fn update(&mut self, rowid: String, row: serde_json::Value) -> UpdateResult {
         log::info!("updating item at {rowid} = {row:?}");
 
-        let Some(row) = row.as_array() else {
-            return UpdateResult::Err("Could not parse row as array".to_string());
+        let Some(row_array) = row.as_array() else {
+            return UpdateResult::Error("Could not parse row as array".to_string());
         };
 
-        let &[
-            Value::Number(rowid),
-            Value::String(name),
-            Value::String(lastname),
-        ] = &row.as_slice()
-        else {
-            return UpdateResult::Err("Could not parse row update".to_string());
-        };
+        if row_array.len() < 2 {
+            return UpdateResult::Error("Row must have at least 2 elements".to_string());
+        }
 
-        let Some(rowid) = rowid.as_u64() else {
-            return UpdateResult::Err("Could not parse rowid as u64".to_string());
+        let Some(name) = row_array[0].as_str() else {
+            return UpdateResult::Error("Name must be a string".to_string());
         };
 
-        self.items.insert(rowid, (name.clone(), lastname.clone()));
+        let Some(lastname) = row_array[1].as_str() else {
+            return UpdateResult::Error("Lastname must be a string".to_string());
+        };
 
-        UpdateResult::Success
+        if self.items.contains_key(&rowid) {
+            self.items
+                .insert(rowid, (name.to_string(), lastname.to_string()));
+            UpdateResult::Ok
+        } else {
+            UpdateResult::NotFound
+        }
     }
 
-    fn delete(&mut self, rowid: u64) -> DeleteResult {
+    fn delete(&mut self, rowid: String) -> DeleteResult {
         log::info!("deleting item: {rowid}");
 
         match self.items.remove(&rowid) {
-            Some(_) => DeleteResult::Success,
-            None => DeleteResult::Err("Could not find rowid".to_string()),
+            Some(_) => DeleteResult::Ok,
+            None => DeleteResult::NotFound,
         }
     }
 
-    fn insert(&mut self, _auto_rowid: bool, row: &Value) -> InsertResult {
+    fn insert(&mut self, row: serde_json::Value) -> InsertResult {
         log::info!("inserting item: {row:?}");
 
-        let Some(row) = row.as_array() else {
-            return InsertResult::Err("Could not parse row as array".to_string());
+        let Some(row_array) = row.as_array() else {
+            return InsertResult::Error("Could not parse row as array".to_string());
+        };
+
+        if row_array.len() < 2 {
+            return InsertResult::Error("Row must have at least 2 elements".to_string());
+        }
+
+        let Some(name) = row_array[0].as_str() else {
+            return InsertResult::Error("Name must be a string".to_string());
         };
 
-        let rowid = match &row.as_slice() {
-            [Value::Null, Value::String(name), Value::String(lastname)] => {
-                // TODO: figure out what auto_rowid means here
-                let rowid = self.items.keys().next_back().unwrap_or(&0u64) + 1;
-                log::info!("rowid: {rowid}");
-
-                self.items.insert(rowid, (name.clone(), lastname.clone()));
-
-                rowid
-            }
-            [
-                Value::Number(rowid),
-                Value::String(name),
-                Value::String(lastname),
-            ] => {
-                let Some(rowid) = rowid.as_u64() else {
-                    return InsertResult::Err("Could not parse rowid as u64".to_string());
-                };
-
-                self.items.insert(rowid, (name.clone(), lastname.clone()));
-
-                rowid
-            }
-            _ => {
-                return InsertResult::Constraint;
-            }
+        let Some(lastname) = row_array[1].as_str() else {
+            return InsertResult::Error("Lastname must be a string".to_string());
         };
 
-        InsertResult::Success(rowid)
+        let rowid = self.next_id.to_string();
+        self.next_id += 1;
+
+        self.items
+            .insert(rowid.clone(), (name.to_string(), lastname.to_string()));
+
+        InsertResult::Ok(rowid)
     }
     fn shutdown(&self) {
         info!("Shutting down");
@@ -183,7 +177,7 @@ mod tests {
 
     #[test]
     fn test_generate_returns_initial_data() {
-        let table = WriteableTable::new();
+        let mut table = WriteableTable::new();
         let response = table.generate(ExtensionPluginRequest::default());
         let rows = response.response.expect("should have rows");
 
@@ -198,14 +192,14 @@ mod tests {
     fn test_insert_with_auto_rowid() {
         let mut table = WriteableTable::new();
 
-        // Insert with null rowid (auto-assign)
-        let row = json!([null, "alice", "smith"]);
-        let result = table.insert(true, &row);
+        // Insert new row (auto-assign rowid)
+        let row = json!(["alice", "smith"]);
+        let result = table.insert(row);
 
-        let InsertResult::Success(rowid) = result else {
-            panic!("Expected InsertResult::Success");
+        let InsertResult::Ok(rowid) = result else {
+            panic!("Expected InsertResult::Ok");
         };
-        assert_eq!(rowid, 3); // Next after 0, 1, 2
+        assert_eq!(rowid, "3"); // Next after 0, 1, 2
 
         // Verify the row was added
         let response = table.generate(ExtensionPluginRequest::default());
@@ -214,39 +208,39 @@ mod tests {
     }
 
     #[test]
-    fn test_insert_with_explicit_rowid() {
+    fn test_insert_another_row() {
         let mut table = WriteableTable::new();
 
-        // Insert with explicit rowid
-        let row = json!([100, "bob", "jones"]);
-        let result = table.insert(false, &row);
+        // Insert another row
+        let row = json!(["bob", "jones"]);
+        let result = table.insert(row);
 
-        let InsertResult::Success(rowid) = result else {
-            panic!("Expected InsertResult::Success");
+        let InsertResult::Ok(rowid) = result else {
+            panic!("Expected InsertResult::Ok");
         };
-        assert_eq!(rowid, 100);
+        assert_eq!(rowid, "3");
     }
 
     #[test]
-    fn test_insert_invalid_row_returns_constraint() {
+    fn test_insert_invalid_row_returns_error() {
         let mut table = WriteableTable::new();
 
         // Invalid row format
         let row = json!(["invalid"]);
-        let result = table.insert(false, &row);
+        let result = table.insert(row);
 
-        assert!(matches!(result, InsertResult::Constraint));
+        assert!(matches!(result, InsertResult::Error(_)));
     }
 
     #[test]
     fn test_update_existing_row() {
         let mut table = WriteableTable::new();
 
         // Update row 0 (foo -> updated)
-        let row = json!([0, "updated_name", "updated_lastname"]);
-        let result = table.update(0, &row);
+        let row = json!(["updated_name", "updated_lastname"]);
+        let result = table.update("0".to_string(), row);
 
-        assert!(matches!(result, UpdateResult::Success));
+        assert!(matches!(result, UpdateResult::Ok));
 
         // Verify the update
         let response = table.generate(ExtensionPluginRequest::default());
@@ -263,18 +257,18 @@ mod tests {
 
         // Invalid row (not an array)
         let row = json!({"name": "test"});
-        let result = table.update(0, &row);
+        let result = table.update("0".to_string(), row);
 
-        assert!(matches!(result, UpdateResult::Err(_)));
+        assert!(matches!(result, UpdateResult::Error(_)));
     }
 
     #[test]
     fn test_delete_existing_row() {
         let mut table = WriteableTable::new();
 
         // Delete row 0
-        let result = table.delete(0);
-        assert!(matches!(result, DeleteResult::Success));
+        let result = table.delete("0".to_string());
+        assert!(matches!(result, DeleteResult::Ok));
 
         // Verify deletion
         let response = table.generate(ExtensionPluginRequest::default());
@@ -287,18 +281,18 @@ mod tests {
         let mut table = WriteableTable::new();
 
         // Try to delete non-existent row
-        let result = table.delete(999);
+        let result = table.delete("999".to_string());
 
-        assert!(matches!(result, DeleteResult::Err(_)));
+        assert!(matches!(result, DeleteResult::NotFound));
     }
 
     #[test]
     fn test_full_crud_workflow() {
         let mut table = WriteableTable::new();
 
         // Create
-        let row = json!([null, "new_user", "new_lastname"]);
-        let InsertResult::Success(new_rowid) = table.insert(true, &row) else {
+        let row = json!(["new_user", "new_lastname"]);
+        let InsertResult::Ok(new_rowid) = table.insert(row) else {
             panic!("Insert failed");
         };
 
@@ -308,14 +302,14 @@ mod tests {
         assert_eq!(rows.len(), 4);
 
         // Update
-        let updated = json!([new_rowid, "modified", "user"]);
+        let updated = json!(["modified", "user"]);
         assert!(matches!(
-            table.update(new_rowid, &updated),
-            UpdateResult::Success
+            table.update(new_rowid.clone(), updated),
+            UpdateResult::Ok
         ));
 
         // Delete
-        assert!(matches!(table.delete(new_rowid), DeleteResult::Success));
+        assert!(matches!(table.delete(new_rowid), DeleteResult::Ok));
 
         // Verify final state
         let response = table.generate(ExtensionPluginRequest::default());

osquery-rust/src/client/mod.rs

@@ -0,0 +1,17 @@
+//! Client module for osquery communication
+//!
+//! This module provides client implementations for communicating with osquery daemon.
+//! The main components are:
+//!
+//! - `trait_def`: Core trait definitions for client communication
+//! - `thrift_client`: Thrift-based client implementation
+
+pub mod thrift_client;
+pub mod trait_def;
+
+// Re-export public items for compatibility
+pub use thrift_client::{Client, ThriftClient};
+pub use trait_def::OsqueryClient;
+
+#[cfg(test)]
+pub use trait_def::MockOsqueryClient;